From 47507ddbdb1aed29bc05705984b4074caae38911 Mon Sep 17 00:00:00 2001
From: David Zane <davizane@amazon.com>
Date: Thu, 25 Jan 2024 18:16:57 -0800
Subject: [PATCH] Can match phase fix

Signed-off-by: David Zane <davizane@amazon.com>
---
 .../opensearch/action/search/AbstractSearchAsyncAction.java  | 5 ++++-
 .../org/opensearch/action/search/SearchRequestContext.java   | 4 ++++
 .../action/search/SearchRequestCoordinatorTrace.java         | 3 ---
 .../org/opensearch/action/search/TransportSearchAction.java  | 2 +-
 4 files changed, 9 insertions(+), 5 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
index 5d912fcbd4c83..90feebab83b4c 100644
--- a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
+++ b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
@@ -443,7 +443,10 @@ private void onPhaseEnd(SearchRequestContext searchRequestContext) {
 
     private void onPhaseStart(SearchPhase phase, SearchRequestContext searchRequestContext) {
         setCurrentPhase(phase);
-        this.searchRequestContext.getSearchRequestOperationsListener().onPhaseStart(this, searchRequestContext);
+        try {
+            SearchPhaseName.valueOf(phase.getSearchPhaseName().toString());
+            this.searchRequestContext.getSearchRequestOperationsListener().onPhaseStart(this, searchRequestContext);
+        } catch (IllegalArgumentException ignored) {}
     }
 
     private void onRequestEnd(SearchRequestContext searchRequestContext) {
diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
index 70598190e4418..c9b980c7c9072 100644
--- a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
+++ b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
@@ -13,6 +13,7 @@
 import org.opensearch.common.annotation.InternalApi;
 import org.opensearch.telemetry.tracing.Span;
 import org.opensearch.telemetry.tracing.Tracer;
+import org.opensearch.telemetry.tracing.noop.NoopSpan;
 import org.opensearch.telemetry.tracing.noop.NoopTracer;
 
 import java.util.EnumMap;
@@ -67,6 +68,9 @@ class SearchRequestContext {
         this.absoluteStartNanos = System.nanoTime();
         this.phaseTookMap = new HashMap<>();
         this.shardStats = new EnumMap<>(ShardStatsFieldNames.class);
+        this.tracer = tracer;
+        this.requestSpan = NoopSpan.INSTANCE;
+        this.phaseSpan = NoopSpan.INSTANCE;
     }
 
     public SearchRequest getSearchRequest() {
diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestCoordinatorTrace.java b/server/src/main/java/org/opensearch/action/search/SearchRequestCoordinatorTrace.java
index 3d934b6b04fa1..658e9858fb3c4 100644
--- a/server/src/main/java/org/opensearch/action/search/SearchRequestCoordinatorTrace.java
+++ b/server/src/main/java/org/opensearch/action/search/SearchRequestCoordinatorTrace.java
@@ -30,9 +30,6 @@ public SearchRequestCoordinatorTrace(Tracer tracer) {
 
     @Override
     void onPhaseStart(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
-        if (searchRequestContext.getPhaseSpan() != null) {
-            searchRequestContext.getPhaseSpan().endSpan();
-        }
         searchRequestContext.setPhaseSpan(
             tracer.startSpan(
                 SpanBuilder.from(
diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
index 93efb54284bb8..2512988f56f01 100644
--- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
+++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -1253,7 +1253,7 @@ private AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction
                         clusters,
                         searchRequestContext
                     );
-                    return new SearchPhase(action.getName()) {
+                    return new SearchPhase("none") {
                         @Override
                         public void run() {
                             action.start();