diff --git a/build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml b/build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml index 05606dc6d574..152507019408 100644 --- a/build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml +++ b/build-tools/src/main/resources/software/amazon/awssdk/spotbugs-suppressions.xml @@ -547,4 +547,14 @@ + + + + + + + + + + diff --git a/core/retries-spi/pom.xml b/core/retries-spi/pom.xml index 1b63df065554..c975ef1005d4 100644 --- a/core/retries-spi/pom.xml +++ b/core/retries-spi/pom.xml @@ -64,5 +64,11 @@ assertj-core test + + + org.mockito + mockito-core + test + diff --git a/core/retries-spi/src/main/java/software/amazon/awssdk/retries/api/RetryStrategy.java b/core/retries-spi/src/main/java/software/amazon/awssdk/retries/api/RetryStrategy.java index 491bab172e65..ce5732977157 100644 --- a/core/retries-spi/src/main/java/software/amazon/awssdk/retries/api/RetryStrategy.java +++ b/core/retries-spi/src/main/java/software/amazon/awssdk/retries/api/RetryStrategy.java @@ -15,9 +15,12 @@ package software.amazon.awssdk.retries.api; +import java.util.concurrent.CompletableFuture; import java.util.function.Predicate; import software.amazon.awssdk.annotations.SdkPublicApi; import software.amazon.awssdk.annotations.ThreadSafe; +import software.amazon.awssdk.utils.CompletableFutureUtils; +import software.amazon.awssdk.utils.SdkAutoCloseable; import software.amazon.awssdk.utils.builder.SdkBuilder; /** @@ -40,9 +43,9 @@ */ @ThreadSafe @SdkPublicApi -public interface RetryStrategy { +public interface RetryStrategy extends SdkAutoCloseable { /** - * Invoked before the first request attempt. + * Acquire a retry token synchronously. Invoked before the first request attempt. * *

Callers MUST wait for the {@code delay} returned by this call before making the first attempt. Callers that wish to * retry a failed attempt MUST call {@link #refreshRetryToken} before doing so. @@ -55,7 +58,27 @@ public interface RetryStrategy { AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request); /** - * Invoked before each subsequent (non-first) request attempt. + * Acquire a retry token asynchronously. Invoked before the first request attempt. + * + *

Callers MUST wait for the {@code delay} returned by this call before making the first attempt. Callers that wish to + * retry a failed attempt MUST call {@link #refreshRetryToken} before doing so. + * + *

If the attempt was successful, callers MUST call {@link #recordSuccess}. + * + *

The future is completed exceptionally with {@link NullPointerException} if a required parameter is {@code null}. + *

The future is completed exceptionally with {@link TokenAcquisitionFailedException} if a token cannot be acquired. + */ + default CompletableFuture acquireInitialTokenAsync(AcquireInitialTokenRequest request) { + try { + AcquireInitialTokenResponse response = acquireInitialToken(request); + return CompletableFuture.completedFuture(response); + } catch (Throwable t) { + return CompletableFutureUtils.failedFuture(t); + } + } + + /** + * Refresh a retry token synchronously. Invoked before each subsequent (non-first) request attempt. * *

Callers MUST wait for the {@code delay} returned by this call before making the next attempt. If the next attempt * fails, callers MUST re-call {@link #refreshRetryToken} before attempting another retry. This call invalidates the provided @@ -70,6 +93,29 @@ public interface RetryStrategy { */ RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request); + /** + * Refresh a retry token asynchronously. Invoked before each subsequent (non-first) request attempt. + * + *

Callers MUST wait for the {@code delay} returned by this call before making the next attempt. If the next attempt + * fails, callers MUST re-call {@link #refreshRetryToken} before attempting another retry. This call invalidates the provided + * token, and returns a new one. Callers MUST use the new token. + * + *

If the attempt was successful, callers MUST call {@link #recordSuccess}. + * + *

The future is completed exceptionally with {@link NullPointerException} if a required parameter is not specified. + *

The future is completed exceptionally with {@link IllegalArgumentException} if the provided token was not issued by + * this strategy or the provided token was already used for a previous refresh or success call. + *

The future is completed exceptionally with {@link TokenAcquisitionFailedException} if a token cannot be acquired. + */ + default CompletableFuture refreshRetryTokenAsync(RefreshRetryTokenRequest request) { + try { + RefreshRetryTokenResponse response = refreshRetryToken(request); + return CompletableFuture.completedFuture(response); + } catch (Throwable t) { + return CompletableFutureUtils.failedFuture(t); + } + } + /** * Invoked after an attempt succeeds. * @@ -102,6 +148,10 @@ default boolean useClientDefaults() { */ Builder toBuilder(); + @Override + default void close() { + } + /** * Builder to create immutable instances of {@link RetryStrategy}. */ diff --git a/core/retries-spi/src/test/java/software/amazon/awssdk/retries/api/RetryStrategyTest.java b/core/retries-spi/src/test/java/software/amazon/awssdk/retries/api/RetryStrategyTest.java new file mode 100644 index 000000000000..3e419f1ab8aa --- /dev/null +++ b/core/retries-spi/src/test/java/software/amazon/awssdk/retries/api/RetryStrategyTest.java @@ -0,0 +1,61 @@ +/* + * 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.retries.api; + +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import org.junit.jupiter.api.Test; + +public class RetryStrategyTest { + @Test + void acquireInitialTokenAsync_syncThrows_reportedThroughFuture() { + RetryStrategy retryStrategy = mock(RetryStrategy.class); + + when(retryStrategy.acquireInitialTokenAsync(any(AcquireInitialTokenRequest.class))) + .thenCallRealMethod(); + + RuntimeException t = new RuntimeException("oops"); + when(retryStrategy.acquireInitialToken(any(AcquireInitialTokenRequest.class))) + .thenThrow(t); + + assertThatThrownBy(() -> retryStrategy.acquireInitialTokenAsync(AcquireInitialTokenRequest.create("test")).join()) + .hasRootCause(t); + } + + @Test + void refreshRetryTokenAsync_syncThrows_reportedThroughFuture() { + RetryStrategy retryStrategy = mock(RetryStrategy.class); + + when(retryStrategy.refreshRetryTokenAsync(any(RefreshRetryTokenRequest.class))) + .thenCallRealMethod(); + + RuntimeException t = new RuntimeException("oops"); + when(retryStrategy.refreshRetryToken(any(RefreshRetryTokenRequest.class))) + .thenThrow(t); + + RetryToken token = mock(RetryToken.class); + RefreshRetryTokenRequest request = RefreshRetryTokenRequest.builder() + .token(token) + .failure(new RuntimeException("failure")) + .build(); + + assertThatThrownBy(() -> retryStrategy.refreshRetryTokenAsync(request).join()) + .hasRootCause(t); + } +} diff --git a/core/retries/pom.xml b/core/retries/pom.xml index 4cf5d1bc08a4..e9ab029e34d8 100644 --- a/core/retries/pom.xml +++ b/core/retries/pom.xml @@ -70,5 +70,10 @@ assertj-core test + + org.mockito + mockito-core + test + diff --git a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/BaseRetryStrategy.java b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/BaseRetryStrategy.java index dfda06093ab3..6c6551dc68e2 100644 --- a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/BaseRetryStrategy.java +++ b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/BaseRetryStrategy.java @@ -39,6 +39,7 @@ import software.amazon.awssdk.retries.internal.circuitbreaker.TokenBucket; import software.amazon.awssdk.retries.internal.circuitbreaker.TokenBucketStore; import software.amazon.awssdk.utils.Logger; +import software.amazon.awssdk.utils.Pair; import software.amazon.awssdk.utils.ToString; import software.amazon.awssdk.utils.Validate; @@ -86,7 +87,7 @@ public abstract class BaseRetryStrategy implements DefaultAwareRetryStrategy { * @see RetryStrategy#acquireInitialToken(AcquireInitialTokenRequest) */ @Override - public final AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request) { + public AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request) { logAcquireInitialToken(request); DefaultRetryToken token = DefaultRetryToken.builder().scope(request.scope()).build(); return AcquireInitialTokenResponse.create(token, computeInitialBackoff(request)); @@ -98,7 +99,20 @@ public final AcquireInitialTokenResponse acquireInitialToken(AcquireInitialToken * @see RetryStrategy#refreshRetryToken(RefreshRetryTokenRequest) */ @Override - public final RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request) { + public RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request) { + Pair refreshedToken = refreshTokenOrThrow(request); + Duration backoff = computeBackoff(request, refreshedToken.left()); + + logRefreshTokenSuccess(refreshedToken.left(), refreshedToken.right(), backoff); + return RefreshRetryTokenResponseImpl.create(refreshedToken.left(), backoff); + } + + /** + * Attempt to refresh the token for a retry or throws {@link TokenAcquisitionFailedException} if unable to do so. + * + * @return A pair of the refreshed token and the successful acquire response from the token bucket. + */ + protected Pair refreshTokenOrThrow(RefreshRetryTokenRequest request) { DefaultRetryToken token = asDefaultRetryToken(request.token()); // Check if we meet the preconditions needed for retrying. These will throw if the expected condition is not meet. @@ -115,12 +129,8 @@ public final RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenReques // All the conditions required to retry were meet, update the internal state before retrying. updateStateForRetry(request); - // Refresh the retry token and compute the backoff delay. - DefaultRetryToken refreshedToken = refreshToken(request, acquireResponse); - Duration backoff = computeBackoff(request, refreshedToken); - - logRefreshTokenSuccess(refreshedToken, acquireResponse, backoff); - return RefreshRetryTokenResponseImpl.create(refreshedToken, backoff); + // Refresh the retry token + return Pair.of(refreshToken(request, acquireResponse), acquireResponse); } /** @@ -335,7 +345,7 @@ private String acquisitionFailedMessage(AcquireResponse response) { response.maxCapacity()); } - private void logAcquireInitialToken(AcquireInitialTokenRequest request) { + protected void logAcquireInitialToken(AcquireInitialTokenRequest request) { // Request attempt 1 token acquired (backoff: 0ms, cost: 0, capacity: 500/500) TokenBucket tokenBucket = tokenBucketStore.tokenBucketForScope(request.scope()); log.debug(() -> String.format("Request attempt 1 token acquired " @@ -343,7 +353,7 @@ private void logAcquireInitialToken(AcquireInitialTokenRequest request) { tokenBucket.currentCapacity(), tokenBucket.maxCapacity())); } - private void logRefreshTokenSuccess(DefaultRetryToken token, AcquireResponse acquireResponse, Duration delay) { + protected void logRefreshTokenSuccess(DefaultRetryToken token, AcquireResponse acquireResponse, Duration delay) { log.debug(() -> String.format("Request attempt %d token acquired " + "(backoff: %dms, cost: %d, capacity: %d/%d)", token.attempt(), delay.toMillis(), diff --git a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java index c52c7859526f..685bba2d5e46 100644 --- a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java +++ b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/DefaultAdaptiveRetryStrategy.java @@ -16,22 +16,27 @@ package software.amazon.awssdk.retries.internal; import java.time.Duration; +import java.util.concurrent.CompletableFuture; import java.util.function.Predicate; import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.retries.AdaptiveRetryStrategy; import software.amazon.awssdk.retries.api.AcquireInitialTokenRequest; +import software.amazon.awssdk.retries.api.AcquireInitialTokenResponse; import software.amazon.awssdk.retries.api.BackoffStrategy; import software.amazon.awssdk.retries.api.RefreshRetryTokenRequest; +import software.amazon.awssdk.retries.api.RefreshRetryTokenResponse; +import software.amazon.awssdk.retries.internal.circuitbreaker.AcquireResponse; import software.amazon.awssdk.retries.internal.circuitbreaker.TokenBucketStore; import software.amazon.awssdk.retries.internal.ratelimiter.RateLimiterTokenBucket; import software.amazon.awssdk.retries.internal.ratelimiter.RateLimiterTokenBucketStore; +import software.amazon.awssdk.utils.CompletableFutureUtils; import software.amazon.awssdk.utils.Logger; +import software.amazon.awssdk.utils.Pair; import software.amazon.awssdk.utils.Validate; @SdkInternalApi public final class DefaultAdaptiveRetryStrategy extends BaseRetryStrategy implements AdaptiveRetryStrategy { - private static final Logger LOG = Logger.loggerFor(DefaultAdaptiveRetryStrategy.class); private final RateLimiterTokenBucketStore rateLimiterTokenBucketStore; @@ -42,16 +47,53 @@ public final class DefaultAdaptiveRetryStrategy } @Override - protected Duration computeInitialBackoff(AcquireInitialTokenRequest request) { + public AcquireInitialTokenResponse acquireInitialToken(AcquireInitialTokenRequest request) { + return CompletableFutureUtils.joinLikeSync(acquireInitialTokenAsync(request)); + } + + @Override + public RefreshRetryTokenResponse refreshRetryToken(RefreshRetryTokenRequest request) { + return CompletableFutureUtils.joinLikeSync(refreshRetryTokenAsync(request)); + } + + @Override + public CompletableFuture acquireInitialTokenAsync(AcquireInitialTokenRequest request) { + logAcquireInitialToken(request); RateLimiterTokenBucket bucket = rateLimiterTokenBucketStore.tokenBucketForScope(request.scope()); - return bucket.tryAcquire().delay(); + CompletableFuture acquireResult = bucket.acquireAsync(); + + return acquireResult.thenApply(r -> { + DefaultRetryToken token = DefaultRetryToken.builder().scope(request.scope()).build(); + return AcquireInitialTokenResponse.create(token, Duration.ZERO); + }); } @Override - protected Duration computeBackoff(RefreshRetryTokenRequest request, DefaultRetryToken token) { - Duration backoff = super.computeBackoff(request, token); + public CompletableFuture refreshRetryTokenAsync(RefreshRetryTokenRequest request) { + DefaultRetryToken token = (DefaultRetryToken) request.token(); + Pair refreshResult; + try { + refreshResult = refreshTokenOrThrow(request); + } catch (Throwable t) { + return CompletableFutureUtils.failedFuture(t); + } + + DefaultRetryToken refreshedToken = refreshResult.left(); + AcquireResponse acquireResponse = refreshResult.right(); RateLimiterTokenBucket bucket = rateLimiterTokenBucketStore.tokenBucketForScope(token.scope()); - return backoff.plus(bucket.tryAcquire().delay()); + CompletableFuture acquireResult = bucket.acquireAsync(); + return acquireResult.thenApply(r -> { + // Note: This is the backoff imposed standard retry strategy, *not* the rate limiter. This must still be honored by + // the caller before sending the request. + Duration backoff = computeBackoff(request, refreshedToken); + logRefreshTokenSuccess(refreshedToken, acquireResponse, backoff); + return RefreshRetryTokenResponse.create(refreshedToken, backoff); + }); + } + + @Override + public void close() { + rateLimiterTokenBucketStore.close(); } @Override diff --git a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java index 2d8743af9924..86a718f18906 100644 --- a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java +++ b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucket.java @@ -16,42 +16,150 @@ package software.amazon.awssdk.retries.internal.ratelimiter; import java.time.Duration; -import java.util.concurrent.atomic.AtomicReference; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import java.util.function.Function; import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.annotations.SdkTestInternalApi; +import software.amazon.awssdk.annotations.ThreadSafe; +import software.amazon.awssdk.utils.CompletableFutureUtils; +import software.amazon.awssdk.utils.SdkAutoCloseable; /** * The {@link RateLimiterTokenBucket} keeps track of past throttling responses and adapts to slow down the send rate to adapt to - * the service. It does this by suggesting a delay amount as result of a {@link #tryAcquire()} call. Callers must update its - * internal state by calling {@link #updateRateAfterThrottling()} when getting a throttling response or - * {@link #updateRateAfterSuccess()} when getting successful response. + * the service. It does this by delaying the completion of the future returned by {@link #acquireAsync()} until the requested + * capacity is available. Callers must update its internal state by calling {@link #updateRateAfterThrottling()} when getting a + * throttling response or {@link #updateRateAfterSuccess()} when getting successful response. * - *

This class is thread-safe, its internal current state is kept in the inner class {@link PersistentState} which is stored - * using an {@link AtomicReference}. This class is converted to {@link TransientState} when the state needs to be mutated and - * converted back to a {@link PersistentState} and stored using {@link AtomicReference#compareAndSet(Object, Object)}. + *

This class is thread-safe. * *

The algorithm used is adapted from the network congestion avoidance algorithm * CUBIC. */ @SdkInternalApi -public class RateLimiterTokenBucket { - private final AtomicReference stateReference; +@ThreadSafe +public class RateLimiterTokenBucket implements SdkAutoCloseable { + // Thread used for capacity waiting and notifying. + private final ScheduledExecutorService scheduler; + + // Protect access to other members below + private final Object lock = new Object(); + private final RateLimiterClock clock; + // the collection of futures returned to threads currently waiting for capacity. + // Futures are completed in FIFO order. + // The size of this equal to the number of threads concurrently accessing this bucket. + private final Deque> waiting = new ArrayDeque<>(); + private PersistentState state; + private boolean open = true; + private boolean notifierRunning = false; - RateLimiterTokenBucket(RateLimiterClock clock) { + RateLimiterTokenBucket(RateLimiterClock clock, ScheduledExecutorService scheduler) { this.clock = clock; - this.stateReference = new AtomicReference<>(new PersistentState()); + this.scheduler = scheduler; + this.state = new PersistentState(); + } + + @Override + public void close() { + doClose(null); + } + + private void doClose(Throwable cause) { + synchronized (lock) { + open = false; + IllegalStateException closedException = new IllegalStateException("Rate limiter bucket is closed", cause); + while (true) { + CompletableFuture w = waiting.poll(); + if (w == null) { + break; + } + w.completeExceptionally(closedException); + } + } } /** - * Acquire tokens from the bucket. If the bucket contains enough capacity to satisfy the request, this method will return in - * {@link RateLimiterAcquireResponse#delay()} a {@link Duration#ZERO} value, otherwise it will return the amount of time the - * callers need to wait until enough tokens are refilled. + * Acquire a token from the bucket. + * + * @return A future that is completed when the requested amount is acquired from this bucket. */ - public RateLimiterAcquireResponse tryAcquire() { - StateUpdate update = updateState(ts -> ts.tokenBucketAcquire(clock, 1.0)); - return RateLimiterAcquireResponse.create(update.result); + public CompletableFuture acquireAsync() { + synchronized (lock) { + if (!open) { + return CompletableFutureUtils.failedFuture(new IllegalStateException("Rate limiter bucket is closed")); + } + + // fast path and avoid scheduling in the executor if throttling isn't enabled. + if (!state.enabled) { + return CompletableFuture.completedFuture(null); + } + + CompletableFuture future = new CompletableFuture<>(); + waiting.add(future); + if (!notifierRunning) { + notifierRunning = scheduleOrClose(this::doNotify, Duration.ZERO); + } + return future; + } + } + + + @SdkTestInternalApi + Deque> waiting() { + return waiting; + } + + @SdkTestInternalApi + boolean isClosed() { + return !open; + } + + private void doNotify() { + while (true) { + CompletableFuture w; + synchronized (lock) { + w = waiting.poll(); + if (w == null) { + notifierRunning = false; + return; + } + + TransientState.AcquireResult acquireResult = updateState(t -> t.tokenBucketAcquire(clock, 1.0)).result; + + // 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; + } + + } + // Acquire was successful, signal the waiting thread. + w.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); + } + return false; } /** @@ -88,24 +196,15 @@ private StateUpdate consumeState(Consumer mutator) { } /** - * Converts the stored persistent state into a transient one and transforms it using the provided function. The provided - * function is expected to update the transient state in-place and return a value that will be returned to the caller in the - * {@link StateUpdate#result} field. The mutated transient value is converted back to a persistent one and stored in the - * atomic reference if no changes were made in-between. If another thread changes the value in-between, the operation is - * retried until succeeded. + * Converts the stored persistent state into a transient one and transforms it using the provided function. */ private StateUpdate updateState(Function mutator) { - PersistentState current; - PersistentState updated; - T result; - do { - current = stateReference.get(); - TransientState transientState = current.toTransient(); - result = mutator.apply(transientState); - updated = transientState.toPersistent(); - } while (!stateReference.compareAndSet(current, updated)); - - return new StateUpdate<>(updated, result); + synchronized (lock) { + TransientState transientState = state.toTransient(); + T result = mutator.apply(transientState); + state = transientState.toPersistent(); + return new StateUpdate<>(state, result); + } } static class StateUpdate { @@ -163,17 +262,38 @@ PersistentState toPersistent() { * a {@link Duration#ZERO} value, otherwise it will return the amount of time the callers need to wait until enough tokens * are refilled. */ - Duration tokenBucketAcquire(RateLimiterClock clock, double amount) { + AcquireResult tokenBucketAcquire(RateLimiterClock clock, double amount) { if (!this.enabled) { - return Duration.ZERO; + return new AcquireResult(true, Duration.ZERO); } refill(clock); - double waitTime = 0.0; if (this.currentCapacity < amount) { - waitTime = (amount - this.currentCapacity) / this.fillRate; + double diff = amount - currentCapacity; + double waitTime = diff / this.fillRate; + double waitTimeMs = waitTime * 1_000.0; + Duration duration = Duration.ofMillis((long) Math.ceil(waitTimeMs)); + return new AcquireResult(false, duration); } this.currentCapacity -= amount; - return Duration.ofNanos((long) (waitTime * 1_000_000_000.0)); + return new AcquireResult(true, Duration.ZERO); + } + + private static class AcquireResult { + final boolean successful; + final Duration refillWait; + + AcquireResult(boolean successful, Duration refillWait) { + this.successful = successful; + this.refillWait = refillWait; + } + + boolean isSuccessful() { + return successful; + } + + Duration refillWait() { + return refillWait; + } } /** diff --git a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java index 303e6d61d7f5..a2a5f2f966ed 100644 --- a/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java +++ b/core/retries/src/main/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStore.java @@ -15,8 +15,13 @@ package software.amazon.awssdk.retries.internal.ratelimiter; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.annotations.SdkTestInternalApi; import software.amazon.awssdk.annotations.ToBuilderIgnoreField; +import software.amazon.awssdk.utils.SdkAutoCloseable; +import software.amazon.awssdk.utils.ThreadFactoryBuilder; import software.amazon.awssdk.utils.Validate; import software.amazon.awssdk.utils.builder.CopyableBuilder; import software.amazon.awssdk.utils.builder.ToCopyableBuilder; @@ -27,23 +32,59 @@ */ @SdkInternalApi public final class RateLimiterTokenBucketStore - implements ToCopyableBuilder { + implements ToCopyableBuilder, SdkAutoCloseable { private static final int MAX_ENTRIES = 128; + private static final String THREAD_NAME_PREFIX = "sdk-adaptive-rate-limiter-"; + private static final RateLimiterClock DEFAULT_CLOCK = new SystemClock(); private final LruCache scopeToTokenBucket; private final RateLimiterClock clock; + private final ScheduledExecutorService scheduler; + private final boolean closeScheduler; private RateLimiterTokenBucketStore(Builder builder) { - this.clock = Validate.paramNotNull(builder.clock, "clock"); - this.scopeToTokenBucket = LruCache.builder(x -> new RateLimiterTokenBucket(clock)) + this(builder.clock, + resolveScheduler(builder), + builder.scheduler == null); + } + + private RateLimiterTokenBucketStore(RateLimiterClock clock, ScheduledExecutorService scheduler, boolean closeScheduler) { + this.clock = Validate.paramNotNull(clock, "clock"); + this.scheduler = Validate.paramNotNull(scheduler, "scheduler"); + this.closeScheduler = closeScheduler; + this.scopeToTokenBucket = LruCache.builder( + x -> new RateLimiterTokenBucket(clock, scheduler)) .maxSize(MAX_ENTRIES) .build(); } + @Override + public void close() { + scopeToTokenBucket.evictAll(); + if (closeScheduler) { + scheduler.shutdownNow(); + } + } + public RateLimiterTokenBucket tokenBucketForScope(String scope) { return scopeToTokenBucket.get(scope); } + @SdkTestInternalApi + ScheduledExecutorService scheduler() { + return scheduler; + } + + private static ScheduledExecutorService resolveScheduler(Builder b) { + if (b.scheduler != null) { + return b.scheduler; + } + return Executors.newSingleThreadScheduledExecutor(new ThreadFactoryBuilder() + .daemonThreads(true) + .threadNamePrefix(THREAD_NAME_PREFIX) + .build()); + } + @Override @ToBuilderIgnoreField("scopeToTokenBucket") public Builder toBuilder() { @@ -56,6 +97,7 @@ public static RateLimiterTokenBucketStore.Builder builder() { public static class Builder implements CopyableBuilder { private RateLimiterClock clock; + private ScheduledExecutorService scheduler; Builder() { this.clock = DEFAULT_CLOCK; @@ -63,6 +105,7 @@ public static class Builder implements CopyableBuilderwill not be closed when {@link #close() closing} this bucket store. + * + * @return This object for method chaining. + */ + public Builder scheduler(ScheduledExecutorService scheduler) { + this.scheduler = scheduler; + return this; + } + @Override public RateLimiterTokenBucketStore build() { return new RateLimiterTokenBucketStore(this); diff --git a/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java b/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java new file mode 100644 index 000000000000..b915045e1b5a --- /dev/null +++ b/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketStoreTest.java @@ -0,0 +1,88 @@ +/* + * 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.retries.internal.ratelimiter; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ScheduledExecutorService; +import org.junit.jupiter.api.Test; + +public class RateLimiterTokenBucketStoreTest { + @Test + void close_schedulerProvided_schedulerNotClosed() { + ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class); + RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder() + .clock(new SystemClock()) + .scheduler(scheduler) + .build(); + + store.close(); + + verify(scheduler, never()).shutdownNow(); + verify(scheduler, never()).shutdown(); + } + + @Test + void close_schedulerNotProvidedOnBuilder_schedulerClosed() { + RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder() + .clock(new SystemClock()) + .build(); + + store.close(); + + assertThat(store.scheduler().isShutdown()).isTrue(); + } + + @Test + void close_closesAllCacheEntries() { + int entries = 64; + List buckets = new ArrayList<>(entries); + + ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class); + RateLimiterTokenBucketStore store = RateLimiterTokenBucketStore.builder() + .clock(new SystemClock()) + .scheduler(scheduler) + .build(); + + List> futures = new ArrayList<>(entries); + + for (int i = 0; i < entries; ++i) { + RateLimiterTokenBucket bucket = store.tokenBucketForScope(Integer.toString(i)); + buckets.add(bucket); + + // enable throttling so futures from acquireAsync get queued + bucket.updateRateAfterThrottling(); + futures.add(bucket.acquireAsync()); + } + + store.close(); + + // New acquires from the closed bucket should fail + assertThat(buckets).allSatisfy(b -> assertThatThrownBy(b.acquireAsync()::join) + .hasMessageContaining("Rate limiter bucket is closed")); + + // All pending futures should be failed + assertThat(futures).allSatisfy(cf -> assertThatThrownBy(cf::join) + .hasMessageContaining("Rate limiter bucket is closed")); + } +} diff --git a/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java b/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java index 14021ebde812..806826a20d07 100644 --- a/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java +++ b/core/retries/src/test/java/software/amazon/awssdk/retries/internal/ratelimiter/RateLimiterTokenBucketTest.java @@ -16,44 +16,206 @@ package software.amazon.awssdk.retries.internal.ratelimiter; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.AssertionsForClassTypes.within; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; import java.util.Arrays; import java.util.Collection; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.MethodSource; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; class RateLimiterTokenBucketTest { - private static MutableClock clock = null; - private static RateLimiterTokenBucket tokenBucket = null; private static final double EPSILON = 0.0001; + private MutableClock clock = null; + private ScheduledExecutorService scheduler = null; + private RateLimiterTokenBucket tokenBucket = null; - @BeforeAll - static void setup() { + @BeforeEach + void setup() { clock = new MutableClock(); - tokenBucket = new RateLimiterTokenBucket(clock); + scheduler = mock(ScheduledExecutorService.class); + tokenBucket = new RateLimiterTokenBucket(clock, scheduler); } - @ParameterizedTest - @MethodSource("parameters") - void testCase(TestCase testCase) { + @Test + void acquireAsync_bucketClosed_futureCompletedExceptionally() { + tokenBucket.close(); + CompletableFuture f = tokenBucket.acquireAsync(); + assertThatThrownBy(f::join).satisfies(t -> { + Throwable cause = t.getCause(); + assertThat(cause).isExactlyInstanceOf(IllegalStateException.class); + assertThat(cause).hasMessage("Rate limiter bucket is closed"); + }); + } + + @Test + void acquireAsync_notEnabled_doesNotScheduleTask() { + CompletableFuture f = tokenBucket.acquireAsync(); + assertThat(f).isCompleted(); + verifyNoInteractions(scheduler); + } + + @Test + void acquireAsync_enabled_schedulesTask() { + tokenBucket.updateRateAfterThrottling(); + + tokenBucket.acquireAsync(); + verify(scheduler).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + } + + + @Test + void acquireAsync_scheduleFails_completesFutureExceptionally() { + tokenBucket.updateRateAfterThrottling(); + + doThrow(new RejectedExecutionException("no")).when(scheduler).schedule(any(Runnable.class), + anyLong(), + any(TimeUnit.class)); + + CompletableFuture f = tokenBucket.acquireAsync(); + assertThatThrownBy(f::join).satisfies(t -> { + Throwable cause = t.getCause(); + assertThat(cause).hasMessage("Rate limiter bucket is closed"); + }); + } + + @Test + void acquireAsync_scheduleFails_futureNotInWaitingDeque() { + tokenBucket.updateRateAfterThrottling(); + + doThrow(new RejectedExecutionException("no")).when(scheduler).schedule(any(Runnable.class), + anyLong(), + any(TimeUnit.class)); + + CompletableFuture f = tokenBucket.acquireAsync(); + assertThat(f).isCompletedExceptionally(); + assertThat(tokenBucket.waiting()).isEmpty(); + } + + @Test + void acquireAsync_scheduleFails_closesBucket() { + tokenBucket.updateRateAfterThrottling(); + + doThrow(new RejectedExecutionException("no")).when(scheduler).schedule(any(Runnable.class), + anyLong(), + any(TimeUnit.class)); + + CompletableFuture f = tokenBucket.acquireAsync(); + assertThat(f).isCompletedExceptionally(); + assertThat(tokenBucket.isClosed()).isTrue(); + } + + @Test + void close_completesAllPendingFutures() { + // enable throttling so futures actually get queued instead of being completed immediately + tokenBucket.updateRateAfterThrottling(); + + List> futures = IntStream.range(0, 10) + .mapToObj(i -> tokenBucket.acquireAsync()) + .collect(Collectors.toList()); + + tokenBucket.close(); + + assertThat(futures).allSatisfy(f -> { + assertThatThrownBy(f::join).satisfies(t -> { + Throwable cause = t.getCause(); + assertThat(cause).isExactlyInstanceOf(IllegalStateException.class); + assertThat(cause).hasMessage("Rate limiter bucket is closed"); + }); + }); + + assertThat(tokenBucket.waiting()).isEmpty(); + } + + @Test + void close_doesShutDownExecutor() { + tokenBucket.close(); + verifyNoInteractions(scheduler); + } + + @Test + void doNotify_scheduleRejected_failsFuture() { + tokenBucket.updateRateAfterThrottling(); + + // Empty bucket at default rate of 0.5 tokens per second should be 2seconds + when(scheduler.schedule(any(Runnable.class), eq(2000L), eq(TimeUnit.MILLISECONDS))) + .thenThrow(new RejectedExecutionException("no")); + + // 0L is the initial schedule from acquireAsync, capture the doNotify schedule and execute that. + when(scheduler.schedule(any(Runnable.class), eq(0L), any(TimeUnit.class))).thenAnswer(i -> { + Runnable r = i.getArgument(0); + r.run(); + return null; + }); + + tokenBucket.acquireAsync(); + assertThat(tokenBucket.isClosed()).isTrue(); + } + + @Test + void doNotify_scheduleRejected_closesBucket() { + tokenBucket.updateRateAfterThrottling(); + + // Empty bucket at default rate of 0.5 tokens per second should be 2seconds + when(scheduler.schedule(any(Runnable.class), eq(2000L), eq(TimeUnit.MILLISECONDS))) + .thenThrow(new RejectedExecutionException("no")); + + // 0L is the initial schedule from acquireAsync, capture the doNotify schedule and execute that. + when(scheduler.schedule(any(Runnable.class), eq(0L), any(TimeUnit.class))).thenAnswer(i -> { + Runnable r = i.getArgument(0); + r.run(); + return null; + }); + + CompletableFuture future = tokenBucket.acquireAsync(); + assertThatThrownBy(future::join).hasRootCauseInstanceOf(RejectedExecutionException.class); + assertThat(tokenBucket.isClosed()).isTrue(); + } + + @Test + void sendingRateEndToEndTest() { + for (TestCase sendingRateTestCase : sendingRateTestCases()) { + assertSendingRateTestCase(sendingRateTestCase); + } + } + + void assertSendingRateTestCase(TestCase testCase) { clock.setCurrent(testCase.givenTimestamp); RateLimiterUpdateResponse res; - tokenBucket.tryAcquire(); + if (testCase.throttleResponse) { res = tokenBucket.updateRateAfterThrottling(); } else { res = tokenBucket.updateRateAfterSuccess(); } double measuredTxRate = res.measuredTxRate(); - assertThat(measuredTxRate).isCloseTo(testCase.expectMeasuredTxRate, within(EPSILON)); + assertThat(measuredTxRate) + .as("%s: Measured TX rate", testCase) + .isCloseTo(testCase.expectMeasuredTxRate, within(EPSILON)); double fillRate = res.fillRate(); - assertThat(fillRate).isCloseTo(testCase.expectFillRate, within(EPSILON)); + assertThat(fillRate) + .as("%s: Fill rate", testCase) + .isCloseTo(testCase.expectFillRate, within(EPSILON)); } - - static Collection parameters() { + static Collection sendingRateTestCases() { + // Note: Test cases are not independent. Each case depends on the state of the bucket being correctly updated from the + // previous test. return Arrays.asList( new TestCase() .givenSuccessResponse() @@ -174,6 +336,15 @@ TestCase expectFillRate(double expectFillRate) { return this; } + @Override + public String toString() { + return "TestCase{" + + "throttleResponse=" + throttleResponse + + ", givenTimestamp=" + givenTimestamp + + ", expectMeasuredTxRate=" + expectMeasuredTxRate + + ", expectFillRate=" + expectFillRate + + '}'; + } } static class MutableClock implements RateLimiterClock { diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStage.java index 15b81d30d75f..19ba5036ecc7 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStage.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStage.java @@ -36,7 +36,6 @@ import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.http.SdkHttpResponse; import software.amazon.awssdk.utils.CompletableFutureUtils; -import software.amazon.awssdk.utils.Either; import software.amazon.awssdk.utils.Logger; /** @@ -90,14 +89,14 @@ public CompletableFuture> execute() { } public void attemptFirstExecute(CompletableFuture> future) { - Duration backoffDelay = retryableStageHelper.acquireInitialToken(); - if (backoffDelay.isZero()) { - attemptExecute(future); - } else { - retryableStageHelper.logBackingOff(backoffDelay); - long totalDelayMillis = backoffDelay.toMillis(); - scheduledExecutor.schedule(() -> attemptExecute(future), totalDelayMillis, MILLISECONDS); - } + retryableStageHelper.acquireInitialTokenAsync().whenComplete((r, t) -> { + if (t != null) { + future.completeExceptionally(t); + return; + } + scheduledExecutor.schedule(() -> attemptExecute(future), r.toMillis(), MILLISECONDS); + }); + } private void attemptExecute(CompletableFuture> future) { @@ -140,33 +139,38 @@ private void attemptExecute(CompletableFuture> future) { } public void maybeAttemptExecute(CompletableFuture> future) { - Either backoffDelay = retryableStageHelper.tryRefreshToken(suggestedDelay()); - - Optional acquireFailureDelay = backoffDelay.right(); - if (acquireFailureDelay.isPresent()) { - Duration delay = acquireFailureDelay.get(); - retryableStageHelper.logAcquireFailureBackingOff(delay); - SdkException disallowedException = retryableStageHelper.retryPolicyDisallowedRetryException(); - // Avoid needless scheduling if we won't wait - if (delay.isZero()) { - future.completeExceptionally(disallowedException); - } else { - scheduledExecutor.schedule(() -> future.completeExceptionally(disallowedException), - delay.toMillis(), MILLISECONDS); + retryableStageHelper.tryRefreshTokenAsync(suggestedDelay()).whenComplete((backoffDelay, t) -> { + if (t != null) { + future.completeExceptionally(t); + return; } - return; - } - // We failed the last attempt, but will retry. The response handler wants to know when that happens. - responseHandler.onError(retryableStageHelper.getLastException()); - // Reset the request provider to the original one before retries, in case it was modified downstream. - context.requestProvider(originalRequestBody); + Optional acquireFailureDelay = backoffDelay.right(); + if (acquireFailureDelay.isPresent()) { + Duration delay = acquireFailureDelay.get(); + retryableStageHelper.logAcquireFailureBackingOff(delay); + SdkException disallowedException = retryableStageHelper.retryPolicyDisallowedRetryException(); + // Avoid needless scheduling if we won't wait + if (delay.isZero()) { + future.completeExceptionally(disallowedException); + } else { + scheduledExecutor.schedule(() -> future.completeExceptionally(disallowedException), + delay.toMillis(), MILLISECONDS); + } + return; + } + // We failed the last attempt, but will retry. The response handler wants to know when that happens. + responseHandler.onError(retryableStageHelper.getLastException()); + + // Reset the request provider to the original one before retries, in case it was modified downstream. + context.requestProvider(originalRequestBody); - // get() is safe, Either requires left OR right to be present - Duration successDelay = backoffDelay.left().get(); - retryableStageHelper.logBackingOff(successDelay); - long totalDelayMillis = successDelay.toMillis(); - scheduledExecutor.schedule(() -> attemptExecute(future), totalDelayMillis, MILLISECONDS); + // get() is safe, Either requires left OR right to be present + Duration successDelay = backoffDelay.left().get(); + retryableStageHelper.logBackingOff(successDelay); + long totalDelayMillis = successDelay.toMillis(); + scheduledExecutor.schedule(() -> attemptExecute(future), totalDelayMillis, MILLISECONDS); + }); } private void maybeRetryExecute(CompletableFuture> future, Exception exception) { diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/utils/RetryableStageHelper.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/utils/RetryableStageHelper.java index 840df78367fd..39780334c526 100644 --- a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/utils/RetryableStageHelper.java +++ b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/utils/RetryableStageHelper.java @@ -23,6 +23,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Optional; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import software.amazon.awssdk.annotations.SdkInternalApi; import software.amazon.awssdk.core.Response; @@ -60,6 +61,7 @@ public final class RetryableStageHelper { public static final String SDK_RETRY_INFO_HEADER = "amz-sdk-request"; public static final ExecutionAttribute LAST_BACKOFF_DELAY_DURATION = new ExecutionAttribute<>("LastBackoffDuration"); + private static final String RETRY_TOKEN_SCOPE = "GLOBAL"; private final SdkHttpFullRequest request; private final boolean isLongPollingOperation; @@ -72,6 +74,7 @@ public final class RetryableStageHelper { private SdkHttpResponse lastResponse; private SdkException lastException; + public RetryableStageHelper(SdkHttpFullRequest request, RequestExecutionContext context, HttpClientDependencies dependencies) { @@ -106,14 +109,31 @@ public void startingAttempt() { * value is {@link AdaptiveRetryStrategy}. */ public Duration acquireInitialToken() { - String scope = "GLOBAL"; - AcquireInitialTokenRequest acquireRequest = AcquireInitialTokenRequest.create(scope); - AcquireInitialTokenResponse acquireResponse = retryStrategy().acquireInitialToken(acquireRequest); - RetryToken retryToken = acquireResponse.token(); - Duration delay = acquireResponse.delay(); - context.executionAttributes().putAttribute(RETRY_TOKEN, retryToken); - context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, delay); - return delay; + AcquireInitialTokenResponse result = doBlockAcquireInitialToken(); + context.executionAttributes().putAttribute(RETRY_TOKEN, result.token()); + context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, result.delay()); + return result.delay(); + } + + private AcquireInitialTokenResponse doBlockAcquireInitialToken() { + AcquireInitialTokenRequest acquireRequest = AcquireInitialTokenRequest.create(RETRY_TOKEN_SCOPE); + + return retryStrategy().acquireInitialToken(acquireRequest); + } + + public CompletableFuture acquireInitialTokenAsync() { + AcquireInitialTokenRequest acquireRequest = AcquireInitialTokenRequest.create(RETRY_TOKEN_SCOPE); + + return retryStrategy().acquireInitialTokenAsync(acquireRequest).whenComplete( + (result, t) -> { + if (t != null) { + return; + } + + context.executionAttributes().putAttribute(RETRY_TOKEN, result.token()); + context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, result.delay()); + }) + .thenApply(AcquireInitialTokenResponse::delay); } /** @@ -138,16 +158,12 @@ public void recordAttemptSucceeded() { * code should not retry. */ public Either tryRefreshToken(Duration suggestedDelay) { - RetryToken retryToken = context.executionAttributes().getAttribute(RETRY_TOKEN); - RefreshRetryTokenResponse refreshResponse; + RetryToken retryToken; + Duration attemptDelay; try { - RefreshRetryTokenRequest refreshRequest = RefreshRetryTokenRequest.builder() - .failure(this.lastException) - .token(retryToken) - .isLongPolling(isLongPollingOperation) - .suggestedDelay(suggestedDelay) - .build(); - refreshResponse = retryStrategy().refreshRetryToken(refreshRequest); + RefreshRetryTokenResponse tryRefreshResult = doBlockingRefreshRetryToken(suggestedDelay); + attemptDelay = tryRefreshResult.delay(); + retryToken = tryRefreshResult.token(); } catch (TokenAcquisitionFailedException e) { context.executionAttributes().putAttribute(RETRY_TOKEN, e.token()); Optional acquireFailureDelay = e.delay(); @@ -158,12 +174,64 @@ public Either tryRefreshToken(Duration suggestedDelay) { } return Either.right(Duration.ZERO); } - Duration acquireSuccessDelay = refreshResponse.delay(); - context.executionAttributes().putAttribute(RETRY_TOKEN, refreshResponse.token()); + Duration acquireSuccessDelay = attemptDelay; + context.executionAttributes().putAttribute(RETRY_TOKEN, retryToken); context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, acquireSuccessDelay); return Either.left(acquireSuccessDelay); } + private RefreshRetryTokenResponse doBlockingRefreshRetryToken(Duration suggestedDelay) { + RetryToken retryToken = context.executionAttributes().getAttribute(RETRY_TOKEN); + + RefreshRetryTokenRequest refreshRequest = RefreshRetryTokenRequest.builder() + .failure(this.lastException) + .token(retryToken) + .isLongPolling(isLongPollingOperation) + .suggestedDelay(suggestedDelay) + .build(); + + return retryStrategy().refreshRetryToken(refreshRequest); + } + + public CompletableFuture> tryRefreshTokenAsync(Duration suggestedDelay) { + CompletableFuture> cf = new CompletableFuture<>(); + + RetryToken retryToken = context.executionAttributes().getAttribute(RETRY_TOKEN); + + RefreshRetryTokenRequest refreshRequest = RefreshRetryTokenRequest.builder() + .failure(this.lastException) + .token(retryToken) + .isLongPolling(isLongPollingOperation) + .suggestedDelay(suggestedDelay) + .build(); + + retryStrategy().refreshRetryTokenAsync(refreshRequest).whenComplete((r, t) -> { + if (t != null) { + if (t instanceof TokenAcquisitionFailedException) { + TokenAcquisitionFailedException e = (TokenAcquisitionFailedException) t; + context.executionAttributes().putAttribute(RETRY_TOKEN, e.token()); + Optional acquireFailureDelay = e.delay(); + if (acquireFailureDelay.isPresent()) { + Duration acquireDelay = acquireFailureDelay.get(); + context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, acquireDelay); + cf.complete(Either.right(acquireDelay)); + } + cf.complete(Either.right(Duration.ZERO)); + } else { + cf.completeExceptionally(t); + } + return; + } + + Duration acquireSuccessDelay = r.delay(); + context.executionAttributes().putAttribute(RETRY_TOKEN, r.token()); + context.executionAttributes().putAttribute(LAST_BACKOFF_DELAY_DURATION, acquireSuccessDelay); + cf.complete(Either.left(acquireSuccessDelay)); + }); + + return cf; + } + /** * Return the exception that should be thrown, because the retry strategy did not allow the request to be retried. */ diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/AsyncClientRetryStrategyExceptionTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/AsyncClientRetryStrategyExceptionTest.java index 983be14c8451..011df8916320 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/AsyncClientRetryStrategyExceptionTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/AsyncClientRetryStrategyExceptionTest.java @@ -61,6 +61,7 @@ public class AsyncClientRetryStrategyExceptionTest { public void exceptionInInitialTokenReportedInFuture() { Exception exception = new RuntimeException(MESSAGE); when(retryStrategy.acquireInitialToken(any())).thenThrow(exception); + when(retryStrategy.acquireInitialTokenAsync(any())).thenCallRealMethod(); CompletableFuture responseFuture = makeRequest(); @@ -73,8 +74,11 @@ public void exceptionInRefreshTokenReportedInFuture() { AcquireInitialTokenResponse.create(new RetryToken() { }, Duration.ZERO) ); + when(retryStrategy.acquireInitialTokenAsync(any())).thenCallRealMethod(); + Exception exception = new RuntimeException(MESSAGE); when(retryStrategy.refreshRetryToken(any())).thenThrow(exception); + when(retryStrategy.refreshRetryTokenAsync(any())).thenCallRealMethod(); CompletableFuture responseFuture = makeRequest(); diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerSignerValidationTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerSignerValidationTest.java index b0c19a521895..4bceeb8a2d0b 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerSignerValidationTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerSignerValidationTest.java @@ -18,6 +18,8 @@ import static org.assertj.core.api.Assertions.assertThatNoException; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import java.util.Arrays; @@ -26,6 +28,8 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -179,7 +183,6 @@ public void execute_requestHasHttpsEndpoint_usesBearerAuth_succeeds() throws Exc SdkAsyncClientHandler sdkAsyncClientHandler = new SdkAsyncClientHandler(testClientConfiguration()); CompletableFuture execute = sdkAsyncClientHandler.execute(executionParams()); - SdkAsyncHttpResponseHandler capturedHandler = executeRequestCaptor.getValue().responseHandler(); Map> headers = new HashMap<>(); headers.put("foo", Arrays.asList("bar")); @@ -227,10 +230,18 @@ private ClientExecutionParams executionParams() { } private SdkClientConfiguration testClientConfiguration() { + ScheduledExecutorService exec = mock(ScheduledExecutorService.class); + when(exec.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))).thenAnswer(i -> { + Runnable r = i.getArgument(0); + r.run(); + return null; + }); + return HttpTestUtils.testClientConfiguration() .toBuilder() .option(SdkClientOption.ASYNC_HTTP_CLIENT, httpClient) .option(SdkAdvancedClientOption.SIGNER, signer) + .option(SdkClientOption.SCHEDULED_EXECUTOR_SERVICE, exec) .build(); } } diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerTest.java index 6482ab936fc2..fef4caf7f256 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/client/handler/AsyncClientHandlerTest.java @@ -18,6 +18,8 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; @@ -27,6 +29,7 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import org.junit.Before; import org.junit.Test; @@ -143,10 +146,17 @@ private ClientExecutionParams clientExecutionParams() { } public SdkClientConfiguration clientConfiguration() { + ScheduledExecutorService exec = mock(ScheduledExecutorService.class); + when(exec.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))).thenAnswer(i -> { + Runnable r = i.getArgument(0); + r.run(); + return null; + }); return HttpTestUtils.testClientConfiguration().toBuilder() .option(SdkClientOption.ASYNC_HTTP_CLIENT, httpClient) .option(SdkClientOption.RETRY_POLICY, RetryPolicy.none()) .option(SdkClientOption.RETRY_STRATEGY, DefaultRetryStrategy.doNotRetry()) + .option(SdkClientOption.SCHEDULED_EXECUTOR_SERVICE, exec) .build(); } } diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStageTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStageTest.java index 827e1a7d1cef..f4b3d830cd78 100644 --- a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStageTest.java +++ b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncRetryableStageTest.java @@ -60,6 +60,7 @@ import software.amazon.awssdk.retries.api.RetryStrategy; import software.amazon.awssdk.retries.api.RetryToken; import software.amazon.awssdk.retries.api.TokenAcquisitionFailedException; +import software.amazon.awssdk.utils.CompletableFutureUtils; public class AsyncRetryableStageTest extends BaseRetryableStageTest { private RetryStrategy mockRetryStrategy; @@ -89,7 +90,8 @@ void methodSetup() { when(mockAcquireInitialTokenResponse.token()).thenReturn(mockRetryToken); when(mockAcquireInitialTokenResponse.delay()).thenReturn(Duration.ZERO); - when(mockRetryStrategy.acquireInitialToken(any())).thenReturn(mockAcquireInitialTokenResponse); + when(mockRetryStrategy.acquireInitialTokenAsync(any())) + .thenReturn(CompletableFuture.completedFuture(mockAcquireInitialTokenResponse)); mockDelegatePipeline = mock(RequestPipeline.class); } @@ -141,17 +143,22 @@ void execute_acquireDelay_behavesCorrectly(AcquireDelayTestCase testCase) throws when(mockDelegatePipeline.execute(any(), any())).thenReturn(CompletableFuture.completedFuture(response)); if (testCase.isFailure()) { - when(mockRetryStrategy.refreshRetryToken(any())).thenThrow( - new TokenAcquisitionFailedException("Acquire failed", mockRetryToken, null, testCase.failureDelay()) + when(mockRetryStrategy.refreshRetryTokenAsync(any())).thenReturn( + CompletableFutureUtils.failedFuture(new TokenAcquisitionFailedException("Acquire failed", mockRetryToken, null, + testCase.failureDelay())) ); } else { // only retry once, otherwise we'll get into an infinite loop AtomicBoolean first = new AtomicBoolean(); - when(mockRetryStrategy.refreshRetryToken(any())).thenAnswer(i -> { + when(mockRetryStrategy.refreshRetryTokenAsync(any())).thenAnswer(i -> { if (first.compareAndSet(false, true)) { - return RefreshRetryTokenResponse.create(mockRetryToken, testCase.successDelay()); + return CompletableFuture.completedFuture(RefreshRetryTokenResponse.create(mockRetryToken, + testCase.successDelay())); } - throw new TokenAcquisitionFailedException("Acquire failed", mockRetryToken, null, Duration.ZERO); + return CompletableFutureUtils.failedFuture(new TokenAcquisitionFailedException("Acquire failed", + mockRetryToken, + null, + Duration.ZERO)); }); } @@ -229,7 +236,7 @@ void execute_retryableException_treatsRetryAfterCorrectly(RetryAfterTestCase tes ArgumentCaptor refreshRequestCaptor = ArgumentCaptor.forClass(RefreshRetryTokenRequest.class); - verify(mockRetryStrategy).refreshRetryToken(refreshRequestCaptor.capture()); + verify(mockRetryStrategy).refreshRetryTokenAsync(refreshRequestCaptor.capture()); RefreshRetryTokenRequest refreshRequest = refreshRequestCaptor.getValue(); @@ -292,7 +299,7 @@ void execute_delegateThrows_noHttpResponse_uses0SuggestedDelay(boolean newRetrie ArgumentCaptor refreshRequestCaptor = ArgumentCaptor.forClass(RefreshRetryTokenRequest.class); - verify(mockRetryStrategy).refreshRetryToken(refreshRequestCaptor.capture()); + verify(mockRetryStrategy).refreshRetryTokenAsync(refreshRequestCaptor.capture()); RefreshRetryTokenRequest refreshRequest = refreshRequestCaptor.getValue(); diff --git a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crt/S3MetaRequestPauseObservable.java b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crt/S3MetaRequestPauseObservable.java index 5ddb41219d39..9307b38b79eb 100644 --- a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crt/S3MetaRequestPauseObservable.java +++ b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crt/S3MetaRequestPauseObservable.java @@ -43,7 +43,11 @@ public void subscribe(S3MetaRequestWrapper request) { * Pause the request */ public ResumeToken pause() { - return pause.apply(request); + S3MetaRequestWrapper r = this.request; + if (r != null) { + return pause.apply(r); + } + return null; } } diff --git a/test/codegen-generated-classes-test/src/main/resources/codegen-resources/json/service-2.json b/test/codegen-generated-classes-test/src/main/resources/codegen-resources/json/service-2.json new file mode 100644 index 000000000000..c0372d42aa46 --- /dev/null +++ b/test/codegen-generated-classes-test/src/main/resources/codegen-resources/json/service-2.json @@ -0,0 +1,38 @@ +{ + "version":"2.0", + "metadata":{ + "apiVersion":"2010-05-08", + "endpointPrefix":"json-service-endpoint", + "globalEndpoint": "json-service.amazonaws.com", + "jsonVersion":"1.0", + "protocol":"json", + "serviceAbbreviation":"Aws Json Service", + "serviceFullName":"Some Service That Uses AWS JSON", + "serviceId":"Aws Json Service", + "signingName": "aws-json-service", + "signatureVersion":"v4", + "uid":"aws-json-service-2010-05-08", + "awsQueryCompatible":{} + }, + "operations":{ + "AllType": { + "name": "APostOperation", + "http": { + "method": "POST", + "requestUri": "/" + }, + "httpChecksumRequired": true + } + }, + "shapes": { + "OneShape": { + "type": "structure", + "members": { + "StringMember": { + "shape": "String" + } + } + }, + "String":{"type":"string"} + } +} diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/AsyncRequestCompressionTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/AsyncRequestCompressionTest.java index df5b327a5442..b6d0ec381a7d 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/AsyncRequestCompressionTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/AsyncRequestCompressionTest.java @@ -83,7 +83,7 @@ public void asyncNonStreamingOperation_compressionEnabledThresholdOverridden_com c -> c.minimumCompressionThresholdInBytes(1))) .build(); - asyncClient.putOperationWithRequestCompression(request); + asyncClient.putOperationWithRequestCompression(request).join(); SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); @@ -104,7 +104,7 @@ public void asyncNonStreamingOperation_payloadSizeLessThanCompressionThreshold_d .body(SdkBytes.fromUtf8String(UNCOMPRESSED_BODY)) .build(); - asyncClient.putOperationWithRequestCompression(request); + asyncClient.putOperationWithRequestCompression(request).join(); SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); @@ -148,7 +148,7 @@ public void asyncNonStreamingOperation_compressionEnabledThresholdOverriddenWith c -> c.minimumCompressionThresholdInBytes(1))) .build(); - asyncClient.putOperationWithRequestCompression(request); + asyncClient.putOperationWithRequestCompression(request).join(); SdkHttpFullRequest loggedRequest = (SdkHttpFullRequest) mockAsyncHttpClient.getLastRequest(); InputStream loggedStream = loggedRequest.contentStreamProvider().get().newStream(); diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/BusinessMetricsUserAgentTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/BusinessMetricsUserAgentTest.java index ab9394e70ac0..f95312956612 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/BusinessMetricsUserAgentTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/BusinessMetricsUserAgentTest.java @@ -246,7 +246,7 @@ void validate_serviceUserAgent_format() { .overrideConfiguration(c -> c.addExecutionInterceptor(interceptor)) .build(); - client.headOperation(); + assertThatThrownBy(client.headOperation()::join); String userAgent = assertAndGetUserAgentString(); assertThat(userAgent).contains("AmazonProtocolRestJson#" + ServiceVersionInfo.VERSION); diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/HttpChecksumInHeaderTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/HttpChecksumInHeaderTest.java index f24f7f35d1c6..1424545509bc 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/HttpChecksumInHeaderTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/HttpChecksumInHeaderTest.java @@ -105,9 +105,9 @@ public void sync_json_nonStreaming_unsignedPayload_with_Sha1_in_header() { } @Test - public void aync_json_nonStreaming_unsignedPayload_with_Sha1_in_header() { + public void async_json_nonStreaming_unsignedPayload_with_Sha1_in_header() { jsonAsyncClient.operationWithChecksumNonStreaming( - r -> r.checksumAlgorithm(ChecksumAlgorithm.SHA1).stringMember("Hello world").build()); + r -> r.checksumAlgorithm(ChecksumAlgorithm.SHA1).stringMember("Hello world").build()).join(); assertThat(getAsyncRequest().firstMatchingHeader("Content-MD5")).isNotPresent(); //Note that content will be of form "{"stringMember":"Hello world"}" assertThat(getAsyncRequest().firstMatchingHeader("x-amz-checksum-sha1")).hasValue("M68rRwFal7o7B3KEMt3m0w39TaA="); @@ -132,7 +132,7 @@ public void sync_xml_nonStreaming_unsignedEmptyPayload_with_Sha1_in_header() { } @Test - public void aync_xml_nonStreaming_unsignedPayload_with_Sha1_in_header() { + public void async_xml_nonStreaming_unsignedPayload_with_Sha1_in_header() { xmlAsyncClient.operationWithChecksumNonStreaming(r -> r.stringMember("Hello world") .checksumAlgorithm(software.amazon.awssdk.services.protocolrestxml.model.ChecksumAlgorithm.SHA1).build()).join(); assertThat(getAsyncRequest().firstMatchingHeader("Content-MD5")).isNotPresent(); @@ -141,7 +141,7 @@ public void aync_xml_nonStreaming_unsignedPayload_with_Sha1_in_header() { } @Test - public void aync_xml_nonStreaming_unsignedEmptyPayload_with_Sha1_in_header() { + public void async_xml_nonStreaming_unsignedEmptyPayload_with_Sha1_in_header() { xmlAsyncClient.operationWithChecksumNonStreaming(r -> r.checksumAlgorithm(software.amazon.awssdk.services.protocolrestxml.model.ChecksumAlgorithm.SHA1).build()).join(); assertThat(getAsyncRequest().firstMatchingHeader("Content-MD5")).isNotPresent(); diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/NoneAuthTypeRequestTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/NoneAuthTypeRequestTest.java index 20ceeee91815..d193c7087f11 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/NoneAuthTypeRequestTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/NoneAuthTypeRequestTest.java @@ -113,14 +113,14 @@ public void sync_json_authorization_is_present_for_defaultAuth() { @Test public void async_json_authorization_is_absent_for_noneAuthType() { - jsonAsyncClient.operationWithNoneAuthType(o -> o.booleanMember(true)); + jsonAsyncClient.operationWithNoneAuthType(o -> o.booleanMember(true)).join(); assertThat(getAsyncRequest().firstMatchingHeader("Authorization")).isNotPresent(); verify(credentialsProvider, times(0)).resolveIdentity(any(ResolveIdentityRequest.class)); } @Test public void async_json_authorization_is_present_for_defaultAuth() { - jsonAsyncClient.jsonValuesOperation(); + jsonAsyncClient.jsonValuesOperation().join(); assertThat(getAsyncRequest().firstMatchingHeader("Authorization")).isPresent(); verify(credentialsProvider, times(1)).resolveIdentity(any(ResolveIdentityRequest.class)); } @@ -141,14 +141,14 @@ public void sync_xml_authorization_is_present_for_defaultAuth() { @Test public void async_xml_authorization_is_absent_for_noneAuthType() { - xmlAsyncClient.operationWithNoneAuthType(o -> o.booleanMember(true)); + xmlAsyncClient.operationWithNoneAuthType(o -> o.booleanMember(true)).join(); assertThat(getAsyncRequest().firstMatchingHeader("Authorization")).isNotPresent(); verify(credentialsProvider, times(0)).resolveIdentity(any(ResolveIdentityRequest.class)); } @Test public void async_xml_authorization_is_present_for_defaultAuth() { - xmlAsyncClient.jsonValuesOperation(json -> json.jsonValueMember("one")); + xmlAsyncClient.jsonValuesOperation(json -> json.jsonValueMember("one")).join(); assertThat(getAsyncRequest().firstMatchingHeader("Authorization")).isPresent(); verify(credentialsProvider, times(1)).resolveIdentity(any(ResolveIdentityRequest.class)); } diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/ProfileFileConfigurationTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/ProfileFileConfigurationTest.java index d8a11a0b76b2..21df46df84ff 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/ProfileFileConfigurationTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/ProfileFileConfigurationTest.java @@ -41,6 +41,7 @@ import software.amazon.awssdk.http.auth.aws.scheme.AwsV4AuthScheme; import software.amazon.awssdk.http.auth.aws.signer.AwsV4HttpSigner; import software.amazon.awssdk.http.auth.spi.signer.AsyncSignRequest; +import software.amazon.awssdk.http.auth.spi.signer.AsyncSignedRequest; import software.amazon.awssdk.http.auth.spi.signer.SignRequest; import software.amazon.awssdk.http.auth.spi.signer.SignedRequest; import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity; @@ -197,7 +198,7 @@ public void streamingOperation_asyncHttpSigner_profileIsHonoredForCredentialsAnd env.remove(SdkSystemSetting.AWS_ACCESS_KEY_ID); env.remove(SdkSystemSetting.AWS_SECRET_ACCESS_KEY); - Mockito.when(signer.signAsync(any(AsyncSignRequest.class))).thenReturn(CompletableFuture.completedFuture(any(AsyncSignRequest.class))); + Mockito.when(signer.signAsync(any(AsyncSignRequest.class))).thenReturn(CompletableFuture.completedFuture(mock(AsyncSignedRequest.class))); ProtocolRestJsonAsyncClient asyncClient = asyncClientWithHttpSignerOverride(); diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolquery/AsyncOperationCancelTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolquery/AsyncOperationCancelTest.java index e243a489ddbb..4498222bd438 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolquery/AsyncOperationCancelTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolquery/AsyncOperationCancelTest.java @@ -15,6 +15,8 @@ package software.amazon.awssdk.services.protocolquery; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -35,6 +37,8 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; /** @@ -49,13 +53,24 @@ public class AsyncOperationCancelTest { private CompletableFuture executeFuture; + private ScheduledExecutorService scheduledExec; + @Before public void setUp() { + scheduledExec = mock(ScheduledExecutorService.class); + + when(scheduledExec.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))).thenAnswer(i -> { + Runnable r = i.getArgument(0); + r.run(); + return null; + }); + client = ProtocolQueryAsyncClient.builder() .region(Region.US_EAST_1) .credentialsProvider(StaticCredentialsProvider.create( AwsBasicCredentials.create("foo", "bar"))) .httpClient(mockHttpClient) + .overrideConfiguration(o -> o.scheduledExecutorService(scheduledExec)) .build(); executeFuture = new CompletableFuture(); diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolquery/MoveQueryParamsToBodyTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolquery/MoveQueryParamsToBodyTest.java index 559c36a8019b..9532d4bd83c6 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolquery/MoveQueryParamsToBodyTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolquery/MoveQueryParamsToBodyTest.java @@ -129,7 +129,7 @@ public void customInterceptor_asyncClient_additionalQueryParamsAdded_paramsAlsoM .build(); ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(AsyncExecuteRequest.class); - asyncClient.membersInQueryParams(r -> r.stringQueryParam("hello")); + assertThatThrownBy(asyncClient.membersInQueryParams(r -> r.stringQueryParam("hello"))::join); verify(asyncMockHttpClient, atLeast(1)).execute(requestCaptor.capture()); verifyParametersMovedToBody_asyncClient(requestCaptor); @@ -164,7 +164,7 @@ public void requestOverrideConfiguration_asyncClient_additionalQueryParamsAdded_ ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(AsyncExecuteRequest.class); - asyncClient.membersInQueryParams(r -> r.stringQueryParam("hello").overrideConfiguration(createOverrideConfigWithQueryParams())); + assertThatThrownBy(asyncClient.membersInQueryParams(r -> r.stringQueryParam("hello").overrideConfiguration(createOverrideConfigWithQueryParams()))::join); verify(asyncMockHttpClient, atLeast(1)).execute(requestCaptor.capture()); verifyParametersMovedToBody_asyncClient(requestCaptor); } @@ -199,7 +199,7 @@ public void asyncClient_noQueryParamsAdded_onlyContainsOriginalContent() throws .build(); ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(AsyncExecuteRequest.class); - asyncClient.allTypes(r -> r.stringMember("hello")); + assertThatThrownBy(asyncClient.allTypes(r -> r.stringMember("hello"))::join); verify(asyncMockHttpClient, atLeast(1)).execute(requestCaptor.capture()); diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolrestjson/AsyncOperationCancelTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolrestjson/AsyncOperationCancelTest.java index 932f4a9032ac..0fceef6d1a11 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolrestjson/AsyncOperationCancelTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolrestjson/AsyncOperationCancelTest.java @@ -17,10 +17,14 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import io.reactivex.Flowable; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -43,7 +47,8 @@ import software.amazon.awssdk.services.protocolrestjson.model.StreamingOutputOperationResponse; /** - * Test to ensure that cancelling the future returned for an async operation will cancel the future returned by the async HTTP client. + * Test to ensure that cancelling the future returned for an async operation will cancel the future returned by the async HTTP + * client. */ @RunWith(MockitoJUnitRunner.class) public class AsyncOperationCancelTest { @@ -54,30 +59,45 @@ public class AsyncOperationCancelTest { private CompletableFuture executeFuture; + private ScheduledExecutorService scheduledExec; + @Before public void setUp() { + scheduledExec = mock(ScheduledExecutorService.class); + + when(scheduledExec.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))).thenAnswer(i -> { + Runnable r = i.getArgument(0); + r.run(); + return null; + }); + client = ProtocolRestJsonAsyncClient.builder() - .region(Region.US_EAST_1) - .credentialsProvider(StaticCredentialsProvider.create( - AwsBasicCredentials.create("foo", "bar"))) - .httpClient(mockHttpClient) - .build(); + .region(Region.US_EAST_1) + .credentialsProvider(StaticCredentialsProvider.create( + AwsBasicCredentials.create("foo", "bar"))) + .httpClient(mockHttpClient) + .overrideConfiguration(o -> o.scheduledExecutorService(scheduledExec)) + .build(); executeFuture = new CompletableFuture<>(); + when(mockHttpClient.execute(any())).thenReturn(executeFuture); } @Test - public void testNonStreamingOperation() { - CompletableFuture responseFuture = client.allTypes(r -> {}); + public void testNonStreamingOperation() throws InterruptedException { + CompletableFuture responseFuture = client.allTypes(r -> { + }); responseFuture.cancel(true); + // Thread.sleep(1000); assertThat(executeFuture.isCompletedExceptionally()).isTrue(); assertThat(executeFuture.isCancelled()).isTrue(); } @Test public void testStreamingOperation() { - CompletableFuture responseFuture = client.streamingInputOperation(r -> {}, AsyncRequestBody.empty()); + CompletableFuture responseFuture = client.streamingInputOperation(r -> { + }, AsyncRequestBody.empty()); responseFuture.cancel(true); assertThat(executeFuture.isCompletedExceptionally()).isTrue(); assertThat(executeFuture.isCancelled()).isTrue(); @@ -95,25 +115,26 @@ public void testStreamingOutputOperation() { @Test public void testEventStreamingOperation() throws InterruptedException { CompletableFuture responseFuture = - client.eventStreamOperation(r -> {}, + client.eventStreamOperation(r -> { + }, Flowable.just(InputEventStream.inputEventBuilder().build()), - new EventStreamOperationResponseHandler() { - @Override - public void responseReceived(EventStreamOperationResponse response) { - } - - @Override - public void onEventStream(SdkPublisher publisher) { - } - - @Override - public void exceptionOccurred(Throwable throwable) { - } - - @Override - public void complete() { - } - }); + new EventStreamOperationResponseHandler() { + @Override + public void responseReceived(EventStreamOperationResponse response) { + } + + @Override + public void onEventStream(SdkPublisher publisher) { + } + + @Override + public void exceptionOccurred(Throwable throwable) { + } + + @Override + public void complete() { + } + }); responseFuture.cancel(true); assertThat(executeFuture.isCompletedExceptionally()).isTrue(); assertThat(executeFuture.isCancelled()).isTrue(); diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolrestxml/AsyncOperationCancelTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolrestxml/AsyncOperationCancelTest.java index 51deef21c7e5..5de704f296f7 100644 --- a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolrestxml/AsyncOperationCancelTest.java +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/protocolrestxml/AsyncOperationCancelTest.java @@ -15,6 +15,8 @@ package software.amazon.awssdk.services.protocolrestxml; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -39,6 +41,8 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; /** @@ -53,13 +57,24 @@ public class AsyncOperationCancelTest { private CompletableFuture executeFuture; + private ScheduledExecutorService scheduledExec; + @Before public void setUp() { + scheduledExec = mock(ScheduledExecutorService.class); + + when(scheduledExec.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))).thenAnswer(i -> { + Runnable r = i.getArgument(0); + r.run(); + return null; + }); + client = ProtocolRestXmlAsyncClient.builder() .region(Region.US_EAST_1) .credentialsProvider(StaticCredentialsProvider.create( AwsBasicCredentials.create("foo", "bar"))) .httpClient(mockHttpClient) + .overrideConfiguration(o -> o.scheduledExecutorService(scheduledExec)) .build(); executeFuture = new CompletableFuture(); diff --git a/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/retry/AdaptiveRetryRateLimitingTest.java b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/retry/AdaptiveRetryRateLimitingTest.java new file mode 100644 index 000000000000..24bc160e9f9f --- /dev/null +++ b/test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/retry/AdaptiveRetryRateLimitingTest.java @@ -0,0 +1,278 @@ +/* + * 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.services.retry; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.any; +import static com.github.tomakehurst.wiremock.client.WireMock.anyUrl; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options; +import static org.assertj.core.api.Assertions.assertThat; + +import com.github.tomakehurst.wiremock.common.FileSource; +import com.github.tomakehurst.wiremock.extension.Parameters; +import com.github.tomakehurst.wiremock.extension.ResponseDefinitionTransformer; +import com.github.tomakehurst.wiremock.http.Request; +import com.github.tomakehurst.wiremock.http.ResponseDefinition; +import com.github.tomakehurst.wiremock.junit5.WireMockExtension; +import com.github.tomakehurst.wiremock.stubbing.ServeEvent; +import java.net.URI; +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentLinkedDeque; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder; +import software.amazon.awssdk.awscore.client.builder.AwsDefaultClientBuilder; +import software.amazon.awssdk.awscore.retry.AwsRetryStrategy; +import software.amazon.awssdk.core.SdkSystemSetting; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.retries.AdaptiveRetryStrategy; +import software.amazon.awssdk.services.json.JsonAsyncClient; +import software.amazon.awssdk.services.json.JsonAsyncClientBuilder; +import software.amazon.awssdk.services.json.JsonClient; +import software.amazon.awssdk.services.json.JsonClientBuilder; + +/** + * End-to-end testing for testing the correctness of the rate limiting implementation of the adaptive retry strategy. + */ +class AdaptiveRetryRateLimitingTest { + private static final AtomicLong eventCount = new AtomicLong(0L); + + private static final int throttleTps = 20; + private static final double tpsLower = throttleTps * 0.5; + private static final double tpsUpper = throttleTps * 1.5; + + private static final Duration testDuration = Duration.ofSeconds(30); + private static final Duration warmupDuration = Duration.ofSeconds(2); + private static final Duration tailTrim = Duration.ofSeconds(1); + private static final int numWorkerThreads = Runtime.getRuntime().availableProcessors(); + + private static final RollingThrottle throttler = new RollingThrottle(throttleTps); + + private static String newRetriesPropertySave; + + @RegisterExtension + static WireMockExtension wm = WireMockExtension.newInstance() + .options(options().dynamicPort() + .extensions(new ThrottlingTransformer())) + .build(); + + @BeforeAll + static void setup() { + newRetriesPropertySave = System.getProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property()); + System.setProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property(), "true"); + } + + @AfterAll + static void teardown() { + if (newRetriesPropertySave != null) { + System.setProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property(), newRetriesPropertySave); + } else { + System.clearProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property()); + } + } + + @Test + void staticServerThrottling_sync() throws Exception { + JsonClientBuilder builder = JsonClient.builder(); + configure(builder); + JsonClient client = builder.build(); + + Callable c = () -> { + client.allType(r -> {}); + return null; + }; + + try { + testStaticServerThrottling(c); + } finally { + client.close(); + } + } + + @Test + void staticServerThrottling_async() throws Exception { + JsonAsyncClientBuilder builder = JsonAsyncClient.builder(); + configure(builder); + JsonAsyncClient client = builder.build(); + + Callable c = () -> { + client.allType(r -> {}).join(); + return null; + }; + + try { + testStaticServerThrottling(c); + } finally { + client.close(); + } + } + + private void configure(AwsClientBuilder builder) { + AdaptiveRetryStrategy adaptiveRetryStrategy = AwsRetryStrategy.adaptiveRetryStrategy(true); + + builder.endpointOverride(URI.create(wm.baseUrl())) + .region(Region.US_WEST_2) + .credentialsProvider(StaticCredentialsProvider.create( + AwsBasicCredentials.create("akid", "secret"))) + .overrideConfiguration(o -> o.retryStrategy(adaptiveRetryStrategy)) + .build(); + } + + /** + * The server has a static max TPS of 20 before responding with throttling exceptions. The client sending rate should match + * this within an acceptable margin. + */ + void testStaticServerThrottling(Callable operationInvoker) throws Exception { + + wm.stubFor(any(anyUrl()).willReturn(aResponse().withTransformers("throttling"))); + + + Instant startT = Instant.now(); + Instant endAt = startT.plus(testDuration); + AtomicBoolean stop = new AtomicBoolean(false); + + ExecutorService pool = Executors.newFixedThreadPool(numWorkerThreads); + try { + List> futures = IntStream.range(0, numWorkerThreads) + .mapToObj(i -> pool.submit(() -> { + while (!stop.get() && Instant.now().isBefore(endAt)) { + try { + operationInvoker.call(); + } catch (Exception e) { + // We measure wire-level TPS, not call success. + // Throttling exceptions surface here once retries are exhausted. + } + } + })) + .collect(Collectors.toList()); + + for (Future f : futures) { + f.get(testDuration.getSeconds() + 30, TimeUnit.SECONDS); + } + } finally { + stop.set(true); + } + Instant endT = Instant.now(); + + double measuredTps = tpsInWindow( + wm.getAllServeEvents(), + startT.plus(warmupDuration), + endT.minus(tailTrim)); + + assertThat(measuredTps) + .as("Expected send rate >= %.1f TPS, but got %.1f TPS. " + + "Is the SDK not converging to the server throttle rate?", + tpsLower, measuredTps) + .isGreaterThanOrEqualTo(tpsLower); + assertThat(measuredTps) + .as("Expected send rate <= %.1f TPS, but got %.1f TPS. " + + "Is the SDK not respecting the server throttle rate limit?", + tpsUpper, measuredTps) + .isLessThanOrEqualTo(tpsUpper); + } + + private static double tpsInWindow(List events, Instant from, Instant to) { + long count = events.stream() + .map(e -> e.getRequest().getLoggedDate().toInstant()) + .filter(t -> !t.isBefore(from) && t.isBefore(to)) + .count(); + double seconds = Duration.between(from, to).toMillis() / 1000.0; + return seconds <= 0 ? 0 : count / seconds; + } + + /** + * Sliding 1-second window admission control. Once {@code TPS} requests have been admitted within the past second, returns + * false until the window slides. + */ + static final class RollingThrottle { + private final int tps; + private final ConcurrentLinkedDeque hits = new ConcurrentLinkedDeque<>(); + + RollingThrottle(int tps) { + this.tps = tps; + } + + synchronized boolean tryAcquire() { + long now = System.nanoTime(); + long cutoff = now - TimeUnit.SECONDS.toNanos(1); + while (!hits.isEmpty() && hits.peekFirst() < cutoff) { + hits.pollFirst(); + } + if (hits.size() >= tps) { + return false; + } + hits.addLast(now); + return true; + } + } + + /** + * WireMock transformer that delegates the 200-vs-throttle decision to {@link #throttler}. + */ + public static final class ThrottlingTransformer extends ResponseDefinitionTransformer { + + private static final String successBody = "{}"; + private static final String throttleBody = + "{\"__type\":\"ThrottlingException\",\"message\":\"Rate exceeded\"}"; + + @Override + public String getName() { + return "throttling"; + } + + @Override + public boolean applyGlobally() { + return false; + } + + @Override + public ResponseDefinition transform(Request request, + ResponseDefinition responseDefinition, + FileSource files, + Parameters parameters) { + eventCount.incrementAndGet(); + + if (throttler.tryAcquire()) { + return aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/x-amz-json-1.0") + .withBody(successBody) + .build(); + } + return aResponse() + .withStatus(400) + .withHeader("Content-Type", "application/x-amz-json-1.0") + .withHeader("x-amzn-ErrorType", "ThrottlingException") + .withBody(throttleBody) + .build(); + } + } +} diff --git a/utils/src/main/java/software/amazon/awssdk/utils/cache/lru/LruCache.java b/utils/src/main/java/software/amazon/awssdk/utils/cache/lru/LruCache.java index df7bc222d261..eff5892b01ab 100644 --- a/utils/src/main/java/software/amazon/awssdk/utils/cache/lru/LruCache.java +++ b/utils/src/main/java/software/amazon/awssdk/utils/cache/lru/LruCache.java @@ -15,6 +15,8 @@ package software.amazon.awssdk.utils.cache.lru; +import java.util.ArrayList; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; @@ -77,6 +79,20 @@ public V get(K key) { } } + public List evictAll() { + List evicted = new ArrayList<>(cache.size()); + synchronized (listLock) { + while (true) { + CacheEntry evictedEntry = evict(); + if (evictedEntry == null) { + break; + } + evicted.add(evictedEntry.value); + } + return evicted; + } + } + private CacheEntry newEntry(K key) { V value = valueSupplier.apply(key); return new CacheEntry<>(key, value); @@ -147,11 +163,18 @@ private void addToQueue(CacheEntry entry) { /** * Removes the least recently used entry from the cache, marks it as evicted and removes it from the queue. */ - private void evict() { - leastRecentlyUsed.isEvicted(true); - closeEvictedResourcesIfPossible(leastRecentlyUsed.value); - cache.remove(leastRecentlyUsed.key()); - removeFromQueue(leastRecentlyUsed); + private CacheEntry evict() { + if (leastRecentlyUsed == null) { + return null; + } + + CacheEntry entryToEvict = leastRecentlyUsed; + + entryToEvict.isEvicted(true); + closeEvictedResourcesIfPossible(entryToEvict.value); + cache.remove(entryToEvict.key()); + removeFromQueue(entryToEvict); + return entryToEvict; } private void closeEvictedResourcesIfPossible(V value) { diff --git a/utils/src/test/java/software/amazon/awssdk/utils/cache/lru/LruCacheTest.java b/utils/src/test/java/software/amazon/awssdk/utils/cache/lru/LruCacheTest.java index 2ee389b8849a..12b8dc3ef470 100644 --- a/utils/src/test/java/software/amazon/awssdk/utils/cache/lru/LruCacheTest.java +++ b/utils/src/test/java/software/amazon/awssdk/utils/cache/lru/LruCacheTest.java @@ -16,13 +16,16 @@ package software.amazon.awssdk.utils.cache.lru; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static software.amazon.awssdk.utils.FunctionalUtils.invokeSafely; import java.util.ArrayList; import java.util.Collections; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -40,6 +43,7 @@ import org.junit.jupiter.params.provider.MethodSource; import org.mockito.Spy; import org.mockito.junit.jupiter.MockitoExtension; +import software.amazon.awssdk.utils.SdkAutoCloseable; @ExtendWith(MockitoExtension.class) public class LruCacheTest { @@ -246,6 +250,39 @@ void when_multipleThreadsAreCallingCache_WorksAsExpected(Integer numThreads, } } + @Test + void evictAll_noEntries_returnsEmptyList() { + LruCache cache = simpleCache.get(); + assertThat(cache.evictAll()).isEmpty(); + } + + @Test + void evictAll_maxEntries_returnsAllEntries() { + LruCache cache = simpleCache.get(); + List expected = new ArrayList<>(); + for (int i = 0; i < MAX_SIMPLE_CACHE_SIZE; ++i) { + expected.add(cache.get(i)); + } + assertThat(cache.evictAll()).containsExactlyInAnyOrder(expected.toArray(new String[0])); + } + + @Test + void evictAll_closesEvictedEntries() { + LruCache cache = LruCache. + builder(k -> mock(SdkAutoCloseable.class)) + .maxSize(MAX_SIMPLE_CACHE_SIZE) + .build(); + + for (int i = 0; i < MAX_SIMPLE_CACHE_SIZE; ++i) { + cache.get(i); + } + + List evicted = cache.evictAll(); + + assertThat(evicted).isNotEmpty(); + assertThat(evicted).allSatisfy(e -> verify(e).close()); + } + private static Stream concurrencyTestValues() { // numThreads, numGetsPerThreads, sleepDurationMillis, cacheSize return Stream.of(Arguments.of(1000, 5000, false, 5),