refactor with latest changes from main

Author: chrisparrinello

server/src/main/java/org/elasticsearch/index/search/stats/ShardSearchPhaseAPMMetrics.java

@@ -57,18 +57,15 @@ public ShardSearchPhaseAPMMetrics(MeterRegistry meterRegistry) {
     }
 
     @Override
-    public void onCanMatchPhase(ShardSearchRequest shardSearchRequest, long tookInNanos) {
-        recordPhaseLatency(
-            canMatchPhaseMetric,
-            tookInNanos,
-            shardSearchRequest,
-            null // we don't have the time range filter info at this point in the search
-        );
+    public void onCanMatchPhase(Map<String, Object> searchRequestAttributes, long tookInNanos) {
+        canMatchPhaseMetric.record(TimeUnit.NANOSECONDS.toMillis(tookInNanos), searchRequestAttributes);
     }
 
     @Override
     public void onDfsPhase(SearchContext searchContext, long tookInNanos) {
-        recordPhaseLatency(dfsPhaseMetric, tookInNanos);
+        SearchExecutionContext searchExecutionContext = searchContext.getSearchExecutionContext();
+        Long timeRangeFilterFromMillis = searchExecutionContext.getTimeRangeFilterFromMillis();
+        recordPhaseLatency(dfsPhaseMetric, tookInNanos, searchContext.request(), timeRangeFilterFromMillis);
     }
 
     @Override
@@ -85,10 +82,6 @@ public void onFetchPhase(SearchContext searchContext, long tookInNanos) {
         recordPhaseLatency(fetchPhaseMetric, tookInNanos, searchContext.request(), timeRangeFilterFromMillis);
     }
 
-    private static void recordPhaseLatency(LongHistogram histogramMetric, long tookInNanos) {
-        histogramMetric.record(TimeUnit.NANOSECONDS.toMillis(tookInNanos));
-    }
-
     private static void recordPhaseLatency(
         LongHistogram histogramMetric,
         long tookInNanos,

server/src/main/java/org/elasticsearch/index/shard/SearchOperationListener.java

@@ -12,10 +12,10 @@
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.search.internal.ReaderContext;
 import org.elasticsearch.search.internal.SearchContext;
-import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.transport.TransportRequest;
 
 import java.util.List;
+import java.util.Map;
 
 /**
  * An listener for search, fetch and context events.
@@ -92,10 +92,10 @@ default void onFailedDfsPhase(SearchContext searchContext) {}
      * Executed after the can-match phase successfully finished.
      * Note: this is not invoked if the can match phase execution failed.
      *
-     * @param shardSearchRequest the current shard search request
+     * @param searchRequestAttributes the attributes of the search request
      * @param tookInNanos the number of nanoseconds the can-match execution took
      */
-    default void onCanMatchPhase(ShardSearchRequest shardSearchRequest, long tookInNanos) {}
+    default void onCanMatchPhase(Map<String, Object> searchRequestAttributes, long tookInNanos) {}
 
     /**
      * Executed when a new reader context was created
@@ -248,10 +248,10 @@ public void onDfsPhase(SearchContext searchContext, long tookInNanos) {
         }
 
         @Override
-        public void onCanMatchPhase(ShardSearchRequest shardSearchRequest, long tookInNanos) {
+        public void onCanMatchPhase(Map<String, Object> searchRequestAttributes, long tookInNanos) {
             for (SearchOperationListener listener : listeners) {
                 try {
-                    listener.onCanMatchPhase(shardSearchRequest, tookInNanos);
+                    listener.onCanMatchPhase(searchRequestAttributes, tookInNanos);
                 } catch (Exception e) {
                     logger.warn(() -> "onCanMatchPhase listener [" + listener + "] failed", e);
                 }

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.action.search.CanMatchNodeRequest;
 import org.elasticsearch.action.search.CanMatchNodeResponse;
 import org.elasticsearch.action.search.OnlinePrewarmingService;
+import org.elasticsearch.action.search.SearchRequestAttributesExtractor;
 import org.elasticsearch.action.search.SearchShardTask;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.action.support.TransportActions;
@@ -1936,6 +1937,7 @@ public AliasFilter buildAliasFilter(ProjectState state, String index, Set<Resolv
     public void canMatch(CanMatchNodeRequest request, ActionListener<CanMatchNodeResponse> listener) {
         var shardLevelRequests = request.getShardLevelRequests();
         final List<CanMatchNodeResponse.ResponseOrFailure> responses = new ArrayList<>(shardLevelRequests.size());
+        Map<String, Object> searchRequestAttributes = null;
         for (var shardLevelRequest : shardLevelRequests) {
             long shardCanMatchStartTimeInNanos = System.nanoTime();
             ShardSearchRequest shardSearchRequest = request.createShardSearchRequest(shardLevelRequest);
@@ -1946,7 +1948,18 @@ public void canMatch(CanMatchNodeRequest request, ActionListener<CanMatchNodeRes
                 CanMatchContext canMatchContext = createCanMatchContext(shardSearchRequest);
                 CanMatchShardResponse canMatchShardResponse = canMatch(canMatchContext, true);
                 responses.add(new CanMatchNodeResponse.ResponseOrFailure(canMatchShardResponse));
-                indexShard.getSearchOperationListener().onCanMatchPhase(System.nanoTime() - shardCanMatchStartTimeInNanos);
+
+                if (searchRequestAttributes == null) {
+                    // All of the shards in the request are for the same search, so the attributes will be the same for all shards.
+                    searchRequestAttributes = SearchRequestAttributesExtractor.extractAttributes(
+                        shardSearchRequest,
+                        canMatchContext.getTimeRangeFilterFromMillis(),
+                        shardSearchRequest.nowInMillis()
+                    );
+                }
+
+                indexShard.getSearchOperationListener()
+                    .onCanMatchPhase(searchRequestAttributes, System.nanoTime() - shardCanMatchStartTimeInNanos);
             } catch (Exception e) {
                 responses.add(new CanMatchNodeResponse.ResponseOrFailure(e));
             }

server/src/test/java/org/elasticsearch/search/TelemetryMetrics/ShardSearchPhaseAPMMetricsTests.java

@@ -101,6 +101,7 @@ public void testMetricsDfsQueryThenFetch() {
         );
         final List<Measurement> dfsMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(DFS_SEARCH_PHASE_METRIC);
         assertEquals(num_primaries, dfsMeasurements.size());
+        assertAttributes(dfsMeasurements, false, false);
         final List<Measurement> queryMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(QUERY_SEARCH_PHASE_METRIC);
         assertEquals(num_primaries, queryMeasurements.size());
         final List<Measurement> fetchMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(FETCH_SEARCH_PHASE_METRIC);
@@ -117,8 +118,10 @@ public void testMetricsDfsQueryThenFetchSystem() {
         );
         final List<Measurement> dfsMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(DFS_SEARCH_PHASE_METRIC);
         assertEquals(0, dfsMeasurements.size()); // DFS phase not done for index with single shard
+        assertAttributes(dfsMeasurements, true, false);
         final List<Measurement> queryMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(QUERY_SEARCH_PHASE_METRIC);
         assertEquals(1, queryMeasurements.size());
+        assertAttributes(queryMeasurements, true, false);
         final List<Measurement> fetchMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(FETCH_SEARCH_PHASE_METRIC);
         assertEquals(1, fetchMeasurements.size());
         assertAttributes(fetchMeasurements, true, false);
@@ -369,6 +372,17 @@ public void testTimeRangeFilterAllResults() {
         });
         final List<Measurement> canMatchMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(CAN_MATCH_SEARCH_PHASE_METRIC);
         assertEquals(num_primaries, canMatchMeasurements.size());
+        for (Measurement measurement : canMatchMeasurements) {
+            Map<String, Object> attributes = measurement.attributes();
+            assertEquals(5, attributes.size());
+            assertEquals("user", attributes.get("target"));
+            assertEquals("hits_only", attributes.get("query_type"));
+            assertEquals("_score", attributes.get("sort"));
+            assertEquals(false, attributes.get(SearchRequestAttributesExtractor.SYSTEM_THREAD_ATTRIBUTE_NAME));
+            // the range query was rewritten to one without bounds: we do track the time range filter from value but we don't set
+            // the time range filter field because no range query is executed at the shard level.
+            assertEquals("older_than_14_days", attributes.get("time_range_filter_from"));
+        }
         final List<Measurement> queryMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(QUERY_SEARCH_PHASE_METRIC);
         // the two docs are at most spread across two shards, other shards are empty and get filtered out
         assertThat(queryMeasurements.size(), Matchers.lessThanOrEqualTo(2));