[IMDS] Token cache for Async client

core/imds/src/main/java/software/amazon/awssdk/imds/Ec2MetadataClient.java

@@ -68,6 +68,7 @@ interface Builder extends Ec2MetadataClientBuilder<Ec2MetadataClient.Builder, Ec
          * @return a reference to this builder
          */
         Builder httpClient(SdkHttpClient httpClient);
+
     }
 
 }

core/imds/src/main/java/software/amazon/awssdk/imds/Ec2MetadataClientBuilder.java

@@ -82,17 +82,5 @@ public interface Ec2MetadataClientBuilder<B, T> extends SdkBuilder<Ec2MetadataCl
 
     EndpointMode getEndpointMode();
 
-    /**
-     * Define the token caching strategy to use when executing IMDS requests.
-     * <p>
-     * If not specified, defaults to {@link TokenCacheStrategy#NONE}, and a request to fetch a new token will be executed for each
-     * metadata request.
-     *
-     * @param tokenCacheStrategy the strategy to use for token caching
-     * @return a reference to this builder
-     */
-    B tokenCacheStrategy(TokenCacheStrategy tokenCacheStrategy);
-
-    TokenCacheStrategy getTokenCacheStrategy();
 
 }

core/imds/src/main/java/software/amazon/awssdk/imds/internal/AsyncHttpRequestHelper.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.imds.internal;
+
+import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.http.HttpResponseHandler;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
+import software.amazon.awssdk.core.internal.http.async.AsyncResponseHandler;
+import software.amazon.awssdk.core.internal.http.async.SimpleHttpContentPublisher;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.http.async.AsyncExecuteRequest;
+import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
+import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
+import software.amazon.awssdk.utils.CompletableFutureUtils;
+
+@SdkInternalApi
+final class AsyncHttpRequestHelper {
+
+    private AsyncHttpRequestHelper() {
+        // static utility class
+    }
+
+    public static CompletableFuture<String> sendAsyncMetadataRequest(SdkAsyncHttpClient httpClient,
+                                                                     SdkHttpFullRequest baseRequest,
+                                                                     CompletableFuture<?> parentFuture) {
+        StringResponseHandler stringResponseHandler = new StringResponseHandler();
+        return sendAsync(httpClient, baseRequest, stringResponseHandler, stringResponseHandler::setFuture, parentFuture);
+    }
+
+    public static CompletableFuture<Token> sendAsyncTokenRequest(Duration ttlSeconds,
+                                                                 SdkAsyncHttpClient httpClient,
+                                                                 SdkHttpFullRequest baseRequest) {
+        TokenResponseHandler tokenResponseHandler = new TokenResponseHandler(ttlSeconds.getSeconds());
+        return sendAsync(httpClient, baseRequest, tokenResponseHandler, tokenResponseHandler::setFuture, null);
+    }
+
+    static <T> CompletableFuture<T> sendAsync(SdkAsyncHttpClient client,
+                                              SdkHttpFullRequest request,
+                                              HttpResponseHandler<T> handler,
+                                              Consumer<CompletableFuture<T>> withFuture,
+                                              CompletableFuture<?> parentFuture) {
+        SdkHttpContentPublisher requestContentPublisher = new SimpleHttpContentPublisher(request);
+        TransformingAsyncResponseHandler<T> responseHandler = new AsyncResponseHandler<>(handler,
+                                                                                             Function.identity(),
+                                                                                             new ExecutionAttributes());
+        CompletableFuture<T> responseHandlerFuture = responseHandler.prepare();
+        withFuture.accept(responseHandlerFuture);
+        AsyncExecuteRequest metadataRequest = AsyncExecuteRequest.builder()
+                                                                 .request(request)
+                                                                 .requestContentPublisher(requestContentPublisher)
+                                                                 .responseHandler(responseHandler)
+                                                                 .build();
+        CompletableFuture<Void> executeFuture = client.execute(metadataRequest);
+        if (parentFuture != null) {
+            CompletableFutureUtils.forwardExceptionTo(parentFuture, executeFuture);
+            CompletableFutureUtils.forwardExceptionTo(parentFuture, responseHandlerFuture);
+        }
+        return responseHandlerFuture;
+
+    }
+
+}

