Add configurable in-memory caching to the HTTP Credentials Provider

pom.xml

@@ -58,6 +58,7 @@
         <dep.minio.version>8.5.9</dep.minio.version>
         <dep.docker.version>3.3.6</dep.docker.version>
         <dep.awaitility.version>4.1.1</dep.awaitility.version>
+        <dep.caffeine.version>3.1.8</dep.caffeine.version>
     </properties>
 
     <dependencyManagement>
@@ -108,6 +109,18 @@
                 <version>${project.version}</version>
             </dependency>
 
+            <dependency>
+                <groupId>com.github.ben-manes.caffeine</groupId>
+                <artifactId>caffeine</artifactId>
+                <version>${dep.caffeine.version}</version>
+                <exclusions>
+                    <exclusion>
+                        <groupId>org.checkerframework</groupId>
+                        <artifactId>checker-qual</artifactId>
+                    </exclusion>
+                </exclusions>
+            </dependency>
+
             <dependency>
                 <groupId>com.github.docker-java</groupId>
                 <artifactId>docker-java-api</artifactId>

trino-aws-proxy/pom.xml

@@ -44,6 +44,11 @@
             <artifactId>jackson-dataformat-xml</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.github.ben-manes.caffeine</groupId>
+            <artifactId>caffeine</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>com.google.guava</groupId>
             <artifactId>guava</artifactId>

trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/credentials/http/HttpCredentialsModule.java

@@ -37,7 +37,6 @@ protected void setup(Binder binder)
                 HttpCredentialsProvider.class,
                 innerBinder -> {
                     configBinder(innerBinder).bindConfig(HttpCredentialsProviderConfig.class);
-                    innerBinder.bind(HttpCredentialsProvider.class);
                     httpClientBinder(innerBinder).bindHttpClient(HTTP_CREDENTIALS_PROVIDER_HTTP_CLIENT_NAME, ForHttpCredentialsProvider.class);
                     jsonCodecBinder(innerBinder).bindJsonCodec(Credentials.class);
                 }));

trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/credentials/http/HttpCredentialsProvider.java

@@ -13,6 +13,9 @@
  */
 package io.trino.aws.proxy.server.credentials.http;
 
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.LoadingCache;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Multimaps;
 import com.google.inject.Inject;
@@ -28,6 +31,7 @@
 import java.net.URI;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 
 import static io.airlift.http.client.FullJsonResponseHandler.createFullJsonResponseHandler;
 import static io.airlift.http.client.Request.Builder.prepareGet;
