Improve FutureUtils.get exception handling (#50339)

server/src/main/java/org/elasticsearch/action/support/AdapterActionFuture.java

@@ -19,19 +19,25 @@
 
 package org.elasticsearch.action.support;
 
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.action.ActionFuture;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.util.concurrent.BaseFuture;
 import org.elasticsearch.common.util.concurrent.FutureUtils;
+import org.elasticsearch.common.util.concurrent.UncategorizedExecutionException;
 
 import java.util.concurrent.TimeUnit;
 
 public abstract class AdapterActionFuture<T, L> extends BaseFuture<T> implements ActionFuture<T>, ActionListener<L> {
 
     @Override
     public T actionGet() {
-        return FutureUtils.get(this);
+        try {
+            return FutureUtils.get(this);
+        } catch (ElasticsearchException e) {
+            throw unwrapEsException(e);
+        }
     }
 
     @Override
@@ -51,7 +57,11 @@ public T actionGet(TimeValue timeout) {
 
     @Override
     public T actionGet(long timeout, TimeUnit unit) {
-        return FutureUtils.get(this, timeout, unit);
+        try {
+            return FutureUtils.get(this, timeout, unit);
+        } catch (ElasticsearchException e) {
+            throw unwrapEsException(e);
+        }
     }
 
     @Override
@@ -66,4 +76,11 @@ public void onFailure(Exception e) {
 
     protected abstract T convert(L listenerResponse);
 
+    private static RuntimeException unwrapEsException(ElasticsearchException esEx) {
+        Throwable root = esEx.unwrapCause();
+        if (root instanceof RuntimeException) {
+            return (RuntimeException) root;
+        }
+        return new UncategorizedExecutionException("Failed execution", root);
+    }
 }

server/src/main/java/org/elasticsearch/cluster/action/index/MappingUpdatedAction.java

@@ -31,7 +31,7 @@
 import org.elasticsearch.common.settings.Setting.Property;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
-import org.elasticsearch.common.util.concurrent.FutureUtils;
+import org.elasticsearch.common.util.concurrent.UncategorizedExecutionException;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.mapper.MapperService;
@@ -89,7 +89,16 @@ public void onFailure(Exception e) {
             });
     }
 
+    // todo: this explicit unwrap should not be necessary, but is until guessRootCause is fixed to allow wrapped non-es exception.
     private static Exception unwrapException(Exception cause) {
-        return cause instanceof ElasticsearchException ? FutureUtils.unwrapEsException((ElasticsearchException) cause) : cause;
+        return cause instanceof ElasticsearchException ? unwrapEsException((ElasticsearchException) cause) : cause;
+    }
+
+    private static RuntimeException unwrapEsException(ElasticsearchException esEx) {
+        Throwable root = esEx.unwrapCause();
+        if (root instanceof RuntimeException) {
+            return (RuntimeException) root;
+        }
+        return new UncategorizedExecutionException("Failed execution", root);
     }
 }

server/src/main/java/org/elasticsearch/common/util/concurrent/FutureUtils.java

@@ -19,7 +19,6 @@
 
 package org.elasticsearch.common.util.concurrent;
 
-import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ElasticsearchTimeoutException;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.SuppressForbidden;
@@ -86,21 +85,10 @@ public static  <T> T get(Future<T> future, long timeout, TimeUnit unit) {
     }
 
     public static RuntimeException rethrowExecutionException(ExecutionException e) {
-        if (e.getCause() instanceof ElasticsearchException) {
-            ElasticsearchException esEx = (ElasticsearchException) e.getCause();
-            return unwrapEsException(esEx);
-        } else if (e.getCause() instanceof RuntimeException) {
+        if (e.getCause() instanceof RuntimeException) {
             return (RuntimeException) e.getCause();
         } else {
             return new UncategorizedExecutionException("Failed execution", e);
         }
     }
-
-    public static RuntimeException unwrapEsException(ElasticsearchException esEx) {
-        Throwable root = esEx.unwrapCause();
-        if (root instanceof ElasticsearchException || root instanceof RuntimeException) {
-            return (RuntimeException) root;
-        }
-        return new UncategorizedExecutionException("Failed execution", root);
-    }
 }

server/src/test/java/org/elasticsearch/action/StepListenerTests.java

@@ -22,11 +22,13 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.transport.RemoteTransportException;
 import org.junit.After;
 import org.junit.Before;
 
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 
 import static org.hamcrest.Matchers.equalTo;
@@ -110,4 +112,20 @@ private void executeAction(Runnable runnable) {
             runnable.run();
         }
     }
+
+    /**
+     * This test checks that we no longer unwrap exceptions when using StepListener.
+     */
+    public void testNoUnwrap() {
+        StepListener<String> step = new StepListener<>();
+        step.onFailure(new RemoteTransportException("test", new RuntimeException("expected")));
+        AtomicReference<RuntimeException> exception = new AtomicReference<>();
+        step.whenComplete(null, e -> {
+            exception.set((RuntimeException) e);
+        });
+
+        assertEquals(RemoteTransportException.class, exception.get().getClass());
+        RuntimeException e = expectThrows(RuntimeException.class, () -> step.result());
+        assertEquals(RemoteTransportException.class, e.getClass());
+    }
 }

server/src/test/java/org/elasticsearch/action/support/AdapterActionFutureTests.java

@@ -19,12 +19,16 @@
 
 package org.elasticsearch.action.support;
 
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.common.util.concurrent.UncategorizedExecutionException;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.transport.RemoteTransportException;
 
 import java.util.Objects;
 import java.util.concurrent.BrokenBarrierException;
 import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -90,4 +94,28 @@ protected String convert(final Integer listenerResponse) {
         thread.join();
     }
 
+    public void testUnwrapException() {
+        checkUnwrap(new RemoteTransportException("test", new RuntimeException()), RuntimeException.class, RemoteTransportException.class);
+        checkUnwrap(new RemoteTransportException("test", new Exception()),
+            UncategorizedExecutionException.class, RemoteTransportException.class);
+        checkUnwrap(new Exception(), UncategorizedExecutionException.class, Exception.class);
+        checkUnwrap(new ElasticsearchException("test", new Exception()), ElasticsearchException.class, ElasticsearchException.class);
+    }
+
+    private void checkUnwrap(Exception exception, Class<? extends Exception> actionGetException,
+                             Class<? extends Exception> getException) {
+        final AdapterActionFuture<Void, Void> adapter = new AdapterActionFuture<Void, Void>() {
+            @Override
+            protected Void convert(Void listenerResponse) {
+                fail();
+                return null;
+            }
+        };
+
+        adapter.onFailure(exception);
+        assertEquals(actionGetException, expectThrows(RuntimeException.class, adapter::actionGet).getClass());
+        assertEquals(actionGetException, expectThrows(RuntimeException.class, () -> adapter.actionGet(10, TimeUnit.SECONDS)).getClass());
+        assertEquals(getException, expectThrows(ExecutionException.class, () -> adapter.get()).getCause().getClass());
+        assertEquals(getException, expectThrows(ExecutionException.class, () -> adapter.get(10, TimeUnit.SECONDS)).getCause().getClass());
+    }
 }