Add support for extended search usage telemetry (#135306)

* Add extended telemetry

* Cleanup

* Update docs/changelog/135306.yaml

* [CI] Auto commit changes from spotless

* [CI] Update transport version definitions

* Fix serialization

* Add tests

* Update transport

* Tests

* [CI] Auto commit changes from spotless

* Yaml

* Fix error in SearchUsageStatsTests due to randomization

* Ensure mutation in serialization tests

* [CI] Auto commit changes from spotless

* Revert "Yaml"

This reverts commit 9f4e40d.

* [CI] Update transport version definitions

* Remove old debug breakpoint

* PR feedback: Refactor extended stats into objects instead of forcing to Map<String,Long>

* Make ExtendedSeearchUsageMetric a NamedWriteable

* Regenerate transport version

* Fix serialization

* [CI] Auto commit changes from spotless

* PR Feedback

* [CI] Auto commit changes from spotless

* PR feedback

* Regenerate transport version

---------

Co-authored-by: elasticsearchmachine <[email protected]>

Author: kderusso

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/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<NamedWriteableRegistry.Entry> getNamedWriteables() {
+        return List.of(
+            new NamedWriteableRegistry.Entry(
+                ExtendedSearchUsageMetric.class,
+                ExtendedSearchUsageLongCounter.NAME,
+                ExtendedSearchUsageLongCounter::new
+            )
+        );
+    }
 }

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<ExtendedSearchUsageLongCounter> {
+
+    public static final String NAME = "extended_search_usage_long_counter";
+
+    private final Map<String, Long> values;
+
+    public ExtendedSearchUsageLongCounter(Map<String, Long> values) {
+        this.values = values;
+    }
+
+    public ExtendedSearchUsageLongCounter(StreamInput in) throws IOException {
+        this.values = in.readMap(StreamInput::readString, StreamInput::readLong);
+    }
+
+    public Map<String, Long> getValues() {
+        return Collections.unmodifiableMap(values);
+    }
+
+    @Override
+    public ExtendedSearchUsageLongCounter merge(ExtendedSearchUsageMetric<?> other) {
+        assert other instanceof ExtendedSearchUsageLongCounter;
+        ExtendedSearchUsageLongCounter otherLongCounter = (ExtendedSearchUsageLongCounter) other;
+        Map<String, Long> 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;
+    }
+}

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<T extends ExtendedSearchUsageMetric<T>> 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);
+}

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<String, Map<String, ExtendedSearchUsageMetric<?>>> categorizedExtendedData;
+
+    public static final ExtendedSearchUsageStats EMPTY = new ExtendedSearchUsageStats();
+
+    public ExtendedSearchUsageStats() {
+        this.categorizedExtendedData = new HashMap<>();
+    }
+
+    public ExtendedSearchUsageStats(Map<String, Map<String, ExtendedSearchUsageMetric<?>>> 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<String, Map<String, ExtendedSearchUsageMetric<?>>> 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<String, ExtendedSearchUsageMetric<?>> 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<String, ExtendedSearchUsageMetric<?>> 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);
+    }
+}

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 = ExtendedSearchUsageStats.EMPTY;
     }
 
     /**
@@ -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)
+            : 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<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);
     }
 }