From 7470dcd951f35bdeb46368bdeb45a38775566d84 Mon Sep 17 00:00:00 2001
From: Henning Andersen <33268011+henningandersen@users.noreply.github.com>
Date: Mon, 3 Feb 2020 14:02:31 +0100
Subject: [PATCH] Fix completeWith exception handling (#51734)

ActionListener.completeWith would catch exceptions from
listener.onResponse and deliver them to lister.onFailure, essentially
double notifying the listener. Instead we now assert that listeners do
not throw when using ActionListener.completeWith.

Relates #50886
---
 .../elasticsearch/action/ActionListener.java  | 20 ++++++++++++--
 .../action/ActionListenerTests.java           | 27 +++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/action/ActionListener.java b/server/src/main/java/org/elasticsearch/action/ActionListener.java
index 3cb205f9847b3..6e8cc99f7394e 100644
--- a/server/src/main/java/org/elasticsearch/action/ActionListener.java
+++ b/server/src/main/java/org/elasticsearch/action/ActionListener.java
@@ -315,12 +315,28 @@ protected void innerOnFailure(Exception e) {
     /**
      * Completes the given listener with the result from the provided supplier accordingly.
      * This method is mainly used to complete a listener with a block of synchronous code.
+     *
+     * If the supplier fails, the listener's onFailure handler will be called.
+     * It is the responsibility of {@code delegate} to handle its own exceptions inside `onResponse` and `onFailure`.
      */
     static <Response> void completeWith(ActionListener<Response> listener, CheckedSupplier<Response, ? extends Exception> supplier) {
+        Response response;
         try {
-            listener.onResponse(supplier.get());
+            response = supplier.get();
         } catch (Exception e) {
-            listener.onFailure(e);
+            try {
+                listener.onFailure(e);
+            } catch (RuntimeException ex) {
+                assert false : ex;
+                throw ex;
+            }
+            return;
+        }
+        try {
+            listener.onResponse(response);
+        } catch (RuntimeException ex) {
+            assert false : ex;
+            throw ex;
         }
     }
 }
diff --git a/server/src/test/java/org/elasticsearch/action/ActionListenerTests.java b/server/src/test/java/org/elasticsearch/action/ActionListenerTests.java
index 3577577dd9e7a..a399f3c9c67ed 100644
--- a/server/src/test/java/org/elasticsearch/action/ActionListenerTests.java
+++ b/server/src/test/java/org/elasticsearch/action/ActionListenerTests.java
@@ -233,6 +233,33 @@ public void testCompleteWith() {
         ActionListener.completeWith(onFailureListener, () -> { throw new IOException("not found"); });
         assertThat(onFailureListener.isDone(), equalTo(true));
         assertThat(expectThrows(ExecutionException.class, onFailureListener::get).getCause(), instanceOf(IOException.class));
+
+        AtomicReference<Exception> exReference = new AtomicReference<>();
+        ActionListener<String> listener = new ActionListener<>() {
+            @Override
+            public void onResponse(String s) {
+                if (s == null) {
+                    throw new IllegalArgumentException("simulate onResponse exception");
+                }
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                exReference.set(e);
+                if (e instanceof IllegalArgumentException) {
+                    throw (IllegalArgumentException) e;
+                }
+            }
+        };
+
+        AssertionError assertionError = expectThrows(AssertionError.class, () -> ActionListener.completeWith(listener, () -> null));
+        assertThat(assertionError.getCause(), instanceOf(IllegalArgumentException.class));
+        assertNull(exReference.get());
+
+        assertionError = expectThrows(AssertionError.class, () -> ActionListener.completeWith(listener,
+            () -> { throw new IllegalArgumentException(); }));
+        assertThat(assertionError.getCause(), instanceOf(IllegalArgumentException.class));
+        assertThat(exReference.get(), instanceOf(IllegalArgumentException.class));
     }
 
     /**