Reduced resource consumption of async credential providers.

Fixes #3259.

1. Share thread pools across async credential providers (anything using CachedSupplier's NonBlocking prefetch strategy).
2. Log a warning if an extreme number of concurrent refreshes are happening, to help users detect when they're not closing their credential providers.

Even though this is an increase in resource sharing, it should not cause increased availability risks. Because these threads are only used for background refreshes, if one particular type of credential provider has availability problems (e.g. SSO or STS high latency), it only disables background refreshes, not prefetches or synchronous fetches.

Author: millems

.changes/next-release/feature-AWSSDKforJavav2-36b8c47.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Share background refresh threads across async credential providers to reduce base SDK resource consumption."
+}

.changes/next-release/feature-AWSSDKforJavav2-8f6d13b.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Log a warning when an extreme number of async credential providers are running in parallel, because it could indicate that the user is not closing their clients or credential providers when they are done using them."
+}

utils/src/main/java/software/amazon/awssdk/utils/cache/CachedSupplier.java

@@ -37,7 +37,7 @@
  * This should be created using {@link #builder(Supplier)}.
  */
 @SdkProtectedApi
-public final class CachedSupplier<T> implements Supplier<T>, SdkAutoCloseable {
+public class CachedSupplier<T> implements Supplier<T>, SdkAutoCloseable {
     /**
      * Maximum time to wait for a blocking refresh lock before calling refresh again. This is to rate limit how many times we call
      * refresh. In the ideal case, refresh always occurs in a timely fashion and only one thread actually does the refresh.

utils/src/main/java/software/amazon/awssdk/utils/cache/NonBlocking.java

@@ -18,14 +18,17 @@
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 import java.time.Duration;
-import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.ThreadFactoryBuilder;
-import software.amazon.awssdk.utils.Validate;
 
 /**
  * A {@link CachedSupplier.PrefetchStrategy} that will run a single thread in the background to update the value. A call to
@@ -37,40 +40,102 @@
 public class NonBlocking implements CachedSupplier.PrefetchStrategy {
     private static final Logger log = Logger.loggerFor(NonBlocking.class);
 
+    /**
+     * The maximum number of concurrent refreshes allowed across all NonBlocking instances. This is to limit the amount of 
+     * traffic that can be generated by background refreshes. If this is exceeded, a background refresh gets skipped. This just
+     * increases the chance of latency being pushed to the cached supplier caller, which is preferable to running out of memory.
+     */
+    private static final int MAX_CONCURRENT_REFRESHES = 100;
+
+    /**
+     * By default, how often we periodically call get() on the cached supplier. This may not necessarily call the downstream
+     * service if the cache is not stale.
+     */
+    private static final Duration DEFAULT_REFRESH_FREQUENCY = Duration.ofSeconds(60);
+
+    /**
+     * By default, how much we jitter the {@link #DEFAULT_REFRESH_FREQUENCY}. This is done to prevent the case that a large
+     * number of NonBlocking instances are created at once, so they all try to refresh at the same time.
+     */
+    private static final Duration DEFAULT_REFRESH_FREQUENCY_JITTER = Duration.ofSeconds(10);
+
+    /**
+     * Threads used to periodically kick off credential refreshes based on the {@link #asyncRefreshFrequency}.
+     */
+    private static final ScheduledThreadPoolExecutor SCHEDULER =
+        new ScheduledThreadPoolExecutor(3, new ThreadFactoryBuilder().threadNamePrefix("sdk-cached-supplier-scheduler")
+                                                                     .daemonThreads(true)
+                                                                     .build());
+
+    /**
+     * Threads used to do the actual work of refreshing the credentials (because the cached supplier might block, so we don't
+     * want the work to be done by a small thread pool). This executor is created as unbounded, but in reality it is limited by
+     * the {@link #MAX_CONCURRENT_REFRESHES} via {@link #BACKGROUND_REFRESH_LEASE}.
+     */
+    private static final ThreadPoolExecutor EXECUTOR =
+        new ThreadPoolExecutor(0,
+                               Integer.MAX_VALUE,
+                               DEFAULT_REFRESH_FREQUENCY.toMillis() + DEFAULT_REFRESH_FREQUENCY_JITTER.toMillis() + 5_000,
+                               MILLISECONDS,
+                               new SynchronousQueue<>(),
+                               new ThreadFactoryBuilder().daemonThreads(true).build());
+    /**
+     * A set of leases used to prevent concurrent refreshes beyond the limit described in {@link #MAX_CONCURRENT_REFRESHES}. If
+     * a lease cannot be acquired, the refresh is skipped.
+     */
+    private static final Semaphore BACKGROUND_REFRESH_LEASE = new Semaphore(MAX_CONCURRENT_REFRESHES);
+
+    /**
+     * An incrementing number, used to uniquely identify an instance of NonBlocking in the {@link #asyncThreadName}.
+     */
+    private static final AtomicLong THREAD_NUMBER = new AtomicLong(0);
+
     /**
      * Whether we are currently refreshing the supplier. This is used to make sure only one caller is blocking at a time.
      */
     private final AtomicBoolean currentlyRefreshing = new AtomicBoolean(false);
 
+    /**
+     * Name of the thread refreshing the cache for this strategy.
+     */
+    private final String asyncThreadName;
+
     /**
      * How frequently to automatically refresh the supplier in the background.
      */
     private final Duration asyncRefreshFrequency;
 
     /**
-     * Single threaded executor to asynchronous refresh the value.
+     * Whether this strategy has been shutdown (and should stop doing background refreshes)
      */
-    private final ScheduledExecutorService executor;
+    private volatile boolean shutdown = false;
 
     /**
      * Create a non-blocking prefetch strategy that uses the provided value for the name of the background thread that will be
      * performing the update.
      */
     public NonBlocking(String asyncThreadName) {
-        this(asyncThreadName, Duration.ofMinutes(1));
+        this(asyncThreadName, defaultRefreshFrequency());
+    }
+
+    private static Duration defaultRefreshFrequency() {
+        // We jitter the default refresh frequency with each instance, so that objects created at the same time will not all be
+        // refreshing at the exact same times.
+        long jitter = DEFAULT_REFRESH_FREQUENCY_JITTER.toMillis();
+        long asyncRefreshFrequency = DEFAULT_REFRESH_FREQUENCY.toMillis() +
+                                     ThreadLocalRandom.current().nextLong(-jitter, jitter + 1);
+        return Duration.ofMillis(asyncRefreshFrequency);
     }
 
     @SdkTestInternalApi
     NonBlocking(String asyncThreadName, Duration asyncRefreshFrequency) {
-        this.executor = newExecutor(asyncThreadName);
+        this.asyncThreadName = asyncThreadName + THREAD_NUMBER.getAndIncrement();
         this.asyncRefreshFrequency = asyncRefreshFrequency;
     }
 
-    private static ScheduledExecutorService newExecutor(String asyncThreadName) {
-        Validate.paramNotBlank(asyncThreadName, "asyncThreadName");
-        return new ScheduledThreadPoolExecutor(1, new ThreadFactoryBuilder().daemonThreads(true)
-                                                                            .threadNamePrefix(asyncThreadName)
-                                                                            .build());
+    @SdkTestInternalApi
+    static ThreadPoolExecutor executor() {
+        return EXECUTOR;
     }
 
     @Override
@@ -79,10 +144,35 @@ public void initializeCachedSupplier(CachedSupplier<?> cachedSupplier) {
     }
 
     private void scheduleRefresh(CachedSupplier<?> cachedSupplier) {
-        executor.schedule(() -> {
+        SCHEDULER.schedule(() -> {
+            Thread.currentThread().setName(asyncThreadName);
+
+            if (shutdown) {
+                return;
+            }
+
+            if (!BACKGROUND_REFRESH_LEASE.tryAcquire()) {
+                log.warn(() -> "Skipped a background refresh to limit SDK resource consumption. Are you closing your SDK "
+                               + "resources?");
+                scheduleRefresh(cachedSupplier);
+                return;
+            }
+
             try {
-                cachedSupplier.get();
-            } finally {
+                EXECUTOR.execute(() -> {
+                    try {
+                        Thread.currentThread().setName(asyncThreadName);
+                        cachedSupplier.get();
+                    } catch (Exception e) {
+                        log.error(() -> "Background refresh failed: " + e.getMessage(), e);
+                    } finally {
+                        BACKGROUND_REFRESH_LEASE.release();
+                        scheduleRefresh(cachedSupplier);
+                    }
+                });
+            } catch (Throwable t) {
+                BACKGROUND_REFRESH_LEASE.release();
+                log.warn(() -> "Failed to submit a background refresh task.", t);
                 scheduleRefresh(cachedSupplier);
             }
         }, asyncRefreshFrequency.toMillis(), MILLISECONDS);
@@ -93,7 +183,7 @@ public void prefetch(Runnable valueUpdater) {
         // Only run one async refresh at a time.
         if (currentlyRefreshing.compareAndSet(false, true)) {
             try {
-                executor.submit(() -> {
+                EXECUTOR.submit(() -> {
                     try {
                         valueUpdater.run();
                     } catch (RuntimeException e) {
@@ -111,6 +201,6 @@ public void prefetch(Runnable valueUpdater) {
 
     @Override
     public void close() {
-        executor.shutdown();
+        shutdown = true;
     }
 }

utils/src/test/java/software/amazon/awssdk/utils/cache/NonBlockingTest.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file 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 software.amazon.awssdk.utils.cache;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.jupiter.api.MethodOrderer;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestMethodOrder;
+import org.mockito.Mockito;
+import org.mockito.stubbing.Answer;
+
+// These tests assert on the largest pool size, so the test order matters (smaller expected sizes before larger expected sizes)
+@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+public class NonBlockingTest {
+    @Test
+    @Order(1)
+    public void threadsAreSharedBetweenNonBlockingInstances() throws InterruptedException {
+        List<NonBlocking> nbs = new ArrayList<>();
+        try {
+            // Create 99 concurrent non-blocking instances
+            for (int i = 0; i < 99; i++) {
+                NonBlocking nb = new NonBlocking("test", Duration.ofMillis(100));
+                nb.initializeCachedSupplier(Mockito.mock(CachedSupplier.class));
+                nbs.add(nb);
+                Thread.sleep(10);
+            }
+
+            Thread.sleep(1_000);
+
+            // Make sure we used less-than 99 to do the refreshes.
+            assertThat(NonBlocking.executor().getLargestPoolSize()).isLessThan(99);
+        } finally {
+            nbs.forEach(NonBlocking::close);
+        }
+    }
+
+    @Test
+    @Order(2)
+    public void refreshesAreMaxed() throws InterruptedException {
+        CachedSupplier<?> slowSupplier = Mockito.mock(CachedSupplier.class);
+        Mockito.when(slowSupplier.get()).thenAnswer((Answer<Void>) invocation -> {
+            Thread.sleep(1_000);
+            return null;
+        });
+
+        List<NonBlocking> nbs = new ArrayList<>();
+        try {
+            for (int i = 0; i < 1_000; i++) {
+                NonBlocking nb = new NonBlocking("test", Duration.ofMillis(0));
+                nb.initializeCachedSupplier(slowSupplier);
+                nbs.add(nb);
+            }
+
+            Thread.sleep(1_000);
+
+            // In a perfect world this would be capped to 100, but the mechanism we use to limit concurrent refreshes usually
+            // means more than 100 get created. 150 should be a reasonable limit to check for, because without the limiter it
+            // would be ~1000.
+            assertThat(NonBlocking.executor().getLargestPoolSize()).isLessThanOrEqualTo(150);
+        } finally {
+            nbs.forEach(NonBlocking::close);
+        }
+    }
+}