Reduced resource consumption of async credential providers.

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/pom.xml

@@ -94,6 +94,27 @@
             <artifactId>reactive-streams-tck</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-api</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-core</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-slf4j-impl</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>jcl-over-slf4j</artifactId>
+            <scope>test</scope>
+            <version>${slf4j.version}</version>
+        </dependency>
     </dependencies>
 
     <build>

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.
@@ -107,7 +107,7 @@ private boolean cacheIsStale() {
         if (cachedValue.staleTime() == null) {
             return false;
         }
-        return Instant.now().isAfter(cachedValue.staleTime());
+        return !Instant.now().isBefore(cachedValue.staleTime());
     }
 
     /**
@@ -118,7 +118,7 @@ private boolean shouldInitiateCachePrefetch() {
         if (cachedValue.prefetchTime() == null) {
             return false;
         }
-        return Instant.now().isAfter(cachedValue.prefetchTime());
+        return !Instant.now().isBefore(cachedValue.prefetchTime());
     }
 
     /**
@@ -146,7 +146,7 @@ private void refreshCache() {
                     }
 
                     // It wasn't, call the supplier to update it.
-                    cachedValue = valueSupplier.get();
+                    cachedValue = prefetchStrategy.fetch(valueSupplier);
                 }
             } finally {
                 if (lockAcquired) {
@@ -215,6 +215,10 @@ public interface PrefetchStrategy extends SdkAutoCloseable {
          */
         void prefetch(Runnable valueUpdater);
 
+        default <T> RefreshResult<T> fetch(Supplier<RefreshResult<T>> supplier) {
+            return supplier.get();
+        }
+
         /**
          * Invoked when the prefetch strategy is registered with a {@link CachedSupplier}.
          */

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

@@ -15,17 +15,26 @@
 
 package software.amazon.awssdk.utils.cache;
 
-import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.time.temporal.ChronoUnit.HOURS;
+import static java.time.temporal.ChronoUnit.MINUTES;
 
 import java.time.Duration;