@@ -36,10 +40,20 @@
 public class HttpCredentialsProvider
         implements CredentialsProvider
 {
+    private record CredentialsKey(String emulatedAccessKey, Optional<String> session)
+    {
+        private CredentialsKey {
+            requireNonNull(emulatedAccessKey, "emulatedAccessKey is null");
+            requireNonNull(session, "session is null");
+        }
+    }
+
     private final HttpClient httpClient;
     private final JsonCodec<Credentials> jsonCodec;
     private final URI httpCredentialsProviderEndpoint;
     private final Map<String, String> httpHeaders;
+    private final Optional<LoadingCache<CredentialsKey, Optional<Credentials>>> credentialsCache;
+    private final Function<CredentialsKey, Optional<Credentials>> credentialsFetcher;
 
     @Inject
     public HttpCredentialsProvider(@ForHttpCredentialsProvider HttpClient httpClient, HttpCredentialsProviderConfig config, JsonCodec<Credentials> jsonCodec)
@@ -48,13 +62,39 @@ public HttpCredentialsProvider(@ForHttpCredentialsProvider HttpClient httpClient
         this.jsonCodec = requireNonNull(jsonCodec, "jsonCodec is null");
         this.httpCredentialsProviderEndpoint = config.getEndpoint();
         this.httpHeaders = ImmutableMap.copyOf(config.getHttpHeaders());
+        if (config.getCacheSize() > 0 && config.getCacheTtl().toMillis() > 0) {
+            LoadingCache<CredentialsKey, Optional<Credentials>> cache = Caffeine.newBuilder()
+                    .maximumSize(config.getCacheSize())
+                    .expireAfterWrite(config.getCacheTtl().toJavaTime())
+                    .build(this::fetchCredentials);
+            this.credentialsCache = Optional.of(cache);
+            this.credentialsFetcher = cache::get;
+        }
+        else {
+            this.credentialsCache = Optional.empty();
+            this.credentialsFetcher = this::fetchCredentials;
+        }
     }
 
     @Override
     public Optional<Credentials> credentials(String emulatedAccessKey, Optional<String> session)
     {
-        UriBuilder uriBuilder = UriBuilder.fromUri(httpCredentialsProviderEndpoint).path(emulatedAccessKey);
-        session.ifPresent(sessionToken -> uriBuilder.queryParam("sessionToken", sessionToken));
+        return credentialsFetcher.apply(new CredentialsKey(emulatedAccessKey, session));
+    }
+
+    @VisibleForTesting
+    void resetCache()
+    {
+        credentialsCache.ifPresent(instantiatedCache -> {
+            instantiatedCache.invalidateAll();
+            instantiatedCache.cleanUp();
+        });
+    }
+
+    private Optional<Credentials> fetchCredentials(CredentialsKey credentialsKey)
+    {
+        UriBuilder uriBuilder = UriBuilder.fromUri(httpCredentialsProviderEndpoint).path(credentialsKey.emulatedAccessKey());
+        credentialsKey.session().ifPresent(sessionToken -> uriBuilder.queryParam("sessionToken", sessionToken));
         Request.Builder requestBuilder = prepareGet()
                 .addHeaders(Multimaps.forMap(httpHeaders))
                 .setUri(uriBuilder.build());

trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/credentials/http/HttpCredentialsProviderConfig.java

@@ -14,7 +14,12 @@
 package io.trino.aws.proxy.server.credentials.http;
 
 import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableMap;
 import io.airlift.configuration.Config;
+import io.airlift.configuration.ConfigDescription;
+import io.airlift.units.Duration;
+import io.airlift.units.MinDuration;
+import jakarta.validation.constraints.Min;
 import jakarta.validation.constraints.NotNull;
 
 import java.net.URI;
@@ -25,7 +30,9 @@
 public class HttpCredentialsProviderConfig
 {
     private URI endpoint;
-    private Map<String, String> httpHeaders = Map.of();
+    private Map<String, String> httpHeaders = ImmutableMap.of();
+    private long cacheSize;
+    private Duration cacheTtl = Duration.ZERO;
 
     @NotNull
     public URI getEndpoint()
@@ -34,9 +41,10 @@ public URI getEndpoint()
     }
 
     @Config("credentials-provider.http.endpoint")
-    public HttpCredentialsProviderConfig setEndpoint(String endpoint)
+    @ConfigDescription("URL to retrieve credentials from, the username will be passed as a path under this URL")
+    public HttpCredentialsProviderConfig setEndpoint(URI endpoint)
     {
-        this.endpoint = URI.create(endpoint);
+        this.endpoint = endpoint;
         return this;
     }
 
@@ -46,6 +54,9 @@ public Map<String, String> getHttpHeaders()
     }
 
     @Config("credentials-provider.http.headers")
+    @ConfigDescription(
+            "Additional headers to include in requests, in the format header-1-name:header-1-value,header-2-name:header-2-value. " +
+                    "If a header value needs to include a comma, it should be doubled")
     public HttpCredentialsProviderConfig setHttpHeaders(String httpHeadersList)
     {
         try {
@@ -62,4 +73,32 @@ public HttpCredentialsProviderConfig setHttpHeaders(String httpHeadersList)
         }
         return this;
     }
+
+    @Config("credentials-provider.http.cache-size")
+    @ConfigDescription("In-memory cache size for the credentials provider, defaults to 0 (no caching)")
+    public HttpCredentialsProviderConfig setCacheSize(long cacheSize)
+    {
+        this.cacheSize = cacheSize;
+        return this;
+    }
+
+    @Min(0)
+    public long getCacheSize()
+    {
+        return cacheSize;
+    }
+
+    @Config("credentials-provider.http.cache-ttl")
+    @ConfigDescription("In-memory cache TTL for the credentials provider, defaults to 0 seconds (no caching)")
+    public HttpCredentialsProviderConfig setCacheTtl(Duration cacheTtl)
+    {
+        this.cacheTtl = cacheTtl;
+        return this;
+    }
+
+    @MinDuration("0s")
+    public Duration getCacheTtl()
+    {
+        return cacheTtl;
+    }
 }