Add support for async rate limiter waiting

Closes #318

Author: jhalterman

CHANGELOG.md

@@ -3,6 +3,11 @@
 ### Improvements
 
 - Issue #309 - Introduced a `Bulkhead` policy.
+- Issue #318 - Add non-blocking async waiting for rate limiters.
+
+### SPI Changes
+
+- `PolicyExecutor.preExecuteAsync` was introduced to support async pre-execution. This is backwards compatible with `preExecute`.
 
 # 3.1.0
 

src/main/java/dev/failsafe/internal/RateLimiterExecutor.java

@@ -18,9 +18,14 @@
 import dev.failsafe.RateLimitExceededException;
 import dev.failsafe.RateLimiter;
 import dev.failsafe.spi.ExecutionResult;
+import dev.failsafe.spi.FailsafeFuture;
 import dev.failsafe.spi.PolicyExecutor;
+import dev.failsafe.spi.Scheduler;
 
 import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 
 /**
  * A PolicyExecutor that handles failures according to a {@link RateLimiter}.
@@ -50,4 +55,31 @@ protected ExecutionResult<R> preExecute() {
       return ExecutionResult.failure(e);
     }
   }
+
+  @Override
+  protected CompletableFuture<ExecutionResult<R>> preExecuteAsync(Scheduler scheduler, FailsafeFuture<R> future) {
+    CompletableFuture<ExecutionResult<R>> promise = new CompletableFuture<>();
+    long waitNanos = rateLimiter.acquirePermitWaitNanos(maxWaitTime);
+    if (waitNanos == -1)
+      promise.complete(ExecutionResult.failure(new RateLimitExceededException(rateLimiter)));
+    else {
+      try {
+        Future<?> scheduledWait = scheduler.schedule(() -> {
+          // Signal for execution and post-execution to proceed with a non-result
+          return promise.complete(ExecutionResult.none());
+        }, waitNanos, TimeUnit.NANOSECONDS);
+
+        // Propagate outer cancellations to the RateLimiter future and its promise
+        future.setCancelFn(this, (mayInterrupt, cancelResult) -> {
+          scheduledWait.cancel(mayInterrupt);
+          promise.complete(cancelResult);
+        });
+      } catch (Throwable t) {
+        // Hard scheduling failure
+        promise.completeExceptionally(t);
+      }
+    }
+
+    return promise;
+  }
 }

src/main/java/dev/failsafe/internal/RateLimiterImpl.java

@@ -21,8 +21,10 @@
 import dev.failsafe.internal.util.Assert;
 import dev.failsafe.internal.util.Durations;
 import dev.failsafe.spi.PolicyExecutor;
+import dev.failsafe.spi.Scheduler;
 
 import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -81,4 +83,12 @@ public boolean tryAcquirePermits(int permits, Duration maxWaitTime) throws Inter
   public PolicyExecutor<R> toExecutor(int policyIndex) {
     return new RateLimiterExecutor<>(this, policyIndex);
   }
+
+  /**
+   * Returns the wait nanos for an acquired permit which can be used to externally wait.
+   */
+  long acquirePermitWaitNanos(Duration maxWaitTime) {
+    Assert.notNull(maxWaitTime, "maxWaitTime");
+    return stats.acquirePermits(1, Durations.ofSafeNanos(maxWaitTime));
+  }
 }

src/main/java/dev/failsafe/spi/PolicyExecutor.java

@@ -15,9 +15,9 @@
  */
 package dev.failsafe.spi;
 
-import dev.failsafe.internal.EventHandler;
 import dev.failsafe.ExecutionContext;
 import dev.failsafe.Policy;
+import dev.failsafe.internal.EventHandler;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Function;
@@ -61,6 +61,17 @@ protected ExecutionResult<R> preExecute() {
     return null;
   }
 
+  /**
+   * Called before an async execution to return an alternative result or failure such as if execution is not allowed or
+   * needed. Returns {@code null} if pre execution is not performed. If the resulting future is completed with a {@link
+   * ExecutionResult#isNonResult() non-result}, then execution and post-execution should still be performed. If the
+   * resulting future is completed with {@code null}, then the execution is assumed to have been cancelled.
+   */
+  protected CompletableFuture<ExecutionResult<R>> preExecuteAsync(Scheduler scheduler, FailsafeFuture<R> future) {
+    ExecutionResult<R> result = preExecute();
+    return result == null ? null : CompletableFuture.completedFuture(result);
+  }
+
   /**
    * Performs an execution by calling pre-execute else calling the supplier and doing a post-execute.
    */
