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;
@@ -36,10 +39,19 @@
 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 LoadingCache<CredentialsKey, Optional<Credentials>> credentialsCache;
 
  @Inject
  public HttpCredentialsProvider(@ForHttpCredentialsProvider HttpClient httpClient, HttpCredentialsProviderConfig config, JsonCodec<Credentials> jsonCodec)
@@ -48,13 +60,29 @@ public HttpCredentialsProvider(@ForHttpCredentialsProvider HttpClient httpClient
  this.jsonCodec = requireNonNull(jsonCodec, "jsonCodec is null");
  this.httpCredentialsProviderEndpoint = config.getEndpoint();
  this.httpHeaders = ImmutableMap.copyOf(config.getHttpHeaders());
+ this.credentialsCache = Caffeine.newBuilder()
+ .maximumSize(config.getCacheSize())
+ .expireAfterWrite(config.getCacheTtl().toJavaTime())
+ .build(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 credentialsCache.get(new CredentialsKey(emulatedAccessKey, session));
+ }
+
+ @VisibleForTesting
+ void resetCache()
+ {
+ credentialsCache.invalidateAll();
+ credentialsCache.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;
+ }
 }