Ensure that cancelled request returns

I tried cancelling a request and found that the client would hang. This change reports an exception for each remaining request instead. Signed-off-by: Michael Froh <[email protected]>
opensearch-project · Jan 24, 2025 · d3331d9 · d3331d9
1 parent 84a7311
commit d3331d9
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/server/src/main/java/org/opensearch/action/search/TransportMultiSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportMultiSearchAction.java
@@ -44,6 +44,7 @@
 import org.opensearch.common.util.concurrent.OpenSearchExecutors;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.tasks.TaskCancelledException;
 import org.opensearch.core.tasks.TaskId;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
@@ -195,7 +196,20 @@ private void handleResponse(final int responseSlot, final MultiSearchResponse.It
                 if (responseCounter.decrementAndGet() == 0) {
                     assert requests.isEmpty();
                     finish();
-                } else if (isCancelled(request.request.getParentTask()) == false) {
+                } else if (isCancelled(request.request.getParentTask())) {
+                    // Drain the rest of the queue
+                    SearchRequestSlot request;
+                    while ((request = requests.poll()) != null) {
+                        responses.set(
+                            request.responseSlot,
+                            new MultiSearchResponse.Item(null, new TaskCancelledException("Parent task was cancelled"))
+                        );
+                        if (responseCounter.decrementAndGet() == 0) {
+                            assert requests.isEmpty();
+                            finish();
+                        }
+                    }
+                } else {
                     if (thread == Thread.currentThread()) {
                         // we are on the same thread, we need to fork to another thread to avoid recursive stack overflow on a single thread
                         threadPool.generic()