Tracing for deep search path

Signed-off-by: David Zane <[email protected]>
opensearch-project · Feb 7, 2024 · 4b0a3ef · 4b0a3ef
1 parent 52b27f4
commit 4b0a3ef
Show file tree

Hide file tree

Showing 9 changed files with 172 additions and 6 deletions.
diff --git a/CHANGELOG.md b/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

diff --git a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
@@ -220,6 +220,7 @@ public final void start() {
                     null
                 )
             );
+            onRequestEnd(searchRequestContext);
             return;
         }
         executePhase(this);

diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java b/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;
     }
 
@@ -89,7 +89,7 @@ void setShardStats(int total, int successful, int skipped, int failed) {
         this.shardStats.put(ShardStatsFieldNames.SEARCH_REQUEST_SLOWLOG_SHARD_FAILED, failed);
     }
 
-    String formattedShardStats() {
+    public String formattedShardStats() {
         if (shardStats.isEmpty()) {
             return "";
         } else {

diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -88,6 +88,10 @@
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
 import org.opensearch.telemetry.metrics.MetricsRegistry;
+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 +177,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 +200,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 +221,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,13 +438,35 @@ private void executeRequest(
         if (originalSearchRequest.isPhaseTook() == null) {
             originalSearchRequest.setPhaseTook(clusterService.getClusterSettings().get(SEARCH_PHASE_TOOK_ENABLED));
         }
-        SearchRequestOperationsListener.CompositeListener requestOperationsListeners = searchRequestOperationsCompositeListenerFactory
-            .buildCompositeListener(originalSearchRequest, logger);
+        SearchRequestOperationsListener.CompositeListener requestOperationsListeners;
+        TraceableSearchRequestOperationsListener traceableSearchRequestOperationsListener = null;
+        if (tracer.isRecording()) {
+            traceableSearchRequestOperationsListener = new TraceableSearchRequestOperationsListener(tracer);
+            requestOperationsListeners = searchRequestOperationsCompositeListenerFactory.buildCompositeListener(
+                originalSearchRequest,
+                logger,
+                traceableSearchRequestOperationsListener
+            );
+        } else {
+            requestOperationsListeners = searchRequestOperationsCompositeListenerFactory.buildCompositeListener(
+                originalSearchRequest,
+                logger
+            );
+        }
         SearchRequestContext searchRequestContext = new SearchRequestContext(requestOperationsListeners, originalSearchRequest);
         searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
 
         PipelinedRequest searchRequest;
         ActionListener<SearchResponse> listener;
+
+        if (traceableSearchRequestOperationsListener != null) {
+            originalListener = TraceableActionListener.create(
+                originalListener,
+                traceableSearchRequestOperationsListener.getRequestSpan(),
+                tracer
+            );
+        }
+
         try {
             searchRequest = searchPipelineService.resolvePipeline(originalSearchRequest);
             listener = searchRequest.transformResponseListener(originalListener);

diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/AttributeNames.java b/server/src/main/java/org/opensearch/telemetry/tracing/AttributeNames.java
@@ -94,4 +94,19 @@ private AttributeNames() {
      * Refresh Policy
      */
     public static final String REFRESH_POLICY = "refresh_policy";
+
+    /**
+     * Search Request Source
+     */
+    public static final String SOURCE = "source";
+
+    /**
+     * Search Response Shard Stats
+     */
+    public static final String SHARDS = "shards";
+
+    /**
+     * Search Response Total Hits
+     */
+    public static final String TOTAL_HITS = "total_hits";
 }
diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/SpanBuilder.java b/server/src/main/java/org/opensearch/telemetry/tracing/SpanBuilder.java
@@ -42,6 +42,15 @@ private SpanBuilder() {
 
     }
 
+    /**
+     * Creates {@link SpanCreationContext} from String
+     * @param spanName String.
+     * @return context.
+     */
+    public static SpanCreationContext from(String spanName) {
+        return SpanCreationContext.server().name(spanName);
+    }
+
     /**
      * Creates {@link SpanCreationContext} from the {@link HttpRequest}
      * @param request Http request.
@@ -170,4 +179,13 @@ 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);
+    }
 }
diff --git a/...a/org/opensearch/telemetry/tracing/listener/TraceableSearchRequestOperationsListener.java b/...a/org/opensearch/telemetry/tracing/listener/TraceableSearchRequestOperationsListener.java
@@ -0,0 +1,82 @@
+/*
+ * 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 org.opensearch.telemetry.tracing.noop.NoopSpan;
+
+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 Span requestSpan;
+    private Span phaseSpan;
+    private SpanScope phaseSpanScope;
+
+    public TraceableSearchRequestOperationsListener(Tracer tracer) {
+        this.tracer = tracer;
+        this.requestSpan = NoopSpan.INSTANCE;
+        this.phaseSpan = NoopSpan.INSTANCE;
+    }
+
+    @Override
+    protected void onPhaseStart(SearchPhaseContext context) {
+        phaseSpan = tracer.startSpan(
+            SpanBuilder.from("coord" + capitalize(context.getCurrentPhase().getName()), new SpanContext(requestSpan))
+        );
+        phaseSpanScope = tracer.withSpanInScope(phaseSpan);
+    }
+
+    @Override
+    protected void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        phaseSpan.endSpan();
+        phaseSpanScope.close();
+    }
+
+    @Override
+    protected void onPhaseFailure(SearchPhaseContext context) {
+        phaseSpan.endSpan();
+        phaseSpanScope.close();
+    }
+
+    @Override
+    public void onRequestStart(SearchRequestContext searchRequestContext) {
+        requestSpan = tracer.startSpan(SpanBuilder.from("coordReq"));
+    }
+
+    @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
+        );
+        requestSpan.addAttribute(
+            AttributeNames.SHARDS,
+            searchRequestContext.formattedShardStats().isEmpty() ? "null" : searchRequestContext.formattedShardStats()
+        );
+    }
+
+    public Span getRequestSpan() {
+        return requestSpan;
+    }
+}
diff --git a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java
@@ -2310,7 +2310,8 @@ public void onFailure(final Exception e) {
                             client
                         ),
                         NoopMetricsRegistry.INSTANCE,
-                        searchRequestOperationsCompositeListenerFactory
+                        searchRequestOperationsCompositeListenerFactory,
+                        NoopTracer.INSTANCE
                     )
                 );
                 actions.put(

diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java
@@ -19,6 +19,7 @@
 import org.opensearch.http.HttpResponse;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.telemetry.tracing.attributes.Attributes;
+import org.opensearch.telemetry.tracing.noop.NoopSpan;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.transport.Transport;
 import org.opensearch.transport.TransportException;
@@ -32,6 +33,14 @@
 
 public class SpanBuilderTests extends OpenSearchTestCase {
 
+    public void testString() {
+        String spanName = "test-name";
+        SpanCreationContext context = SpanBuilder.from(spanName);
+        Attributes attributes = context.getAttributes();
+        assertEquals(spanName, context.getSpanName());
+        assertNull(attributes);
+    }
+
     public void testHttpRequestContext() {
         HttpRequest httpRequest = createHttpRequest();
         SpanCreationContext context = SpanBuilder.from(httpRequest);
@@ -67,6 +76,16 @@ public void testTransportContext() {
         assertEquals(connection.getNode().getHostAddress(), attributes.getAttributesMap().get(AttributeNames.TRANSPORT_TARGET_HOST));
     }
 
+    public void testParentSpan() {
+        String spanName = "test-name";
+        SpanContext parentSpanContext = new SpanContext(NoopSpan.INSTANCE);
+        SpanCreationContext context = SpanBuilder.from(spanName, parentSpanContext);
+        Attributes attributes = context.getAttributes();
+        assertNull(attributes);
+        assertEquals(spanName, context.getSpanName());
+        assertEquals(parentSpanContext, context.getParent());
+    }
+
     private static Transport.Connection createTransportConnection() {
         return new Transport.Connection() {
             @Override
-Original file line number
+Diff line change
@@ Expand Up / @@ -220,6 +220,7 @@ public final void start() { @@
                         null
                     )
                 );
+                onRequestEnd(searchRequestContext);
                 return;
             }
             executePhase(this);
@@ Expand Down @@