Add async waiting support for bulkheads

Author: jhalterman

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

@@ -20,9 +20,14 @@
 import dev.failsafe.ExecutionContext;
 import dev.failsafe.RateLimitExceededException;
 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 Bulkhead}.
@@ -53,6 +58,39 @@ protected ExecutionResult<R> preExecute() {
     }
   }
 
+  @Override
+  protected CompletableFuture<ExecutionResult<R>> preExecuteAsync(Scheduler scheduler, FailsafeFuture<R> future) {
+    CompletableFuture<ExecutionResult<R>> promise = new CompletableFuture<>();
+    CompletableFuture<Void> acquireFuture = bulkhead.acquirePermitAsync();
+    acquireFuture.whenComplete((result, error) -> {
+      // Signal for execution to proceed
+      promise.complete(ExecutionResult.none());
+    });
+
+    if (!promise.isDone()) {
+      try {
+        // Schedule bulkhead permit timeout
+        Future<?> timeoutFuture = scheduler.schedule(() -> {
+          promise.complete(ExecutionResult.failure(new BulkheadFullException(bulkhead)));
+          acquireFuture.cancel(true);
+          return null;
+        }, maxWaitTime.toNanos(), TimeUnit.NANOSECONDS);
+
+        // Propagate outer cancellations to the promise, bulkhead acquire future, and timeout future
+        future.setCancelFn(this, (mayInterrupt, cancelResult) -> {
+          promise.complete(cancelResult);
+          acquireFuture.cancel(mayInterrupt);
+          timeoutFuture.cancel(mayInterrupt);
+        });
+      } catch (Throwable t) {
+        // Hard scheduling failure
+        promise.completeExceptionally(t);
+      }
+    }
+
+    return promise;
+  }
+
   @Override
   public void onSuccess(ExecutionResult<R> result) {
     bulkhead.releasePermit();

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

@@ -17,26 +17,31 @@
 
 import dev.failsafe.Bulkhead;
 import dev.failsafe.BulkheadConfig;
-import dev.failsafe.internal.util.Durations;
+import dev.failsafe.internal.util.FutureLinkedList;
 import dev.failsafe.spi.PolicyExecutor;
 
 import java.time.Duration;
-import java.util.concurrent.Semaphore;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 
 /**
- * A Bulkhead implementation.
+ * A Bulkhead implementation that supports sync and async waiting.
  *
  * @param <R> result type
  * @author Jonathan Halterman
  */
 public class BulkheadImpl<R> implements Bulkhead<R> {
+  private static final CompletableFuture<Void> NULL_FUTURE = CompletableFuture.completedFuture(null);
   private final BulkheadConfig<R> config;
-  private final Semaphore semaphore;
+  private final int maxPermits;
+
+  // Mutable state
+  private int permits;
+  private final FutureLinkedList futures = new FutureLinkedList();
 
   public BulkheadImpl(BulkheadConfig<R> config) {
     this.config = config;
-    semaphore = new Semaphore(config.getMaxConcurrency(), true);
+    maxPermits = config.getMaxConcurrency();
+    permits = maxPermits;
   }
 
   @Override
@@ -46,22 +51,57 @@ public BulkheadConfig<R> getConfig() {
 
   @Override
   public void acquirePermit() throws InterruptedException {
-    semaphore.acquire();
+    try {
+      acquirePermitAsync().get();
+    } catch (CancellationException | ExecutionException ignore) {
+      // Not possible since the future will always be completed with null
+    }
   }
 
   @Override
-  public boolean tryAcquirePermit() {
-    return semaphore.tryAcquire();
+  public synchronized boolean tryAcquirePermit() {
+    if (permits > 0) {
+      permits -= 1;
+      return true;
+    }
+    return false;
   }
 
   @Override
   public boolean tryAcquirePermit(Duration maxWaitTime) throws InterruptedException {
-    return semaphore.tryAcquire(Durations.ofSafeNanos(maxWaitTime).toNanos(), TimeUnit.NANOSECONDS);
+    CompletableFuture<Void> future = acquirePermitAsync();
+    if (future == NULL_FUTURE)
+      return true;
+
+    try {
+      future.get(maxWaitTime.toNanos(), TimeUnit.NANOSECONDS);
+      return true;
+    } catch (CancellationException | ExecutionException | TimeoutException e) {
+      return false;
+    }
+  }
+
+  /**
+   * Returns a CompletableFuture that is completed when a permit is acquired. Externally completing this future will
+   * remove the waiter from the bulkhead's internal queue.
+   */
+  synchronized CompletableFuture<Void> acquirePermitAsync() {
+    if (permits > 0) {
+      permits -= 1;
+      return NULL_FUTURE;
+    } else {
+      return futures.add();
+    }
   }
 
   @Override
-  public void releasePermit() {
-    semaphore.release();
+  public synchronized void releasePermit() {
+    if (permits < maxPermits) {
+      permits += 1;
+      CompletableFuture<Void> future = futures.pollFirst();
+      if (future != null)
+        future.complete(null);
+    }
   }
 
   @Override

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

@@ -64,15 +64,17 @@ protected CompletableFuture<ExecutionResult<R>> preExecuteAsync(Scheduler schedu
       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());
+        // Wait for the permit
+        Future<?> permitWaitFuture = scheduler.schedule(() -> {
+          // Signal for execution to proceed
+          promise.complete(ExecutionResult.none());
+          return null;
         }, waitNanos, TimeUnit.NANOSECONDS);
 
-        // Propagate outer cancellations to the RateLimiter future and its promise
+        // Propagate outer cancellations to the promise and permit wait future
         future.setCancelFn(this, (mayInterrupt, cancelResult) -> {
-          scheduledWait.cancel(mayInterrupt);
           promise.complete(cancelResult);
+          permitWaitFuture.cancel(mayInterrupt);
         });
       } catch (Throwable t) {
         // Hard scheduling failure

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

@@ -28,7 +28,7 @@
 import java.util.concurrent.TimeUnit;
 
 /**
- * A RateLimiter implementation.
+ * A RateLimiter implementation that supports smooth and bursty rate limiting.
  *
  * @param <R> result type
  */

src/main/java/dev/failsafe/internal/util/FutureLinkedList.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+package dev.failsafe.internal.util;
+
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * A LinkedList of CompletableFutures that removes a future from the list when it's completed.
+ * <p>
+ * This class is threadsafe.
+ * </p>
+ *
+ * @author Jonathan Halterman
+ */
+public final class FutureLinkedList {
+  Node head;
+  Node tail;
+
+  static class Node {
+    Node previous;
+    Node next;
+    CompletableFuture<Void> future;
+  }
+
+  /**
+   * Adds a new CompletableFuture to the list and returns it. The returned future will be removed from the list when
+   * it's completed.
+   */
+  public synchronized CompletableFuture<Void> add() {
+    Node node = new Node();
+    node.future = new CompletableFuture<>();
+    node.future.whenComplete((result, error) -> remove(node));
+
+    if (head == null)
+      head = tail = node;
+    else {
+      tail.next = node;
+      node.previous = tail;
+      tail = node;
+    }
+    return node.future;
+  }
+
+  /**
+   * Returns and removes the first future in the list, else returns {@code null} if the list is empty.
+   */
+  public synchronized CompletableFuture<Void> pollFirst() {
+    Node previousHead = head;
+    if (head != null) {
+      head = head.next;
+      if (head != null)
+        head.previous = null;
+    }
+    return previousHead == null ? null : previousHead.future;
+  }
+
+  private synchronized void remove(Node node) {
+    if (node.previous != null)
+      node.previous.next = node.next;
+    if (node.next != null)
+      node.next.previous = node.previous;
+    if (head == node)
+      head = node.next;
+    if (tail == node)
+      tail = node.previous;
+  }
+}

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

@@ -15,21 +15,37 @@
  */
 package dev.failsafe.functional;
 
-import dev.failsafe.*;
+import dev.failsafe.Bulkhead;
+import dev.failsafe.BulkheadFullException;
+import dev.failsafe.Failsafe;
 import dev.failsafe.testing.Testing;
 import org.testng.annotations.Test;
 
 import java.time.Duration;
 
-import static dev.failsafe.internal.InternalTesting.resetBulkhead;
-import static dev.failsafe.internal.InternalTesting.resetLimiter;
-import static org.testng.Assert.assertEquals;
-
 /**
  * Tests various Bulkhead scenarios.
  */
 @Test
 public class BulkheadTest extends Testing {
+  public void testPermitAcquiredAfterWait() {
+    // Given
+    Bulkhead<Object> bulkhead = Bulkhead.builder(2).withMaxWaitTime(Duration.ofSeconds(1)).build();
+
+    // When / Then
+    testGetSuccess(() -> {
+      bulkhead.tryAcquirePermit();
+      bulkhead.tryAcquirePermit(); // bulkhead should be full
+
+      runInThread(() -> {
+        Thread.sleep(200);
+        bulkhead.releasePermit(); // bulkhead should not be full
+      });
+    }, Failsafe.with(bulkhead), ctx -> {
+      return "test";
+    }, "test");
+  }
+
   public void shouldThrowBulkheadFullExceptionAfterPermitsExceeded() {
     // Given
     Bulkhead<Object> bulkhead = Bulkhead.of(2);
@@ -54,52 +70,4 @@ public void testMaxWaitTimeExceeded() {
     testRunFailure(Failsafe.with(bulkhead), ctx -> {
     }, BulkheadFullException.class);
   }
-
-  /**
-   * Tests a scenario where Bulkhead rejects some retried executions, which prevents the user's Supplier from being
-   * called.
-   */
-  public void testRejectedWithRetries() {
-    Stats rpStats = new Stats();
-    Stats rlStats = new Stats();
-    RetryPolicy<Object> rp = withStatsAndLogs(RetryPolicy.builder().withMaxAttempts(7), rpStats).build();
-    Bulkhead<Object> bh = withStatsAndLogs(Bulkhead.builder(2), rlStats).build();
-    bh.tryAcquirePermit();
-    bh.tryAcquirePermit(); // bulkhead should be full
-
-    testRunFailure(() -> {
-      rpStats.reset();
-      rlStats.reset();
-    }, Failsafe.with(rp, bh), ctx -> {
-      System.out.println("Executing");
-      throw new Exception();
-    }, (f, e) -> {
-      assertEquals(e.getAttemptCount(), 7);
-      assertEquals(e.getExecutionCount(), 0);
-      assertEquals(rpStats.failedAttemptCount, 7);
-      assertEquals(rpStats.retryCount, 6);
-    }, BulkheadFullException.class);
-  }
-
-  /**
-   * Asserts that a bulkhead propagates an InterruptedException.
-   */
-  public void testAcquirePermitWithInterrupt() {
-    Bulkhead<Object> bulkhead = Bulkhead.builder(1).withMaxWaitTime(Duration.ofSeconds(5)).build();
-    bulkhead.tryAcquirePermit(); // Bulkhead should be full
-
-    testRunFailure(() -> {
-      Thread thread = Thread.currentThread();
-      runInThread(() -> {
-        Thread.sleep(100);
-        thread.interrupt();
-      });
-    }, Failsafe.with(bulkhead), ctx -> {
-      System.out.println("Executing");
-      throw new Exception();
-    }, InterruptedException.class);
-
-    // Reset interrupt flag
-    Thread.interrupted();
-  }
 }

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

@@ -105,6 +105,16 @@ public void shouldCancelAsyncRateLimiterWaitingOnPermit() throws Throwable {
     });
   }
 
+  public void shouldCancelBulkheadWaitingOnPermit() throws Throwable {
+    Bulkhead<Void> bulkhead = Bulkhead.<Void>builder(2).withMaxWaitTime(Duration.ofSeconds(1)).build();
+    bulkhead.tryAcquirePermit();
+    bulkhead.tryAcquirePermit(); // bulkhead should be full
+
+    assertCancel(Failsafe.with(bulkhead), ctx -> {
+      fail("Execution should be cancelled during preExecute");
+    });
+  }
+
   /**
    * Asserts that cancelling a FailsafeFuture causes both retry policies to stop.
    */