diff --git a/docs/changelog/135306.yaml b/docs/changelog/135306.yaml new file mode 100644 index 0000000000000..bfc422d309e6a --- /dev/null +++ b/docs/changelog/135306.yaml @@ -0,0 +1,5 @@ +pr: 135306 +summary: Add support for extended search usage telemetry +area: Relevance +type: enhancement +issues: [] diff --git a/server/src/main/java/org/elasticsearch/action/ActionModule.java b/server/src/main/java/org/elasticsearch/action/ActionModule.java index 9173c9f120964..5d4acd0a3f691 100644 --- a/server/src/main/java/org/elasticsearch/action/ActionModule.java +++ b/server/src/main/java/org/elasticsearch/action/ActionModule.java @@ -69,6 +69,8 @@ import org.elasticsearch.action.admin.cluster.snapshots.status.TransportSnapshotsStatusAction; import org.elasticsearch.action.admin.cluster.state.ClusterStateAction; import org.elasticsearch.action.admin.cluster.state.TransportClusterStateAction; +import org.elasticsearch.action.admin.cluster.stats.ExtendedSearchUsageLongCounter; +import org.elasticsearch.action.admin.cluster.stats.ExtendedSearchUsageMetric; import org.elasticsearch.action.admin.cluster.stats.TransportClusterStatsAction; import org.elasticsearch.action.admin.cluster.storedscripts.GetScriptContextAction; import org.elasticsearch.action.admin.cluster.storedscripts.GetScriptLanguageAction; @@ -1078,4 +1080,14 @@ public RestController getRestController() { public ReservedClusterStateService getReservedClusterStateService() { return reservedClusterStateService; } + + public List getNamedWriteables() { + return List.of( + new NamedWriteableRegistry.Entry( + ExtendedSearchUsageMetric.class, + ExtendedSearchUsageLongCounter.NAME, + ExtendedSearchUsageLongCounter::new + ) + ); + } } diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageLongCounter.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageLongCounter.java new file mode 100644 index 0000000000000..d2968cb78d62c --- /dev/null +++ b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageLongCounter.java @@ -0,0 +1,81 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.action.admin.cluster.stats; + +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Collections; +import java.util.Map; +import java.util.Objects; + +/** + * An {@link ExtendedSearchUsageMetric} implementation that holds a map of values to counts. + */ +public class ExtendedSearchUsageLongCounter implements ExtendedSearchUsageMetric { + + public static final String NAME = "extended_search_usage_long_counter"; + + private final Map values; + + public ExtendedSearchUsageLongCounter(Map values) { + this.values = values; + } + + public ExtendedSearchUsageLongCounter(StreamInput in) throws IOException { + this.values = in.readMap(StreamInput::readString, StreamInput::readLong); + } + + public Map getValues() { + return Collections.unmodifiableMap(values); + } + + @Override + public ExtendedSearchUsageLongCounter merge(ExtendedSearchUsageMetric other) { + assert other instanceof ExtendedSearchUsageLongCounter; + ExtendedSearchUsageLongCounter otherLongCounter = (ExtendedSearchUsageLongCounter) other; + Map values = new java.util.HashMap<>(this.values); + otherLongCounter.getValues().forEach((key, otherValue) -> { values.merge(key, otherValue, Long::sum); }); + return new ExtendedSearchUsageLongCounter(values); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeMap(values, StreamOutput::writeString, StreamOutput::writeLong); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + for (String key : values.keySet()) { + builder.field(key, values.get(key)); + } + return builder; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + ExtendedSearchUsageLongCounter that = (ExtendedSearchUsageLongCounter) o; + return Objects.equals(values, that.values); + } + + @Override + public int hashCode() { + return values.hashCode(); + } + + @Override + public String getWriteableName() { + return NAME; + } +} diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageMetric.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageMetric.java new file mode 100644 index 0000000000000..354dd306b08dd --- /dev/null +++ b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageMetric.java @@ -0,0 +1,26 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.action.admin.cluster.stats; + +import org.elasticsearch.common.io.stream.NamedWriteable; +import org.elasticsearch.xcontent.ToXContentFragment; + +/** + * Represents a metric colleged as part of {@link ExtendedSearchUsageStats}. + */ +public interface ExtendedSearchUsageMetric> extends NamedWriteable, ToXContentFragment { + + /** + * Merges two equivalent metrics together for statistical reporting. + * @param other Another {@link ExtendedSearchUsageMetric}. + * @return ExtendedSearchUsageMetric The merged metric. + */ + T merge(ExtendedSearchUsageMetric other); +} diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageStats.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageStats.java new file mode 100644 index 0000000000000..d39b46bea3049 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageStats.java @@ -0,0 +1,116 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.action.admin.cluster.stats; + +import org.elasticsearch.common.Strings; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.io.stream.Writeable; +import org.elasticsearch.usage.SearchUsage; +import org.elasticsearch.xcontent.ToXContent; +import org.elasticsearch.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +/** + * Provides extended statistics for {@link SearchUsage} beyond the basic counts provided in {@link SearchUsageStats}. + */ +public class ExtendedSearchUsageStats implements Writeable, ToXContent { + + /** + * A map of categories to extended data. Categories correspond to a high-level search usage statistic, + * e.g. `queries`, `rescorers`, `sections`, `retrievers`. + * + * Extended data is further segmented by name, e.g., collecting specific statistics for certain retrievers only. + */ + private final Map>> categorizedExtendedData; + + public static final ExtendedSearchUsageStats EMPTY = new ExtendedSearchUsageStats(); + + public ExtendedSearchUsageStats() { + this.categorizedExtendedData = new HashMap<>(); + } + + public ExtendedSearchUsageStats(Map>> categorizedExtendedData) { + this.categorizedExtendedData = categorizedExtendedData; + } + + public ExtendedSearchUsageStats(StreamInput in) throws IOException { + this.categorizedExtendedData = in.readMap( + StreamInput::readString, + i -> i.readMap(StreamInput::readString, p -> p.readNamedWriteable(ExtendedSearchUsageMetric.class)) + ); + } + + public Map>> getCategorizedExtendedData() { + return Collections.unmodifiableMap(categorizedExtendedData); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeMap( + categorizedExtendedData, + StreamOutput::writeString, + (o, v) -> o.writeMap(v, StreamOutput::writeString, (p, q) -> out.writeNamedWriteable(q)) + ); + } + + public void merge(ExtendedSearchUsageStats other) { + other.categorizedExtendedData.forEach((key, otherMap) -> { + categorizedExtendedData.merge(key, otherMap, (existingMap, newMap) -> { + Map> mergedMap = new HashMap<>(existingMap); + newMap.forEach( + (innerKey, innerValue) -> { mergedMap.merge(innerKey, innerValue, (existing, incoming) -> (existing).merge(incoming)); } + ); + return mergedMap; + }); + }); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + + builder.startObject(); + for (String category : categorizedExtendedData.keySet()) { + builder.startObject(category); + Map> names = categorizedExtendedData.get(category); + for (String name : names.keySet()) { + builder.startObject(name); + names.get(name).toXContent(builder, params); + builder.endObject(); + } + builder.endObject(); + } + builder.endObject(); + return builder; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + ExtendedSearchUsageStats that = (ExtendedSearchUsageStats) o; + return Objects.equals(categorizedExtendedData, that.categorizedExtendedData); + } + + @Override + public int hashCode() { + return Objects.hash(categorizedExtendedData); + } + + @Override + public String toString() { + return Strings.toString(this, true, true); + } +} diff --git a/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStats.java b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStats.java index a6e80b5efd08c..081a21e3742ea 100644 --- a/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStats.java +++ b/server/src/main/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStats.java @@ -9,6 +9,7 @@ package org.elasticsearch.action.admin.cluster.stats; +import org.elasticsearch.TransportVersion; import org.elasticsearch.common.Strings; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; @@ -31,11 +32,15 @@ * accumulate stats for the entire cluster and return them as part of the {@link ClusterStatsResponse}. */ public final class SearchUsageStats implements Writeable, ToXContentFragment { + + static final TransportVersion EXTENDED_SEARCH_USAGE_TELEMETRY = TransportVersion.fromName("extended_search_usage_telemetry"); + private long totalSearchCount; private final Map queries; private final Map rescorers; private final Map sections; private final Map retrievers; + private final ExtendedSearchUsageStats extendedSearchUsageStats; /** * Creates a new empty stats instance, that will get additional stats added through {@link #add(SearchUsageStats)} @@ -46,6 +51,7 @@ public SearchUsageStats() { this.sections = new HashMap<>(); this.rescorers = new HashMap<>(); this.retrievers = new HashMap<>(); + this.extendedSearchUsageStats = ExtendedSearchUsageStats.EMPTY; } /** @@ -57,6 +63,7 @@ public SearchUsageStats( Map rescorers, Map sections, Map retrievers, + ExtendedSearchUsageStats extendedSearchUsageStats, long totalSearchCount ) { this.totalSearchCount = totalSearchCount; @@ -64,6 +71,7 @@ public SearchUsageStats( this.sections = sections; this.rescorers = rescorers; this.retrievers = retrievers; + this.extendedSearchUsageStats = extendedSearchUsageStats; } public SearchUsageStats(StreamInput in) throws IOException { @@ -72,6 +80,9 @@ public SearchUsageStats(StreamInput in) throws IOException { this.totalSearchCount = in.readVLong(); this.rescorers = in.getTransportVersion().onOrAfter(V_8_12_0) ? in.readMap(StreamInput::readLong) : Map.of(); this.retrievers = in.getTransportVersion().onOrAfter(V_8_16_0) ? in.readMap(StreamInput::readLong) : Map.of(); + this.extendedSearchUsageStats = in.getTransportVersion().supports(EXTENDED_SEARCH_USAGE_TELEMETRY) + ? new ExtendedSearchUsageStats(in) + : ExtendedSearchUsageStats.EMPTY; } @Override @@ -86,6 +97,9 @@ public void writeTo(StreamOutput out) throws IOException { if (out.getTransportVersion().onOrAfter(V_8_16_0)) { out.writeMap(retrievers, StreamOutput::writeLong); } + if (out.getTransportVersion().supports(EXTENDED_SEARCH_USAGE_TELEMETRY)) { + extendedSearchUsageStats.writeTo(out); + } } /** @@ -96,6 +110,7 @@ public void add(SearchUsageStats stats) { stats.rescorers.forEach((rescorer, count) -> rescorers.merge(rescorer, count, Long::sum)); stats.sections.forEach((query, count) -> sections.merge(query, count, Long::sum)); stats.retrievers.forEach((query, count) -> retrievers.merge(query, count, Long::sum)); + this.extendedSearchUsageStats.merge(stats.extendedSearchUsageStats); this.totalSearchCount += stats.totalSearchCount; } @@ -112,6 +127,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.map(sections); builder.field("retrievers"); builder.map(retrievers); + builder.field("extended"); + extendedSearchUsageStats.toXContent(builder, params); } builder.endObject(); return builder; @@ -133,6 +150,10 @@ public Map getRetrieversUsage() { return Collections.unmodifiableMap(retrievers); } + public ExtendedSearchUsageStats getExtendedSearchUsage() { + return extendedSearchUsageStats; + } + public long getTotalSearchCount() { return totalSearchCount; } @@ -150,12 +171,13 @@ public boolean equals(Object o) { && queries.equals(that.queries) && rescorers.equals(that.rescorers) && sections.equals(that.sections) - && retrievers.equals(that.retrievers); + && retrievers.equals(that.retrievers) + && extendedSearchUsageStats.equals(that.extendedSearchUsageStats); } @Override public int hashCode() { - return Objects.hash(totalSearchCount, queries, rescorers, sections, retrievers); + return Objects.hash(totalSearchCount, queries, rescorers, sections, retrievers, extendedSearchUsageStats); } @Override diff --git a/server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java b/server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java index 690f3155971ca..8c0384abd7197 100644 --- a/server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java +++ b/server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestClusterStatsAction.java @@ -34,7 +34,8 @@ public class RestClusterStatsAction extends BaseRestHandler { "verbose-dense-vector-mapping-stats", "ccs-stats", "retrievers-usage-stats", - "esql-stats" + "esql-stats", + "extended-search-usage-stats" ); private static final Set SUPPORTED_QUERY_PARAMETERS = Set.of("include_remotes", "nodeId", REST_TIMEOUT_PARAM); diff --git a/server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java b/server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java index e9b414aec86f2..f35ecb0b369fb 100644 --- a/server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java +++ b/server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java @@ -59,6 +59,7 @@ private SearchCapabilities() {} private static final String KNN_FILTER_ON_NESTED_FIELDS_CAPABILITY = "knn_filter_on_nested_fields"; private static final String BUCKET_SCRIPT_PARENT_MULTI_BUCKET_ERROR = "bucket_script_parent_multi_bucket_error"; private static final String EXCLUDE_SOURCE_VECTORS_SETTING = "exclude_source_vectors_setting"; + private static final String CLUSTER_STATS_EXTENDED_USAGE = "extended-search-usage-stats"; public static final Set CAPABILITIES; static { @@ -88,6 +89,7 @@ private SearchCapabilities() {} capabilities.add(KNN_FILTER_ON_NESTED_FIELDS_CAPABILITY); capabilities.add(BUCKET_SCRIPT_PARENT_MULTI_BUCKET_ERROR); capabilities.add(EXCLUDE_SOURCE_VECTORS_SETTING); + capabilities.add(CLUSTER_STATS_EXTENDED_USAGE); CAPABILITIES = Set.copyOf(capabilities); } } diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RescorerRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RescorerRetrieverBuilder.java index c2bda5587e1bb..aa016d7909c4a 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/RescorerRetrieverBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/RescorerRetrieverBuilder.java @@ -46,7 +46,7 @@ public final class RescorerRetrieverBuilder extends CompoundRetrieverBuilder { RetrieverBuilder innerRetriever = parser.namedObject(RetrieverBuilder.class, n, context); - context.trackRetrieverUsage(innerRetriever.getName()); + context.trackRetrieverUsage(innerRetriever); return innerRetriever; }, RETRIEVER_FIELD); PARSER.declareField(constructorArg(), (parser, context) -> { diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java index 0a62e9f968e4f..bfdd9a341fb53 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java @@ -33,6 +33,7 @@ import java.util.List; import java.util.Locale; import java.util.Objects; +import java.util.Set; /** * A retriever represents an API element that returns an ordered list of top @@ -132,7 +133,7 @@ protected static RetrieverBuilder parseInnerRetrieverBuilder(XContentParser pars throw new ParsingException(new XContentLocation(nonfe.getLineNumber(), nonfe.getColumnNumber()), message, nonfe); } - context.trackRetrieverUsage(retrieverName); + context.trackRetrieverUsage(retrieverBuilder); if (parser.currentToken() != XContentParser.Token.END_OBJECT) { throw new ParsingException( @@ -243,6 +244,16 @@ public ActionRequestValidationException validate( return validationException; } + /** + * @return Additional fields associated with this retriever that we want to track in + * {@link org.elasticsearch.action.admin.cluster.stats.SearchUsageStats}. + * + * Individual retrievers should override this to add their own specific custom fields. + */ + public Set getExtendedUsageFields() { + return Set.of(); + } + // ---- FOR TESTING XCONTENT PARSING ---- public abstract String getName(); diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverParserContext.java b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverParserContext.java index bdf3f8a194546..61f019f9f6e6d 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverParserContext.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverParserContext.java @@ -37,8 +37,9 @@ public void trackRescorerUsage(String name) { searchUsage.trackRescorerUsage(name); } - public void trackRetrieverUsage(String name) { - searchUsage.trackRetrieverUsage(name); + public void trackRetrieverUsage(RetrieverBuilder retriever) { + searchUsage.trackRetrieverUsage(retriever.getName()); + searchUsage.trackRetrieverExtendedDataUsage(retriever.getName(), retriever.getExtendedUsageFields()); } public boolean clusterSupportsFeature(NodeFeature nodeFeature) { diff --git a/server/src/main/java/org/elasticsearch/usage/SearchUsage.java b/server/src/main/java/org/elasticsearch/usage/SearchUsage.java index e35594fb161ac..8e5ce86fb05f3 100644 --- a/server/src/main/java/org/elasticsearch/usage/SearchUsage.java +++ b/server/src/main/java/org/elasticsearch/usage/SearchUsage.java @@ -10,17 +10,23 @@ package org.elasticsearch.usage; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; +import java.util.Map; import java.util.Set; /** * Holds usage statistics for an incoming search request */ public final class SearchUsage { + + public final String RETRIEVERS_NAME = "retrievers"; + private final Set queries = new HashSet<>(); private final Set rescorers = new HashSet<>(); private final Set sections = new HashSet<>(); private final Set retrievers = new HashSet<>(); + private final ExtendedUsageTracker extendedUsage = new ExtendedUsageTracker(); /** * Track the usage of the provided query @@ -43,11 +49,20 @@ public void trackRescorerUsage(String name) { rescorers.add(name); } - /** - * Track retrieve usage - */ public void trackRetrieverUsage(String retriever) { retrievers.add(retriever); + extendedUsage.initialize(RETRIEVERS_NAME, retriever); + } + + /** + * Track the usage of extended data for a specific category + */ + private void trackExtendedDataUsage(String category, String name, Set values) { + extendedUsage.track(category, name, values); + } + + public void trackRetrieverExtendedDataUsage(String name, Set values) { + trackExtendedDataUsage(RETRIEVERS_NAME, name, values); } /** @@ -77,4 +92,43 @@ public Set getSectionsUsage() { public Set getRetrieverUsage() { return Collections.unmodifiableSet(retrievers); } + + /** + * Returns the extended data that has been tracked for the search request + */ + public Map>> getExtendedDataUsage() { + return extendedUsage.getUsage(); + } + + private static final class ExtendedUsageTracker { + + /** + * A map of categories to extended data. Categories correspond to a high-level search usage statistic, + * e.g. `queries`, `rescorers`, `sections`, `retrievers`. + * + * Extended data is further segmented by name, for example collecting specific statistics for certain retrievers only. + * Finally we keep track of the set of values we are tracking for each category and name. + */ + private final Map>> categoriesToExtendedUsage = new HashMap<>(); + + public void initialize(String category, String name) { + categoriesToExtendedUsage.computeIfAbsent(category, k -> new HashMap<>()).computeIfAbsent(name, k -> new HashSet<>()); + } + + public void track(String category, String name, String value) { + categoriesToExtendedUsage.computeIfAbsent(category, k -> new HashMap<>()) + .computeIfAbsent(name, k -> new HashSet<>()) + .add(value); + } + + public void track(String category, String name, Set values) { + categoriesToExtendedUsage.computeIfAbsent(category, k -> new HashMap<>()) + .computeIfAbsent(name, k -> new HashSet<>()) + .addAll(values); + } + + public Map>> getUsage() { + return Collections.unmodifiableMap(categoriesToExtendedUsage); + } + } } diff --git a/server/src/main/java/org/elasticsearch/usage/SearchUsageHolder.java b/server/src/main/java/org/elasticsearch/usage/SearchUsageHolder.java index ef802723cf164..68b9b55bc6d1c 100644 --- a/server/src/main/java/org/elasticsearch/usage/SearchUsageHolder.java +++ b/server/src/main/java/org/elasticsearch/usage/SearchUsageHolder.java @@ -9,13 +9,19 @@ package org.elasticsearch.usage; +import org.elasticsearch.action.admin.cluster.stats.ExtendedSearchUsageLongCounter; +import org.elasticsearch.action.admin.cluster.stats.ExtendedSearchUsageMetric; +import org.elasticsearch.action.admin.cluster.stats.ExtendedSearchUsageStats; import org.elasticsearch.action.admin.cluster.stats.SearchUsageStats; import org.elasticsearch.common.util.Maps; import java.util.Collections; +import java.util.HashMap; import java.util.Map; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.LongAdder; +import java.util.function.Function; /** * Service responsible for holding search usage statistics, like the number of used search sections and queries. @@ -28,6 +34,7 @@ public final class SearchUsageHolder { private final Map rescorersUsage = new ConcurrentHashMap<>(); private final Map sectionsUsage = new ConcurrentHashMap<>(); private final Map retrieversUsage = new ConcurrentHashMap<>(); + private final Map>> extendedSearchUsage = new ConcurrentHashMap<>(); SearchUsageHolder() {} @@ -48,6 +55,7 @@ public void updateUsage(SearchUsage searchUsage) { for (String retriever : searchUsage.getRetrieverUsage()) { retrieversUsage.computeIfAbsent(retriever, q -> new LongAdder()).increment(); } + updateExtendedUsage(searchUsage.getExtendedDataUsage()); } /** @@ -62,12 +70,45 @@ public SearchUsageStats getSearchUsageStats() { rescorersUsage.forEach((query, adder) -> rescorersUsageMap.put(query, adder.longValue())); Map retrieversUsageMap = Maps.newMapWithExpectedSize(retrieversUsage.size()); retrieversUsage.forEach((retriever, adder) -> retrieversUsageMap.put(retriever, adder.longValue())); + ExtendedSearchUsageStats extendedSearchUsageStats = new ExtendedSearchUsageStats(getExtendedSearchUsage()); + return new SearchUsageStats( Collections.unmodifiableMap(queriesUsageMap), Collections.unmodifiableMap(rescorersUsageMap), Collections.unmodifiableMap(sectionsUsageMap), Collections.unmodifiableMap(retrieversUsageMap), + extendedSearchUsageStats, totalSearchCount.longValue() ); } + + private Map>> getExtendedSearchUsage() { + return unmodifiableMap( + extendedSearchUsage, + nameMap -> unmodifiableMap( + nameMap, + valueMap -> new ExtendedSearchUsageLongCounter(unmodifiableMap(valueMap, LongAdder::longValue)) + ) + ); + } + + private void updateExtendedUsage(Map>> extendedDataUsage) { + extendedDataUsage.forEach( + (category, nameMap) -> nameMap.forEach( + (name, valueSet) -> valueSet.forEach( + value -> extendedSearchUsage.computeIfAbsent(category, k -> new ConcurrentHashMap<>()) + .computeIfAbsent(name, k -> new ConcurrentHashMap<>()) + .computeIfAbsent(value, k -> new LongAdder()) + .increment() + ) + ) + ); + } + + private static Map unmodifiableMap(Map in, Function valueMapper) { + Map map = new HashMap<>(in.size()); + in.forEach((k, v) -> map.put(k, valueMapper.apply(v))); + return Collections.unmodifiableMap(map); + } + } diff --git a/server/src/main/resources/transport/definitions/referable/extended_search_usage_telemetry.csv b/server/src/main/resources/transport/definitions/referable/extended_search_usage_telemetry.csv new file mode 100644 index 0000000000000..4b012195cac35 --- /dev/null +++ b/server/src/main/resources/transport/definitions/referable/extended_search_usage_telemetry.csv @@ -0,0 +1 @@ +9177000 diff --git a/server/src/main/resources/transport/upper_bounds/9.2.csv b/server/src/main/resources/transport/upper_bounds/9.2.csv index 78180d915cd67..980b8e87c6329 100644 --- a/server/src/main/resources/transport/upper_bounds/9.2.csv +++ b/server/src/main/resources/transport/upper_bounds/9.2.csv @@ -1 +1 @@ -roles_security_stats,9176000 +extended_search_usage_telemetry,9177000 diff --git a/server/src/test/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageStatsTests.java b/server/src/test/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageStatsTests.java new file mode 100644 index 0000000000000..24d3ffdf411c2 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageStatsTests.java @@ -0,0 +1,92 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +package org.elasticsearch.action.admin.cluster.stats; + +import org.elasticsearch.common.io.stream.NamedWriteableRegistry; +import org.elasticsearch.common.io.stream.Writeable.Reader; +import org.elasticsearch.test.AbstractWireSerializingTestCase; + +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.action.admin.cluster.stats.ExtendedSearchUsageStats.EMPTY; + +public class ExtendedSearchUsageStatsTests extends AbstractWireSerializingTestCase { + + @Override + protected Reader instanceReader() { + return ExtendedSearchUsageStats::new; + } + + @Override + protected NamedWriteableRegistry getNamedWriteableRegistry() { + return new NamedWriteableRegistry( + List.of( + new NamedWriteableRegistry.Entry( + ExtendedSearchUsageMetric.class, + ExtendedSearchUsageLongCounter.NAME, + ExtendedSearchUsageLongCounter::new + ) + ) + ); + } + + public static ExtendedSearchUsageStats randomExtendedSearchUsage() { + return randomExtendedSearchUsage(randomBoolean()); + } + + public static ExtendedSearchUsageStats randomExtendedSearchUsage(boolean empty) { + if (empty) { + return EMPTY; + } + Map>> categoriesToExtendedData = new HashMap<>(); + + // TODO: Gate this behind a randomBoolean() in the future when we have other categories to add. + categoriesToExtendedData.put("retrievers", randomExtendedRetrieversData()); + + return new ExtendedSearchUsageStats(categoriesToExtendedData); + } + + private static Map> randomExtendedRetrieversData() { + Map> retrieversData = new HashMap<>(); + + // TODO: Gate this behind a randomBoolean() in the future when we have other values to add. + ExtendedSearchUsageMetric values = new ExtendedSearchUsageLongCounter(Map.of("chunk_rescorer", randomLongBetween(1, 10))); + retrieversData.put("text_similarity_reranker", values); + + return retrieversData; + } + + @Override + protected ExtendedSearchUsageStats createTestInstance() { + return randomExtendedSearchUsage(); + } + + @Override + protected ExtendedSearchUsageStats mutateInstance(ExtendedSearchUsageStats instance) throws IOException { + Map>> current = instance.getCategorizedExtendedData(); + Map>> modified = new HashMap<>(); + if (current.isEmpty()) { + modified.put( + "retrievers", + Map.of("text_similarity_reranker", new ExtendedSearchUsageLongCounter(Map.of("chunk_rescorer", randomLongBetween(1, 10)))) + ); + } else if (randomBoolean()) { + modified.put( + "retrivers", + Map.of("text_similarity_reranker", new ExtendedSearchUsageLongCounter(Map.of("chunk_rescorer", randomLongBetween(11, 20)))) + ); + } + return new ExtendedSearchUsageStats(modified); + } + +} diff --git a/server/src/test/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStatsTests.java b/server/src/test/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStatsTests.java index 46b757407e6a9..5eec89d6e4402 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStatsTests.java +++ b/server/src/test/java/org/elasticsearch/action/admin/cluster/stats/SearchUsageStatsTests.java @@ -12,6 +12,7 @@ import org.elasticsearch.TransportVersion; import org.elasticsearch.TransportVersions; import org.elasticsearch.common.Strings; +import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.io.stream.Writeable.Reader; import org.elasticsearch.test.AbstractWireSerializingTestCase; import org.elasticsearch.test.TransportVersionUtils; @@ -21,8 +22,25 @@ import java.util.List; import java.util.Map; +import static org.elasticsearch.action.admin.cluster.stats.ExtendedSearchUsageStats.EMPTY; +import static org.elasticsearch.action.admin.cluster.stats.ExtendedSearchUsageStatsTests.randomExtendedSearchUsage; +import static org.elasticsearch.action.admin.cluster.stats.SearchUsageStats.EXTENDED_SEARCH_USAGE_TELEMETRY; + public class SearchUsageStatsTests extends AbstractWireSerializingTestCase { + @Override + protected NamedWriteableRegistry getNamedWriteableRegistry() { + return new NamedWriteableRegistry( + List.of( + new NamedWriteableRegistry.Entry( + ExtendedSearchUsageMetric.class, + ExtendedSearchUsageLongCounter.NAME, + ExtendedSearchUsageLongCounter::new + ) + ) + ); + } + private static final List QUERY_TYPES = List.of( "match", "bool", @@ -97,6 +115,7 @@ protected SearchUsageStats createTestInstance() { randomRescorerUsage(randomIntBetween(0, RESCORER_TYPES.size())), randomSectionsUsage(randomIntBetween(0, SECTIONS.size())), randomRetrieversUsage(randomIntBetween(0, RETRIEVERS.size())), + randomExtendedSearchUsage(), randomLongBetween(10, Long.MAX_VALUE) ); } @@ -110,6 +129,7 @@ protected SearchUsageStats mutateInstance(SearchUsageStats instance) { instance.getRescorerUsage(), instance.getSectionsUsage(), instance.getRetrieversUsage(), + instance.getExtendedSearchUsage(), instance.getTotalSearchCount() ); case 1 -> new SearchUsageStats( @@ -117,6 +137,7 @@ protected SearchUsageStats mutateInstance(SearchUsageStats instance) { randomValueOtherThan(instance.getRescorerUsage(), () -> randomRescorerUsage(randomIntBetween(0, RESCORER_TYPES.size()))), instance.getSectionsUsage(), instance.getRetrieversUsage(), + instance.getExtendedSearchUsage(), instance.getTotalSearchCount() ); case 2 -> new SearchUsageStats( @@ -124,6 +145,7 @@ protected SearchUsageStats mutateInstance(SearchUsageStats instance) { instance.getRescorerUsage(), randomValueOtherThan(instance.getSectionsUsage(), () -> randomSectionsUsage(randomIntBetween(0, SECTIONS.size()))), instance.getRetrieversUsage(), + instance.getExtendedSearchUsage(), instance.getTotalSearchCount() ); case 3 -> new SearchUsageStats( @@ -131,6 +153,7 @@ protected SearchUsageStats mutateInstance(SearchUsageStats instance) { instance.getRescorerUsage(), instance.getSectionsUsage(), randomValueOtherThan(instance.getRetrieversUsage(), () -> randomSectionsUsage(randomIntBetween(0, SECTIONS.size()))), + instance.getExtendedSearchUsage(), instance.getTotalSearchCount() ); case 4 -> new SearchUsageStats( @@ -138,6 +161,7 @@ protected SearchUsageStats mutateInstance(SearchUsageStats instance) { instance.getRescorerUsage(), instance.getSectionsUsage(), instance.getRetrieversUsage(), + instance.getExtendedSearchUsage(), randomValueOtherThan(instance.getTotalSearchCount(), () -> randomLongBetween(10, Long.MAX_VALUE)) ); default -> throw new IllegalStateException("Unexpected value: " + i); @@ -150,14 +174,33 @@ public void testAdd() { assertEquals(Map.of(), searchUsageStats.getRescorerUsage()); assertEquals(Map.of(), searchUsageStats.getSectionsUsage()); assertEquals(0, searchUsageStats.getTotalSearchCount()); + assertEquals(EMPTY, searchUsageStats.getExtendedSearchUsage()); + + ExtendedSearchUsageStats extendedSearchUsageStats = new ExtendedSearchUsageStats( + Map.of("retrievers", Map.of("text_similarity_reranker", new ExtendedSearchUsageLongCounter(Map.of("chunk_rescorer", 10L)))) + ); + ExtendedSearchUsageStats anotherExtendedSearchUsageStats = new ExtendedSearchUsageStats( + Map.of("retrievers", Map.of("text_similarity_reranker", new ExtendedSearchUsageLongCounter(Map.of("chunk_rescorer", 5L)))) + ); + ExtendedSearchUsageStats combinedExtendedSearchUsageStats = new ExtendedSearchUsageStats( + Map.of("retrievers", Map.of("text_similarity_reranker", new ExtendedSearchUsageLongCounter(Map.of("chunk_rescorer", 15L)))) + ); searchUsageStats.add( - new SearchUsageStats(Map.of("match", 10L), Map.of("query", 5L), Map.of("query", 10L), Map.of("knn", 10L), 10L) + new SearchUsageStats( + Map.of("match", 10L), + Map.of("query", 5L), + Map.of("query", 10L), + Map.of("knn", 10L), + extendedSearchUsageStats, + 10L + ) ); assertEquals(Map.of("match", 10L), searchUsageStats.getQueryUsage()); assertEquals(Map.of("query", 10L), searchUsageStats.getSectionsUsage()); assertEquals(Map.of("query", 5L), searchUsageStats.getRescorerUsage()); assertEquals(10L, searchUsageStats.getTotalSearchCount()); + assertEquals(extendedSearchUsageStats, searchUsageStats.getExtendedSearchUsage()); searchUsageStats.add( new SearchUsageStats( @@ -165,6 +208,7 @@ public void testAdd() { Map.of("query", 5L, "learning_to_rank", 2L), Map.of("query", 10L, "knn", 1L), Map.of("knn", 10L, "rrf", 2L), + anotherExtendedSearchUsageStats, 10L ) ); @@ -172,6 +216,7 @@ public void testAdd() { assertEquals(Map.of("query", 20L, "knn", 1L), searchUsageStats.getSectionsUsage()); assertEquals(Map.of("query", 10L, "learning_to_rank", 2L), searchUsageStats.getRescorerUsage()); assertEquals(Map.of("knn", 20L, "rrf", 2L), searchUsageStats.getRetrieversUsage()); + assertEquals(combinedExtendedSearchUsageStats, searchUsageStats.getExtendedSearchUsage()); assertEquals(20L, searchUsageStats.getTotalSearchCount()); } @@ -181,11 +226,12 @@ public void testToXContent() throws IOException { Map.of("query", 2L), Map.of("query", 10L), Map.of("knn", 10L), + EMPTY, 10L ); assertEquals( "{\"search\":{\"total\":10,\"queries\":{\"term\":1},\"rescorers\":{\"query\":2}," - + "\"sections\":{\"query\":10},\"retrievers\":{\"knn\":10}}}", + + "\"sections\":{\"query\":10},\"retrievers\":{\"knn\":10},\"extended\":{}}}", Strings.toString(searchUsageStats) ); } @@ -200,6 +246,7 @@ public void testSerializationBWC() throws IOException { version.onOrAfter(TransportVersions.V_8_12_0) ? randomRescorerUsage(RESCORER_TYPES.size()) : Map.of(), randomSectionsUsage(SECTIONS.size()), version.onOrAfter(TransportVersions.V_8_16_0) ? randomRetrieversUsage(RETRIEVERS.size()) : Map.of(), + version.supports(EXTENDED_SEARCH_USAGE_TELEMETRY) ? randomExtendedSearchUsage() : EMPTY, randomLongBetween(0, Long.MAX_VALUE) ); assertSerialization(testInstance, version); diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java index b724cd133f0c5..491e937337db2 100644 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java @@ -66,7 +66,7 @@ public final class QueryRuleRetrieverBuilder extends CompoundRetrieverBuilder p.map(), MATCH_CRITERIA_FIELD); PARSER.declareNamedObject(constructorArg(), (p, c, n) -> { RetrieverBuilder innerRetriever = p.namedObject(RetrieverBuilder.class, n, c); - c.trackRetrieverUsage(innerRetriever.getName()); + c.trackRetrieverUsage(innerRetriever); return innerRetriever; }, RETRIEVER_FIELD); PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java index 2f664f6c4e103..d45a93c1444e3 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java @@ -25,9 +25,11 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import static org.elasticsearch.search.rank.RankBuilder.DEFAULT_RANK_WINDOW_SIZE; import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; @@ -87,7 +89,7 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder static { PARSER.declareNamedObject(constructorArg(), (p, c, n) -> { RetrieverBuilder innerRetriever = p.namedObject(RetrieverBuilder.class, n, c); - c.trackRetrieverUsage(innerRetriever.getName()); + c.trackRetrieverUsage(innerRetriever); return innerRetriever; }, RETRIEVER_FIELD); PARSER.declareString(optionalConstructorArg(), INFERENCE_ID_FIELD); @@ -223,6 +225,17 @@ protected SearchSourceBuilder finalizeSourceBuilder(SearchSourceBuilder sourceBu return sourceBuilder; } + @Override + public Set getExtendedUsageFields() { + Set extendedFields = new HashSet<>(); + + if (chunkScorerConfig != null) { + extendedFields.add(CHUNK_RESCORER_FIELD.getPreferredName()); + } + + return extendedFields; + } + @Override public String getName() { return TextSimilarityRankBuilder.NAME; diff --git a/x-pack/plugin/monitoring/src/test/java/org/elasticsearch/xpack/monitoring/collector/cluster/ClusterStatsMonitoringDocTests.java b/x-pack/plugin/monitoring/src/test/java/org/elasticsearch/xpack/monitoring/collector/cluster/ClusterStatsMonitoringDocTests.java index c1c618412e110..806b82088101d 100644 --- a/x-pack/plugin/monitoring/src/test/java/org/elasticsearch/xpack/monitoring/collector/cluster/ClusterStatsMonitoringDocTests.java +++ b/x-pack/plugin/monitoring/src/test/java/org/elasticsearch/xpack/monitoring/collector/cluster/ClusterStatsMonitoringDocTests.java @@ -592,7 +592,8 @@ public void testToXContent() throws IOException { "queries": {}, "rescorers": {}, "sections": {}, - "retrievers": {} + "retrievers": {}, + "extended": {} }, "dense_vector": { "value_count": 0 diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverComponent.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverComponent.java index 963ba6883e7c9..3fe03237125e1 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverComponent.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverComponent.java @@ -67,7 +67,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws static { PARSER.declareNamedObject(constructorArg(), (p, c, n) -> { RetrieverBuilder innerRetriever = p.namedObject(RetrieverBuilder.class, n, c); - c.trackRetrieverUsage(innerRetriever.getName()); + c.trackRetrieverUsage(innerRetriever); return innerRetriever; }, RETRIEVER_FIELD); PARSER.declareFloat(optionalConstructorArg(), WEIGHT_FIELD); diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverComponent.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverComponent.java index 4946407fb19fb..15dfee43976eb 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverComponent.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverComponent.java @@ -88,7 +88,7 @@ public static RRFRetrieverComponent fromXContent(XContentParser parser, Retrieve parser.nextToken(); String retrieverType = parser.currentName(); retriever = parser.namedObject(RetrieverBuilder.class, retrieverType, context); - context.trackRetrieverUsage(retriever.getName()); + context.trackRetrieverUsage(retriever); parser.nextToken(); } else if (WEIGHT_FIELD.match(fieldName, parser.getDeprecationHandler())) { if (weight != null) { @@ -114,7 +114,7 @@ public static RRFRetrieverComponent fromXContent(XContentParser parser, Retrieve return new RRFRetrieverComponent(retriever, weight); } else { RetrieverBuilder retriever = parser.namedObject(RetrieverBuilder.class, firstFieldName, context); - context.trackRetrieverUsage(retriever.getName()); + context.trackRetrieverUsage(retriever); if (parser.nextToken() != XContentParser.Token.END_OBJECT) { throw new ParsingException(parser.getTokenLocation(), "unknown field [{}] after retriever", parser.currentName()); } diff --git a/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilder.java b/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilder.java index 0e254b4fcd5b3..a876d520e4b6b 100644 --- a/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilder.java +++ b/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilder.java @@ -68,7 +68,7 @@ public final class PinnedRetrieverBuilder extends CompoundRetrieverBuilder SpecifiedDocument.PARSER.apply(p, null), DOCS_FIELD); PARSER.declareNamedObject(constructorArg(), (p, c, n) -> { RetrieverBuilder innerRetriever = p.namedObject(RetrieverBuilder.class, n, c); - c.trackRetrieverUsage(innerRetriever.getName()); + c.trackRetrieverUsage(innerRetriever); return innerRetriever; }, RETRIEVER_FIELD); PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);