Tracing for deep search path (opensearch-project#12103)

* Tracing for deep search path

Signed-off-by: David Zane <davizane@amazon.com>

* Refactor search phase tracing instrumentation to detach span creation from the SearchRequestOperationsListener

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>

* Close WrappedPhase span

Signed-off-by: David Zane <davizane@amazon.com>

* Add asserting listener class

Signed-off-by: David Zane <davizane@amazon.com>

* Cleanup AbstractSearchAsyncActionTests tests

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>

* Cleanup SearchAsyncActionTests tests

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>

* Cleanup CanMatchPreFilterSearchPhaseTests tests

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>

---------

Signed-off-by: David Zane <davizane@amazon.com>
Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
Co-authored-by: Andriy Redko <andriy.redko@aiven.io>

Author: dzane17

CHANGELOG.md

@@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add Remote Store Migration Experimental flag and allow mixed mode clusters under same ([#11986](https://github.com/opensearch-project/OpenSearch/pull/11986))
 - Remote reindex: Add support for configurable retry mechanism ([#12561](https://github.com/opensearch-project/OpenSearch/pull/12561))
 - [Admission Control] Integrate IO Usage Tracker to the Resource Usage Collector Service and Emit IO Usage Stats ([#11880](https://github.com/opensearch-project/OpenSearch/pull/11880))
+- 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

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanContext.java

@@ -31,4 +31,19 @@ public SpanContext(Span span) {
     Span getSpan() {
         return span;
     }
+
+    /**
+     * Sets the error for the current span behind this context
+     * @param cause error
+     */
+    public void setError(final Exception cause) {
+        span.setError(cause);
+    }
+
+    /**
+     * Ends current span
+     */
+    public void endSpan() {
+        span.endSpan();
+    }
 }

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanCreationContext.java

@@ -79,8 +79,8 @@ public SpanCreationContext attributes(Attributes attributes) {
     }
 
     /**
-     * Sets the parent for spann
-     * @param parent parent
+     * Sets the parent for span
+     * @param parent parent span context
      * @return spanCreationContext
      */
     public SpanCreationContext parent(SpanContext parent) {

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

@@ -58,6 +58,10 @@
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.internal.ShardSearchRequest;
 import org.opensearch.search.pipeline.PipelinedRequest;
+import org.opensearch.telemetry.tracing.Span;
+import org.opensearch.telemetry.tracing.SpanCreationContext;
+import org.opensearch.telemetry.tracing.SpanScope;
+import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.transport.Transport;
 
 import java.util.ArrayDeque;
@@ -116,6 +120,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
     private final Map<String, PendingExecutions> pendingExecutionsPerNode = new ConcurrentHashMap<>();
     private final boolean throttleConcurrentRequests;
     private final SearchRequestContext searchRequestContext;
+    private final Tracer tracer;
 
     private SearchPhase currentPhase;
     private boolean currentPhaseHasLifecycle;
@@ -140,7 +145,8 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
         SearchPhaseResults<Result> resultConsumer,
         int maxConcurrentRequestsPerNode,
         SearchResponse.Clusters clusters,
-        SearchRequestContext searchRequestContext
+        SearchRequestContext searchRequestContext,
+        Tracer tracer
     ) {
         super(name);
         final List<SearchShardIterator> toSkipIterators = new ArrayList<>();
@@ -177,6 +183,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
         this.results = resultConsumer;
         this.clusters = clusters;
         this.searchRequestContext = searchRequestContext;
+        this.tracer = tracer;
     }
 
     @Override
@@ -221,6 +228,7 @@ public final void start() {
                     null
                 )
             );
+            onRequestEnd(searchRequestContext);
             return;
         }
         executePhase(this);
@@ -460,15 +468,24 @@ private void onRequestEnd(SearchRequestContext searchRequestContext) {
     }
 
     private void executePhase(SearchPhase phase) {
-        try {
+        Span phaseSpan = tracer.startSpan(SpanCreationContext.server().name("[phase/" + phase.getName() + "]"));
+        try (final SpanScope scope = tracer.withSpanInScope(phaseSpan)) {
             onPhaseStart(phase);
             phase.recordAndRun();
         } catch (Exception e) {
             if (logger.isDebugEnabled()) {
                 logger.debug(new ParameterizedMessage("Failed to execute [{}] while moving to [{}] phase", request, phase.getName()), e);
             }
 
+            if (currentPhaseHasLifecycle == false) {
+                phaseSpan.setError(e);
+            }
+
             onPhaseFailure(phase, "", e);
+        } finally {
+            if (currentPhaseHasLifecycle == false) {
+                phaseSpan.endSpan();
+            }
         }
     }
 
@@ -733,7 +750,7 @@ public void sendSearchResponse(InternalSearchResponse internalSearchResponse, At
     @Override
     public final void onPhaseFailure(SearchPhase phase, String msg, Throwable cause) {
         if (currentPhaseHasLifecycle) {
-            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/CanMatchPreFilterSearchPhase.java

@@ -44,6 +44,7 @@
 import org.opensearch.search.sort.FieldSortBuilder;
 import org.opensearch.search.sort.MinAndMax;
 import org.opensearch.search.sort.SortOrder;
+import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.transport.Transport;
 
 import java.util.Comparator;
@@ -91,7 +92,8 @@ final class CanMatchPreFilterSearchPhase extends AbstractSearchAsyncAction<CanMa
         SearchTask task,
         Function<GroupShardsIterator<SearchShardIterator>, SearchPhase> phaseFactory,
         SearchResponse.Clusters clusters,
-        SearchRequestContext searchRequestContext
+        SearchRequestContext searchRequestContext,
+        Tracer tracer
     ) {
         // We set max concurrent shard requests to the number of shards so no throttling happens for can_match requests
         super(
@@ -112,7 +114,8 @@ final class CanMatchPreFilterSearchPhase extends AbstractSearchAsyncAction<CanMa
             new CanMatchSearchPhaseResults(shardsIts.size()),
             shardsIts.size(),
             clusters,
-            searchRequestContext
+            searchRequestContext,
+            tracer
         );
         this.phaseFactory = phaseFactory;
         this.shardsIts = shardsIts;

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

@@ -41,6 +41,7 @@
 import org.opensearch.search.dfs.AggregatedDfs;
 import org.opensearch.search.dfs.DfsSearchResult;
 import org.opensearch.search.internal.AliasFilter;
+import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.transport.Transport;
 
 import java.util.List;
@@ -77,7 +78,8 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
         final ClusterState clusterState,
         final SearchTask task,
         SearchResponse.Clusters clusters,
-        SearchRequestContext searchRequestContext
+        SearchRequestContext searchRequestContext,
+        final Tracer tracer
     ) {
         super(
             SearchPhaseName.DFS_PRE_QUERY.getName(),
@@ -97,7 +99,8 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
             new ArraySearchPhaseResults<>(shardsIts.size()),
             request.getMaxConcurrentShardRequests(),
             clusters,
-            searchRequestContext
+            searchRequestContext,
+            tracer
         );
         this.queryPhaseResultConsumer = queryPhaseResultConsumer;
         this.searchPhaseController = searchPhaseController;

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

@@ -43,6 +43,7 @@
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.internal.ShardSearchRequest;
 import org.opensearch.search.query.QuerySearchResult;
+import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.transport.Transport;
 
 import java.util.Map;
@@ -82,7 +83,8 @@ class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<SearchPh
         ClusterState clusterState,
         SearchTask task,
         SearchResponse.Clusters clusters,
-        SearchRequestContext searchRequestContext
+        SearchRequestContext searchRequestContext,
+        final Tracer tracer
     ) {
         super(
             SearchPhaseName.QUERY.getName(),
@@ -102,7 +104,8 @@ class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<SearchPh
             resultConsumer,
             request.getMaxConcurrentShardRequests(),
             clusters,
-            searchRequestContext
+            searchRequestContext,
+            tracer
         );
         this.topDocsSize = SearchPhaseController.getTopDocsSize(request);
         this.trackTotalHitsUpTo = request.resolveTrackTotalHitsUpTo();

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