@@ -69,7 +80,7 @@ public Function<SyncExecutionInternal<R>, ExecutionResult<R>> apply(
     return execution -> {
       ExecutionResult<R> result = preExecute();
       if (result != null) {
-        // Still need to preExecute when returning an alternative result before making it to the terminal Supplier
+        // Still need to preExecute when short-circuiting an execution with an alternative result
         execution.preExecute();
         return result;
       }
@@ -78,23 +89,6 @@ public Function<SyncExecutionInternal<R>, ExecutionResult<R>> apply(
     };
   }
 
-  /**
-   * Performs synchronous post-execution handling for a {@code result}.
-   */
-  public ExecutionResult<R> postExecute(ExecutionInternal<R> execution, ExecutionResult<R> result) {
-    execution.recordAttempt();
-    if (isFailure(result)) {
-      result = onFailure(execution, result.withFailure());
-      handleFailure(result, execution);
-    } else {
-      result = result.withSuccess();
-      onSuccess(result);
-      handleSuccess(result, execution);
-    }
-
-    return result;
-  }
-
   /**
    * Performs an async execution by calling pre-execute else calling the supplier and doing a post-execute. Implementors
    * must handle a null result from a supplier, which indicates that an async execution has occurred, that a result will
@@ -105,21 +99,65 @@ public Function<AsyncExecutionInternal<R>, CompletableFuture<ExecutionResult<R>>
     FailsafeFuture<R> future) {
 
     return execution -> {
+      CompletableFuture<ExecutionResult<R>> promise = new CompletableFuture<>();
+      Runnable runnable = () -> innerFn.apply(execution)
+        .thenCompose(r -> r == null ? ExecutionResult.nullFuture() : postExecuteAsync(execution, r, scheduler, future))
+        .whenComplete((postResult, postError) -> {
+          if (postError != null)
+            promise.completeExceptionally(postError);
+          else
+            promise.complete(postResult);
+        });
+
       if (!execution.isRecorded()) {
-        ExecutionResult<R> result = preExecute();
-        if (result != null) {
-          // Still need to preExecute when returning an alternative result before making it to the terminal Supplier
-          execution.preExecute();
-          return CompletableFuture.completedFuture(result);
+        CompletableFuture<ExecutionResult<R>> preResult = preExecuteAsync(scheduler, future);
+        if (preResult != null) {
+          preResult.whenComplete((result, error) -> {
+            if (error != null)
+              promise.completeExceptionally(error);
+            else if (result != null) {
+              // Check for non-result, which occurs after a rate limiter preExecute delay
+              if (result.isNonResult()) {
+                // Execute and post-execute
+                runnable.run();
+              } else {
+                // Still need to preExecute when short-circuiting an execution with an alternative result
+                execution.preExecute();
+                promise.complete(result);
+              }
+            } else {
+              // If result is null, the execution is assumed to have been cancelled
+              promise.complete(null);
+            }
+          });
+
+          return promise;
         }
       }
 
-      return innerFn.apply(execution).thenCompose(r -> {
-        return r == null ? ExecutionResult.nullFuture() : postExecuteAsync(execution, r, scheduler, future);
-      });
+      // Execute and post-execute
+      runnable.run();
+      return promise;
     };
   }
 
+  /**
+   * Performs synchronous post-execution handling for a {@code result}.
+   */
+  public ExecutionResult<R> postExecute(ExecutionInternal<R> execution, ExecutionResult<R> result) {
+    execution.recordAttempt();
+    if (isFailure(result)) {
+      result = onFailure(execution, result.withFailure());
+      handleFailure(result, execution);
+    } else {
+      result = result.withSuccess();
+      onSuccess(result);
+      handleSuccess(result, execution);
+    }
+
+    return result;
+  }
+
   /**
    * Performs potentially asynchronous post-execution handling for a {@code result}.
    */

src/test/java/dev/failsafe/functional/FailsafeFutureCancellationTest.java

@@ -16,6 +16,7 @@
 package dev.failsafe.functional;
 
 import dev.failsafe.*;
+import dev.failsafe.function.ContextualRunnable;
 import dev.failsafe.testing.Testing;
 import net.jodah.concurrentunit.Waiter;
 import dev.failsafe.event.ExecutionCompletedEvent;
@@ -41,13 +42,13 @@ void beforeMethod() {
     waiter = new Waiter();
   }
 
-  private <R> void assertCancel(FailsafeExecutor<R> executor, ContextualSupplier<R, R> supplier) throws Throwable {
+  private void assertCancel(FailsafeExecutor<Void> executor, ContextualRunnable<Void> runnable) throws Throwable {
     // Given
-    CompletableFuture<R> future = executor.onComplete(e -> {
+    CompletableFuture<Void> future = executor.onComplete(e -> {
       waiter.assertNull(e.getResult());
       waiter.assertTrue(e.getFailure() instanceof CancellationException);
       waiter.resume();
-    }).getAsync(supplier);
+    }).runAsync(runnable);
 
     Testing.sleep(300);
 
@@ -62,23 +63,45 @@ private <R> void assertCancel(FailsafeExecutor<R> executor, ContextualSupplier<R
     assertThrows(future::get, CancellationException.class);
   }
 
-  public void shouldCancelOnGetAsyncWithRetries() throws Throwable {
-    assertCancel(Failsafe.with(retryAlways), ctx -> {
+  public void shouldCancelAsyncRetriesWithPendingDelay() throws Throwable {
+    RetryPolicy<Void> retryPolicy = RetryPolicy.<Void>builder().withDelay(Duration.ofMinutes(1)).build();
+    assertCancel(Failsafe.with(retryPolicy), ctx -> {
+      throw new IllegalStateException();
+    });
+  }
+
+  public void shouldCancelAsyncRetriesWithBlockedExecution() throws Throwable {
+    assertCancel(Failsafe.with(RetryPolicy.ofDefaults()), ctx -> {
       try {
         waiter.assertFalse(ctx.isCancelled());
         Thread.sleep(1000);
       } catch (InterruptedException e) {
         waiter.assertTrue(ctx.isCancelled());
         throw e;
       }
-      return false;
     });
   }
 
-  public void shouldCancelOnGetAsyncWithTimeout() throws Throwable {
+  public void shouldCancelAsyncTimeoutWithBlockedExecution() throws Throwable {
     assertCancel(Failsafe.with(Timeout.of(Duration.ofMinutes(1))), ctx -> {
-      Thread.sleep(1000);
-      return "test";
+      try {
+        waiter.assertFalse(ctx.isCancelled());
+        Thread.sleep(1000);
+      } catch (InterruptedException e) {
+        waiter.assertTrue(ctx.isCancelled());
+        throw e;
+      }
+    });
+  }
+
+  public void shouldCancelAsyncRateLimiterWaitingOnPermit() throws Throwable {
+    RateLimiter<Void> limiter = RateLimiter.<Void>smoothBuilder(1, Duration.ofSeconds(1))
+      .withMaxWaitTime(Duration.ofMinutes(1))
+      .build();
+    limiter.tryAcquirePermit(); // All permits should be used now
+
+    assertCancel(Failsafe.with(limiter), ctx -> {
+      fail("Execution should be cancelled during preExecute");
     });
   }
 

src/test/java/dev/failsafe/functional/RateLimiterTest.java

@@ -15,10 +15,7 @@
  */
 package dev.failsafe.functional;
 
-import dev.failsafe.Failsafe;
-import dev.failsafe.RateLimitExceededException;
-import dev.failsafe.RateLimiter;
-import dev.failsafe.RetryPolicy;
+import dev.failsafe.*;
 import dev.failsafe.testing.Testing;
 import org.testng.annotations.Test;
 
@@ -49,9 +46,7 @@ public void shouldThrowRateLimitExceededExceptionAfterPermitsExceeded() {
    */
   public void testMaxWaitTimeExceeded() {
     // Given
-    RateLimiter<Object> limiter = RateLimiter.smoothBuilder(Duration.ofMillis(10))
-      .withMaxWaitTime(Duration.ofMillis(20))
-      .build();
+    RateLimiter<Object> limiter = RateLimiter.smoothBuilder(Duration.ofMillis(10)).build();
 
     // When / Then
     testRunFailure(() -> {
@@ -90,24 +85,25 @@ public void testRejectedWithRetries() {
   }
 
   /**
-   * Asserts that a rate limiter propagates an InterruptedException.
+   * Asserts that a rate limiter propagates a sync InterruptedException.
    */
   public void testAcquirePermitWithInterrupt() {
     RateLimiter<Object> limiter = RateLimiter.smoothBuilder(Duration.ofSeconds(1))
       .withMaxWaitTime(Duration.ofSeconds(5))
       .build();
+    limiter.tryAcquirePermit();
 
-    testRunFailure(() -> {
-      resetLimiter(limiter);
-      limiter.tryAcquirePermit();
-      Thread thread = Thread.currentThread();
-      runInThread(() -> {
-        Thread.sleep(100);
-        thread.interrupt();
-      });
-    }, Failsafe.with(limiter), ctx -> {
+    Thread thread = Thread.currentThread();
+    runInThread(() -> {
+      Thread.sleep(100);
+      thread.interrupt();
+    });
+    assertThrows(() -> Failsafe.with(limiter).run(() -> {
       System.out.println("Executing");
       throw new Exception();
-    }, InterruptedException.class);
+    }), FailsafeException.class, InterruptedException.class);
+
+    // Reset interrupt flag
+    Thread.interrupted();
   }
 }