Refactor search phase tracing instrumentation to detach span creation…

… from the SearchRequestOperationsListener Signed-off-by: Andriy Redko <[email protected]>
opensearch-project · Feb 19, 2024 · ec73bef · ec73bef
1 parent fcde9cb
commit ec73bef
Show file tree

Hide file tree

Showing 12 changed files with 101 additions and 50 deletions.
diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanContext.java b/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();
+    }
 }
diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanCreationContext.java b/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) {

diff --git a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java b/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
@@ -455,7 +462,8 @@ private void onRequestEnd(SearchRequestContext searchRequestContext) {
     }
 
     private void executePhase(SearchPhase phase) {
-        try {
+        final Span phaseSpan = tracer.startSpan(SpanCreationContext.server().name("[phase/" + phase.getName() + "]"));
+        try (final SpanScope scope = tracer.withSpanInScope(phaseSpan)) {
             onPhaseStart(phase);
             phase.recordAndRun();
         } catch (Exception e) {

diff --git a/server/src/main/java/org/opensearch/action/search/CanMatchPreFilterSearchPhase.java b/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;

diff --git a/server/src/main/java/org/opensearch/action/search/SearchDfsQueryThenFetchAsyncAction.java b/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;

diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryThenFetchAsyncAction.java b/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();

diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -393,7 +393,8 @@ public AbstractSearchAsyncAction<? extends SearchPhaseResult> asyncSearchAction(
                     new ArraySearchPhaseResults<>(shardsIts.size()),
                     searchRequest.getMaxConcurrentShardRequests(),
                     clusters,
-                    searchRequestContext
+                    searchRequestContext,
+                    tracer
                 ) {
                     @Override
                     protected void executePhaseOnShard(
@@ -1258,7 +1259,8 @@ private AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction
                     )
                 ),
                 clusters,
-                searchRequestContext
+                searchRequestContext,
+                tracer
             );
         } else {
             final QueryPhaseResultConsumer queryResultConsumer = searchPhaseController.newSearchPhaseResults(
@@ -1289,7 +1291,8 @@ private AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction
                         clusterState,
                         task,
                         clusters,
-                        searchRequestContext
+                        searchRequestContext,
+                        tracer
                     );
                     break;
                 case QUERY_THEN_FETCH:
@@ -1310,7 +1313,8 @@ private AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction
                         clusterState,
                         task,
                         clusters,
-                        searchRequestContext
+                        searchRequestContext,
+                        tracer
                     );
                     break;
                 default:

diff --git a/...a/org/opensearch/telemetry/tracing/listener/TraceableSearchRequestOperationsListener.java b/...a/org/opensearch/telemetry/tracing/listener/TraceableSearchRequestOperationsListener.java
@@ -13,13 +13,9 @@
 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
  *
@@ -28,14 +24,12 @@
 public final class TraceableSearchRequestOperationsListener extends SearchRequestOperationsListener {
     private final Tracer tracer;
     private final Span requestSpan;
-    private Span phaseSpan;
-    private SpanScope phaseSpanScope;
+    private SpanContext phaseSpanContext;
 
     public TraceableSearchRequestOperationsListener(final Tracer tracer, final Span requestSpan) {
         this.tracer = tracer;
         this.requestSpan = requestSpan;
-        this.phaseSpan = null;
-        this.phaseSpanScope = null;
+        this.phaseSpanContext = null;
     }
 
     public static SearchRequestOperationsListener create(final Tracer tracer, final Span requestSpan) {
@@ -48,23 +42,23 @@ public static SearchRequestOperationsListener create(final Tracer tracer, final
 
     @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)));
+        assert phaseSpanContext == null : "There should be only one search phase active at a time";
+        phaseSpanContext = tracer.getCurrentSpan();
     }
 
     @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;
+        assert phaseSpanContext != null : "There should be a search phase active at that time";
+        phaseSpanContext.endSpan();
+        phaseSpanContext = 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;
+        assert phaseSpanContext != null : "There should be a search phase active at that time";
+        phaseSpanContext.setError((Exception) cause);
+        phaseSpanContext.endSpan();
+        phaseSpanContext = null;
     }
 
     @Override

diff --git a/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java
@@ -58,6 +58,7 @@
 import org.opensearch.search.internal.ShardSearchContextId;
 import org.opensearch.search.internal.ShardSearchRequest;
 import org.opensearch.search.query.QuerySearchResult;
+import org.opensearch.telemetry.tracing.noop.NoopTracer;
 import org.opensearch.test.InternalAggregationTestCase;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.TestThreadPool;
@@ -122,7 +123,7 @@ protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searc
             }
 
             @Override
-            protected void onPhaseFailure(SearchPhaseContext context) {
+            protected void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
                 assertThat(phase, is(context.getCurrentPhase()));
                 phase = null;
             }
@@ -205,7 +206,8 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
             new SearchRequestContext(
                 new SearchRequestOperationsListener.CompositeListener(List.of(assertingListener), LogManager.getLogger()),
                 request
-            )
+            ),
+            NoopTracer.INSTANCE
         ) {
             @Override
             protected SearchPhase getNextPhase(final SearchPhaseResults<SearchPhaseResult> results, SearchPhaseContext context) {
@@ -746,7 +748,8 @@ private SearchDfsQueryThenFetchAsyncAction createSearchDfsQueryThenFetchAsyncAct
             new SearchRequestContext(
                 new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger),
                 searchRequest
-            )
+            ),
+            NoopTracer.INSTANCE
         );
     }
 
@@ -799,7 +802,8 @@ private SearchQueryThenFetchAsyncAction createSearchQueryThenFetchAsyncAction(
             new SearchRequestContext(
                 new SearchRequestOperationsListener.CompositeListener(searchRequestOperationsListeners, logger),
                 searchRequest
-            )
+            ),
+            NoopTracer.INSTANCE
         ) {
             @Override
             ShardSearchFailure[] buildShardFailures() {