Support prefetch & stale time for S3 web identity provider

Author: jirislav

lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemConfig.java

@@ -159,6 +159,8 @@ public static RetryStrategy getRetryStrategy(RetryMode retryMode)
     private String sseKmsKeyId;
     private String sseCustomerKey;
     private boolean useWebIdentityTokenCredentialsProvider;
+    private int webIdentityTokenCredentialsPrefetchTimeSeconds = 300;
+    private int webIdentityTokenCredentialsStaleTimeSeconds = 60;
     private SignerType signerType;
     private DataSize streamingPartSize = DataSize.of(32, MEGABYTE);
     private boolean requesterPays;
@@ -397,6 +399,32 @@ public S3FileSystemConfig setUseWebIdentityTokenCredentialsProvider(boolean useW
         return this;
     }
 
+    public int getWebIdentityTokenCredentialsPrefetchTimeSeconds()
+    {
+        return webIdentityTokenCredentialsPrefetchTimeSeconds;
+    }
+
+    @Config("s3.web-identity-token-credentials-prefetch-time-seconds")
+    @ConfigDescription("Configure the amount of time, relative to STS token expiration, that the cached credentials are considered close to stale and should be updated. Prefetch updates will occur between the specified time and the stale time of the provider. Prefetch updates are asynchronous.")
+    public S3FileSystemConfig setWebIdentityTokenCredentialsPrefetchTimeSeconds(int webIdentityTokenCredentialsPrefetchTimeSeconds)
+    {
+        this.webIdentityTokenCredentialsPrefetchTimeSeconds = webIdentityTokenCredentialsPrefetchTimeSeconds;
+        return this;
+    }
+
+    public int getWebIdentityTokenCredentialsStaleTimeSeconds()
+    {
+        return webIdentityTokenCredentialsStaleTimeSeconds;
+    }
+
+    @Config("s3.web-identity-token-credentials-stale-time-seconds")
+    @ConfigDescription("Configure the amount of time, relative to STS token expiration, that the cached credentials are considered stale and must be updated. All threads using S3 client will block until the value is updated.")
+    public S3FileSystemConfig setWebIdentityTokenCredentialsStaleTimeSeconds(int webIdentityTokenCredentialsStaleTimeSeconds)
+    {
+        this.webIdentityTokenCredentialsStaleTimeSeconds = webIdentityTokenCredentialsStaleTimeSeconds;
+        return this;
+    }
+
     public String getSseCustomerKey()
     {
         return sseCustomerKey;

lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemLoader.java

@@ -41,6 +41,8 @@
 import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
 
 import java.net.URI;
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
@@ -162,6 +164,8 @@ private static S3ClientFactory s3ClientFactory(SdkHttpClient httpClient, OpenTel
         Optional<String> staticEndpoint = Optional.ofNullable(config.getEndpoint());
         boolean pathStyleAccess = config.isPathStyleAccess();
         boolean useWebIdentityTokenCredentialsProvider = config.isUseWebIdentityTokenCredentialsProvider();
+        Duration webIdentityTokenCredentialsPrefetchTime = Duration.of(config.getWebIdentityTokenCredentialsPrefetchTimeSeconds(), ChronoUnit.SECONDS);
+        Duration webIdentityTokenCredentialsStaleTime = Duration.of(config.getWebIdentityTokenCredentialsStaleTimeSeconds(), ChronoUnit.SECONDS);
         Optional<String> staticIamRole = Optional.ofNullable(config.getIamRole());
         String staticRoleSessionName = config.getRoleSessionName();
         String externalId = config.getExternalId();
@@ -192,6 +196,8 @@ private static S3ClientFactory s3ClientFactory(SdkHttpClient httpClient, OpenTel
             if (useWebIdentityTokenCredentialsProvider) {
                 s3.credentialsProvider(WebIdentityTokenFileCredentialsProvider.builder()
                         .asyncCredentialUpdateEnabled(true)
+                        .prefetchTime(webIdentityTokenCredentialsPrefetchTime)
+                        .staleTime(webIdentityTokenCredentialsStaleTime)
                         .build());
             }
             else if (iamRole.isPresent()) {
@@ -219,6 +225,8 @@ private static S3Presigner s3PreSigner(SdkHttpClient httpClient, OpenTelemetry o
         Optional<String> staticEndpoint = Optional.ofNullable(config.getEndpoint());
         boolean pathStyleAccess = config.isPathStyleAccess();
         boolean useWebIdentityTokenCredentialsProvider = config.isUseWebIdentityTokenCredentialsProvider();
+        Duration webIdentityTokenCredentialsPrefetchTime = Duration.of(config.getWebIdentityTokenCredentialsPrefetchTimeSeconds(), ChronoUnit.SECONDS);
+        Duration webIdentityTokenCredentialsStaleTime = Duration.of(config.getWebIdentityTokenCredentialsStaleTimeSeconds(), ChronoUnit.SECONDS);
         Optional<String> staticIamRole = Optional.ofNullable(config.getIamRole());
         String staticRoleSessionName = config.getRoleSessionName();
         String externalId = config.getExternalId();
@@ -236,6 +244,8 @@ private static S3Presigner s3PreSigner(SdkHttpClient httpClient, OpenTelemetry o
         if (useWebIdentityTokenCredentialsProvider) {
             s3.credentialsProvider(WebIdentityTokenFileCredentialsProvider.builder()
                     .asyncCredentialUpdateEnabled(true)
+                    .prefetchTime(webIdentityTokenCredentialsPrefetchTime)
+                    .staleTime(webIdentityTokenCredentialsStaleTime)
                     .build());
         }
         else if (staticIamRole.isPresent()) {