onShardResult and onShardFailure are executed on one shard causes ope…

…nsearch jvm crashed
opensearch-project · Feb 4, 2024 · 4a957d5 · 4a957d5
1 parent a8dd6a0
commit 4a957d5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 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
@@ -285,6 +285,7 @@ private void performPhaseOnShard(final int shardIndex, final SearchShardIterator
             Runnable r = () -> {
                 final Thread thread = Thread.currentThread();
                 try {
+                    final SearchPhase phase = this;
                     executePhaseOnShard(shardIt, shard, new SearchActionListener<Result>(shard, shardIndex) {
                         @Override
                         public void innerOnResponse(Result result) {
@@ -298,7 +299,13 @@ public void innerOnResponse(Result result) {
                         @Override
                         public void onFailure(Exception t) {
                             try {
-                                onShardFailure(shardIndex, shard, shardIt, t);
+                                // It only happens when onPhaseDone() is called and executePhaseOnShard() fails hard with an exception.
+                                if (totalOps.get() == expectedTotalOps) {
+                                    onPhaseFailure(phase, "The phase has failed", t);
+                                } else {
+                                    onShardFailure(shardIndex, shard, shardIt, t);
+                                }
+
                             } finally {
                                 executeNext(pendingExecutions, thread);
                             }

diff --git a/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java
@@ -194,7 +194,11 @@ protected void executePhaseOnShard(
                 if (failExecutePhaseOnShard) {
                     listener.onFailure(new ShardNotFoundException(shardIt.shardId()));
                 } else {
-                    listener.onResponse(new QuerySearchResult());
+                    try {
+                        listener.onResponse(new QuerySearchResult());
+                    } catch (Exception e) {
+                        listener.onFailure(e);
+                    }
                 }
             }