Support retry cancel feature fully for AsyncFeign

OpenFeign · Oct 31, 2022 · 17dc1cc · 17dc1cc
1 parent 73f3309
commit 17dc1cc
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 14 deletions.
diff --git a/core/src/main/java/feign/AsynchronousMethodHandler.java b/core/src/main/java/feign/AsynchronousMethodHandler.java
@@ -88,7 +88,7 @@ public Object invoke(Object[] argv) throws Throwable {
   private CompletableFuture<Object> executeAndDecode(RequestTemplate template,
                                                      Options options,
                                                      Retryer retryer) {
-    CompletableFuture<Object> resultFuture = new CompletableFuture<>();
+    CancellableFuture<Object> resultFuture = new CancellableFuture<>();
 
     executeAndDecode(template, options)
         .whenComplete((response, throwable) -> {
@@ -98,8 +98,8 @@ private CompletableFuture<Object> executeAndDecode(RequestTemplate template,
                 logger.logRetry(metadata.configKey(), logLevel);
               }
 
-              executeAndDecode(template, options, retryer)
-                  .whenComplete(pipeTo(resultFuture));
+              resultFuture.setInner(
+                  executeAndDecode(template, options, retryer));
             }
           } else {
             resultFuture.complete(response);
@@ -109,6 +109,38 @@ private CompletableFuture<Object> executeAndDecode(RequestTemplate template,
     return resultFuture;
   }
 
+  private static class CancellableFuture<T> extends CompletableFuture<T> {
+    private CompletableFuture<T> inner = null;
+
+    public void setInner(CompletableFuture<T> value) {
+      inner = value;
+      inner.whenComplete(pipeTo(this));
+    }
+
+    @Override
+    public boolean cancel(boolean mayInterruptIfRunning) {
+      final boolean result = super.cancel(mayInterruptIfRunning);
+      if (inner != null) {
+        inner.cancel(mayInterruptIfRunning);
+      }
+      return result;
+    }
+
+    private static <T> BiConsumer<? super T, ? super Throwable> pipeTo(CompletableFuture<T> completableFuture) {
+      return (value, throwable) -> {
+        if (completableFuture.isDone()) {
+          return;
+        }
+
+        if (throwable != null) {
+          completableFuture.completeExceptionally(throwable);
+        } else {
+          completableFuture.complete(value);
+        }
+      };
+    }
+  }
+
   private boolean shouldRetry(Retryer retryer,
                               Throwable throwable,
                               CompletableFuture<Object> resultFuture) {
@@ -136,16 +168,6 @@ private boolean shouldRetry(Retryer retryer,
     }
   }
 
-  private static <T> BiConsumer<? super T, ? super Throwable> pipeTo(CompletableFuture<T> completableFuture) {
-    return (value, throwable) -> {
-      if (throwable != null) {
-        completableFuture.completeExceptionally(throwable);
-      } else {
-        completableFuture.complete(value);
-      }
-    };
-  }
-
   private CompletableFuture<Object> executeAndDecode(RequestTemplate template, Options options) {
     Request request = targetRequest(template);
 

diff --git a/core/src/test/java/feign/AsyncFeignTest.java b/core/src/test/java/feign/AsyncFeignTest.java
@@ -662,7 +662,7 @@ public void whenReturnTypeIsResponseNoErrorHandling() throws Throwable {
   }
 
   @ParameterizedTest
-  @ValueSource(ints = {1}) // TODO: Modify it to work even if it is more than 2 tries
+  @ValueSource(ints = {1, 5, 10, 100, 1000})
   public void cancelRetry(final int expectedTryCount) throws Throwable {
     // Arrange
     final CompletableFuture<Boolean> maximumTryCompleted = new CompletableFuture<>();