Add relevant attributes to search took time APM metrics

docs/changelog/134232.yaml

@@ -0,0 +1,5 @@
+pr: 134232
+summary: Add relevant attributes to search took time APM metrics
+area: Search
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/action/search/SearchRequestIntrospector.java

@@ -0,0 +1,245 @@
+/*
+ * 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.search;
+
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.index.query.BoolQueryBuilder;
+import org.elasticsearch.index.query.BoostingQueryBuilder;
+import org.elasticsearch.index.query.ConstantScoreQueryBuilder;
+import org.elasticsearch.index.query.NestedQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.RangeQueryBuilder;
+import org.elasticsearch.search.SearchService;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.sort.FieldSortBuilder;
+import org.elasticsearch.search.sort.ScoreSortBuilder;
+import org.elasticsearch.search.sort.SortBuilder;
+import org.elasticsearch.search.vectors.KnnVectorQueryBuilder;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Used to introspect a search request and extract metadata from it around the features it uses.
+ */
+public class SearchRequestIntrospector {
+
+    /**
+     * Introspects the provided search request and extracts metadata from it about some of its characteristics.
+     *
+     */
+    public static QueryMetadata introspectSearchRequest(SearchRequest searchRequest) {
+
+        String target = introspectIndices(searchRequest.indices());
+
+        String pitOrScroll = null;
+        if (searchRequest.scroll() != null) {
+            pitOrScroll = SCROLL;
+        }
+
+        SearchSourceBuilder searchSourceBuilder = searchRequest.source();
+        if (searchSourceBuilder == null) {
+            return new QueryMetadata(target, ScoreSortBuilder.NAME, HITS_ONLY, false, Strings.EMPTY_ARRAY, pitOrScroll);
+        }
+
+        if (searchSourceBuilder.pointInTimeBuilder() != null) {
+            pitOrScroll = PIT;
+        }
+
+        final String primarySort;
+        if (searchSourceBuilder.sorts() == null || searchSourceBuilder.sorts().isEmpty()) {
+            primarySort = ScoreSortBuilder.NAME;
+        } else {
+            primarySort = introspectPrimarySort(searchSourceBuilder.sorts().getFirst());
+        }
+
+        final String queryType = introspectQueryType(searchSourceBuilder);
+
+        QueryMetadataBuilder queryMetadataBuilder = new QueryMetadataBuilder();
+        if (searchSourceBuilder.query() != null) {
+            introspectQueryBuilder(searchSourceBuilder.query(), queryMetadataBuilder);
+        }
+
+        final boolean hasKnn = searchSourceBuilder.knnSearch().isEmpty() == false || queryMetadataBuilder.knnQuery;
+
+        return new QueryMetadata(
+            target,
+            primarySort,
+            queryType,
+            hasKnn,
+            queryMetadataBuilder.rangeFields.toArray(new String[0]),
+            pitOrScroll
+        );
+    }
+
+    private static final class QueryMetadataBuilder {
+        private boolean knnQuery = false;
+        private final List<String> rangeFields = new ArrayList<>();
+    }
+
+    public record QueryMetadata(
+        String target,
+        String primarySort,
+        String queryType,
+        boolean knn,
+        String[] rangeFields,
+        String pitOrScroll
+    ) {
+
+        public Map<String, Object> toAttributesMap() {
+            Map<String, Object> attributes = new HashMap<>();
+            attributes.put(TARGET_ATTRIBUTE, target);
+            attributes.put(SORT_ATTRIBUTE, primarySort);
+            if (pitOrScroll == null) {
+                attributes.put(QUERY_TYPE_ATTRIBUTE, queryType);
+            } else {
+                attributes.put(QUERY_TYPE_ATTRIBUTE, Arrays.asList(queryType, pitOrScroll));
+            }
+
+            attributes.put(KNN_ATTRIBUTE, knn);
+            attributes.put(RANGES_ATTRIBUTE, rangeFields);
+            return attributes;
+        }
+    }
+
+    private static final String TARGET_ATTRIBUTE = "target";
+    private static final String SORT_ATTRIBUTE = "sort";
+    private static final String QUERY_TYPE_ATTRIBUTE = "query_type";
+    private static final String KNN_ATTRIBUTE = "knn";
+    private static final String RANGES_ATTRIBUTE = "ranges";
+
+    private static final String TARGET_KIBANA = ".kibana";
+    private static final String TARGET_ML = ".ml";
+    private static final String TARGET_FLEET = ".fleet";
+    private static final String TARGET_SLO = ".slo";
+    private static final String TARGET_ALERTS = ".alerts";
+    private static final String TARGET_ELASTIC = ".elastic";
+    private static final String TARGET_DS = ".ds-";
+    private static final String TARGET_OTHERS = ".others";
+    private static final String TARGET_USER = "user";
+
+    static String introspectIndices(String[] indices) {
+        // Note that indices are expected to be resolved, meaning wildcards are not handled on purpose
+        // If indices resolve to data streams, the name of the data stream is returned as opposed to its backing indices
+        if (indices.length == 1) {
+            String index = indices[0];
+            if (index.startsWith(".")) {
+                if (index.startsWith(TARGET_KIBANA)) {
+                    return TARGET_KIBANA;
+                }
+                if (index.startsWith(TARGET_ML)) {
+                    return TARGET_ML;
+                }
+                if (index.startsWith(TARGET_FLEET)) {
+                    return TARGET_FLEET;
+                }
+                if (index.startsWith(TARGET_SLO)) {
+                    return TARGET_SLO;
+                }
+                if (index.startsWith(TARGET_ALERTS)) {
+                    return TARGET_ALERTS;
+                }
+                if (index.startsWith(TARGET_ELASTIC)) {
+                    return TARGET_ELASTIC;
+                }
+                // this happens only when data streams backing indices are searched explicitly
+                if (index.startsWith(TARGET_DS)) {
+                    return TARGET_DS;
+                }
+                return TARGET_OTHERS;
+            }
+        }
+        return TARGET_USER;
+    }
+
+    private static final String TIMESTAMP = "@timestamp";
+    private static final String EVENT_INGESTED = "event.ingested";
+    private static final String _DOC = "_doc";
+    private static final String FIELD = "field";
+
+    static String introspectPrimarySort(SortBuilder<?> primarySortBuilder) {
+        if (primarySortBuilder instanceof FieldSortBuilder fieldSort) {
+            return switch (fieldSort.getFieldName()) {
+                case TIMESTAMP -> TIMESTAMP;
+                case EVENT_INGESTED -> EVENT_INGESTED;
+                case _DOC -> _DOC;
+                default -> FIELD;
+            };
+        }
+        return primarySortBuilder.getWriteableName();
+    }
+
+    private static final String HITS_AND_AGGS = "hits_and_aggs";
+    private static final String HITS_ONLY = "hits_only";
+    private static final String AGGS_ONLY = "aggs_only";
+    private static final String COUNT_ONLY = "count_only";
+    private static final String PIT = "pit";
+    private static final String SCROLL = "scroll";
+
+    public static final Map<String, Object> SEARCH_SCROLL_ATTRIBUTES = Map.of(QUERY_TYPE_ATTRIBUTE, SCROLL);
+
+    static String introspectQueryType(SearchSourceBuilder searchSourceBuilder) {
+        int size = searchSourceBuilder.size();
+        if (size == -1) {
+            size = SearchService.DEFAULT_SIZE;
+        }
+        if (searchSourceBuilder.aggregations() != null && size > 0) {
+            return HITS_AND_AGGS;
+        }
+        if (searchSourceBuilder.aggregations() != null) {
+            return AGGS_ONLY;
+        }
+        if (size > 0) {
+            return HITS_ONLY;
+        }
+        return COUNT_ONLY;
+    }
+
+    static void introspectQueryBuilder(QueryBuilder queryBuilder, QueryMetadataBuilder queryMetadataBuilder) {
+        switch (queryBuilder) {
+            case BoolQueryBuilder bool:
+                for (QueryBuilder must : bool.must()) {
+                    introspectQueryBuilder(must, queryMetadataBuilder);
+                }
+                for (QueryBuilder filter : bool.filter()) {
+                    introspectQueryBuilder(filter, queryMetadataBuilder);
+                }
+                if (bool.must().isEmpty() && bool.filter().isEmpty() && bool.mustNot().isEmpty() && bool.should().size() == 1) {
+                    introspectQueryBuilder(bool.should().getFirst(), queryMetadataBuilder);
+                }
+                // Note that should clauses are ignored unless there's only one that becomes mandatory
+                // must_not clauses are also ignored for now
+                break;
+            case ConstantScoreQueryBuilder constantScore:
+                introspectQueryBuilder(constantScore.innerQuery(), queryMetadataBuilder);
+                break;
+            case BoostingQueryBuilder boosting:
+                introspectQueryBuilder(boosting.positiveQuery(), queryMetadataBuilder);
+                break;
+            case NestedQueryBuilder nested:
+                introspectQueryBuilder(nested.query(), queryMetadataBuilder);
+                break;
+            case RangeQueryBuilder range:
+                switch (range.fieldName()) {
+                    case TIMESTAMP -> queryMetadataBuilder.rangeFields.add(TIMESTAMP);
+                    case EVENT_INGESTED -> queryMetadataBuilder.rangeFields.add(EVENT_INGESTED);
+                    default -> queryMetadataBuilder.rangeFields.add(FIELD);
+                }
+                break;
+            case KnnVectorQueryBuilder knn:
+                queryMetadataBuilder.knnQuery = true;
+                break;
+            default:
+        }
+    }
+}

