From 868dc3c4e4e643f2cdbe03f1b6a94fba3ba3c51a Mon Sep 17 00:00:00 2001
From: Donghyeon Kim <wplong11@users.noreply.github.com>
Date: Sun, 23 Oct 2022 10:16:04 +0900
Subject: [PATCH] Support retry cancel feature fully for AsyncFeign

---
 .../java/feign/AsynchronousMethodHandler.java | 49 ++++++++++++++-----
 core/src/test/java/feign/AsyncFeignTest.java  |  2 +-
 2 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/core/src/main/java/feign/AsynchronousMethodHandler.java b/core/src/main/java/feign/AsynchronousMethodHandler.java
index 2376d00a21..99225c0b35 100644
--- 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,9 @@ 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 +110,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 +169,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
index da10da8619..1db75e9748 100644
--- 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<>();