onShardResult and onShardFailure are executed on one shard causes ope…

…nsearch jvm crashed Signed-off-by: kkewwei <[email protected]>
opensearch-project · Mar 10, 2024 · dd46b25 · dd46b25
1 parent 2b17902
commit dd46b25
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
@@ -286,6 +286,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) {
@@ -299,7 +300,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
@@ -221,7 +221,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);
+                    }
                 }
             }