core/imds/src/main/java/software/amazon/awssdk/imds/internal/AsyncTokenCache.java

@@ -0,0 +1,131 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.imds.internal;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.utils.Logger;
+
+/**
+ * A cache for the IMDS {@link Token}, which can be refreshed through an asynchronous operation.
+ * A call to the {@link AsyncTokenCache#get()} method returns an already completed future if the cached token is still fresh.
+ * If the cached token was expired, the returned future will be completed once the refresh process has been
+ * completed.
+ * In the case where multiple call to <pre>get</pre> are made while the token is expired, all CompletableFuture returned
+ * will be completed once the single refresh process completes.
+ *
+ */
+@SdkInternalApi
+final class AsyncTokenCache implements Supplier<CompletableFuture<Token>> {
+
+    private static final Logger log = Logger.loggerFor(AsyncTokenCache.class);
+
+    /**
+     * The currently cached value.
+     * Only modified through synchronized block, under the refreshLock.
+     */
+    private volatile Token cachedToken;
+
+    /**
+     * The asynchronous operation that is used to refresh the token.
+     * The Supplier must not block the current thread and is responsible to start the process that will complete the future.
+     * A call the {@link AsyncTokenCache#get} method does not join or wait for the supplied future to finish, it only refreshes
+     * the token once it finishes.
+     */
+    private final Supplier<CompletableFuture<Token>> supplier;
+
+    /**
+     * The collection of future that are waiting for the refresh call to complete. If a call to {@link AsyncTokenCache#get()}
+     * is made while the token request is happening, a future will be added to this collection. All future in this collection
+     * are completed once the token request is done.
+     * Should only be modified while holding the lock on {@link AsyncTokenCache#refreshLock}.
+     */
+    private Collection<CompletableFuture<Token>> waitingFutures = new ArrayList<>();
+
+    /**
+     * Indicates if the token refresh request is currently running or not.
+     */
+    private final AtomicBoolean refreshRunning = new AtomicBoolean(false);
+
+    private final Object refreshLock = new Object();
+
+    AsyncTokenCache(Supplier<CompletableFuture<Token>> supplier) {
+        this.supplier = supplier;
+    }
+
+    @Override
+    public CompletableFuture<Token> get() {
+        Token currentValue = cachedToken;
+        if (!needsRefresh(currentValue)) {
+            log.debug(() -> "IMDS Token is not expired");
+            return CompletableFuture.completedFuture(currentValue);
+        }
+        synchronized (refreshLock) {
+            // Make sure the value wasn't refreshed while we were waiting for the lock.
+            currentValue = cachedToken;
+            if (!needsRefresh(currentValue)) {
+                return CompletableFuture.completedFuture(currentValue);
+            }
+            CompletableFuture<Token> result = new CompletableFuture<>();
+            waitingFutures.add(result);
+            if (!refreshRunning.get()) {
+                startRefresh();
+            }
+            return result;
+        }
+    }
+
+    private void startRefresh() {
+        log.debug(() -> "IMDS token expired or null, starting asynchronous refresh.");
+        CompletableFuture<Token> tokenRequest = supplier.get();
+        refreshRunning.set(true); // After supplier.get(), in case that throws an exception
+        tokenRequest.whenComplete((token, throwable) -> {
+            Collection<CompletableFuture<Token>> toComplete;
+            synchronized (refreshLock) {
+                // Instead of completing the waiting future while holding the lock, we copy the list reference and
+                // release the lock before completing them. This is just in case someone (naughty) is doing something
+                // blocking on the complete calls. It's not good that they're doing that, but at least
+                // it won't block other threads trying to acquire the lock.
+                toComplete = waitingFutures;
+                waitingFutures = new ArrayList<>();
+                refreshRunning.set(false);
+                if (token != null) {
+                    log.debug(() -> "IMDS token refresh completed. Token value: " + token.value());
+                    cachedToken = token;
+                } else {
+                    log.error(() -> "IMDS token refresh completed with error.", throwable);
+                }
+            }
+
+            toComplete.forEach(future -> {
+                if (throwable == null) {
+                    future.complete(token);
+                } else {
+                    future.completeExceptionally(throwable);
+                }
+            });
+        });
+    }
+
+    private boolean needsRefresh(Token token) {
+        return token == null || token.isExpired();
+    }
+
+}