From 79034613d8dd687701e31db25d78d74a05d49005 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   | 5 ++++-
 .../action/search/SearchRequestCoordinatorTrace.java         | 3 ---
 .../org/opensearch/action/search/TransportSearchAction.java  | 2 +-
 4 files changed, 9 insertions(+), 6 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 abd9f962939f4..a19a60a918cca 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 4d36347810530..5c32b61232046 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;
@@ -63,10 +64,12 @@ class SearchRequestContext {
     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);
+        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 5b4aabe6d5c2f..e9bcea26453e9 100644
--- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
+++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -1346,7 +1346,7 @@ private AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction
                         clusters,
                         searchRequestContext
                     );
-                    return new SearchPhase(action.getName()) {
+                    return new SearchPhase("none") {
                         @Override
                         public void run() {
                             action.start();