Tracing for deep search path

Signed-off-by: David Zane <[email protected]>

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Allow to pass the list settings through environment variables (like [], ["a", "b", "c"], ...) ([#10625](https://github.com/opensearch-project/OpenSearch/pull/10625))
 - [Admission Control] Integrate CPU AC with ResourceUsageCollector and add CPU AC stats to nodes/stats ([#10887](https://github.com/opensearch-project/OpenSearch/pull/10887))
 - [S3 Repository] Add setting to control connection count for sync client ([#12028](https://github.com/opensearch-project/OpenSearch/pull/12028))
+- Tracing for deep search path ([#12103](https://github.com/opensearch-project/OpenSearch/pull/12103))
 
 ### Dependencies
 - Bump `log4j-core` from 2.18.0 to 2.19.0

plugins/query-insights/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java

@@ -116,7 +116,7 @@ public void onPhaseStart(SearchPhaseContext context) {}
     public void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {}
 
     @Override
-    public void onPhaseFailure(SearchPhaseContext context) {}
+    public void onPhaseFailure(SearchPhaseContext context, Throwable cause) {}
 
     @Override
     public void onRequestStart(SearchRequestContext searchRequestContext) {}

server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java

@@ -220,6 +220,7 @@ public final void start() {
                     null
                 )
             );
+            onRequestEnd(searchRequestContext);
             return;
         }
         executePhase(this);
@@ -717,7 +718,7 @@ public void sendSearchResponse(InternalSearchResponse internalSearchResponse, At
     @Override
     public final void onPhaseFailure(SearchPhase phase, String msg, Throwable cause) {
         if (SearchPhaseName.isValidName(phase.getName())) {
-            this.searchRequestContext.getSearchRequestOperationsListener().onPhaseFailure(this);
+            this.searchRequestContext.getSearchRequestOperationsListener().onPhaseFailure(this, cause);
         }
         raisePhaseFailure(new SearchPhaseExecutionException(phase.getName(), msg, cause, buildShardFailures()));
     }

server/src/main/java/org/opensearch/action/search/SearchRequestContext.java

@@ -78,7 +78,7 @@ void setTotalHits(TotalHits totalHits) {
         this.totalHits = totalHits;
     }
 
-    TotalHits totalHits() {
+    public TotalHits totalHits() {
         return totalHits;
     }
 

server/src/main/java/org/opensearch/action/search/SearchRequestOperationsListener.java

@@ -22,6 +22,16 @@
 @InternalApi
 public abstract class SearchRequestOperationsListener {
     private volatile boolean enabled;
+    public static final SearchRequestOperationsListener NOOP = new SearchRequestOperationsListener(false) {
+        @Override
+        protected void onPhaseStart(SearchPhaseContext context) {}
+
+        @Override
+        protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {}
+
+        @Override
+        protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {}
+    };
 
     protected SearchRequestOperationsListener() {
         this.enabled = true;
@@ -35,7 +45,7 @@ protected SearchRequestOperationsListener(final boolean enabled) {
 
     protected abstract void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext);
 
-    protected abstract void onPhaseFailure(SearchPhaseContext context);
+    protected abstract void onPhaseFailure(SearchPhaseContext context, Throwable cause);
 
     protected void onRequestStart(SearchRequestContext searchRequestContext) {}
 
@@ -91,10 +101,10 @@ protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searc
         }
 
         @Override
-        protected void onPhaseFailure(SearchPhaseContext context) {
+        protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
             for (SearchRequestOperationsListener listener : listeners) {
                 try {
-                    listener.onPhaseFailure(context);
+                    listener.onPhaseFailure(context, cause);
                 } catch (Exception e) {
                     logger.warn(() -> new ParameterizedMessage("onPhaseFailure listener [{}] failed", listener), e);
                 }

server/src/main/java/org/opensearch/action/search/SearchRequestSlowLog.java

@@ -140,7 +140,7 @@ protected void onPhaseStart(SearchPhaseContext context) {}
     protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {}
 
     @Override
-    protected void onPhaseFailure(SearchPhaseContext context) {}
+    protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {}
 
     @Override
     protected void onRequestStart(SearchRequestContext searchRequestContext) {}

server/src/main/java/org/opensearch/action/search/SearchRequestStats.java

@@ -71,7 +71,7 @@ protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searc
     }
 
     @Override
-    protected void onPhaseFailure(SearchPhaseContext context) {
+    protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
         phaseStatsMap.get(context.getCurrentPhase().getSearchPhaseName()).current.dec();
     }
 

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

@@ -88,6 +88,12 @@
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
 import org.opensearch.telemetry.metrics.MetricsRegistry;
+import org.opensearch.telemetry.tracing.Span;
+import org.opensearch.telemetry.tracing.SpanBuilder;
+import org.opensearch.telemetry.tracing.SpanScope;
+import org.opensearch.telemetry.tracing.Tracer;
+import org.opensearch.telemetry.tracing.listener.TraceableActionListener;
+import org.opensearch.telemetry.tracing.listener.TraceableSearchRequestOperationsListener;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.RemoteClusterAware;
 import org.opensearch.transport.RemoteClusterService;
@@ -173,6 +179,7 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
     private final CircuitBreaker circuitBreaker;
     private final SearchPipelineService searchPipelineService;
     private final SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory;
+    private final Tracer tracer;
 
     private volatile boolean searchQueryMetricsEnabled;
 
@@ -195,7 +202,8 @@ public TransportSearchAction(
         NamedWriteableRegistry namedWriteableRegistry,
         SearchPipelineService searchPipelineService,
         MetricsRegistry metricsRegistry,
-        SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory
+        SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory,
+        Tracer tracer
     ) {
         super(SearchAction.NAME, transportService, actionFilters, (Writeable.Reader<SearchRequest>) SearchRequest::new);
         this.client = client;
@@ -215,6 +223,7 @@ public TransportSearchAction(
         this.searchRequestOperationsCompositeListenerFactory = searchRequestOperationsCompositeListenerFactory;
         clusterService.getClusterSettings()
             .addSettingsUpdateConsumer(SEARCH_QUERY_METRICS_ENABLED_SETTING, this::setSearchQueryMetricsEnabled);
+        this.tracer = tracer;
     }
 
     private void setSearchQueryMetricsEnabled(boolean searchQueryMetricsEnabled) {
@@ -431,49 +440,58 @@ private void executeRequest(
         if (originalSearchRequest.isPhaseTook() == null) {
             originalSearchRequest.setPhaseTook(clusterService.getClusterSettings().get(SEARCH_PHASE_TOOK_ENABLED));
         }
-        SearchRequestOperationsListener.CompositeListener requestOperationsListeners = searchRequestOperationsCompositeListenerFactory
-            .buildCompositeListener(originalSearchRequest, logger);
-        SearchRequestContext searchRequestContext = new SearchRequestContext(requestOperationsListeners, originalSearchRequest);
-        searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
-
-        PipelinedRequest searchRequest;
-        ActionListener<SearchResponse> listener;
-        try {
-            searchRequest = searchPipelineService.resolvePipeline(originalSearchRequest);
-            listener = searchRequest.transformResponseListener(originalListener);
-        } catch (Exception e) {
-            originalListener.onFailure(e);
-            return;
-        }
 
-        ActionListener<SearchRequest> requestTransformListener = ActionListener.wrap(sr -> {
-            if (searchQueryMetricsEnabled) {
-                try {
-                    searchQueryCategorizer.categorize(sr.source());
-                } catch (Exception e) {
-                    logger.error("Error while trying to categorize the query.", e);
-                }
+        final Span requestSpan = tracer.startSpan(SpanBuilder.from(task, actionName));
+        try (final SpanScope spanScope = tracer.withSpanInScope(requestSpan)) {
+            SearchRequestOperationsListener.CompositeListener requestOperationsListeners;
+            final ActionListener<SearchResponse> updatedListener = TraceableActionListener.create(originalListener, requestSpan, tracer);
+            requestOperationsListeners = searchRequestOperationsCompositeListenerFactory.buildCompositeListener(
+                originalSearchRequest,
+                logger,
+                TraceableSearchRequestOperationsListener.create(tracer, requestSpan)
+            );
+            SearchRequestContext searchRequestContext = new SearchRequestContext(requestOperationsListeners, originalSearchRequest);
+            searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
+
+            PipelinedRequest searchRequest;
+            ActionListener<SearchResponse> listener;
+            try {
+                searchRequest = searchPipelineService.resolvePipeline(originalSearchRequest);
+                listener = searchRequest.transformResponseListener(updatedListener);
+            } catch (Exception e) {
+                updatedListener.onFailure(e);
+                return;
             }
 
-            ActionListener<SearchSourceBuilder> rewriteListener = buildRewriteListener(
-                sr,
-                task,
-                timeProvider,
-                searchAsyncActionProvider,
-                listener,
-                searchRequestContext
-            );
-            if (sr.source() == null) {
-                rewriteListener.onResponse(sr.source());
-            } else {
-                Rewriteable.rewriteAndFetch(
-                    sr.source(),
-                    searchService.getRewriteContext(timeProvider::getAbsoluteStartMillis),
-                    rewriteListener
+            ActionListener<SearchRequest> requestTransformListener = ActionListener.wrap(sr -> {
+                if (searchQueryMetricsEnabled) {
+                    try {
+                        searchQueryCategorizer.categorize(sr.source());
+                    } catch (Exception e) {
+                        logger.error("Error while trying to categorize the query.", e);
+                    }
+                }
+
+                ActionListener<SearchSourceBuilder> rewriteListener = buildRewriteListener(
+                    sr,
+                    task,
+                    timeProvider,
+                    searchAsyncActionProvider,
+                    listener,
+                    searchRequestContext
                 );
-            }
-        }, listener::onFailure);
-        searchRequest.transformRequest(requestTransformListener);
+                if (sr.source() == null) {
+                    rewriteListener.onResponse(sr.source());
+                } else {
+                    Rewriteable.rewriteAndFetch(
+                        sr.source(),
+                        searchService.getRewriteContext(timeProvider::getAbsoluteStartMillis),
+                        rewriteListener
+                    );
+                }
+            }, listener::onFailure);
+            searchRequest.transformRequest(requestTransformListener);
+        }
     }
 
     private ActionListener<SearchSourceBuilder> buildRewriteListener(

server/src/main/java/org/opensearch/telemetry/tracing/AttributeNames.java

@@ -75,6 +75,16 @@ private AttributeNames() {
      */
     public static final String TRANSPORT_ACTION = "action";
 
+    /**
+     * Task id
+     */
+    public static final String TASK_ID = "task_id";
+
+    /**
+     * Parent task id
+     */
+    public static final String PARENT_TASK_ID = "parent_task_id";
+
     /**
      * Index Name
      */
@@ -99,4 +109,9 @@ private AttributeNames() {
      * Refresh Policy
      */
     public static final String REFRESH_POLICY = "refresh_policy";
+
+    /**
+     * Search Response Total Hits
+     */
+    public static final String TOTAL_HITS = "total_hits";
 }

server/src/main/java/org/opensearch/telemetry/tracing/SpanBuilder.java

@@ -15,6 +15,7 @@
 import org.opensearch.core.common.Strings;
 import org.opensearch.http.HttpRequest;
 import org.opensearch.rest.RestRequest;
+import org.opensearch.tasks.Task;
 import org.opensearch.telemetry.tracing.attributes.Attributes;
 import org.opensearch.transport.TcpChannel;
 import org.opensearch.transport.Transport;
@@ -196,4 +197,43 @@ private static Attributes buildSpanAttributes(String nodeId, ReplicatedWriteRequ
         return attributes;
     }
 
+    /**
+     * Creates {@link SpanCreationContext} with parent set to specified SpanContext.
+     * @param spanName name of span.
+     * @param parentSpan target parent span.
+     * @return context
+     */
+    public static SpanCreationContext from(String spanName, SpanContext parentSpan) {
+        return SpanCreationContext.server().name(spanName).parent(parentSpan);
+    }
+
+    /**
+     * Creates {@link SpanCreationContext} with parent set to specified SpanContext.
+     * @param task search task.
+     * @param actionName action.
+     * @return context
+     */
+    public static SpanCreationContext from(Task task, String actionName) {
+        return SpanCreationContext.server().name(createSpanName(task, actionName)).attributes(buildSpanAttributes(task, actionName));
+    }
+
+    private static Attributes buildSpanAttributes(Task task, String actionName) {
+        Attributes attributes = Attributes.create().addAttribute(AttributeNames.TRANSPORT_ACTION, actionName);
+        if (task != null) {
+            attributes.addAttribute(AttributeNames.TASK_ID, task.getId());
+            if (task.getParentTaskId() != null && task.getParentTaskId().isSet()) {
+                attributes.addAttribute(AttributeNames.PARENT_TASK_ID, task.getParentTaskId().getId());
+            }
+        }
+        return attributes;
+
+    }
+
+    private static String createSpanName(Task task, String actionName) {
+        if (task != null) {
+            return task.getType() + SEPARATOR + task.getAction();
+        } else {
+            return actionName;
+        }
+    }
 }

server/src/main/java/org/opensearch/telemetry/tracing/listener/TraceableSearchRequestOperationsListener.java

@@ -0,0 +1,81 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.tracing.listener;
+
+import org.opensearch.action.search.SearchPhaseContext;
+import org.opensearch.action.search.SearchRequestContext;
+import org.opensearch.action.search.SearchRequestOperationsListener;
+import org.opensearch.telemetry.tracing.AttributeNames;
+import org.opensearch.telemetry.tracing.Span;
+import org.opensearch.telemetry.tracing.SpanBuilder;
+import org.opensearch.telemetry.tracing.SpanContext;
+import org.opensearch.telemetry.tracing.SpanScope;
+import org.opensearch.telemetry.tracing.Tracer;
+
+import static org.opensearch.core.common.Strings.capitalize;
+
+/**
+ * SearchRequestOperationsListener subscriber for search request tracing
+ *
+ * @opensearch.internal
+ */
+public final class TraceableSearchRequestOperationsListener extends SearchRequestOperationsListener {
+    private final Tracer tracer;
+    private final Span requestSpan;
+    private Span phaseSpan;
+    private SpanScope phaseSpanScope;
+
+    public TraceableSearchRequestOperationsListener(final Tracer tracer, final Span requestSpan) {
+        this.tracer = tracer;
+        this.requestSpan = requestSpan;
+        this.phaseSpan = null;
+        this.phaseSpanScope = null;
+    }
+
+    public static SearchRequestOperationsListener create(final Tracer tracer, final Span requestSpan) {
+        if (tracer.isRecording()) {
+            return new TraceableSearchRequestOperationsListener(tracer, requestSpan);
+        } else {
+            return SearchRequestOperationsListener.NOOP;
+        }
+    }
+
+    @Override
+    protected void onPhaseStart(SearchPhaseContext context) {
+        assert phaseSpan == null : "There should be only one search phase active at a time";
+        phaseSpan = tracer.startSpan(SpanBuilder.from(capitalize(context.getCurrentPhase().getName()), new SpanContext(requestSpan)));
+    }
+
+    @Override
+    protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        assert phaseSpan != null : "There should be a search phase active at that time";
+        phaseSpan.endSpan();
+        phaseSpan = null;
+    }
+
+    @Override
+    protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
+        assert phaseSpan != null : "There should be a search phase active at that time";
+        phaseSpan.setError(new Exception(cause));
+        phaseSpan.endSpan();
+        phaseSpan = null;
+    }
+
+    @Override
+    public void onRequestStart(SearchRequestContext searchRequestContext) {}
+
+    @Override
+    public void onRequestEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        // add response-related attributes on request end
+        requestSpan.addAttribute(
+            AttributeNames.TOTAL_HITS,
+            searchRequestContext.totalHits() == null ? 0 : searchRequestContext.totalHits().value
+        );
+    }
+}