Coordinator req & phase tracing

dzane17 · Jan 5, 2024 · 4973fdb · 4973fdb
1 parent 3a3da4f
commit 4973fdb
Show file tree

Hide file tree

Showing 17 changed files with 324 additions and 97 deletions.
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
                 )
             );
+            searchRequestContext.getSearchRequestOperationsListener().onRequestEnd(searchRequestContext);
             return;
         }
         executePhase(this);
@@ -439,18 +440,18 @@ private void onPhaseEnd(SearchRequestContext searchRequestContext) {
         this.searchRequestContext.getSearchRequestOperationsListener().onPhaseEnd(this, searchRequestContext);
     }
 
-    private void onPhaseStart(SearchPhase phase) {
+    private void onPhaseStart(SearchPhase phase, SearchRequestContext searchRequestContext) {
         setCurrentPhase(phase);
-        this.searchRequestContext.getSearchRequestOperationsListener().onPhaseStart(this);
+        this.searchRequestContext.getSearchRequestOperationsListener().onPhaseStart(this, searchRequestContext);
     }
 
     private void onRequestEnd(SearchRequestContext searchRequestContext) {
-        this.searchRequestContext.getSearchRequestOperationsListener().onRequestEnd(this, searchRequestContext);
+        this.searchRequestContext.getSearchRequestOperationsListener().onRequestEnd(searchRequestContext);
     }
 
     private void executePhase(SearchPhase phase) {
         try {
-            onPhaseStart(phase);
+            onPhaseStart(phase, searchRequestContext);
             phase.recordAndRun();
         } catch (Exception e) {
             if (logger.isDebugEnabled()) {
@@ -703,6 +704,7 @@ public void sendSearchResponse(InternalSearchResponse internalSearchResponse, At
                     searchContextId = null;
                 }
             }
+            searchRequestContext.setSearchTask(getTask());
             searchRequestContext.setTotalHits(internalSearchResponse.hits().getTotalHits());
             searchRequestContext.setShardStats(results.getNumShards(), successfulOps.get(), skippedOps.get(), failures.length);
             onPhaseEnd(searchRequestContext);
@@ -714,7 +716,7 @@ public void sendSearchResponse(InternalSearchResponse internalSearchResponse, At
 
     @Override
     public final void onPhaseFailure(SearchPhase phase, String msg, Throwable cause) {
-        this.searchRequestContext.getSearchRequestOperationsListener().onPhaseFailure(this);
+        this.searchRequestContext.getSearchRequestOperationsListener().onPhaseFailure(this, searchRequestContext);
         raisePhaseFailure(new SearchPhaseExecutionException(phase.getName(), msg, cause, buildShardFailures()));
     }
 

diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
@@ -11,6 +11,9 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.lucene.search.TotalHits;
 import org.opensearch.common.annotation.InternalApi;
+import org.opensearch.telemetry.tracing.Span;
+import org.opensearch.telemetry.tracing.Tracer;
+import org.opensearch.telemetry.tracing.noop.NoopTracer;
 
 import java.util.EnumMap;
 import java.util.HashMap;
@@ -25,26 +28,59 @@
  */
 @InternalApi
 class SearchRequestContext {
+    private final SearchRequest searchRequest;
+    private SearchTask searchTask;
     private final SearchRequestOperationsListener searchRequestOperationsListener;
     private long absoluteStartNanos;
     private final Map<String, Long> phaseTookMap;
     private TotalHits totalHits;
     private final EnumMap<ShardStatsFieldNames, Integer> shardStats;
+    private final Tracer tracer;
+    private Span requestSpan;
+    private Span phaseSpan;
 
     /**
      * This constructor is for testing only
      */
     SearchRequestContext() {
-        this(new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()));
+        this(new SearchRequest());
     }
 
-    SearchRequestContext(SearchRequestOperationsListener searchRequestOperationsListener) {
+    /**
+     * This constructor is for testing only
+     */
+    SearchRequestContext(SearchRequest searchRequest) {
+        this(searchRequest, new SearchRequestOperationsListener.CompositeListener(List.of(), LogManager.getLogger()), NoopTracer.INSTANCE);
+    }
+
+    /**
+     * This constructor is for testing only
+     */
+    SearchRequestContext(SearchRequest searchRequest, SearchRequestOperationsListener searchRequestOperationsListener) {
+        this(searchRequest, searchRequestOperationsListener, NoopTracer.INSTANCE);
+    }
+
+    SearchRequestContext(SearchRequest searchRequest, SearchRequestOperationsListener searchRequestOperationsListener, Tracer tracer) {
+        this.searchRequest = searchRequest;
         this.searchRequestOperationsListener = searchRequestOperationsListener;
+        this.tracer = tracer;
         this.absoluteStartNanos = System.nanoTime();
         this.phaseTookMap = new HashMap<>();
         this.shardStats = new EnumMap<>(ShardStatsFieldNames.class);
     }
 
+    public SearchRequest getSearchRequest() {
+        return searchRequest;
+    }
+
+    public void setSearchTask(SearchTask searchTask) {
+        this.searchTask = searchTask;
+    }
+
+    public SearchTask getSearchTask() {
+        return searchTask;
+    }
+
     SearchRequestOperationsListener getSearchRequestOperationsListener() {
         return searchRequestOperationsListener;
     }
@@ -76,7 +112,7 @@ void setTotalHits(TotalHits totalHits) {
         this.totalHits = totalHits;
     }
 
-    TotalHits totalHits() {
+    public TotalHits totalHits() {
         return totalHits;
     }
 
@@ -87,7 +123,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 {
@@ -105,6 +141,26 @@ String formattedShardStats() {
             );
         }
     }
+
+    public Tracer getTracer() {
+        return tracer;
+    }
+
+    public void setRequestSpan(Span requestSpan) {
+        this.requestSpan = requestSpan;
+    }
+
+    public Span getRequestSpan() {
+        return requestSpan;
+    }
+
+    public void setPhaseSpan(Span phaseSpan) {
+        this.phaseSpan = phaseSpan;
+    }
+
+    public Span getPhaseSpan() {
+        return phaseSpan;
+    }
 }
 
 enum ShardStatsFieldNames {

diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestCoordinatorTrace.java b/server/src/main/java/org/opensearch/action/search/SearchRequestCoordinatorTrace.java
@@ -0,0 +1,87 @@
+/*
+ * 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.action.search;
+
+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;
+
+/**
+ * Listener for search request tracing on the coordinator node
+ *
+ * @opensearch.internal
+ */
+public final class SearchRequestCoordinatorTrace extends SearchRequestOperationsListener {
+    private final Tracer tracer;
+
+    public SearchRequestCoordinatorTrace(Tracer tracer) {
+        this.tracer = tracer;
+    }
+
+    @Override
+    void onPhaseStart(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        if (searchRequestContext.getPhaseSpan() != null) {
+            searchRequestContext.getPhaseSpan().endSpan();
+        }
+        searchRequestContext.setPhaseSpan(
+            tracer.startSpan(
+                SpanBuilder.from(
+                    "coordinator" + capitalize(context.getCurrentPhase().getName()),
+                    new SpanContext(searchRequestContext.getRequestSpan())
+                )
+            )
+        );
+        SpanScope spanScope = tracer.withSpanInScope(searchRequestContext.getPhaseSpan());
+    }
+
+    @Override
+    void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        searchRequestContext.getPhaseSpan().endSpan();
+        searchRequestContext.setPhaseSpan(null);
+    }
+
+    @Override
+    void onPhaseFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        searchRequestContext.getPhaseSpan().endSpan();
+    }
+
+    @Override
+    void onRequestStart(SearchRequestContext searchRequestContext) {}
+
+    @Override
+    void onRequestEnd(SearchRequestContext searchRequestContext) {
+        Span requestSpan = searchRequestContext.getRequestSpan();
+
+        // add response-related attributes on request end
+        requestSpan.addAttribute(
+            AttributeNames.TOTAL_HITS,
+            searchRequestContext.totalHits() == null ? "0" : searchRequestContext.totalHits().toString()
+        );
+        requestSpan.addAttribute(
+            AttributeNames.SHARDS,
+            searchRequestContext.formattedShardStats().isEmpty() ? "no data" : searchRequestContext.formattedShardStats()
+        );
+        requestSpan.addAttribute(
+            AttributeNames.SOURCE,
+            searchRequestContext.getSearchRequest().source() == null
+                ? "no source"
+                : searchRequestContext.getSearchRequest().source().toString()
+        );
+    }
+
+    @Override
+    void onRequestFailure(SearchRequestContext searchRequestContext) {
+        onRequestEnd(searchRequestContext);
+    }
+}
diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsListener.java b/server/src/main/java/org/opensearch/action/search/SearchRequestOperationsListener.java
@@ -22,15 +22,19 @@
 @InternalApi
 abstract class SearchRequestOperationsListener {
 
-    abstract void onPhaseStart(SearchPhaseContext context);
+    abstract void onPhaseStart(SearchPhaseContext context, SearchRequestContext searchRequestContext);
 
     abstract void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext);
 
-    abstract void onPhaseFailure(SearchPhaseContext context);
+    abstract void onPhaseFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext);
 
     void onRequestStart(SearchRequestContext searchRequestContext) {}
 
-    void onRequestEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {}
+    void onRequestEnd(SearchRequestContext searchRequestContext) {}
+
+    void onRequestFailure(SearchRequestContext searchRequestContext) {
+        onRequestEnd(searchRequestContext);
+    }
 
     /**
      * Holder of Composite Listeners
@@ -48,10 +52,10 @@ static final class CompositeListener extends SearchRequestOperationsListener {
         }
 
         @Override
-        void onPhaseStart(SearchPhaseContext context) {
+        void onPhaseStart(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
             for (SearchRequestOperationsListener listener : listeners) {
                 try {
-                    listener.onPhaseStart(context);
+                    listener.onPhaseStart(context, searchRequestContext);
                 } catch (Exception e) {
                     logger.warn(() -> new ParameterizedMessage("onPhaseStart listener [{}] failed", listener), e);
                 }
@@ -70,10 +74,10 @@ void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestCo
         }
 
         @Override
-        void onPhaseFailure(SearchPhaseContext context) {
+        void onPhaseFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
             for (SearchRequestOperationsListener listener : listeners) {
                 try {
-                    listener.onPhaseFailure(context);
+                    listener.onPhaseFailure(context, searchRequestContext);
                 } catch (Exception e) {
                     logger.warn(() -> new ParameterizedMessage("onPhaseFailure listener [{}] failed", listener), e);
                 }
@@ -92,14 +96,25 @@ void onRequestStart(SearchRequestContext searchRequestContext) {
         }
 
         @Override
-        public void onRequestEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        public void onRequestEnd(SearchRequestContext searchRequestContext) {
             for (SearchRequestOperationsListener listener : listeners) {
                 try {
-                    listener.onRequestEnd(context, searchRequestContext);
+                    listener.onRequestEnd(searchRequestContext);
                 } catch (Exception e) {
                     logger.warn(() -> new ParameterizedMessage("onRequestEnd listener [{}] failed", listener), e);
                 }
             }
         }
+
+        @Override
+        public void onRequestFailure(SearchRequestContext searchRequestContext) {
+            for (SearchRequestOperationsListener listener : listeners) {
+                try {
+                    listener.onRequestFailure(searchRequestContext);
+                } catch (Exception e) {
+                    logger.warn(() -> new ParameterizedMessage("onRequestFailure listener [{}] failed", listener), e);
+                }
+            }
+        }
     }
 }