server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java

@@ -333,7 +333,12 @@ public long buildTookInMillis() {
 
     @Override
     protected void doExecute(Task task, SearchRequest searchRequest, ActionListener<SearchResponse> listener) {
-        executeRequest((SearchTask) task, searchRequest, new SearchResponseActionListener(listener), AsyncSearchActionProvider::new);
+        executeRequest(
+            (SearchTask) task,
+            searchRequest,
+            new SearchResponseActionListener(listener, searchRequest),
+            AsyncSearchActionProvider::new
+        );
     }
 
     void executeRequest(
@@ -526,7 +531,7 @@ void executeRequest(
                 // We set the keep alive to -1 to indicate that we don't need the pit id in the response.
                 // This is needed since we delete the pit prior to sending the response so the id doesn't exist anymore.
                 source.pointInTimeBuilder(new PointInTimeBuilder(resp.getPointInTimeId()).setKeepAlive(TimeValue.MINUS_ONE));
-                var pitListener = new SearchResponseActionListener(delegate) {
+                var pitListener = new SearchResponseActionListener(delegate, original) {
                     @Override
                     public void onResponse(SearchResponse response) {
                         // we need to close the PIT first so we delay the release of the response to after the closing
@@ -2012,9 +2017,11 @@ private class SearchResponseActionListener extends DelegatingActionListener<Sear
         implements
             TelemetryListener {
         private final CCSUsage.Builder usageBuilder;
+        private final SearchRequest searchRequest;
 
-        SearchResponseActionListener(ActionListener<SearchResponse> listener) {
+        SearchResponseActionListener(ActionListener<SearchResponse> listener, SearchRequest searchRequest) {
             super(listener);
+            this.searchRequest = searchRequest;
             if (listener instanceof SearchResponseActionListener srListener) {
                 usageBuilder = srListener.usageBuilder;
             } else {
@@ -2046,7 +2053,7 @@ public void setClient(Task task) {
         @Override
         public void onResponse(SearchResponse searchResponse) {
             try {
-                searchResponseMetrics.recordTookTime(searchResponse.getTookInMillis());
+                searchResponseMetrics.recordTookTime(searchResponse.getTookInMillis(), searchRequest);
                 SearchResponseMetrics.ResponseCountTotalStatus responseCountTotalStatus =
                     SearchResponseMetrics.ResponseCountTotalStatus.SUCCESS;
                 if (searchResponse.getShardFailures() != null && searchResponse.getShardFailures().length > 0) {

server/src/main/java/org/elasticsearch/action/search/TransportSearchScrollAction.java

@@ -60,7 +60,7 @@ protected void doExecute(Task task, SearchScrollRequest request, ActionListener<
             @Override
             public void onResponse(SearchResponse searchResponse) {
                 try {
-                    searchResponseMetrics.recordTookTime(searchResponse.getTookInMillis());
+                    searchResponseMetrics.recordTookTimeForSearchScroll(searchResponse.getTookInMillis());
                     SearchResponseMetrics.ResponseCountTotalStatus responseCountTotalStatus =
                         SearchResponseMetrics.ResponseCountTotalStatus.SUCCESS;
                     if (searchResponse.getShardFailures() != null && searchResponse.getShardFailures().length > 0) {

server/src/main/java/org/elasticsearch/rest/action/search/SearchResponseMetrics.java

@@ -9,6 +9,8 @@
 
 package org.elasticsearch.rest.action.search;
 
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.search.SearchRequestIntrospector;
 import org.elasticsearch.telemetry.metric.LongCounter;
 import org.elasticsearch.telemetry.metric.LongHistogram;
 import org.elasticsearch.telemetry.metric.MeterRegistry;
@@ -66,8 +68,14 @@ private SearchResponseMetrics(LongHistogram tookDurationTotalMillisHistogram, Lo
         this.responseCountTotalCounter = responseCountTotalCounter;
     }
 
-    public long recordTookTime(long tookTime) {
-        tookDurationTotalMillisHistogram.record(tookTime);
+    public long recordTookTimeForSearchScroll(long tookTime) {
+        tookDurationTotalMillisHistogram.record(tookTime, SearchRequestIntrospector.SEARCH_SCROLL_ATTRIBUTES);
+        return tookTime;
+    }
+
+    public long recordTookTime(long tookTime, SearchRequest searchRequest) {
+        SearchRequestIntrospector.QueryMetadata queryMetadata = SearchRequestIntrospector.introspectSearchRequest(searchRequest);
+        tookDurationTotalMillisHistogram.record(tookTime, queryMetadata.toAttributesMap());
         return tookTime;
     }