From e6c14ef8c0d2bcc7130f484be1671b3cccba11ed Mon Sep 17 00:00:00 2001 From: "Alexandre [C] Luc Rehnby-Martin" Date: Wed, 16 Sep 2026 15:53:53 -0700 Subject: [PATCH 1/4] feat(java): add token auto-refresh --- README.md | 2 +- java/CHANGELOG.md | 5 + java/README.md | 66 +- .../examples/ConnectWithLettuce.java | 13 +- .../integration/PublicApiTest.java | 4 + .../auth/ElastiCacheIamAuthTokenManager.java | 604 ++++++++++++++++++ .../auth/ElastiCacheIamAuthTokenProvider.java | 2 +- .../auth/TokenRefreshException.java | 35 + .../amazon/elasticache/auth/package-info.java | 2 +- .../ElastiCacheIamAuthTokenManagerTest.java | 452 +++++++++++++ 10 files changed, 1174 insertions(+), 11 deletions(-) create mode 100644 java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java create mode 100644 java/src/main/java/software/amazon/elasticache/auth/TokenRefreshException.java create mode 100644 java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java 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..4b59d40 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,75 @@ 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, the 15-minute token lifetime. + +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..0fd0680 --- /dev/null +++ b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java @@ -0,0 +1,604 @@ +/* + * 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.regions.Region; +import software.amazon.awssdk.regions.providers.AwsRegionProvider; + +/** + * 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 String CLOSED_MESSAGE = "The token manager is 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 CachedToken cached; + private CompletableFuture inFlight; + private Cancellable refreshTask; + private Cancellable retryTask; + 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 + ? software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider + .builder() + .build() + : builder.credentialsProvider) + .awsRegionEnvironmentProvider(builder.awsRegionEnvironmentProvider == null + ? () -> System.getenv("AWS_REGION") + : builder.awsRegionEnvironmentProvider) + .regionProvider(builder.regionProvider == null + ? software.amazon.awssdk.regions.providers + .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; + 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) { + 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(CLOSED_MESSAGE)); + } + if (ownsScheduler) { + scheduler.close(); + } + } + + private CompletableFuture refresh() { + CompletableFuture refresh; + synchronized (lock) { + ensureOpen(); + if (inFlight != null) { + 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( + CLOSED_MESSAGE, lastError)); + return; + } + retryTask = null; + } + + try { + long issuedAt = currentTimeMillis.getAsLong(); + String token = provider.getToken(); + synchronized (lock) { + if (closed) { + refresh.completeExceptionally(new TokenRefreshException(CLOSED_MESSAGE)); + return; + } + } + install(token, issuedAt); + refresh.complete(token); + } catch (RuntimeException exception) { + synchronized (lock) { + if (closed) { + refresh.completeExceptionally( + new TokenRefreshException(CLOSED_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(CLOSED_MESSAGE, lastError)); + } else { + retryTask = task; + } + } + } catch (RejectedExecutionException exception) { + refresh.completeExceptionally( + new TokenRefreshException(CLOSED_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; + 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 install(String token, long issuedAt) { + synchronized (lock) { + cached = new CachedToken( + token, + issuedAt + refreshAfterMillis, + issuedAt + ElastiCacheIamAuthTokenProvider.TOKEN_TTL.toMillis()); + scheduleRefreshLocked(refreshAfterMillis); + } + notifyTokenChanged(token); + } + + private void notifyTokenChanged(String token) { + if (onTokenChanged == null) { + return; + } + try { + onTokenChanged.accept(token); + } catch (Throwable ignored) { + // A consumer callback cannot invalidate a successfully 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; + } + 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(ElastiCacheIamAuthTokenProvider.TOKEN_TTL) >= 0) { + throw new InvalidParameterException( + "Invalid value for parameter 'refreshAfter': must be at least " + + "1 millisecond and less than 15 minutes."); + } + 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 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 15-minute token lifetime. + * + * @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 after each newly generated token is installed. + * + *

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; + } + + /** + * 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 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 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..525ea6d --- /dev/null +++ b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java @@ -0,0 +1,452 @@ +/* + * 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.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.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +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() throws InterruptedException { + TestScheduler scheduler = new TestScheduler(); + AtomicInteger credentialCalls = new AtomicInteger(); + CountDownLatch entered = 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 { + release.await(); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(exception); + } + return FIRST_CREDENTIALS; + }) + .signingClock(SIGNING_CLOCK) + .currentTimeMillis(scheduler::now) + .scheduler(scheduler) + .build(); + try { + CompletableFuture first = + CompletableFuture.supplyAsync(manager::getToken); + entered.await(); + CompletableFuture second = + CompletableFuture.supplyAsync(manager::getToken); + 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 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); + assertEquals(token, harness.manager.getToken()); + assertEquals(9, harness.credentialCalls.get()); + + harness.failing.set(false); + harness.scheduler.advance(0); + assertNotEquals(token, harness.manager.getToken()); + } finally { + harness.manager.close(); + } + } + + @Test + void throwsRefreshExceptionAfterCachedTokenExpires() { + Harness harness = new Harness(Duration.ofSeconds(899), null); + try { + harness.manager.getToken(); + harness.failing.set(true); + harness.scheduler.advance(911_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 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() throws InterruptedException { + 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(); + try { + release.await(); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + } + 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; + } + }); + entered.await(); + manager.close(); + release.countDown(); + + assertTrue(result.join() instanceof TokenRefreshException); + assertEquals(0, scheduler.createdCount()); + } + + @Test + void validatesRefreshInterval() { + for (Duration value : java.util.Arrays.asList( + Duration.ZERO, + Duration.ofMillis(-1), + Duration.ofNanos(1), + Duration.ofMinutes(15), + Duration.ofMinutes(16))) { + assertThrows( + InvalidParameterException.class, + () -> baseBuilder(new TestScheduler(), new AtomicInteger(), new AtomicBoolean()) + .refreshAfter(value) + .build()); + } + } + + @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) { + for (int attempt = 0; attempt < 10_000 + && scheduler.createdCount() < expected; attempt++) { + Thread.yield(); + } + assertEquals(expected, scheduler.createdCount()); + } + + 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); + return scheduled; + } + + @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 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(); + } + } + } +} From 233e23c34d46fe596abc947f3cde69c518bea966 Mon Sep 17 00:00:00 2001 From: "Alexandre [C] Luc Rehnby-Martin" Date: Fri, 18 Sep 2026 10:32:52 -0700 Subject: [PATCH 2/4] Handle fatal token refresh failures --- .../auth/ElastiCacheIamAuthTokenManager.java | 9 ++- .../ElastiCacheIamAuthTokenManagerTest.java | 66 +++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java index 0fd0680..840b076 100644 --- a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java +++ b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java @@ -226,7 +226,14 @@ private void runRefreshAttempt( } install(token, issuedAt); refresh.complete(token); - } catch (RuntimeException exception) { + } 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( diff --git a/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java index 525ea6d..30a666a 100644 --- a/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java +++ b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java @@ -178,6 +178,72 @@ void retriesInitialRefreshWithCappedExponentialBackoff() { 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 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(); From db4f2d945c7bfb751dd3a5e146c43d75dddaaef5 Mon Sep 17 00:00:00 2001 From: "Alexandre [C] Luc Rehnby-Martin" Date: Fri, 18 Sep 2026 11:41:21 -0700 Subject: [PATCH 3/4] Address token manager concurrency review --- java/README.md | 5 +- .../auth/ElastiCacheIamAuthTokenManager.java | 111 +++++++--- .../ElastiCacheIamAuthTokenManagerTest.java | 204 ++++++++++++++++-- 3 files changed, 272 insertions(+), 48 deletions(-) diff --git a/java/README.md b/java/README.md index 4b59d40..044fa66 100644 --- a/java/README.md +++ b/java/README.md @@ -107,7 +107,10 @@ try (ElastiCacheIamAuthTokenManager auth = 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, the 15-minute token lifetime. +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 asynchronously after the new +token is available to callers. 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 diff --git a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java index 840b076..f645987 100644 --- a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java +++ b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java @@ -9,7 +9,9 @@ import java.util.Objects; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.Executor; import java.util.concurrent.Executors; +import java.util.concurrent.ForkJoinPool; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; @@ -21,8 +23,10 @@ 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. @@ -39,7 +43,12 @@ public final class ElastiCacheIamAuthTokenManager implements AutoCloseable { 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."; @@ -48,15 +57,19 @@ public final class ElastiCacheIamAuthTokenManager implements AutoCloseable { private final ElastiCacheIamAuthTokenProvider provider; private final long refreshAfterMillis; private final Consumer onTokenChanged; + private final Executor callbackExecutor; 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) { @@ -66,16 +79,13 @@ private ElastiCacheIamAuthTokenManager(Builder builder) { .serverlessCacheName(builder.serverlessCacheName) .replicationGroupId(builder.replicationGroupId) .credentialsProvider(builder.credentialsProvider == null - ? software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider - .builder() - .build() + ? DefaultCredentialsProvider.builder().build() : builder.credentialsProvider) .awsRegionEnvironmentProvider(builder.awsRegionEnvironmentProvider == null ? () -> System.getenv("AWS_REGION") : builder.awsRegionEnvironmentProvider) .regionProvider(builder.regionProvider == null - ? software.amazon.awssdk.regions.providers - .DefaultAwsRegionProviderChain.builder().build() + ? DefaultAwsRegionProviderChain.builder().build() : builder.regionProvider) .signingClock(builder.signingClock == null ? Clock.systemUTC() @@ -83,10 +93,15 @@ private ElastiCacheIamAuthTokenManager(Builder builder) { .build(); this.refreshAfterMillis = resolveRefreshAfterMillis(builder.refreshAfter); this.onTokenChanged = builder.onTokenChanged; + this.callbackExecutor = builder.callbackExecutor == null + ? ForkJoinPool.commonPool() + : builder.callbackExecutor; 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; @@ -139,7 +154,7 @@ public String getToken() { long now = currentTimeMillis.getAsLong(); if (cached != null && now < cached.expiresAt) { String token = cached.token; - if (now >= cached.refreshAt) { + if (now >= cached.refreshAt && now >= retryNotBefore) { scheduleRefreshLocked(0); } return token; @@ -168,7 +183,8 @@ public void close() { pending = inFlight; } if (pending != null) { - pending.completeExceptionally(new TokenRefreshException(CLOSED_MESSAGE)); + pending.completeExceptionally( + new TokenRefreshException(REFRESH_CANCELLED_MESSAGE)); } if (ownsScheduler) { scheduler.close(); @@ -180,6 +196,9 @@ private CompletableFuture refresh() { synchronized (lock) { ensureOpen(); if (inFlight != null) { + if (onSharedRefresh != null) { + onSharedRefresh.run(); + } return inFlight; } refresh = new CompletableFuture<>(); @@ -209,7 +228,7 @@ private void runRefreshAttempt( synchronized (lock) { if (closed) { refresh.completeExceptionally(new TokenRefreshException( - CLOSED_MESSAGE, lastError)); + REFRESH_CANCELLED_MESSAGE, lastError)); return; } retryTask = null; @@ -218,14 +237,10 @@ private void runRefreshAttempt( try { long issuedAt = currentTimeMillis.getAsLong(); String token = provider.getToken(); - synchronized (lock) { - if (closed) { - refresh.completeExceptionally(new TokenRefreshException(CLOSED_MESSAGE)); - return; - } + if (beforeInstall != null) { + beforeInstall.run(); } - install(token, issuedAt); - refresh.complete(token); + installAndComplete(refresh, token, issuedAt); } catch (Throwable failure) { if (failure instanceof Error) { refresh.completeExceptionally(failure); @@ -237,7 +252,8 @@ private void runRefreshAttempt( synchronized (lock) { if (closed) { refresh.completeExceptionally( - new TokenRefreshException(CLOSED_MESSAGE, exception)); + new TokenRefreshException( + REFRESH_CANCELLED_MESSAGE, exception)); return; } } @@ -261,14 +277,15 @@ private void scheduleRetry( if (closed) { task.cancel(); refresh.completeExceptionally( - new TokenRefreshException(CLOSED_MESSAGE, lastError)); + new TokenRefreshException( + REFRESH_CANCELLED_MESSAGE, lastError)); } else { retryTask = task; } } } catch (RejectedExecutionException exception) { refresh.completeExceptionally( - new TokenRefreshException(CLOSED_MESSAGE, lastError)); + new TokenRefreshException(REFRESH_CANCELLED_MESSAGE, lastError)); } } @@ -279,6 +296,7 @@ private void finishFailedRefresh( long now = currentTimeMillis.getAsLong(); if (cached != null && now < cached.expiresAt) { validToken = cached.token; + retryNotBefore = now + MAX_DELAY_MILLIS; scheduleRefreshLocked(MAX_DELAY_MILLIS); } } @@ -296,15 +314,23 @@ private void finishFailedRefresh( } } - private void install(String token, long issuedAt) { + 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 + ElastiCacheIamAuthTokenProvider.TOKEN_TTL.toMillis()); + issuedAt + EFFECTIVE_TOKEN_LIFETIME.toMillis()); + retryNotBefore = 0; scheduleRefreshLocked(refreshAfterMillis); + refresh.complete(token); + notifyTokenChanged(token); } - notifyTokenChanged(token); } private void notifyTokenChanged(String token) { @@ -312,9 +338,15 @@ private void notifyTokenChanged(String token) { return; } try { - onTokenChanged.accept(token); + callbackExecutor.execute(() -> { + try { + onTokenChanged.accept(token); + } catch (Throwable ignored) { + // A consumer callback cannot invalidate a successfully signed token. + } + }); } catch (Throwable ignored) { - // A consumer callback cannot invalidate a successfully signed token. + // Failure to dispatch a callback cannot invalidate a signed token. } } @@ -363,6 +395,9 @@ private static String await(CompletableFuture future) { if (cause instanceof RuntimeException) { throw (RuntimeException) cause; } + if (cause instanceof Error) { + throw (Error) cause; + } throw exception; } } @@ -378,10 +413,10 @@ private static long resolveRefreshAfterMillis(Duration refreshAfter) { if (resolved.isNegative() || resolved.isZero() || millis <= 0 - || resolved.compareTo(ElastiCacheIamAuthTokenProvider.TOKEN_TTL) >= 0) { + || resolved.compareTo(EFFECTIVE_TOKEN_LIFETIME) >= 0) { throw new InvalidParameterException( "Invalid value for parameter 'refreshAfter': must be at least " - + "1 millisecond and less than 15 minutes."); + + "1 millisecond and less than 14 minutes 40 seconds."); } return millis; } @@ -428,12 +463,15 @@ public static final class Builder { private AwsCredentialsProvider credentialsProvider; private Duration refreshAfter; private Consumer onTokenChanged; + private Executor callbackExecutor; 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() {} @@ -497,7 +535,8 @@ public Builder credentialsProvider(AwsCredentialsProvider credentialsProvider) { * 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 15-minute token lifetime. + * 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 @@ -508,7 +547,8 @@ public Builder refreshAfter(Duration refreshAfter) { } /** - * Sets a callback invoked after each newly generated token is installed. + * Sets a callback invoked asynchronously after each newly generated token + * is installed and made available to token callers. * *

Callback failures are ignored and do not discard the new token. * @@ -551,6 +591,23 @@ Builder scheduler(Scheduler scheduler) { return this; } + Builder callbackExecutor(Executor callbackExecutor) { + this.callbackExecutor = + Objects.requireNonNull(callbackExecutor, "callbackExecutor"); + 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. * diff --git a/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java index 30a666a..21cabf1 100644 --- a/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java +++ b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java @@ -7,6 +7,7 @@ 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; @@ -19,6 +20,9 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -114,10 +118,11 @@ void returnsNormalizedCredentialsAndCachedToken() { } @Test - void sharesInitialTokenGenerationAcrossConcurrentCallers() throws InterruptedException { + 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) @@ -127,7 +132,9 @@ void sharesInitialTokenGenerationAcrossConcurrentCallers() throws InterruptedExc credentialCalls.incrementAndGet(); entered.countDown(); try { - release.await(); + 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); @@ -137,13 +144,15 @@ void sharesInitialTokenGenerationAcrossConcurrentCallers() throws InterruptedExc .signingClock(SIGNING_CLOCK) .currentTimeMillis(scheduler::now) .scheduler(scheduler) + .onSharedRefresh(shared::countDown) .build(); try { CompletableFuture first = CompletableFuture.supplyAsync(manager::getToken); - entered.await(); + await(entered); CompletableFuture second = CompletableFuture.supplyAsync(manager::getToken); + await(shared); release.countDown(); assertEquals(first.join(), second.join()); @@ -208,6 +217,51 @@ void initialErrorClearsInFlightAndAllowsRetry() { } } + @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(); @@ -253,11 +307,14 @@ void servesValidTokenWhileRefreshFailsAndRecovers() { 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(0); + harness.scheduler.advance(5_000); assertNotEquals(token, harness.manager.getToken()); } finally { harness.manager.close(); @@ -265,12 +322,14 @@ void servesValidTokenWhileRefreshFailsAndRecovers() { } @Test - void throwsRefreshExceptionAfterCachedTokenExpires() { - Harness harness = new Harness(Duration.ofSeconds(899), null); + void treatsTokenAsExpiredAtEffectiveLifetime() { + Harness harness = new Harness(Duration.ofSeconds(879), null); try { harness.manager.getToken(); harness.failing.set(true); - harness.scheduler.advance(911_000); + // 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 { @@ -307,6 +366,37 @@ void ignoresTokenChangedCallbackFailures() { } } + @Test + void blockingTokenChangedCallbackDoesNotDelayTokenCaller() throws Exception { + TestScheduler scheduler = new TestScheduler(); + CountDownLatch callbackEntered = new CountDownLatch(1); + CountDownLatch releaseCallback = new CountDownLatch(1); + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ElastiCacheIamAuthTokenManager manager = + baseBuilder(scheduler, new AtomicInteger(), new AtomicBoolean()) + .callbackExecutor(callbackExecutor) + .onTokenChanged(token -> { + 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()); + } finally { + releaseCallback.countDown(); + manager.close(); + callbackExecutor.shutdownNow(); + } + } + @Test void closeCancelsBackgroundWorkAndRejectsRequests() { Harness harness = new Harness(); @@ -319,7 +409,7 @@ void closeCancelsBackgroundWorkAndRejectsRequests() { } @Test - void closePreventsRetryWhenInFlightGenerationFails() throws InterruptedException { + void closePreventsRetryWhenInFlightGenerationFails() { TestScheduler scheduler = new TestScheduler(); CountDownLatch entered = new CountDownLatch(1); CountDownLatch release = new CountDownLatch(1); @@ -329,11 +419,7 @@ void closePreventsRetryWhenInFlightGenerationFails() throws InterruptedException .region(Region.US_EAST_1) .credentialsProvider(() -> { entered.countDown(); - try { - release.await(); - } catch (InterruptedException exception) { - Thread.currentThread().interrupt(); - } + await(release); throw new IllegalStateException("credential chain failed"); }) .signingClock(SIGNING_CLOCK) @@ -349,11 +435,46 @@ void closePreventsRetryWhenInFlightGenerationFails() throws InterruptedException return exception; } }); - entered.await(); + await(entered); manager.close(); release.countDown(); - assertTrue(result.join() instanceof TokenRefreshException); + 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()) + .callbackExecutor(Runnable::run) + .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()); } @@ -363,6 +484,8 @@ void validatesRefreshInterval() { Duration.ZERO, Duration.ofMillis(-1), Duration.ofNanos(1), + Duration.ofSeconds(880), + Duration.ofSeconds(899), Duration.ofMinutes(15), Duration.ofMinutes(16))) { assertThrows( @@ -371,6 +494,12 @@ void validatesRefreshInterval() { .refreshAfter(value) .build()); } + + ElastiCacheIamAuthTokenManager manager = + baseBuilder(new TestScheduler(), new AtomicInteger(), new AtomicBoolean()) + .refreshAfter(Duration.ofSeconds(879)) + .build(); + manager.close(); } @Test @@ -396,17 +525,35 @@ private static ElastiCacheIamAuthTokenManager.Builder baseBuilder( .signingClock(SIGNING_CLOCK) .currentTimeMillis(scheduler::now) .random(() -> 0.5) - .scheduler(scheduler); + .scheduler(scheduler) + .callbackExecutor(Runnable::run); } private static void awaitCreatedTasks(TestScheduler scheduler, int expected) { - for (int attempt = 0; attempt < 10_000 - && scheduler.createdCount() < expected; attempt++) { - Thread.yield(); - } + 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(); @@ -451,6 +598,7 @@ public synchronized ElastiCacheIamAuthTokenManager.Cancellable schedule( TestTask scheduled = new TestTask(time + delayMillis, delayMillis, task); pending.add(scheduled); created.add(scheduled); + notifyAll(); return scheduled; } @@ -489,6 +637,22 @@ private synchronized TestTask lastCreated() { 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 From c2d041c0eade0ba13f1a55a14aee7f5113242353 Mon Sep 17 00:00:00 2001 From: "Alexandre [C] Luc Rehnby-Martin" Date: Tue, 22 Sep 2026 07:49:45 -0700 Subject: [PATCH 4/4] Run token callbacks on refresh thread --- java/README.md | 5 ++- .../auth/ElastiCacheIamAuthTokenManager.java | 32 ++++++------- .../ElastiCacheIamAuthTokenManagerTest.java | 45 ++++++++++--------- 3 files changed, 42 insertions(+), 40 deletions(-) diff --git a/java/README.md b/java/README.md index 044fa66..0f3793f 100644 --- a/java/README.md +++ b/java/README.md @@ -109,8 +109,9 @@ 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 asynchronously after the new -token is available to callers. +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 diff --git a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java index f645987..bb289ba 100644 --- a/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java +++ b/java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java @@ -9,9 +9,7 @@ import java.util.Objects; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; -import java.util.concurrent.Executor; import java.util.concurrent.Executors; -import java.util.concurrent.ForkJoinPool; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; @@ -57,7 +55,6 @@ public final class ElastiCacheIamAuthTokenManager implements AutoCloseable { private final ElastiCacheIamAuthTokenProvider provider; private final long refreshAfterMillis; private final Consumer onTokenChanged; - private final Executor callbackExecutor; private final LongSupplier currentTimeMillis; private final DoubleSupplier random; private final Scheduler scheduler; @@ -93,9 +90,6 @@ private ElastiCacheIamAuthTokenManager(Builder builder) { .build(); this.refreshAfterMillis = resolveRefreshAfterMillis(builder.refreshAfter); this.onTokenChanged = builder.onTokenChanged; - this.callbackExecutor = builder.callbackExecutor == null - ? ForkJoinPool.commonPool() - : builder.callbackExecutor; this.currentTimeMillis = builder.currentTimeMillis == null ? System::currentTimeMillis : builder.currentTimeMillis; @@ -329,8 +323,8 @@ private void installAndComplete( retryNotBefore = 0; scheduleRefreshLocked(refreshAfterMillis); refresh.complete(token); - notifyTokenChanged(token); } + notifyTokenChanged(token); } private void notifyTokenChanged(String token) { @@ -338,7 +332,7 @@ private void notifyTokenChanged(String token) { return; } try { - callbackExecutor.execute(() -> { + scheduler.execute(() -> { try { onTokenChanged.accept(token); } catch (Throwable ignored) { @@ -463,7 +457,6 @@ public static final class Builder { private AwsCredentialsProvider credentialsProvider; private Duration refreshAfter; private Consumer onTokenChanged; - private Executor callbackExecutor; private Supplier awsRegionEnvironmentProvider; private AwsRegionProvider regionProvider; private Clock signingClock; @@ -547,10 +540,12 @@ public Builder refreshAfter(Duration refreshAfter) { } /** - * Sets a callback invoked asynchronously after each newly generated token - * is installed and made available to token callers. + * Sets a callback invoked on the manager's refresh thread after each newly + * generated token is installed and made available to token callers. * - *

Callback failures are ignored and do not discard the new token. + *

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 @@ -591,12 +586,6 @@ Builder scheduler(Scheduler scheduler) { return this; } - Builder callbackExecutor(Executor callbackExecutor) { - this.callbackExecutor = - Objects.requireNonNull(callbackExecutor, "callbackExecutor"); - return this; - } - Builder beforeInstall(Runnable beforeInstall) { this.beforeInstall = Objects.requireNonNull(beforeInstall, "beforeInstall"); return this; @@ -625,6 +614,8 @@ interface Cancellable { interface Scheduler { Cancellable schedule(Runnable task, long delayMillis); + void execute(Runnable task); + void close(); } @@ -639,6 +630,11 @@ public Cancellable schedule(Runnable task, long delayMillis) { return () -> future.cancel(false); } + @Override + public void execute(Runnable task) { + executor.execute(task); + } + @Override public void close() { executor.shutdownNow(); diff --git a/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java index 21cabf1..1e0022f 100644 --- a/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java +++ b/java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java @@ -20,11 +20,10 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; 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; @@ -368,32 +367,35 @@ void ignoresTokenChangedCallbackFailures() { @Test void blockingTokenChangedCallbackDoesNotDelayTokenCaller() throws Exception { - TestScheduler scheduler = new TestScheduler(); CountDownLatch callbackEntered = new CountDownLatch(1); CountDownLatch releaseCallback = new CountDownLatch(1); - ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); - ElastiCacheIamAuthTokenManager manager = - baseBuilder(scheduler, new AtomicInteger(), new AtomicBoolean()) - .callbackExecutor(callbackExecutor) - .onTokenChanged(token -> { - callbackEntered.countDown(); - try { - releaseCallback.await(); - } catch (InterruptedException exception) { - Thread.currentThread().interrupt(); - } - }) - .build(); + 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(); - callbackExecutor.shutdownNow(); } } @@ -455,7 +457,6 @@ void closeAfterSigningPreventsTokenInstallationAndNotification() { CountDownLatch install = new CountDownLatch(1); ElastiCacheIamAuthTokenManager manager = baseBuilder(scheduler, new AtomicInteger(), new AtomicBoolean()) - .callbackExecutor(Runnable::run) .onTokenChanged(changedTokens::add) .beforeInstall(() -> { signed.countDown(); @@ -525,8 +526,7 @@ private static ElastiCacheIamAuthTokenManager.Builder baseBuilder( .signingClock(SIGNING_CLOCK) .currentTimeMillis(scheduler::now) .random(() -> 0.5) - .scheduler(scheduler) - .callbackExecutor(Runnable::run); + .scheduler(scheduler); } private static void awaitCreatedTasks(TestScheduler scheduler, int expected) { @@ -602,6 +602,11 @@ public synchronized ElastiCacheIamAuthTokenManager.Cancellable schedule( return scheduled; } + @Override + public void execute(Runnable task) { + task.run(); + } + @Override public void close() {}