Fix ADAPTIVE_V2 client close() hang when a queued request's callback …

.changes/next-release/bugfix-AWSSDKforJavav2-84282b4.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Fix an issue where closing an async client using adaptive retries could hang when requests were rate-limited."
+}

core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java

@@ -17,7 +17,9 @@
 
 import java.time.Duration;
 import java.util.ArrayDeque;
+import java.util.ArrayList;
 import java.util.Deque;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -66,29 +68,32 @@ public class RateLimiterTokenBucket implements SdkAutoCloseable {
 
     @Override
     public void close() {
-        doClose(null);
+        completeClose(closeAndDrain(null));
     }
 
-    private void doClose(Throwable cause) {
+    private CloseResult closeAndDrain(Throwable cause) {
         synchronized (lock) {
             open = false;
-            IllegalStateException closedException = new IllegalStateException("Rate limiter bucket is closed", cause);
-            while (true) {
-                CompletableFuture<Void> w = waiting.poll();
-                if (w == null) {
-                    break;
-                }
-                w.completeExceptionally(closedException);
-            }
+            notifierRunning = false;
+            List<CompletableFuture<Void>> waiters = new ArrayList<>(waiting);
+            waiting.clear();
+            return new CloseResult(waiters, cause);
         }
     }
 
+    private void completeClose(CloseResult closeResult) {
+        IllegalStateException closedException = new IllegalStateException("Rate limiter bucket is closed", closeResult.cause);
+        closeResult.waiters.forEach(waiter -> waiter.completeExceptionally(closedException));
+    }
+
     /**
      * Acquire a token from the bucket.
      *
      * @return A future that is completed when the requested amount is acquired from this bucket.
      */
     public CompletableFuture<Void> acquireAsync() {
+        CompletableFuture<Void> future;
+        CloseResult closeResult = null;
         synchronized (lock) {
             if (!open) {
                 return CompletableFutureUtils.failedFuture(new IllegalStateException("Rate limiter bucket is closed"));
@@ -99,13 +104,21 @@ public CompletableFuture<Void> acquireAsync() {
                 return CompletableFuture.completedFuture(null);
             }
 
-            CompletableFuture<Void> future = new CompletableFuture<>();
+            future = new CompletableFuture<>();
             waiting.add(future);
             if (!notifierRunning) {
-                notifierRunning = scheduleOrClose(this::doNotify, Duration.ZERO);
+                notifierRunning = true;
+                try {
+                    schedule(this::doNotify, Duration.ZERO);
+                } catch (Throwable t) {
+                    closeResult = closeAndDrain(t);
+                }
             }
-            return future;
         }
+        if (closeResult != null) {
+            completeClose(closeResult);
+        }
+        return future;
     }
 
 
@@ -121,10 +134,12 @@ boolean isClosed() {
 
     private void doNotify() {
         while (true) {
-            CompletableFuture<Void> w;
+            CompletableFuture<Void> waiter;
+            boolean acquired = false;
+            CloseResult closeResult = null;
             synchronized (lock) {
-                w = waiting.poll();
-                if (w == null) {
+                waiter = waiting.poll();
+                if (waiter == null) {
                     notifierRunning = false;
                     return;
                 }
@@ -134,32 +149,41 @@ private void doNotify() {
                 // Not enough capacity. Try again later when enough time has
                 // passed to refill the bucket at the current rate.
                 if (!acquireResult.isSuccessful()) {
-                    waiting.push(w);
-                    notifierRunning = scheduleOrClose(this::doNotify, acquireResult.refillWait());
-                    return;
+                    waiting.push(waiter);
+                    notifierRunning = true;
+                    try {
+                        schedule(this::doNotify, acquireResult.refillWait());
+                    } catch (Throwable t) {
+                        closeResult = closeAndDrain(t);
+                    }
+                } else {
+                    acquired = true;
                 }
-
+            }
+            if (closeResult != null) {
+                completeClose(closeResult);
+                return;
+            }
+            if (!acquired) {
+                return;
             }
             // Acquire was successful, signal the waiting thread.
-            w.complete(null);
+            waiter.complete(null);
         }
     }
 
     private void schedule(Runnable command, Duration d) {
         scheduler.schedule(command, d.toMillis(), TimeUnit.MILLISECONDS);
     }
 
-    /**
-     * @return true if schedule was successful, false otherwise. If the schedule failed, this bucket will be closed.
-     */
-    private boolean scheduleOrClose(Runnable command, Duration d) {
-        try {
-            schedule(command, d);
-            return true;
-        } catch (Throwable t) {
-            doClose(t);
+    private static final class CloseResult {
+        private final List<CompletableFuture<Void>> waiters;
+        private final Throwable cause;
+
+        private CloseResult(List<CompletableFuture<Void>> waiters, Throwable cause) {
+            this.waiters = waiters;
+            this.cause = cause;
         }
-        return false;
     }
 
     /**

core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java

@@ -17,6 +17,7 @@
 
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.annotations.ToBuilderIgnoreField;
@@ -41,6 +42,8 @@ public final class RateLimiterTokenBucketStore
     private final RateLimiterClock clock;
     private final ScheduledExecutorService scheduler;
     private final boolean closeScheduler;
+    private final AtomicBoolean closed = new AtomicBoolean();
+    private final RateLimiterTokenBucket closedTokenBucket;
 
     private RateLimiterTokenBucketStore(Builder builder) {
         this(builder.clock,
@@ -52,6 +55,8 @@ private RateLimiterTokenBucketStore(RateLimiterClock clock, ScheduledExecutorSer
         this.clock = Validate.paramNotNull(clock, "clock");
         this.scheduler = Validate.paramNotNull(scheduler, "scheduler");
         this.closeScheduler = closeScheduler;
+        this.closedTokenBucket = new RateLimiterTokenBucket(clock, scheduler);
+        this.closedTokenBucket.close();
         this.scopeToTokenBucket = LruCache.<String, RateLimiterTokenBucket>builder(
                                               x -> new RateLimiterTokenBucket(clock, scheduler))
                                           .maxSize(MAX_ENTRIES)
@@ -60,14 +65,25 @@ private RateLimiterTokenBucketStore(RateLimiterClock clock, ScheduledExecutorSer
 
     @Override
     public void close() {
-        scopeToTokenBucket.evictAll();
-        if (closeScheduler) {
-            scheduler.shutdownNow();
+        if (closed.compareAndSet(false, true)) {
+            scopeToTokenBucket.evictAll();
+            if (closeScheduler) {
+                scheduler.shutdownNow();
+            }
         }
     }
 
     public RateLimiterTokenBucket tokenBucketForScope(String scope) {
-        return scopeToTokenBucket.get(scope);
+        if (closed.get()) {
+            return closedTokenBucket;
+        }
+
+        RateLimiterTokenBucket tokenBucket = scopeToTokenBucket.get(scope);
+        if (closed.get()) {
+            tokenBucket.close();
+            return closedTokenBucket;
+        }
+        return tokenBucket;
     }
 
     @SdkTestInternalApi

core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java

@@ -25,6 +25,7 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.Test;
 
 public class RateLimiterTokenBucketStoreTest {
@@ -85,4 +86,48 @@ void close_closesAllCacheEntries() {
         assertThat(futures).allSatisfy(cf -> assertThatThrownBy(cf::join)
             .hasMessageContaining("Rate limiter bucket is closed"));
     }
+
+    @Test
+    void close_whenPendingWaiterReentersStore_completes() throws InterruptedException {
+        ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class);
+        RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder()
+                                                                       .clock(new SystemClock())
+                                                                       .scheduler(scheduler)
+                                                                       .build();
+        RateLimiterTokenBucket tokenBucket = store.tokenBucketForScope("scope");
+        tokenBucket.updateRateAfterThrottling();
+        CompletableFuture<Void> waiter = tokenBucket.acquireAsync();
+        CompletableFuture<RateLimiterTokenBucket> dependent = waiter.handle((r, t) -> store.tokenBucketForScope("scope"));
+
+        AtomicReference<Throwable> failure = new AtomicReference<>();
+        Thread closingThread = new Thread(() -> {
+            try {
+                store.close();
+            } catch (Throwable t) {
+                failure.set(t);
+            }
+        }, "rate-limiter-store-close");
+        closingThread.setDaemon(true);
+        closingThread.start();
+        closingThread.join(5_000L);
+
+        assertThat(closingThread.isAlive()).isFalse();
+        assertThat(failure.get()).isNull();
+        assertThatThrownBy(dependent.join().acquireAsync()::join)
+            .hasMessageContaining("Rate limiter bucket is closed");
+    }
+
+    @Test
+    void tokenBucketForScope_afterClose_returnsClosedBucket() {
+        ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class);
+        RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder()
+                                                                       .clock(new SystemClock())
+                                                                       .scheduler(scheduler)
+                                                                       .build();
+
+        store.close();
+
+        assertThatThrownBy(store.tokenBucketForScope("scope").acquireAsync()::join)
+            .hasMessageContaining("Rate limiter bucket is closed");
+    }
 }

core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java

@@ -31,9 +31,13 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.junit.jupiter.api.BeforeEach;
@@ -148,6 +152,38 @@ void close_doesShutDownExecutor() {
         verifyNoInteractions(scheduler);
     }
 
+    @Test
+    void doNotify_whenRescheduleFails_completesCallbacksOutsideLock() {
+        tokenBucket.updateRateAfterThrottling();
+        AtomicReference<Runnable> notifier = new AtomicReference<>();
+        when(scheduler.schedule(any(Runnable.class), eq(0L), any(TimeUnit.class))).thenAnswer(invocation -> {
+            notifier.set(invocation.getArgument(0));
+            return null;
+        });
+        when(scheduler.schedule(any(Runnable.class), eq(2000L), eq(TimeUnit.MILLISECONDS)))
+            .thenThrow(new RejectedExecutionException("no"));
+
+        CompletableFuture<Void> waiter = tokenBucket.acquireAsync();
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        try {
+            CompletableFuture<Void> callback = waiter.handle((r, t) -> {
+                Future<?> update = executor.submit(tokenBucket::updateRateAfterSuccess);
+                try {
+                    update.get(1, TimeUnit.SECONDS);
+                } catch (Exception e) {
+                    throw new AssertionError(e);
+                }
+                return null;
+            });
+
+            notifier.get().run();
+
+            callback.join();
+        } finally {
+            executor.shutdownNow();
+        }
+    }
+
     @Test
     void doNotify_scheduleRejected_failsFuture() {
         tokenBucket.updateRateAfterThrottling();