Moved jitter to the cached supplier and addressed other comments.

millems · millems · commit cabb263604da · 2022-07-07T16:19:29.000-07:00
diff --git a/utils/src/main/java/software/amazon/awssdk/utils/cache/CachedSupplier.java b/utils/src/main/java/software/amazon/awssdk/utils/cache/CachedSupplier.java
@@ -15,14 +15,19 @@
 
 package software.amazon.awssdk.utils.cache;
 
+import static java.time.temporal.ChronoUnit.MINUTES;
+
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Random;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Supplier;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
+import software.amazon.awssdk.utils.ComparableUtils;
 import software.amazon.awssdk.utils.SdkAutoCloseable;
 import software.amazon.awssdk.utils.Validate;
 
@@ -44,6 +49,16 @@ public class CachedSupplier<T> implements Supplier<T>, SdkAutoCloseable {
      */
     private static final Duration BLOCKING_REFRESH_MAX_WAIT = Duration.ofSeconds(5);
 
+    /**
+     * Maximum amount of jitter to apply to the refresh result's prefetch time.
+     */
+    private static final Duration MAX_PREFETCH_JITTER = Duration.ofMinutes(5);
+
+    /**
+     * Random instance used for jittering refresh results.
+     */
+    private static final Random JITTER_RANDOM = new Random();
+
     /**
      * Used as a primitive form of rate limiting for the speed of our refreshes. This will make sure that the backing supplier has
      * a period of time to update the value when the {@link RefreshResult#staleTime()} arrives without getting called by every
@@ -62,6 +77,11 @@ public class CachedSupplier<T> implements Supplier<T>, SdkAutoCloseable {
      */
     private final AtomicBoolean prefetchStrategyInitialized = new AtomicBoolean(false);
 
+    /**
+     * The maximum amount of prefetch jitter allowed on this instance. Only differs for {@link #MAX_PREFETCH_JITTER} in testing.
+     */
+    private final Duration maxPrefetchJitter;
+
     /**
      * The value currently stored in this cache.
      */
@@ -76,8 +96,9 @@ public class CachedSupplier<T> implements Supplier<T>, SdkAutoCloseable {
     private final Supplier<RefreshResult<T>> valueSupplier;
 
     private CachedSupplier(Builder<T> builder) {
-        this.valueSupplier = Validate.notNull(builder.supplier, "builder.supplier");
+        this.valueSupplier = jitteredValueSupplier(Validate.notNull(builder.supplier, "builder.supplier"));
         this.prefetchStrategy = Validate.notNull(builder.prefetchStrategy, "builder.prefetchStrategy");
+        this.maxPrefetchJitter = Validate.notNull(builder.maxPrefetchJitter, "builder.maxPrefetchJitter");
     }
 
     /**
@@ -163,6 +184,45 @@ private void handleInterruptedException(String message, InterruptedException cau
         throw new IllegalStateException(message, cause);
     }
 
+    /**
+     * Wrap a value supplier with one that jitters its prefetch time.
+     */
+    private Supplier<RefreshResult<T>> jitteredValueSupplier(Supplier<RefreshResult<T>> supplier) {
+        return () -> {
+            RefreshResult<T> result = supplier.get();
+
+            if (result.prefetchTime() == null) {
+                return result;
+            }
+
+            Duration maxJitter = getMaxJitter(result);
+            if (maxJitter.isZero()) {
+                return result;
+            }
+
+            long jitter = Math.abs(JITTER_RANDOM.nextLong() % maxJitter.toMillis());
+            Instant newPrefetchTime = result.prefetchTime().plusMillis(jitter);
+            return RefreshResult.builder(result.value())
+                                .prefetchTime(newPrefetchTime)
+                                .staleTime(result.staleTime())
+                                .build();
+        };
+    }
+
+    private Duration getMaxJitter(RefreshResult<T> result) {
+        if (result.staleTime() == null) {
+            return maxPrefetchJitter;
+        }
+
+        Instant oneMinuteBeforeStale = result.staleTime().minus(1, MINUTES);
+        if (!result.prefetchTime().isBefore(oneMinuteBeforeStale)) {
+            return Duration.ZERO;
+        }
+
+        Duration oneMinuteBeforeStaleDuration = Duration.between(result.prefetchTime(), oneMinuteBeforeStale);
+        return ComparableUtils.minimum(oneMinuteBeforeStaleDuration, maxPrefetchJitter);
+    }
+
     /**
      * Free any resources consumed by the prefetch strategy this supplier is using.
      */
@@ -177,6 +237,7 @@ public void close() {
     public static final class Builder<T> {
         private final Supplier<RefreshResult<T>> supplier;
         private PrefetchStrategy prefetchStrategy = new OneCallerBlocks();
+        private Duration maxPrefetchJitter = MAX_PREFETCH_JITTER;
 
         private Builder(Supplier<RefreshResult<T>> supplier) {
             this.supplier = supplier;
@@ -194,6 +255,15 @@ public Builder<T> prefetchStrategy(PrefetchStrategy prefetchStrategy) {
             return this;
         }
 
+        /**
+         * The maximum amount of time the prefetch time from the configured supplier will be jittered.
+         */
+        @SdkTestInternalApi
+        Builder<T> maxPrefetchJitter(Duration maxPrefetchJitter) {
+            this.maxPrefetchJitter = maxPrefetchJitter;
+            return this;
+        }
+
         /**
          * Create a {@link CachedSupplier} using the current configuration of this builder.
          */
diff --git a/utils/src/main/java/software/amazon/awssdk/utils/cache/NonBlocking.java b/utils/src/main/java/software/amazon/awssdk/utils/cache/NonBlocking.java
@@ -15,13 +15,8 @@
 
 package software.amazon.awssdk.utils.cache;
 
-import static java.time.temporal.ChronoUnit.HOURS;
-import static java.time.temporal.ChronoUnit.MINUTES;
-
 import java.time.Duration;
 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.Semaphore;
@@ -57,11 +52,6 @@ public class NonBlocking implements CachedSupplier.PrefetchStrategy {
      */
     private static final Semaphore CONCURRENT_REFRESH_LEASES = new Semaphore(MAX_CONCURRENT_REFRESHES);
 
-    /**
-     * 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}.
@@ -105,11 +95,6 @@ public class NonBlocking implements CachedSupplier.PrefetchStrategy {
      */
     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)
      */
@@ -130,13 +115,7 @@ public class NonBlocking implements CachedSupplier.PrefetchStrategy {
      * performing the update.
      */
     public NonBlocking(String asyncThreadName) {
-        this(asyncThreadName, Duration.ofSeconds(60));
-    }
-
-    @SdkTestInternalApi
-    NonBlocking(String asyncThreadName, Duration minimumRefreshFrequency) {
         this.asyncThreadName = asyncThreadName + "-" + INSTANCE_NUMBER.getAndIncrement();
-        this.minimumRefreshFrequency = minimumRefreshFrequency;
     }
 
     @SdkTestInternalApi
@@ -160,54 +139,36 @@ public void prefetch(Runnable valueUpdater) {
     @Override
     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);
+        schedulePrefetch(result);
         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);
+    private void schedulePrefetch(RefreshResult<?> result) {
+        if (shutdown || result.staleTime() == null || result.prefetchTime() == null) {
+            return;
         }
 
-        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) {
+        Duration timeUntilPrefetch = Duration.between(Instant.now(), result.prefetchTime());
+        if (timeUntilPrefetch.isNegative() || timeUntilPrefetch.toDays() > 7) {
+            log.debug(() -> "Skipping background refresh because the prefetch time is in the past or too far in the future: " +
+                            result.prefetchTime());
             return;
         }
 
-        Duration waitTime = Duration.between(Instant.now(), refreshTime);
-        log.debug(() -> "Scheduling refresh attempt for " + refreshTime + " (in " + waitTime.toMillis() + " ms)");
+        Instant backgroundRefreshTime = result.prefetchTime().plusSeconds(1);
+        Duration timeUntilBackgroundRefresh = timeUntilPrefetch.plusSeconds(1);
+
+        log.debug(() -> "Scheduling refresh attempt for " + backgroundRefreshTime + " (in " +
+                        timeUntilBackgroundRefresh.toMillis() + " ms)");
 
         ScheduledFuture<?> scheduledTask = SCHEDULER.schedule(() -> {
             runWithInstanceThreadName(() -> {
-                log.debug(() -> "Executing refresh attempt scheduled for " + refreshTime);
+                log.debug(() -> "Executing refresh attempt scheduled for " + backgroundRefreshTime);
 
                 // If the supplier has already been prefetched, this will just be a cache hit.
                 tryRunBackgroundTask(cachedSupplier::get);
             });
-        }, waitTime.toMillis(), TimeUnit.MILLISECONDS);
+        }, timeUntilBackgroundRefresh.toMillis(), TimeUnit.MILLISECONDS);
 
         updateTask(scheduledTask);
 
diff --git a/utils/src/test/java/software/amazon/awssdk/utils/cache/CachedSupplierTest.java b/utils/src/test/java/software/amazon/awssdk/utils/cache/CachedSupplierTest.java
@@ -198,6 +198,7 @@ public void oneCallerBlocksPrefetchStrategyWorks() throws InterruptedException {
         try (WaitingSupplier waitingSupplier = new WaitingSupplier(future(), past())) {
             CachedSupplier<String> cachedSupplier = CachedSupplier.builder(waitingSupplier)
                                                                   .prefetchStrategy(new OneCallerBlocks())
+                                                                  .maxPrefetchJitter(Duration.ZERO)
                                                                   .build();
 
             // Perform one successful "get" to prime the cache.
@@ -225,6 +226,7 @@ public void nonBlockingPrefetchStrategyWorks() {
         try (WaitingSupplier waitingSupplier = new WaitingSupplier(future(), past());
              CachedSupplier<String> cachedSupplier = CachedSupplier.builder(waitingSupplier)
                                                                    .prefetchStrategy(new NonBlocking("test-%s"))
+                                                                   .maxPrefetchJitter(Duration.ZERO)
                                                                    .build()) {
             // Perform one successful "get" to prime the cache.
             waitingSupplier.permits.release(1);
@@ -245,7 +247,8 @@ public void nonBlockingPrefetchStrategyWorks() {
     public void nonBlockingPrefetchStrategyRefreshesInBackground() {
         try (WaitingSupplier waitingSupplier = new WaitingSupplier(now().plusSeconds(62), now());
              CachedSupplier<String> cachedSupplier = CachedSupplier.builder(waitingSupplier)
-                                                                   .prefetchStrategy(new NonBlocking("test-%s", Duration.ZERO))
+                                                                   .prefetchStrategy(new NonBlocking("test-%s"))
+                                                                   .maxPrefetchJitter(Duration.ZERO)
                                                                    .build()) {
             waitingSupplier.permits.release(2);
             cachedSupplier.get();
@@ -309,23 +312,23 @@ public void threadsAreSharedBetweenNonBlockingInstances() throws InterruptedExce
                 CachedSupplier<String> supplier =
                     CachedSupplier.builder(() -> RefreshResult.builder("foo")
                                                               .prefetchTime(now())
-                                                              .staleTime(now())
+                                                              .staleTime(future())
                                                               .build())
-                                  .prefetchStrategy(new NonBlocking("test", Duration.ZERO))
+                                  .prefetchStrategy(new NonBlocking("test"))
+                                  .maxPrefetchJitter(Duration.ofMillis(10))
                                   .build();
                 supplier.get();
                 css.add(supplier);
-                Thread.sleep(10);
             }
 
             int maxActive = 0;
-            for (int i = 0; i < 100; i++) {
+            for (int i = 0; i < 1000; i++) {
                 maxActive = Math.max(maxActive, NonBlocking.executor().getActiveCount());
-                Thread.sleep(10);
+                Thread.sleep(1);
             }
 
             // Make sure we used less-than 99 to do the refreshes.
-            assertThat(maxActive).isBetween(2, 99);
+            assertThat(maxActive).isBetween(1, 99);
         } finally {
             css.forEach(CachedSupplier::close);
         }
@@ -345,7 +348,8 @@ public void activeThreadsHaveMaxCount() throws InterruptedException {
                                             .prefetchTime(now())
                                             .staleTime(now())
                                             .build();
-                    }).prefetchStrategy(new NonBlocking("test", Duration.ZERO))
+                    }).prefetchStrategy(new NonBlocking("test"))
+                      .maxPrefetchJitter(Duration.ZERO)
                       .build();
                 executor.submit(supplier::get);
                 css.add(supplier);
