Refresh IMDS credentials more aggressively. (#3147)

This improves the behavior introduced in #2989 if the customer is using a mock IMDS endpoint with aggressively short session durations.

Author: millems

.changes/next-release/bugfix-AWSSDKforJavav2-d839c10.json

@@ -0,0 +1,6 @@
+{
+    "category": "AWS SDK for Java v2", 
+    "contributor": "", 
+    "type": "bugfix", 
+    "description": "Refresh IMDS credentials more aggressively."
+}

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/InstanceProfileCredentialsProvider.java

@@ -15,8 +15,9 @@
 
 package software.amazon.awssdk.auth.credentials;
 
-import static java.time.temporal.ChronoUnit.HOURS;
 import static java.time.temporal.ChronoUnit.MINUTES;
+import static java.time.temporal.ChronoUnit.SECONDS;
+import static software.amazon.awssdk.utils.ComparableUtils.minimum;
 
 import java.io.IOException;
 import java.net.URI;
@@ -156,24 +157,39 @@ private boolean isLocalCredentialLoadingDisabled() {
 
     private Instant prefetchTime(Instant expiration) {
         Instant now = clock.instant();
-        Instant oneHourFromNow = now.plus(1, HOURS);
 
-        // If expiration time is infinite or farther out than an hour, wait an hour before refreshing
-        if (expiration == null || expiration.isAfter(oneHourFromNow)) {
-            return oneHourFromNow;
+        // If expiration time doesn't exist, refresh in 60 minutes
+        if (expiration == null) {
+            return now.plus(60, MINUTES);
         }
 
-        // If expiration time is within 15 minutes (or in the past), wait 15 minutes and warn the customer that they'll be using
-        // expired credentials.
-        Instant fifteenMinutesFromNow = now.plus(15, MINUTES);
-        if (expiration.isBefore(fifteenMinutesFromNow)) {
-            log.warn(() -> "IMDS credential expiration has been extended due to an IMDS availability outage. A refresh"
-                           + " of these credentials will be attempted again in 15 minutes.");
-            return fifteenMinutesFromNow;
+        // If expiration time is 60+ minutes from now, refresh in 60 minutes or 60 minutes before expiration, whichever is
+        // sooner. This is the average case, where customers are using IMDS and there is no IMDS outage.
+        Instant sixtyMinutesBeforeExpiration = expiration.minus(60, MINUTES);
+        if (now.isBefore(sixtyMinutesBeforeExpiration)) {
+            return minimum(sixtyMinutesBeforeExpiration, now.plus(60, MINUTES));
         }
 
-        // Otherwise, just refresh 15 minutes before the credentials expire.
-        return expiration.minus(15, MINUTES);
+        // If expiration time is 5-60 minutes from now, refresh in 30 minutes or 5 minutes before expiration, whatever is
+        // sooner. This is an unusual case: IMDS is either having an outage or the customer is using a mock IMDS with shorter
+        // default session durations.
+        Instant fiveMinutesBeforeExpiration = expiration.minus(5, MINUTES);
+        if (now.isBefore(fiveMinutesBeforeExpiration)) {
+            return minimum(fiveMinutesBeforeExpiration, now.plus(30, MINUTES));
+        }
+
+        // If expiration time is 0.25-5 minutes from now, refresh 15 seconds before expiration. This is an unusual case: IMDS is
+        // either having an outage or the customer is using a mock IMDS with very aggressive session durations.
+        Instant fifteenSecondsBeforeExpiration = expiration.minus(15, SECONDS);
+        if (now.isBefore(fifteenSecondsBeforeExpiration)) {
+            return fifteenSecondsBeforeExpiration;
+        }
+
+        // These credentials are expired. Try refreshing again in 5 minutes. We can't be more aggressive than that, because we
+        // don't want to overload the IMDS endpoint.
+        log.warn(() -> "IMDS credential expiration has been extended due to an IMDS availability outage. A refresh "
+                       + "of these credentials will be attempted again in 5 minutes.");
+        return now.plus(5, MINUTES);
     }
 
     @Override

core/auth/src/test/java/software/amazon/awssdk/auth/credentials/InstanceProfileCredentialsProviderTest.java

@@ -24,6 +24,8 @@
 import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
 import static java.time.temporal.ChronoUnit.HOURS;
+import static java.time.temporal.ChronoUnit.MINUTES;
+import static java.time.temporal.ChronoUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
@@ -335,6 +337,46 @@ public void resolveCredentials_usesCacheIfImdsFailsOnSecondCall() {
         assertThat(credentialsBefore).isEqualTo(credentialsAfter);
     }
 
+    @Test
+    public void resolveCredentials_callsImdsIfCredentialsWithin5MinutesOfExpiration() {
+        AdjustableClock clock = new AdjustableClock();
+        AwsCredentialsProvider credentialsProvider = credentialsProviderWithClock(clock);
+        Instant now = Instant.now();
+        String successfulCredentialsResponse1 =
+            "{"
+            + "\"AccessKeyId\":\"ACCESS_KEY_ID\","
+            + "\"SecretAccessKey\":\"SECRET_ACCESS_KEY\","
+            + "\"Expiration\":\"" + DateUtils.formatIso8601Date(now) + '"'
+            + "}";
+
+        String successfulCredentialsResponse2 =
+            "{"
+            + "\"AccessKeyId\":\"ACCESS_KEY_ID\","
+            + "\"SecretAccessKey\":\"SECRET_ACCESS_KEY2\","
+            + "\"Expiration\":\"" + DateUtils.formatIso8601Date(now.plus(6, HOURS)) + '"'
+            + "}";
+
+        // Set the time to the past and call IMDS to prime the cache
+        clock.time = now.minus(24, HOURS);
+        stubCredentialsResponse(aResponse().withBody(successfulCredentialsResponse1));
+        AwsCredentials credentials24HoursAgo = credentialsProvider.resolveCredentials();
+
+        // Set the time to 3 minutes before expiration, and fail to call IMDS
+        clock.time = now.minus(3, MINUTES);
+        stubCredentialsResponse(aResponse().withStatus(500));
+        AwsCredentials credentials3MinutesAgo = credentialsProvider.resolveCredentials();
+
+        // Set the time to 10 seconds before expiration, and verify that we still call IMDS to try to get credentials in at the
+        // last moment before expiration
+        clock.time = now.minus(10, SECONDS);
+        stubCredentialsResponse(aResponse().withBody(successfulCredentialsResponse2));
+        AwsCredentials credentials10SecondsAgo = credentialsProvider.resolveCredentials();
+
+        assertThat(credentials24HoursAgo).isEqualTo(credentials3MinutesAgo);
+        assertThat(credentials24HoursAgo.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY");
+        assertThat(credentials10SecondsAgo.secretAccessKey()).isEqualTo("SECRET_ACCESS_KEY2");
+    }
+
     private AwsCredentialsProvider credentialsProviderWithClock(Clock clock) {
         InstanceProfileCredentialsProvider.BuilderImpl builder =
             (InstanceProfileCredentialsProvider.BuilderImpl) InstanceProfileCredentialsProvider.builder();