Add support for extended search usage telemetry

docs/changelog/135306.yaml

@@ -0,0 +1,5 @@
+pr: 135306
+summary: Add support for extended search usage telemetry
+area: Relevance
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/action/admin/cluster/stats/ExtendedSearchUsageStats.java

@@ -0,0 +1,124 @@
+/*
+ * 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, for example collecting specific statistics for certain retrievers only.
+     *
+     * Finally, we have string:count pairs that track each individual metric we wish to track.
+     */
+    private final Map<String, Map<String, Map<String, Long>>> categoriesToExtendedData;
+
+    public ExtendedSearchUsageStats() {
+        this.categoriesToExtendedData = new HashMap<>();
+    }
+
+    public ExtendedSearchUsageStats(Map<String, Map<String, Map<String, Long>>> categoriesToExtendedData) {
+        this.categoriesToExtendedData = categoriesToExtendedData;
+    }
+
+    public ExtendedSearchUsageStats(StreamInput in) throws IOException {
+        this.categoriesToExtendedData = in.readMap(
+            StreamInput::readString,
+            i -> i.readMap(StreamInput::readString, j -> j.readMap(StreamInput::readString, StreamInput::readLong))
+        );
+    }
+
+    public Map<String, Map<String, Map<String, Long>>> getCategoriesToExtendedData() {
+        return Collections.unmodifiableMap(categoriesToExtendedData);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeMap(
+            categoriesToExtendedData,
+            StreamOutput::writeString,
+            (o, v) -> o.writeMap(v, StreamOutput::writeString, (p, q) -> p.writeMap(q, StreamOutput::writeString, StreamOutput::writeLong))
+        );
+    }
+
+    public void merge(ExtendedSearchUsageStats other) {
+        other.categoriesToExtendedData.forEach((key, otherMap) -> {
+            categoriesToExtendedData.merge(key, otherMap, (existingMap, newMap) -> {
+                Map<String, Map<String, Long>> mergedMap = new HashMap<>(existingMap);
+                newMap.forEach((innerKey, innerValue) -> {
+                    mergedMap.merge(innerKey, innerValue, (existingInnerMap, newInnerMap) -> {
+                        Map<String, Long> mergedInnerMap = new HashMap<>(existingInnerMap);
+                        newInnerMap.forEach(
+                            (propertyKey, propertyValue) -> { mergedInnerMap.merge(propertyKey, propertyValue, Long::sum); }
+                        );
+                        return mergedInnerMap;
+                    });
+                });
+                return mergedMap;
+            });
+        });
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+
+        builder.startObject();
+        for (String category : categoriesToExtendedData.keySet()) {
+            builder.startObject(category);
+            Map<String, Map<String, Long>> names = categoriesToExtendedData.get(category);
+            for (String name : names.keySet()) {
+                builder.startObject(name);
+                for (String property : names.get(name).keySet()) {
+                    builder.field(property, names.get(name).get(property));
+                }
+                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(categoriesToExtendedData, that.categoriesToExtendedData);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(categoriesToExtendedData);
+    }
+
+    @Override
+    public String toString() {
+        return Strings.toString(this, true, true);
+    }
+}

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<String, Long> queries;
     private final Map<String, Long> rescorers;
     private final Map<String, Long> sections;
     private final Map<String, Long> 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 = new ExtendedSearchUsageStats();
     }
 
     /**
@@ -57,13 +63,15 @@ public SearchUsageStats(
         Map<String, Long> rescorers,
         Map<String, Long> sections,
         Map<String, Long> retrievers,
+        ExtendedSearchUsageStats extendedSearchUsageStats,
         long totalSearchCount
     ) {
         this.totalSearchCount = totalSearchCount;
         this.queries = queries;
         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)
+            : new ExtendedSearchUsageStats();
     }
 
     @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<String, Long> 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

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<String> SUPPORTED_QUERY_PARAMETERS = Set.of("include_remotes", "nodeId", REST_TIMEOUT_PARAM);
 

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<String> 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);
     }
 }

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -129,6 +129,7 @@ public final class SearchSourceBuilder implements Writeable, ToXContentObject, R
     public static final ParseField POINT_IN_TIME = new ParseField("pit");
     public static final ParseField RUNTIME_MAPPINGS_FIELD = new ParseField("runtime_mappings");
     public static final ParseField RETRIEVER = new ParseField("retriever");
+    public static final ParseField EXTENDED = new ParseField("extended");
 
     private static final boolean RANK_SUPPORTED = Booleans.parseBoolean(System.getProperty("es.search.rank_supported"), true);
 

server/src/main/java/org/elasticsearch/search/retriever/RescorerRetrieverBuilder.java

@@ -46,7 +46,7 @@ public final class RescorerRetrieverBuilder extends CompoundRetrieverBuilder<Res
     static {
         PARSER.declareNamedObject(constructorArg(), (parser, context, n) -> {
             RetrieverBuilder innerRetriever = parser.namedObject(RetrieverBuilder.class, n, context);
-            context.trackRetrieverUsage(innerRetriever.getName());
+            context.trackRetrieverUsage(innerRetriever);
             return innerRetriever;
         }, RETRIEVER_FIELD);
         PARSER.declareField(constructorArg(), (parser, context) -> {

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<String> getExtendedUsageFields() {
+        return Set.of();
+    }
+
     // ---- FOR TESTING XCONTENT PARSING ----
 
     public abstract String getName();

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) {