diff --git a/README.md b/README.md index 466ec61..ff71530 100644 --- a/README.md +++ b/README.md @@ -53,7 +53,7 @@ String password = auth.getToken(); ``` See the [Java README](java/README.md) for Maven setup, region and credential -resolution, and a standalone TLS connection example. +resolution, token caching and background refresh, and connection examples. ### Node.js diff --git a/java/CHANGELOG.md b/java/CHANGELOG.md index 18230a4..f92b138 100644 --- a/java/CHANGELOG.md +++ b/java/CHANGELOG.md @@ -15,6 +15,11 @@ and this package uses [Semantic Versioning](https://semver.org/). typed user-correctable errors, and deterministic SigV4 signing tests. - Credential providers that fail or return unusable credentials are normalized to an actionable `ConfigurationException`. +- `ElastiCacheIamAuthTokenManager` for caching tokens, refreshing them in the + background before expiry, retrying transient failures, and notifying a + token-change callback. +- `TokenRefreshException` for closed managers and refreshes that cannot replace + an expired token. - Maven build, Java version CI matrix, documentation, and dependency updates. ### Notes diff --git a/java/README.md b/java/README.md index e2321e4..0f3793f 100644 --- a/java/README.md +++ b/java/README.md @@ -2,7 +2,8 @@ This Java library generates the IAM authentication token that Amazon ElastiCache requires as the connection password for an IAM-enabled user. It supports both -serverless caches and node-based replication groups. +serverless caches and node-based replication groups, including optional token +caching and background refresh. ## Installation @@ -69,18 +70,79 @@ Missing credentials are reported when a token is requested. Expired or otherwise invalid credentials cannot be detected during local signing: token generation can succeed, but ElastiCache will reject the token when the client connects. +## Caching and background refresh + +`ElastiCacheIamAuthTokenProvider` is stateless and signs a new token on every +`getToken()` call. `ElastiCacheIamAuthTokenManager` caches a token and refreshes +it in the background: + +| | `ElastiCacheIamAuthTokenProvider` | `ElastiCacheIamAuthTokenManager` | +|---|---|---| +| Token per `getToken()` | Freshly signed every call | Cached while valid | +| Refresh | Caller-driven | Background after `refreshAfter` | +| Refresh failures | Surfaced to the caller | Retried while the current token is valid | +| Change notification | None | `onTokenChanged` | +| Cleanup | None needed | `close()` | + +```java +import java.time.Duration; +import software.amazon.awssdk.regions.Region; +import software.amazon.elasticache.auth.ElastiCacheIamAuthTokenManager; + +try (ElastiCacheIamAuthTokenManager auth = + ElastiCacheIamAuthTokenManager.builder() + .serverlessCacheName("my-cache") + .userId("iam-user") + .region(Region.US_EAST_1) + .refreshAfter(Duration.ofMinutes(5)) + .onTokenChanged(token -> + System.out.println("Installed a new IAM auth token")) + .build()) { + ElastiCacheIamAuthTokenManager.Credentials credentials = auth.getCredentials(); + String username = credentials.getUserId(); + String password = credentials.getToken(); +} +``` + +The first token is generated lazily. Concurrent callers share one initial token +generation, and subsequent calls return the cached token while it is valid. +Refresh starts after 5 minutes by default; `refreshAfter(...)` accepts values from +1 millisecond up to, but not including, 14 minutes 40 seconds. The manager reserves +the final 20 seconds of the service's 15-minute token lifetime for clock skew and +in-transit reconnects. Token-change callbacks run on the manager's refresh thread +after the new token is available to callers and should return quickly so they do +not delay that manager's next refresh. + +A refresh makes up to 8 attempts with jittered exponential backoff capped at 5 +seconds. If refresh fails while the cached token is still valid, callers continue +to receive that token and another refresh is scheduled. If the token expires +before it can be replaced, `getToken()` throws `TokenRefreshException` with the +last refresh failure as its cause. + +The refresh thread is a daemon and does not keep the JVM alive. Close the manager +to cancel background work; token requests after close throw +`TokenRefreshException`. + +Use the manager when a client reconnects often or requests credentials more +frequently than the token lifetime. Use the provider when every request should +sign a new token. + ## Reconnecting clients Clients should request fresh credentials whenever they connect or reconnect. The Lettuce example implements `RedisCredentialsProvider` so each credentials request -calls `getToken()`: +uses the manager's current token: ```java RedisCredentialsProvider credentialsProvider = new RedisCredentialsProvider() { @Override public Mono resolveCredentials() { - return Mono.fromSupplier( - () -> RedisCredentials.just(auth.getUserId(), auth.getToken())); + return Mono.fromSupplier(() -> { + ElastiCacheIamAuthTokenManager.Credentials credentials = + auth.getCredentials(); + return RedisCredentials.just( + credentials.getUserId(), credentials.getToken()); + }); } }; ``` diff --git a/java/examples/src/main/java/software/amazon/elasticache/examples/ConnectWithLettuce.java b/java/examples/src/main/java/software/amazon/elasticache/examples/ConnectWithLettuce.java index 9438bef..b2f26a4 100644 --- a/java/examples/src/main/java/software/amazon/elasticache/examples/ConnectWithLettuce.java +++ b/java/examples/src/main/java/software/amazon/elasticache/examples/ConnectWithLettuce.java @@ -11,7 +11,7 @@ import io.lettuce.core.api.StatefulRedisConnection; import reactor.core.publisher.Mono; import software.amazon.awssdk.regions.Region; -import software.amazon.elasticache.auth.ElastiCacheIamAuthTokenProvider; +import software.amazon.elasticache.auth.ElastiCacheIamAuthTokenManager; /** * Connects to an ElastiCache serverless cache with Lettuce and IAM authentication. @@ -39,7 +39,7 @@ public static void main(String[] args) { Region region = Region.of(args[3]); int port = args.length == 5 ? Integer.parseInt(args[4]) : DEFAULT_PORT; - ElastiCacheIamAuthTokenProvider auth = ElastiCacheIamAuthTokenProvider.builder() + ElastiCacheIamAuthTokenManager auth = ElastiCacheIamAuthTokenManager.builder() .serverlessCacheName(cacheName) .userId(userId) .region(region) @@ -48,8 +48,12 @@ public static void main(String[] args) { RedisCredentialsProvider credentialsProvider = new RedisCredentialsProvider() { @Override public Mono resolveCredentials() { - return Mono.fromSupplier( - () -> RedisCredentials.just(auth.getUserId(), auth.getToken())); + return Mono.fromSupplier(() -> { + ElastiCacheIamAuthTokenManager.Credentials credentials = + auth.getCredentials(); + return RedisCredentials.just( + credentials.getUserId(), credentials.getToken()); + }); } }; RedisURI redisUri = RedisURI.Builder.redis(endpoint, port) @@ -68,6 +72,7 @@ public Mono resolveCredentials() { connection.close(); } client.shutdown(); + auth.close(); } } } diff --git a/java/integration-tests/src/test/java/software/amazon/elasticache/integration/PublicApiTest.java b/java/integration-tests/src/test/java/software/amazon/elasticache/integration/PublicApiTest.java index 09397e2..d2e7bcb 100644 --- a/java/integration-tests/src/test/java/software/amazon/elasticache/integration/PublicApiTest.java +++ b/java/integration-tests/src/test/java/software/amazon/elasticache/integration/PublicApiTest.java @@ -10,16 +10,20 @@ import org.junit.jupiter.api.Test; import software.amazon.elasticache.auth.ConfigurationException; +import software.amazon.elasticache.auth.ElastiCacheIamAuthTokenManager; import software.amazon.elasticache.auth.ElastiCacheIamAuthTokenProvider; import software.amazon.elasticache.auth.InvalidParameterException; +import software.amazon.elasticache.auth.TokenRefreshException; import software.amazon.elasticache.auth.ToolkitInputException; class PublicApiTest { @Test void installedArtifactExposesPublicApi() { assertNotNull(ElastiCacheIamAuthTokenProvider.builder()); + assertNotNull(ElastiCacheIamAuthTokenManager.builder()); assertTrue(ToolkitInputException.class.isAssignableFrom(ConfigurationException.class)); assertTrue(ToolkitInputException.class.isAssignableFrom(InvalidParameterException.class)); + assertFalse(ToolkitInputException.class.isAssignableFrom(TokenRefreshException.class)); assertFalse( InvalidParameterException.class.isAssignableFrom(ConfigurationException.class)); } diff --git a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java new file mode 100644 index 0000000..bb289ba --- /dev/null +++ b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java @@ -0,0 +1,664 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package software.amazon.elasticache.auth; + +import java.time.Clock; +import java.time.Duration; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.Executors; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import java.util.function.DoubleSupplier; +import java.util.function.LongSupplier; +import java.util.function.Supplier; + +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.regions.providers.AwsRegionProvider; +import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; + +/** + * Caches an ElastiCache IAM authentication token and refreshes it before expiry. + * + *

The first token is generated lazily. Calls reuse the cached token while it + * remains valid, and concurrent callers share a single refresh. Refresh failures + * are retried in the background while a valid cached token continues to be served. + * + *

Call {@link #close()} when the manager is no longer needed. + */ +public final class ElastiCacheIamAuthTokenManager implements AutoCloseable { + private static final Duration DEFAULT_REFRESH_AFTER = Duration.ofMinutes(5); + private static final int MAX_ATTEMPTS = 8; + private static final long BASE_DELAY_MILLIS = 100; + private static final long MAX_DELAY_MILLIS = 5_000; + private static final double JITTER_RATIO = 0.2; + private static final Duration TOKEN_SERVE_MARGIN = Duration.ofSeconds(20); + private static final Duration EFFECTIVE_TOKEN_LIFETIME = + ElastiCacheIamAuthTokenProvider.TOKEN_TTL.minus(TOKEN_SERVE_MARGIN); + private static final String CLOSED_MESSAGE = "The token manager is closed."; + private static final String REFRESH_CANCELLED_MESSAGE = + "Token refresh was cancelled because the token manager was closed."; + private static final String REFRESH_FAILED_MESSAGE = + "Could not refresh the ElastiCache IAM authentication token before " + + "the previous token expired."; + + private final Object lock = new Object(); + private final ElastiCacheIamAuthTokenProvider provider; + private final long refreshAfterMillis; + private final Consumer onTokenChanged; + private final LongSupplier currentTimeMillis; + private final DoubleSupplier random; + private final Scheduler scheduler; + private final boolean ownsScheduler; + private final Runnable beforeInstall; + private final Runnable onSharedRefresh; + + private CachedToken cached; + private CompletableFuture inFlight; + private Cancellable refreshTask; + private Cancellable retryTask; + private long retryNotBefore; + private boolean closed; + + private ElastiCacheIamAuthTokenManager(Builder builder) { + this.provider = ElastiCacheIamAuthTokenProvider.builder() + .userId(builder.userId) + .region(builder.region) + .serverlessCacheName(builder.serverlessCacheName) + .replicationGroupId(builder.replicationGroupId) + .credentialsProvider(builder.credentialsProvider == null + ? DefaultCredentialsProvider.builder().build() + : builder.credentialsProvider) + .awsRegionEnvironmentProvider(builder.awsRegionEnvironmentProvider == null + ? () -> System.getenv("AWS_REGION") + : builder.awsRegionEnvironmentProvider) + .regionProvider(builder.regionProvider == null + ? DefaultAwsRegionProviderChain.builder().build() + : builder.regionProvider) + .signingClock(builder.signingClock == null + ? Clock.systemUTC() + : builder.signingClock) + .build(); + this.refreshAfterMillis = resolveRefreshAfterMillis(builder.refreshAfter); + this.onTokenChanged = builder.onTokenChanged; + this.currentTimeMillis = builder.currentTimeMillis == null + ? System::currentTimeMillis + : builder.currentTimeMillis; + this.random = builder.random == null ? Math::random : builder.random; + this.beforeInstall = builder.beforeInstall; + this.onSharedRefresh = builder.onSharedRefresh; + if (builder.scheduler == null) { + this.scheduler = new DefaultScheduler(); + this.ownsScheduler = true; + } else { + this.scheduler = builder.scheduler; + this.ownsScheduler = false; + } + } + + /** + * Creates a builder for a caching IAM authentication token manager. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the normalized ElastiCache user ID. + * + * @return lowercase ElastiCache user ID + */ + public String getUserId() { + return provider.getUserId(); + } + + /** + * Returns the normalized user ID and a valid token together. + * + * @return credentials suitable for a client credential-provider adapter + */ + public Credentials getCredentials() { + return new Credentials(getUserId(), getToken()); + } + + /** + * Returns a valid token, generating one on the first call. + * + *

While a cached token remains valid it is returned immediately. Once its + * refresh point is reached, refresh proceeds in the background. + * + * @return current IAM authentication token + * @throws TokenRefreshException if the manager is closed or an expired token + * could not be replaced + */ + public String getToken() { + synchronized (lock) { + ensureOpen(); + long now = currentTimeMillis.getAsLong(); + if (cached != null && now < cached.expiresAt) { + String token = cached.token; + if (now >= cached.refreshAt && now >= retryNotBefore) { + scheduleRefreshLocked(0); + } + return token; + } + } + return await(refresh()); + } + + /** + * Stops all background work and releases the manager's scheduler. + * + *

Subsequent token requests throw {@link TokenRefreshException}. + */ + @Override + public void close() { + CompletableFuture pending; + synchronized (lock) { + if (closed) { + return; + } + closed = true; + cancel(refreshTask); + refreshTask = null; + cancel(retryTask); + retryTask = null; + pending = inFlight; + } + if (pending != null) { + pending.completeExceptionally( + new TokenRefreshException(REFRESH_CANCELLED_MESSAGE)); + } + if (ownsScheduler) { + scheduler.close(); + } + } + + private CompletableFuture refresh() { + CompletableFuture refresh; + synchronized (lock) { + ensureOpen(); + if (inFlight != null) { + if (onSharedRefresh != null) { + onSharedRefresh.run(); + } + return inFlight; + } + refresh = new CompletableFuture<>(); + inFlight = refresh; + refresh.whenComplete((result, error) -> { + synchronized (lock) { + if (inFlight == refresh) { + inFlight = null; + } + } + }); + } + runRefreshAttempt(refresh, 1, null); + return refresh; + } + + private void startBackgroundRefresh() { + try { + refresh(); + } catch (TokenRefreshException ignored) { + // Closing the manager races safely with a scheduled refresh. + } + } + + private void runRefreshAttempt( + CompletableFuture refresh, int attempt, Throwable lastError) { + synchronized (lock) { + if (closed) { + refresh.completeExceptionally(new TokenRefreshException( + REFRESH_CANCELLED_MESSAGE, lastError)); + return; + } + retryTask = null; + } + + try { + long issuedAt = currentTimeMillis.getAsLong(); + String token = provider.getToken(); + if (beforeInstall != null) { + beforeInstall.run(); + } + installAndComplete(refresh, token, issuedAt); + } catch (Throwable failure) { + if (failure instanceof Error) { + refresh.completeExceptionally(failure); + throw (Error) failure; + } + RuntimeException exception = failure instanceof RuntimeException + ? (RuntimeException) failure + : new RuntimeException(failure); + synchronized (lock) { + if (closed) { + refresh.completeExceptionally( + new TokenRefreshException( + REFRESH_CANCELLED_MESSAGE, exception)); + return; + } + } + if (attempt == MAX_ATTEMPTS) { + finishFailedRefresh(refresh, exception); + return; + } + scheduleRetry(refresh, attempt + 1, exception, backoffDelay(attempt)); + } + } + + private void scheduleRetry( + CompletableFuture refresh, + int nextAttempt, + Throwable lastError, + long delayMillis) { + try { + Cancellable task = scheduler.schedule( + () -> runRefreshAttempt(refresh, nextAttempt, lastError), delayMillis); + synchronized (lock) { + if (closed) { + task.cancel(); + refresh.completeExceptionally( + new TokenRefreshException( + REFRESH_CANCELLED_MESSAGE, lastError)); + } else { + retryTask = task; + } + } + } catch (RejectedExecutionException exception) { + refresh.completeExceptionally( + new TokenRefreshException(REFRESH_CANCELLED_MESSAGE, lastError)); + } + } + + private void finishFailedRefresh( + CompletableFuture refresh, RuntimeException lastError) { + String validToken = null; + synchronized (lock) { + long now = currentTimeMillis.getAsLong(); + if (cached != null && now < cached.expiresAt) { + validToken = cached.token; + retryNotBefore = now + MAX_DELAY_MILLIS; + scheduleRefreshLocked(MAX_DELAY_MILLIS); + } + } + if (validToken != null) { + refresh.complete(validToken); + return; + } + synchronized (lock) { + if (cached == null) { + refresh.completeExceptionally(lastError); + } else { + refresh.completeExceptionally( + new TokenRefreshException(REFRESH_FAILED_MESSAGE, lastError)); + } + } + } + + private void installAndComplete( + CompletableFuture refresh, String token, long issuedAt) { + synchronized (lock) { + if (closed) { + refresh.completeExceptionally( + new TokenRefreshException(REFRESH_CANCELLED_MESSAGE)); + return; + } + cached = new CachedToken( + token, + issuedAt + refreshAfterMillis, + issuedAt + EFFECTIVE_TOKEN_LIFETIME.toMillis()); + retryNotBefore = 0; + scheduleRefreshLocked(refreshAfterMillis); + refresh.complete(token); + } + notifyTokenChanged(token); + } + + private void notifyTokenChanged(String token) { + if (onTokenChanged == null) { + return; + } + try { + scheduler.execute(() -> { + try { + onTokenChanged.accept(token); + } catch (Throwable ignored) { + // A consumer callback cannot invalidate a successfully signed token. + } + }); + } catch (Throwable ignored) { + // Failure to dispatch a callback cannot invalidate a signed token. + } + } + + private void scheduleRefreshLocked(long delayMillis) { + cancel(refreshTask); + refreshTask = null; + if (closed) { + return; + } + Cancellable[] scheduled = new Cancellable[1]; + scheduled[0] = scheduler.schedule(() -> { + synchronized (lock) { + if (refreshTask == scheduled[0]) { + refreshTask = null; + } + } + startBackgroundRefresh(); + }, delayMillis); + refreshTask = scheduled[0]; + } + + private long backoffDelay(int failedAttempt) { + long exponential = BASE_DELAY_MILLIS << Math.min(failedAttempt - 1, 30); + long delay = Math.min(exponential, MAX_DELAY_MILLIS); + double jitter = 1 - JITTER_RATIO + 2 * JITTER_RATIO * random.getAsDouble(); + return Math.round(delay * jitter); + } + + private void ensureOpen() { + if (closed) { + throw new TokenRefreshException(CLOSED_MESSAGE); + } + } + + private static void cancel(Cancellable task) { + if (task != null) { + task.cancel(); + } + } + + private static String await(CompletableFuture future) { + try { + return future.join(); + } catch (CompletionException exception) { + Throwable cause = exception.getCause(); + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw exception; + } + } + + private static long resolveRefreshAfterMillis(Duration refreshAfter) { + Duration resolved = refreshAfter == null ? DEFAULT_REFRESH_AFTER : refreshAfter; + long millis; + try { + millis = resolved.toMillis(); + } catch (ArithmeticException exception) { + millis = -1; + } + if (resolved.isNegative() + || resolved.isZero() + || millis <= 0 + || resolved.compareTo(EFFECTIVE_TOKEN_LIFETIME) >= 0) { + throw new InvalidParameterException( + "Invalid value for parameter 'refreshAfter': must be at least " + + "1 millisecond and less than 14 minutes 40 seconds."); + } + return millis; + } + + /** + * Immutable username and token pair returned by {@link #getCredentials()}. + */ + public static final class Credentials { + private final String userId; + private final String token; + + private Credentials(String userId, String token) { + this.userId = userId; + this.token = token; + } + + /** + * Returns the normalized ElastiCache user ID. + * + * @return lowercase user ID + */ + public String getUserId() { + return userId; + } + + /** + * Returns the cached IAM authentication token. + * + * @return IAM authentication token + */ + public String getToken() { + return token; + } + } + + /** + * Builder for {@link ElastiCacheIamAuthTokenManager}. + */ + public static final class Builder { + private String userId; + private Region region; + private String serverlessCacheName; + private String replicationGroupId; + private AwsCredentialsProvider credentialsProvider; + private Duration refreshAfter; + private Consumer onTokenChanged; + private Supplier awsRegionEnvironmentProvider; + private AwsRegionProvider regionProvider; + private Clock signingClock; + private LongSupplier currentTimeMillis; + private DoubleSupplier random; + private Scheduler scheduler; + private Runnable beforeInstall; + private Runnable onSharedRefresh; + + private Builder() {} + + /** + * Sets the IAM-enabled ElastiCache user ID. + * + * @param userId ElastiCache user ID + * @return this builder + */ + public Builder userId(String userId) { + this.userId = userId; + return this; + } + + /** + * Sets the signing region. + * + * @param region AWS region + * @return this builder + */ + public Builder region(Region region) { + this.region = region; + return this; + } + + /** + * Sets the serverless cache name. + * + * @param serverlessCacheName serverless cache name + * @return this builder + */ + public Builder serverlessCacheName(String serverlessCacheName) { + this.serverlessCacheName = serverlessCacheName; + return this; + } + + /** + * Sets the node-based replication group ID. + * + * @param replicationGroupId replication group ID + * @return this builder + */ + public Builder replicationGroupId(String replicationGroupId) { + this.replicationGroupId = replicationGroupId; + return this; + } + + /** + * Sets a credentials provider. The standard AWS provider chain is the default. + * + * @param credentialsProvider AWS credentials provider + * @return this builder + */ + public Builder credentialsProvider(AwsCredentialsProvider credentialsProvider) { + this.credentialsProvider = Objects.requireNonNull( + credentialsProvider, "credentialsProvider"); + return this; + } + + /** + * Sets how long after issuance the token is refreshed in the background. + * + *

The default is 5 minutes. The value must be at least 1 millisecond + * and less than the effective 14-minute-40-second token lifetime. The + * effective lifetime reserves 20 seconds for clock skew and transit time. + * + * @param refreshAfter refresh delay after token issuance + * @return this builder + */ + public Builder refreshAfter(Duration refreshAfter) { + this.refreshAfter = Objects.requireNonNull(refreshAfter, "refreshAfter"); + return this; + } + + /** + * Sets a callback invoked on the manager's refresh thread after each newly + * generated token is installed and made available to token callers. + * + *

The callback should return quickly because it delays this manager's + * next background refresh. Callback failures are ignored and do not discard + * the new token. + * + * @param onTokenChanged token-change callback + * @return this builder + */ + public Builder onTokenChanged(Consumer onTokenChanged) { + this.onTokenChanged = Objects.requireNonNull(onTokenChanged, "onTokenChanged"); + return this; + } + + Builder regionProvider(AwsRegionProvider regionProvider) { + this.regionProvider = Objects.requireNonNull(regionProvider, "regionProvider"); + return this; + } + + Builder awsRegionEnvironmentProvider(Supplier provider) { + this.awsRegionEnvironmentProvider = Objects.requireNonNull(provider, "provider"); + return this; + } + + Builder signingClock(Clock signingClock) { + this.signingClock = Objects.requireNonNull(signingClock, "signingClock"); + return this; + } + + Builder currentTimeMillis(LongSupplier currentTimeMillis) { + this.currentTimeMillis = + Objects.requireNonNull(currentTimeMillis, "currentTimeMillis"); + return this; + } + + Builder random(DoubleSupplier random) { + this.random = Objects.requireNonNull(random, "random"); + return this; + } + + Builder scheduler(Scheduler scheduler) { + this.scheduler = Objects.requireNonNull(scheduler, "scheduler"); + return this; + } + + Builder beforeInstall(Runnable beforeInstall) { + this.beforeInstall = Objects.requireNonNull(beforeInstall, "beforeInstall"); + return this; + } + + Builder onSharedRefresh(Runnable onSharedRefresh) { + this.onSharedRefresh = + Objects.requireNonNull(onSharedRefresh, "onSharedRefresh"); + return this; + } + + /** + * Validates the configuration and creates the manager. + * + * @return configured caching token manager + */ + public ElastiCacheIamAuthTokenManager build() { + return new ElastiCacheIamAuthTokenManager(this); + } + } + + interface Cancellable { + void cancel(); + } + + interface Scheduler { + Cancellable schedule(Runnable task, long delayMillis); + + void execute(Runnable task); + + void close(); + } + + private static final class DefaultScheduler implements Scheduler { + private final ScheduledExecutorService executor = + Executors.newSingleThreadScheduledExecutor(new DaemonThreadFactory()); + + @Override + public Cancellable schedule(Runnable task, long delayMillis) { + ScheduledFuture future = + executor.schedule(task, delayMillis, TimeUnit.MILLISECONDS); + return () -> future.cancel(false); + } + + @Override + public void execute(Runnable task) { + executor.execute(task); + } + + @Override + public void close() { + executor.shutdownNow(); + } + } + + private static final class DaemonThreadFactory implements ThreadFactory { + @Override + public Thread newThread(Runnable runnable) { + Thread thread = new Thread(runnable, "elasticache-iam-token-refresh"); + thread.setDaemon(true); + return thread; + } + } + + private static final class CachedToken { + private final String token; + private final long refreshAt; + private final long expiresAt; + + private CachedToken(String token, long refreshAt, long expiresAt) { + this.token = token; + this.refreshAt = refreshAt; + this.expiresAt = expiresAt; + } + } +} diff --git a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenProvider.java b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenProvider.java index c774a00..d2f1d1b 100644 --- a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenProvider.java +++ b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenProvider.java @@ -39,7 +39,7 @@ public final class ElastiCacheIamAuthTokenProvider { private static final String URL_SCHEME_PREFIX = "https://"; private static final String SERVICE_NAME = "elasticache"; private static final String AWS_REGION_ENV_VAR = "AWS_REGION"; - private static final Duration TOKEN_TTL = Duration.ofSeconds(900); + static final Duration TOKEN_TTL = Duration.ofSeconds(900); private static final Pattern CACHE_NAME_PATTERN = Pattern.compile("^[a-zA-Z][a-zA-Z0-9]*(-[a-zA-Z0-9]+)*$"); private static final Pattern USER_ID_PATTERN = diff --git a/java/src/main/java/software/amazon/elasticache/auth/TokenRefreshException.java b/java/src/main/java/software/amazon/elasticache/auth/TokenRefreshException.java new file mode 100644 index 0000000..46d40d8 --- /dev/null +++ b/java/src/main/java/software/amazon/elasticache/auth/TokenRefreshException.java @@ -0,0 +1,35 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package software.amazon.elasticache.auth; + +/** + * Indicates that a cached token could not be refreshed before it expired, or + * that its token manager was closed. + * + *

This is an operational failure rather than an invalid-input error, so it + * does not extend {@link ToolkitInputException}. + */ +public final class TokenRefreshException extends RuntimeException { + private static final long serialVersionUID = 1L; + + /** + * Creates a token refresh exception. + * + * @param message actionable error message + */ + public TokenRefreshException(String message) { + super(message); + } + + /** + * Creates a token refresh exception with its underlying refresh failure. + * + * @param message actionable error message + * @param cause last refresh failure + */ + public TokenRefreshException(String message, Throwable cause) { + super(message, cause); + } +} diff --git a/java/src/main/java/software/amazon/elasticache/auth/package-info.java b/java/src/main/java/software/amazon/elasticache/auth/package-info.java index 25886ab..a681517 100644 --- a/java/src/main/java/software/amazon/elasticache/auth/package-info.java +++ b/java/src/main/java/software/amazon/elasticache/auth/package-info.java @@ -4,6 +4,6 @@ */ /** - * IAM authentication token generation for Amazon ElastiCache. + * IAM authentication token generation, caching, and refresh for Amazon ElastiCache. */ package software.amazon.elasticache.auth; diff --git a/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java new file mode 100644 index 0000000..1e0022f --- /dev/null +++ b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java @@ -0,0 +1,687 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package software.amazon.elasticache.auth; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.regions.Region; + +class ElastiCacheIamAuthTokenManagerTest { + private static final String CACHE = "my-cache"; + private static final String USER = "testuser"; + private static final AwsCredentials FIRST_CREDENTIALS = AwsBasicCredentials.create( + "AKIAIOSFODNN7EXAMPLE", "first-secret"); + private static final AwsCredentials ROTATED_CREDENTIALS = AwsBasicCredentials.create( + "AKIAI44QH8DHBEXAMPLE", "second-secret"); + private static final Clock SIGNING_CLOCK = + Clock.fixed(Instant.parse("2025-01-01T00:00:00Z"), ZoneOffset.UTC); + + @Test + void cachesTokenUntilBackgroundRefresh() { + Harness harness = new Harness(); + try { + String first = harness.manager.getToken(); + String second = harness.manager.getToken(); + + assertEquals(first, second); + assertEquals(1, harness.credentialCalls.get()); + assertEquals(300_000, harness.scheduler.created.get(0).delayMillis); + } finally { + harness.manager.close(); + } + } + + @Test + void generatesFirstTokenLazily() { + Harness harness = new Harness(); + try { + assertEquals(0, harness.credentialCalls.get()); + assertEquals(0, harness.scheduler.createdCount()); + + harness.manager.getToken(); + + assertEquals(1, harness.credentialCalls.get()); + assertEquals(1, harness.scheduler.createdCount()); + } finally { + harness.manager.close(); + } + } + + @Test + void refreshesInBackgroundWithRotatedCredentials() { + Harness harness = new Harness(); + try { + String first = harness.manager.getToken(); + harness.scheduler.advance(300_000); + String second = harness.manager.getToken(); + + assertEquals(2, harness.credentialCalls.get()); + assertNotEquals(first, second); + assertEquals(2, harness.changedTokens.size()); + } finally { + harness.manager.close(); + } + } + + @Test + void honorsCustomRefreshInterval() { + Harness harness = new Harness(Duration.ofSeconds(30), null); + try { + harness.manager.getToken(); + harness.scheduler.advance(29_999); + assertEquals(1, harness.credentialCalls.get()); + harness.scheduler.advance(1); + assertEquals(2, harness.credentialCalls.get()); + } finally { + harness.manager.close(); + } + } + + @Test + void returnsNormalizedCredentialsAndCachedToken() { + Harness harness = new Harness(null, "TestUser"); + try { + ElastiCacheIamAuthTokenManager.Credentials credentials = + harness.manager.getCredentials(); + + assertEquals(USER, credentials.getUserId()); + assertEquals(credentials.getToken(), harness.manager.getToken()); + assertEquals(1, harness.credentialCalls.get()); + } finally { + harness.manager.close(); + } + } + + @Test + void sharesInitialTokenGenerationAcrossConcurrentCallers() { + TestScheduler scheduler = new TestScheduler(); + AtomicInteger credentialCalls = new AtomicInteger(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch shared = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder() + .serverlessCacheName(CACHE) + .userId(USER) + .region(Region.US_EAST_1) + .credentialsProvider(() -> { + credentialCalls.incrementAndGet(); + entered.countDown(); + try { + if (!release.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting to release provider"); + } + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(exception); + } + return FIRST_CREDENTIALS; + }) + .signingClock(SIGNING_CLOCK) + .currentTimeMillis(scheduler::now) + .scheduler(scheduler) + .onSharedRefresh(shared::countDown) + .build(); + try { + CompletableFuture first = + CompletableFuture.supplyAsync(manager::getToken); + await(entered); + CompletableFuture second = + CompletableFuture.supplyAsync(manager::getToken); + await(shared); + release.countDown(); + + assertEquals(first.join(), second.join()); + assertEquals(1, credentialCalls.get()); + } finally { + release.countDown(); + manager.close(); + } + } + + @Test + void retriesInitialRefreshWithCappedExponentialBackoff() { + Harness harness = new Harness(); + harness.failing.set(true); + CompletableFuture result = CompletableFuture.supplyAsync(() -> { + try { + harness.manager.getToken(); + return null; + } catch (RuntimeException exception) { + return exception; + } + }); + + awaitCreatedTasks(harness.scheduler, 1); + harness.scheduler.advance(12_000); + + assertTrue(result.join() instanceof ConfigurationException); + assertEquals(8, harness.credentialCalls.get()); + assertEquals( + java.util.Arrays.asList(100L, 200L, 400L, 800L, 1_600L, 3_200L, 5_000L), + harness.scheduler.delays()); + harness.manager.close(); + } + + @Test + void initialErrorClearsInFlightAndAllowsRetry() { + TestScheduler scheduler = new TestScheduler(); + AtomicInteger credentialCalls = new AtomicInteger(); + AtomicBoolean throwError = new AtomicBoolean(true); + ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder() + .serverlessCacheName(CACHE) + .userId(USER) + .region(Region.US_EAST_1) + .credentialsProvider(() -> { + credentialCalls.incrementAndGet(); + if (throwError.getAndSet(false)) { + throw new AssertionError("fatal credential provider failure"); + } + return FIRST_CREDENTIALS; + }) + .signingClock(SIGNING_CLOCK) + .currentTimeMillis(scheduler::now) + .scheduler(scheduler) + .build(); + try { + assertThrows(AssertionError.class, manager::getToken); + + assertFalse(manager.getToken().isEmpty()); + assertEquals(2, credentialCalls.get()); + } finally { + manager.close(); + } + } + + @Test + void waitingCallerReceivesOriginalErrorFromSharedRefresh() { + TestScheduler scheduler = new TestScheduler(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch shared = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + AssertionError failure = new AssertionError("fatal credential provider failure"); + ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder() + .serverlessCacheName(CACHE) + .userId(USER) + .region(Region.US_EAST_1) + .credentialsProvider(() -> { + entered.countDown(); + try { + if (!release.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("timed out waiting to release provider"); + } + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(exception); + } + throw failure; + }) + .signingClock(SIGNING_CLOCK) + .currentTimeMillis(scheduler::now) + .scheduler(scheduler) + .onSharedRefresh(shared::countDown) + .build(); + try { + CompletableFuture first = + CompletableFuture.supplyAsync(() -> captureFailure(manager::getToken)); + await(entered); + CompletableFuture second = + CompletableFuture.supplyAsync(() -> captureFailure(manager::getToken)); + await(shared); + release.countDown(); + + assertSame(failure, first.join()); + assertSame(failure, second.join()); + } finally { + release.countDown(); + manager.close(); + } + } + + @Test + void backgroundErrorClearsInFlightAndAllowsLaterRefresh() { + TestScheduler scheduler = new TestScheduler(); + AtomicInteger credentialCalls = new AtomicInteger(); + AtomicBoolean throwError = new AtomicBoolean(); + ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder() + .serverlessCacheName(CACHE) + .userId(USER) + .region(Region.US_EAST_1) + .credentialsProvider(() -> { + int call = credentialCalls.incrementAndGet(); + if (throwError.get()) { + throw new AssertionError("fatal credential provider failure"); + } + return call == 1 ? FIRST_CREDENTIALS : ROTATED_CREDENTIALS; + }) + .signingClock(SIGNING_CLOCK) + .currentTimeMillis(scheduler::now) + .scheduler(scheduler) + .build(); + try { + String first = manager.getToken(); + throwError.set(true); + + assertThrows(AssertionError.class, () -> scheduler.advance(300_000)); + + throwError.set(false); + assertEquals(first, manager.getToken()); + scheduler.advance(0); + assertNotEquals(first, manager.getToken()); + assertEquals(3, credentialCalls.get()); + } finally { + manager.close(); + } + } + + @Test + void servesValidTokenWhileRefreshFailsAndRecovers() { + Harness harness = new Harness(); + try { + String token = harness.manager.getToken(); + harness.failing.set(true); + harness.scheduler.advance(312_000); + + assertEquals(5_000, harness.scheduler.lastCreated().delayMillis); + int tasksAfterFailure = harness.scheduler.createdCount(); + assertEquals(token, harness.manager.getToken()); + assertEquals(token, harness.manager.getToken()); + assertEquals(9, harness.credentialCalls.get()); + assertEquals(tasksAfterFailure, harness.scheduler.createdCount()); + + harness.failing.set(false); + harness.scheduler.advance(5_000); + assertNotEquals(token, harness.manager.getToken()); + } finally { + harness.manager.close(); + } + } + + @Test + void treatsTokenAsExpiredAtEffectiveLifetime() { + Harness harness = new Harness(Duration.ofSeconds(879), null); + try { + harness.manager.getToken(); + harness.failing.set(true); + // The first failed cycle ends after the 880-second effective expiry + // but before the raw 900-second token expiry. + harness.scheduler.advance(891_000); + + CompletableFuture result = CompletableFuture.supplyAsync(() -> { + try { + harness.manager.getToken(); + return null; + } catch (RuntimeException exception) { + return exception; + } + }); + awaitCreatedTasks(harness.scheduler, 9); + harness.scheduler.advance(12_000); + + Throwable error = result.join(); + assertTrue(error instanceof TokenRefreshException); + assertTrue(error.getCause() instanceof ConfigurationException); + } finally { + harness.manager.close(); + } + } + + @Test + void ignoresTokenChangedCallbackFailures() { + Harness harness = new Harness( + null, + null, + token -> { + throw new AssertionError("consumer callback failed"); + }); + try { + assertFalse(harness.manager.getToken().isEmpty()); + assertEquals(1, harness.credentialCalls.get()); + } finally { + harness.manager.close(); + } + } + + @Test + void blockingTokenChangedCallbackDoesNotDelayTokenCaller() throws Exception { + CountDownLatch callbackEntered = new CountDownLatch(1); + CountDownLatch releaseCallback = new CountDownLatch(1); + AtomicReference callbackThread = new AtomicReference<>(); + ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder() + .serverlessCacheName(CACHE) + .userId(USER) + .region(Region.US_EAST_1) + .credentialsProvider(() -> FIRST_CREDENTIALS) + .signingClock(SIGNING_CLOCK) + .onTokenChanged(token -> { + callbackThread.set(Thread.currentThread().getName()); + callbackEntered.countDown(); + try { + releaseCallback.await(); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + } + }) + .build(); + try { + CompletableFuture token = + CompletableFuture.supplyAsync(manager::getToken); + + await(callbackEntered); + assertFalse(token.get(5, TimeUnit.SECONDS).isEmpty()); + assertEquals("elasticache-iam-token-refresh", callbackThread.get()); + } finally { + releaseCallback.countDown(); + manager.close(); + } + } + + @Test + void closeCancelsBackgroundWorkAndRejectsRequests() { + Harness harness = new Harness(); + harness.manager.getToken(); + harness.manager.close(); + harness.scheduler.advance(300_000); + + assertEquals(1, harness.credentialCalls.get()); + assertThrows(TokenRefreshException.class, harness.manager::getToken); + } + + @Test + void closePreventsRetryWhenInFlightGenerationFails() { + TestScheduler scheduler = new TestScheduler(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder() + .serverlessCacheName(CACHE) + .userId(USER) + .region(Region.US_EAST_1) + .credentialsProvider(() -> { + entered.countDown(); + await(release); + throw new IllegalStateException("credential chain failed"); + }) + .signingClock(SIGNING_CLOCK) + .currentTimeMillis(scheduler::now) + .scheduler(scheduler) + .build(); + + CompletableFuture result = CompletableFuture.supplyAsync(() -> { + try { + manager.getToken(); + return null; + } catch (RuntimeException exception) { + return exception; + } + }); + await(entered); + manager.close(); + release.countDown(); + + Throwable failure = result.join(); + assertTrue(failure instanceof TokenRefreshException); + assertEquals( + "Token refresh was cancelled because the token manager was closed.", + failure.getMessage()); + assertEquals(0, scheduler.createdCount()); + } + + @Test + void closeAfterSigningPreventsTokenInstallationAndNotification() { + TestScheduler scheduler = new TestScheduler(); + List changedTokens = new ArrayList<>(); + CountDownLatch signed = new CountDownLatch(1); + CountDownLatch install = new CountDownLatch(1); + ElastiCacheIamAuthTokenManager manager = + baseBuilder(scheduler, new AtomicInteger(), new AtomicBoolean()) + .onTokenChanged(changedTokens::add) + .beforeInstall(() -> { + signed.countDown(); + await(install); + }) + .build(); + CompletableFuture result = + CompletableFuture.supplyAsync(() -> captureFailure(manager::getToken)); + + await(signed); + manager.close(); + install.countDown(); + + Throwable failure = result.join(); + assertTrue(failure instanceof TokenRefreshException); + assertEquals( + "Token refresh was cancelled because the token manager was closed.", + failure.getMessage()); + assertTrue(changedTokens.isEmpty()); + assertEquals(0, scheduler.createdCount()); + } + + @Test + void validatesRefreshInterval() { + for (Duration value : java.util.Arrays.asList( + Duration.ZERO, + Duration.ofMillis(-1), + Duration.ofNanos(1), + Duration.ofSeconds(880), + Duration.ofSeconds(899), + Duration.ofMinutes(15), + Duration.ofMinutes(16))) { + assertThrows( + InvalidParameterException.class, + () -> baseBuilder(new TestScheduler(), new AtomicInteger(), new AtomicBoolean()) + .refreshAfter(value) + .build()); + } + + ElastiCacheIamAuthTokenManager manager = + baseBuilder(new TestScheduler(), new AtomicInteger(), new AtomicBoolean()) + .refreshAfter(Duration.ofSeconds(879)) + .build(); + manager.close(); + } + + @Test + void refreshExceptionIsNotAnInputException() { + assertFalse(ToolkitInputException.class.isAssignableFrom(TokenRefreshException.class)); + } + + private static ElastiCacheIamAuthTokenManager.Builder baseBuilder( + TestScheduler scheduler, + AtomicInteger credentialCalls, + AtomicBoolean failing) { + return ElastiCacheIamAuthTokenManager.builder() + .serverlessCacheName(CACHE) + .userId(USER) + .region(Region.US_EAST_1) + .credentialsProvider(() -> { + int call = credentialCalls.incrementAndGet(); + if (failing.get()) { + throw new IllegalStateException("credential chain failed"); + } + return call == 1 ? FIRST_CREDENTIALS : ROTATED_CREDENTIALS; + }) + .signingClock(SIGNING_CLOCK) + .currentTimeMillis(scheduler::now) + .random(() -> 0.5) + .scheduler(scheduler); + } + + private static void awaitCreatedTasks(TestScheduler scheduler, int expected) { + assertTrue( + scheduler.awaitCreatedCount(expected, 5, TimeUnit.SECONDS), + "timed out waiting for scheduled tasks"); + assertEquals(expected, scheduler.createdCount()); + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS), "timed out waiting for latch"); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new AssertionError("interrupted while waiting for latch", exception); + } + } + + private static Throwable captureFailure(Runnable action) { + try { + action.run(); + return null; + } catch (Throwable failure) { + return failure; + } + } + + private static final class Harness { + private final TestScheduler scheduler = new TestScheduler(); + private final AtomicInteger credentialCalls = new AtomicInteger(); + private final AtomicBoolean failing = new AtomicBoolean(); + private final List changedTokens = new ArrayList<>(); + private final ElastiCacheIamAuthTokenManager manager; + + private Harness() { + this(null, null); + } + + private Harness(Duration refreshAfter, String userId) { + this(refreshAfter, userId, null); + } + + private Harness( + Duration refreshAfter, String userId, java.util.function.Consumer callback) { + ElastiCacheIamAuthTokenManager.Builder builder = + baseBuilder(scheduler, credentialCalls, failing) + .userId(userId == null ? USER : userId) + .onTokenChanged(callback == null ? changedTokens::add : callback); + if (refreshAfter != null) { + builder.refreshAfter(refreshAfter); + } + manager = builder.build(); + } + } + + private static final class TestScheduler + implements ElastiCacheIamAuthTokenManager.Scheduler { + private long time; + private final List pending = new ArrayList<>(); + private final List created = new ArrayList<>(); + + private synchronized long now() { + return time; + } + + @Override + public synchronized ElastiCacheIamAuthTokenManager.Cancellable schedule( + Runnable task, long delayMillis) { + TestTask scheduled = new TestTask(time + delayMillis, delayMillis, task); + pending.add(scheduled); + created.add(scheduled); + notifyAll(); + return scheduled; + } + + @Override + public void execute(Runnable task) { + task.run(); + } + + @Override + public void close() {} + + private synchronized void advance(long milliseconds) { + long target = time + milliseconds; + while (true) { + TestTask next = pending.stream() + .filter(task -> !task.cancelled && task.at <= target) + .min(Comparator.comparingLong(task -> task.at)) + .orElse(null); + if (next == null) { + break; + } + pending.remove(next); + time = next.at; + next.run(); + } + time = target; + } + + private synchronized List delays() { + List delays = new ArrayList<>(); + for (TestTask task : created) { + delays.add(task.delayMillis); + } + return delays; + } + + private synchronized TestTask lastCreated() { + return created.get(created.size() - 1); + } + + private synchronized int createdCount() { + return created.size(); + } + + private synchronized boolean awaitCreatedCount( + int expected, long timeout, TimeUnit unit) { + long remainingNanos = unit.toNanos(timeout); + long deadline = System.nanoTime() + remainingNanos; + while (created.size() < expected && remainingNanos > 0) { + try { + TimeUnit.NANOSECONDS.timedWait(this, remainingNanos); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + return false; + } + remainingNanos = deadline - System.nanoTime(); + } + return created.size() >= expected; + } + } + + private static final class TestTask + implements ElastiCacheIamAuthTokenManager.Cancellable { + private final long at; + private final long delayMillis; + private final Runnable task; + private boolean cancelled; + + private TestTask(long at, long delayMillis, Runnable task) { + this.at = at; + this.delayMillis = delayMillis; + this.task = task; + } + + @Override + public void cancel() { + cancelled = true; + } + + private void run() { + if (!cancelled) { + task.run(); + } + } + } +}