-import java.util.concurrent.ScheduledExecutorService;
+import java.time.Instant;
+import java.util.Optional;
+import java.util.Random;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 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,80 +46,204 @@
 public class NonBlocking implements CachedSupplier.PrefetchStrategy {
     private static final Logger log = Logger.loggerFor(NonBlocking.class);
 
+    /**
+     * The maximum number of concurrent refreshes before we start logging warnings and skipping refreshes.
+     */
+    private static final int MAX_CONCURRENT_REFRESHES = 100;
+
+    /**
+     * The {@link Random} instance used for calculating jitter of the background prefetches.
+     */
+    private static final Random JITTER_RANDOM = new Random();
+
+    /**
+     * Thread used to kick off refreshes during the prefetch window. This does not do the actual refreshing. That's left for
+     * the {@link #EXECUTOR}.
+     */
+    private static final ScheduledThreadPoolExecutor SCHEDULER =
+        new ScheduledThreadPoolExecutor(1, new ThreadFactoryBuilder().daemonThreads(true).build());
+
+    /**
+     * Threads used to do the actual work of refreshing the values (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 we start complaining and
+     * skipping refreshes when there are more than {@link #MAX_CONCURRENT_REFRESHES} running.
+     */
+    private static final ThreadPoolExecutor EXECUTOR =
+        new ThreadPoolExecutor(0, Integer.MAX_VALUE,
+                               60L, TimeUnit.SECONDS,
+                               new SynchronousQueue<>(),
+                               new ThreadFactoryBuilder().daemonThreads(true).build());
+
+    /**
+     * 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);
+    private final AtomicBoolean currentlyPrefetching = new AtomicBoolean(false);
 
     /**
-     * How frequently to automatically refresh the supplier in the background.
+     * Name of the thread refreshing the cache for this strategy.
      */
-    private final Duration asyncRefreshFrequency;
+    private final String asyncThreadName;
 
     /**
-     * Single threaded executor to asynchronous refresh the value.
+     * The refresh task currently scheduled for this non-blocking instance. We ensure that no more than one task is scheduled
+     * per instance.
      */
-    private final ScheduledExecutorService executor;
+    private final AtomicReference<ScheduledFuture<?>> refreshTask = new AtomicReference<>();
+
+    /**
+     * The minimum amount of time allowed between async refreshes, primarily adjustable for testing purposes.
+     */
+    private final Duration minimumRefreshFrequency;
+
+    /**
+     * Whether this strategy has been shutdown (and should stop doing background refreshes)
+     */
+    private volatile boolean shutdown = false;
+
+    /**
+     * The cached supplier using this non-blocking instance.
+     */
+    private volatile CachedSupplier<?> cachedSupplier;
+
+    static {
+        // Ensure that cancelling a task actually removes it from the queue.
+        SCHEDULER.setRemoveOnCancelPolicy(true);
+    }
 
     /**
      * 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, Duration.ofSeconds(60));
     }
 
     @SdkTestInternalApi
-    NonBlocking(String asyncThreadName, Duration asyncRefreshFrequency) {
-        this.executor = newExecutor(asyncThreadName);
-        this.asyncRefreshFrequency = asyncRefreshFrequency;
+    NonBlocking(String asyncThreadName, Duration minimumRefreshFrequency) {
+        this.asyncThreadName = asyncThreadName + "-" + THREAD_NUMBER.getAndIncrement();
+        this.minimumRefreshFrequency = minimumRefreshFrequency;
     }
 
-    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
     public void initializeCachedSupplier(CachedSupplier<?> cachedSupplier) {
-        scheduleRefresh(cachedSupplier);
+        this.cachedSupplier = cachedSupplier;
     }
 
-    private void scheduleRefresh(CachedSupplier<?> cachedSupplier) {
-        executor.schedule(() -> {
-            try {
-                cachedSupplier.get();
-            } finally {
-                scheduleRefresh(cachedSupplier);
-            }
-        }, asyncRefreshFrequency.toMillis(), MILLISECONDS);
+    @Override
+    public void prefetch(Runnable valueUpdater) {
+        // Only run one async prefetch at a time.
+        if (currentlyPrefetching.compareAndSet(false, true)) {
+            tryRunBackgroundTask(valueUpdater, () -> currentlyPrefetching.set(false));
+        }
     }
 
     @Override
-    public void prefetch(Runnable valueUpdater) {
-        // Only run one async refresh at a time.
-        if (currentlyRefreshing.compareAndSet(false, true)) {
-            try {
-                executor.submit(() -> {
-                    try {
-                        valueUpdater.run();
-                    } catch (RuntimeException e) {
-                        log.warn(() -> "Exception occurred in AWS SDK background task.", e);
-                    } finally {
-                        currentlyRefreshing.set(false);
-                    }
-                });
-            } catch (Throwable e) {
-                currentlyRefreshing.set(false);
-                throw e;
-            }
+    public <T> RefreshResult<T> fetch(Supplier<RefreshResult<T>> supplier) {
+        RefreshResult<T> result = supplier.get();
+        if (result.staleTime() == null || result.prefetchTime() == null) {
+            return result;
+        }
+
+        getRefreshTime(result).ifPresent(this::schedulePrefetch);
+        return result;
+    }
+
+    private Optional<Instant> getRefreshTime(RefreshResult<?> result) {
+        Instant minStart = Instant.now().plus(minimumRefreshFrequency);
+        Instant rangeStart = result.prefetchTime().isBefore(minStart) ? minStart : result.prefetchTime();
+
+        if (Duration.between(Instant.now(), rangeStart).toDays() > 7) {
+            log.debug(() -> "Skipping background refresh because the prefetch time is too far in the future: " + rangeStart);
+            return Optional.empty();
+        }
+
+        Instant maxEnd = rangeStart.plus(1, HOURS);
+        Instant rangeEnd = result.staleTime().isAfter(maxEnd) ? maxEnd : result.staleTime().minus(1, MINUTES);
+
+        if (rangeEnd.isBefore(rangeStart)) {
+            return Optional.of(rangeStart);
+        }
+
+        return Optional.of(randomTimeBetween(rangeStart, rangeEnd));
+    }
+
+    private Instant randomTimeBetween(Instant rangeStart, Instant rangeEnd) {
+        Duration timeBetween = Duration.between(rangeStart, rangeEnd);
+        return rangeStart.plusMillis(Math.abs(JITTER_RANDOM.nextLong() % timeBetween.toMillis()));
+    }
+
+    private void schedulePrefetch(Instant refreshTime) {
+        if (shutdown) {
+            return;
+        }
+
+        Duration waitTime = Duration.between(Instant.now(), refreshTime);
+        log.debug(() -> "Scheduling refresh attempt for " + refreshTime + " (in " + waitTime.toMillis() + " ms)");
+        updateTask(SCHEDULER.schedule(() -> {
+            Thread.currentThread().setName(asyncThreadName + "-scheduler");
+            log.debug(() -> "Executing refresh attempt scheduled for " + refreshTime);
+
+            // If the supplier has already been prefetched, this will just be a cache hit.
+            tryRunBackgroundTask(cachedSupplier::get);
+        }, waitTime.toMillis(), TimeUnit.MILLISECONDS));
+
+        if (shutdown) {
+            updateTask(null);
         }
     }
 
     @Override
     public void close() {
-        executor.shutdown();
+        shutdown = true;
+        updateTask(null);
+    }
+
+    public void updateTask(ScheduledFuture<?> newTask) {
+        ScheduledFuture<?> currentTask;
+        do {
+            currentTask = refreshTask.get();
+            if (currentTask != null && !currentTask.isDone()) {
+                currentTask.cancel(false);
+            }
+        } while (!refreshTask.compareAndSet(currentTask, newTask));
+    }
+
+    public void tryRunBackgroundTask(Runnable runnable) {
+        tryRunBackgroundTask(runnable, () -> {
+        });
+    }
+
+    public void tryRunBackgroundTask(Runnable runnable, Runnable finallyRunnable) {
+        try {
+            if (EXECUTOR.getActiveCount() > MAX_CONCURRENT_REFRESHES) {
+                log.warn(() -> "Skipping a background refresh task because there are too many other tasks running.");
+                return;
+            }
+
+            EXECUTOR.submit(() -> {
+                try {
+                    Thread.currentThread().setName(asyncThreadName);
+                    runnable.run();
+                } catch (Throwable t) {
+                    log.warn(() -> "Exception occurred in AWS SDK background task.", t);
+                } finally {
+                    finallyRunnable.run();
+                }
+            });
+        } catch (Throwable t) {
+            log.warn(() -> "Exception occurred when submitting AWS SDK background task.", t);
+        } finally {
+            finallyRunnable.run();
+        }
     }
 }