Add forced IAM token refresh API to the Java toolkit

java/CHANGELOG.md

@@ -18,6 +18,9 @@ and this package uses [Semantic Versioning](https://semver.org/).
 - `ElastiCacheIamAuthTokenManager` for caching tokens, refreshing them in the
   background before expiry, retrying transient failures, and notifying a
   token-change callback.
+- `ElastiCacheIamAuthTokenManager.refreshToken()` for replacing a rejected cached
+  token without waiting for scheduled refresh while sharing concurrent and
+  background refresh work.
 - `TokenRefreshException` for closed managers and refreshes that cannot replace
   an expired token.
 - Maven build, Java version CI matrix, documentation, and dependency updates.

java/README.md

@@ -119,6 +119,14 @@ 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.
 
+Call `refreshToken()` after a client rejects the cached token. It invalidates that
+token and refreshes without waiting for the scheduled refresh, resolving AWS
+credentials again before signing and installing a replacement. Concurrent
+`refreshToken()`, `getToken()`, and `getCredentials()` calls share the same
+replacement operation, including a background refresh already in progress or in
+retry backoff. If the replacement cannot be signed, the rejected token remains
+invalidated and the failure is reported to the caller.
+
 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`.
@@ -147,6 +155,30 @@ RedisCredentialsProvider credentialsProvider = new RedisCredentialsProvider() {
 };
 ```
 
+If the client reports an AUTH rejection, discard that connection before forcing a
+refresh and beginning the next bounded connection attempt:
+
+```java
+try {
+    return openConnection(credentialsProvider);
+} catch (RedisAuthenticationException authFailure) {
+    // Never return the rejected connection to the pool.
+    closeRejectedConnection();
+
+    // Resolve AWS credentials again and install a replacement token. The next
+    // openConnection call obtains it through credentialsProvider.
+    auth.refreshToken();
+    throw authFailure; // Let the existing bounded reconnect policy retry.
+}
+```
+
+`RedisAuthenticationException`, `openConnection`, and
+`closeRejectedConnection` represent the corresponding hooks in your client or
+connection pool. Forced refresh is recovery from a rejected cached token, not
+proof that the IAM user, policy, cache, or region is configured correctly.
+Persistent configuration errors should still surface after the client's normal
+connection-attempt limit.
+
 Run the complete
 [`ConnectWithLettuce`](examples/src/main/java/software/amazon/elasticache/examples/ConnectWithLettuce.java)
 example after installing the toolkit:

java/integration-tests/src/test/java/software/amazon/elasticache/integration/PublicApiTest.java

@@ -9,6 +9,8 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.regions.Region;
 import software.amazon.elasticache.auth.ConfigurationException;
 import software.amazon.elasticache.auth.ElastiCacheIamAuthTokenManager;
 import software.amazon.elasticache.auth.ElastiCacheIamAuthTokenProvider;
@@ -21,6 +23,17 @@ class PublicApiTest {
     void installedArtifactExposesPublicApi() {
         assertNotNull(ElastiCacheIamAuthTokenProvider.builder());
         assertNotNull(ElastiCacheIamAuthTokenManager.builder());
+        try (ElastiCacheIamAuthTokenManager manager =
+                ElastiCacheIamAuthTokenManager.builder()
+                        .serverlessCacheName("test-cache")
+                        .userId("test-user")
+                        .region(Region.US_EAST_1)
+                        .credentialsProvider(() ->
+                                AwsBasicCredentials.create("test-access-key", "test-secret-key"))
+                        .build()) {
+            String token = manager.refreshToken();
+            assertNotNull(token);
+        }
         assertTrue(ToolkitInputException.class.isAssignableFrom(ConfigurationException.class));
         assertTrue(ToolkitInputException.class.isAssignableFrom(InvalidParameterException.class));
         assertFalse(ToolkitInputException.class.isAssignableFrom(TokenRefreshException.class));

java/src/main/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManager.java

@@ -157,6 +157,33 @@ public String getToken() {
         return await(refresh());
     }
 
+    /**
+     * Forces a token refresh without waiting for the scheduled refresh.
+     *
+     * <p>The currently cached token is invalidated before new AWS credentials are
+     * resolved and a replacement token is signed. Concurrent token callers share
+     * the same refresh, including a background refresh already in progress or in
+     * retry backoff.
+     *
+     * <p>Use this operation after a client rejects the cached token. Discard the
+     * rejected connection before opening a replacement connection with the
+     * returned token.
+     *
+     * @return newly installed IAM authentication token
+     * @throws TokenRefreshException if the manager is closed
+     * @throws ConfigurationException if usable AWS credentials cannot be resolved
+     */
+    public String refreshToken() {
+        synchronized (lock) {
+            ensureOpen();
+            cached = null;
+            cancel(refreshTask);
+            refreshTask = null;
+            retryNotBefore = 0;
+        }
+        return await(refresh());
+    }
+
     /**
      * Stops all background work and releases the manager's scheduler.
      *
@@ -285,20 +312,14 @@ private void scheduleRetry(
 
     private void finishFailedRefresh(
             CompletableFuture<String> 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);
+                refresh.complete(cached.token);
+                return;
             }
-        }
-        if (validToken != null) {
-            refresh.complete(validToken);
-            return;
-        }
-        synchronized (lock) {
             if (cached == null) {
                 refresh.completeExceptionally(lastError);
             } else {

java/src/test/java/software/amazon/elasticache/auth/ElastiCacheIamAuthTokenManagerTest.java

@@ -16,10 +16,13 @@
 import java.time.Instant;
 import java.time.ZoneOffset;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Comparator;
 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;
@@ -116,6 +119,164 @@ void returnsNormalizedCredentialsAndCachedToken() {
         }
     }
 
+    @Test
+    void forcedRefreshUsesRotatedCredentialsAndResetsScheduling() {
+        Harness harness = new Harness();
+        try {
+            String first = harness.manager.getToken();
+            TestTask originalRefresh = harness.scheduler.lastCreated();
+
+            String replacement = harness.manager.refreshToken();
+
+            assertNotEquals(first, replacement);
+            assertEquals(replacement, harness.manager.getToken());
+            assertEquals(2, harness.credentialCalls.get());
+            assertEquals(Arrays.asList(first, replacement), harness.changedTokens);
+            assertTrue(originalRefresh.cancelled);
+            assertEquals(300_000, harness.scheduler.lastCreated().delayMillis);
+            harness.scheduler.advance(299_999);
+            assertEquals(2, harness.credentialCalls.get());
+            harness.scheduler.advance(1);
+            assertEquals(3, harness.credentialCalls.get());
+        } finally {
+            harness.manager.close();
+        }
+    }
+
+    @Test
+    void concurrentForcedRefreshesAndTokenCallersShareReplacement() {
+        TestScheduler scheduler = new TestScheduler();
+        AtomicInteger credentialCalls = new AtomicInteger();
+        CountDownLatch replacementStarted = new CountDownLatch(1);
+        CountDownLatch shared = new CountDownLatch(3);
+        CountDownLatch releaseReplacement = new CountDownLatch(1);
+        ExecutorService callers = Executors.newFixedThreadPool(4);
+        ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder()
+                .serverlessCacheName(CACHE)
+                .userId(USER)
+                .region(Region.US_EAST_1)
+                .credentialsProvider(() -> {
+                    int call = credentialCalls.incrementAndGet();
+                    if (call > 1) {
+                        replacementStarted.countDown();
+                        await(releaseReplacement);
+                    }
+                    return call == 1 ? FIRST_CREDENTIALS : ROTATED_CREDENTIALS;
+                })
+                .signingClock(SIGNING_CLOCK)
+                .currentTimeMillis(scheduler::now)
+                .scheduler(scheduler)
+                .onSharedRefresh(shared::countDown)
+                .build();
+        try {
+            String rejected = manager.getToken();
+            CompletableFuture<String> first =
+                    CompletableFuture.supplyAsync(manager::refreshToken, callers);
+            await(replacementStarted);
+            CompletableFuture<String> second =
+                    CompletableFuture.supplyAsync(manager::refreshToken, callers);
+            CompletableFuture<String> token =
+                    CompletableFuture.supplyAsync(manager::getToken, callers);
+            CompletableFuture<ElastiCacheIamAuthTokenManager.Credentials> credentials =
+                    CompletableFuture.supplyAsync(manager::getCredentials, callers);
+            await(shared);
+            releaseReplacement.countDown();
+
+            String replacement = first.join();
+            assertNotEquals(rejected, replacement);
+            assertEquals(replacement, second.join());
+            assertEquals(replacement, token.join());
+            assertEquals(replacement, credentials.join().getToken());
+            assertEquals(2, credentialCalls.get());
+        } finally {
+            releaseReplacement.countDown();
+            manager.close();
+            callers.shutdownNow();
+        }
+    }
+
+    @Test
+    void forcedRefreshJoinsBackgroundRefreshWithoutStaleOverwrite() {
+        TestScheduler scheduler = new TestScheduler();
+        AtomicInteger credentialCalls = new AtomicInteger();
+        CountDownLatch backgroundStarted = new CountDownLatch(1);
+        CountDownLatch shared = new CountDownLatch(1);
+        CountDownLatch releaseBackground = new CountDownLatch(1);
+        ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder()
+                .serverlessCacheName(CACHE)
+                .userId(USER)
+                .region(Region.US_EAST_1)
+                .credentialsProvider(() -> {
+                    int call = credentialCalls.incrementAndGet();
+                    if (call > 1) {
+                        backgroundStarted.countDown();
+                        await(releaseBackground);
+                    }
+                    return call == 1 ? FIRST_CREDENTIALS : ROTATED_CREDENTIALS;
+                })
+                .signingClock(SIGNING_CLOCK)
+                .currentTimeMillis(scheduler::now)
+                .scheduler(scheduler)
+                .onSharedRefresh(shared::countDown)
+                .build();
+        try {
+            String rejected = manager.getToken();
+            CompletableFuture<Void> background =
+                    CompletableFuture.runAsync(() -> scheduler.advance(300_000));
+            await(backgroundStarted);
+            CompletableFuture<String> forced =
+                    CompletableFuture.supplyAsync(manager::refreshToken);
+            await(shared);
+            releaseBackground.countDown();
+
+            String replacement = forced.join();
+            background.join();
+            assertNotEquals(rejected, replacement);
+            assertEquals(replacement, manager.getToken());
+            assertEquals(2, credentialCalls.get());
+            assertEquals(2, scheduler.createdCount());
+        } finally {
+            releaseBackground.countDown();
+            manager.close();
+        }
+    }
+
+    @Test
+    void failedForcedRefreshDoesNotReturnInvalidatedToken() {
+        TestScheduler scheduler = new TestScheduler();
+        AtomicInteger credentialCalls = new AtomicInteger();
+        AtomicBoolean failing = new AtomicBoolean();
+        CountDownLatch shared = new CountDownLatch(1);
+        ElastiCacheIamAuthTokenManager manager =
+                baseBuilder(scheduler, credentialCalls, failing)
+                        .onSharedRefresh(shared::countDown)
+                        .build();
+        try {
+            String rejected = manager.getToken();
+            failing.set(true);
+            CompletableFuture<Throwable> forced = CompletableFuture.supplyAsync(
+                    () -> captureFailure(manager::refreshToken));
+            awaitCreatedTasks(scheduler, 2);
+            CompletableFuture<Throwable> concurrent = CompletableFuture.supplyAsync(
+                    () -> captureFailure(manager::getToken));
+            await(shared);
+            scheduler.advance(12_000);
+
+            Throwable forcedFailure = forced.join();
+            Throwable concurrentFailure = concurrent.join();
+            assertTrue(forcedFailure instanceof ConfigurationException);
+            assertSame(forcedFailure, concurrentFailure);
+            assertFalse(forcedFailure.getMessage().contains(rejected));
+            assertEquals(9, credentialCalls.get());
+
+            failing.set(false);
+            assertNotEquals(rejected, manager.getToken());
+            assertEquals(10, credentialCalls.get());
+        } finally {
+            manager.close();
+        }
+    }
+
     @Test
     void sharesInitialTokenGenerationAcrossConcurrentCallers() {
         TestScheduler scheduler = new TestScheduler();
@@ -181,7 +342,7 @@ void retriesInitialRefreshWithCappedExponentialBackoff() {
         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),
+                Arrays.asList(100L, 200L, 400L, 800L, 1_600L, 3_200L, 5_000L),
                 harness.scheduler.delays());
         harness.manager.close();
     }
@@ -408,6 +569,49 @@ void closeCancelsBackgroundWorkAndRejectsRequests() {
 
         assertEquals(1, harness.credentialCalls.get());
         assertThrows(TokenRefreshException.class, harness.manager::getToken);
+        assertThrows(TokenRefreshException.class, harness.manager::refreshToken);
+    }
+
+    @Test
+    void closeCancelsForcedRefreshInFlight() {
+        TestScheduler scheduler = new TestScheduler();
+        CountDownLatch replacementStarted = new CountDownLatch(1);
+        CountDownLatch releaseReplacement = new CountDownLatch(1);
+        AtomicInteger credentialCalls = new AtomicInteger();
+        ElastiCacheIamAuthTokenManager manager = ElastiCacheIamAuthTokenManager.builder()
+                .serverlessCacheName(CACHE)
+                .userId(USER)
+                .region(Region.US_EAST_1)
+                .credentialsProvider(() -> {
+                    int call = credentialCalls.incrementAndGet();
+                    if (call > 1) {
+                        replacementStarted.countDown();
+                        await(releaseReplacement);
+                    }
+                    return call == 1 ? FIRST_CREDENTIALS : ROTATED_CREDENTIALS;
+                })
+                .signingClock(SIGNING_CLOCK)
+                .currentTimeMillis(scheduler::now)
+                .scheduler(scheduler)
+                .build();
+        try {
+            manager.getToken();
+            CompletableFuture<Throwable> result = CompletableFuture.supplyAsync(
+                    () -> captureFailure(manager::refreshToken));
+            await(replacementStarted);
+            manager.close();
+            releaseReplacement.countDown();
+
+            Throwable failure = result.join();
+            assertTrue(failure instanceof TokenRefreshException);
+            assertEquals(
+                    "Token refresh was cancelled because the token manager was closed.",
+                    failure.getMessage());
+            assertEquals(1, scheduler.createdCount());
+        } finally {
+            releaseReplacement.countDown();
+            manager.close();
+        }
     }
 
     @Test
@@ -481,7 +685,7 @@ void closeAfterSigningPreventsTokenInstallationAndNotification() {
 
     @Test
     void validatesRefreshInterval() {
-        for (Duration value : java.util.Arrays.asList(
+        for (Duration value : Arrays.asList(
                 Duration.ZERO,
                 Duration.ofMillis(-1),
                 Duration.ofNanos(1),