Support prefetch & stale time for S3 web identity provider

Author: jirislav

docs/src/main/sphinx/object-storage/file-system-s3.md

@@ -110,6 +110,21 @@ support:
     Trino on Amazon EKS and using [IAM roles for service accounts
     (IRSA)](https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts.html)
     Defaults to `false`.
+* - `s3.web-identity-token-credentials-prefetch-time`
+  - 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.
+    By default, [AWS SDK v2
+    defaults](https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/sts/auth/StsCredentialsProvider.BaseBuilder.html)
+    are used.
+* - `s3.web-identity-token-credentials-stale-time`
+  - 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.
+    By default, [AWS SDK v2
+    defaults](https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/sts/auth/StsCredentialsProvider.BaseBuilder.html)
+    are used.
 * - `s3.application-id`
   - Specify the application identifier appended to the `User-Agent` header 
     for all requests sent to S3. Defaults to `Trino`.

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 Duration webIdentityTokenCredentialsPrefetchTime;
+    private Duration webIdentityTokenCredentialsStaleTime;
     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 Optional<Duration> getWebIdentityTokenCredentialsPrefetchTime()
+    {
+        return Optional.ofNullable(webIdentityTokenCredentialsPrefetchTime);
+    }
+
+    @Config("s3.web-identity-token-credentials-prefetch-time")
+    @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 setWebIdentityTokenCredentialsPrefetchTime(Duration webIdentityTokenCredentialsPrefetchTime)
+    {
+        this.webIdentityTokenCredentialsPrefetchTime = webIdentityTokenCredentialsPrefetchTime;
+        return this;
+    }
+
+    public Optional<Duration> getWebIdentityTokenCredentialsStaleTime()
+    {
+        return Optional.ofNullable(webIdentityTokenCredentialsStaleTime);
+    }
+
+    @Config("s3.web-identity-token-credentials-stale-time")
+    @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 setWebIdentityTokenCredentialsStaleTime(Duration webIdentityTokenCredentialsStaleTime)
+    {
+        this.webIdentityTokenCredentialsStaleTime = webIdentityTokenCredentialsStaleTime;
+        return this;
+    }
+
     public String getSseCustomerKey()
     {
         return sseCustomerKey;

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

@@ -14,6 +14,7 @@
 package io.trino.filesystem.s3;
 
 import com.google.inject.Inject;
+import io.airlift.units.Duration;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkTelemetry;
 import io.trino.filesystem.Location;
@@ -162,6 +163,8 @@ private static S3ClientFactory s3ClientFactory(SdkHttpClient httpClient, OpenTel
         Optional<String> staticEndpoint = Optional.ofNullable(config.getEndpoint());
         boolean pathStyleAccess = config.isPathStyleAccess();
         boolean useWebIdentityTokenCredentialsProvider = config.isUseWebIdentityTokenCredentialsProvider();
+        Optional<Duration> webIdentityTokenCredentialsPrefetchTime = config.getWebIdentityTokenCredentialsPrefetchTime();
+        Optional<Duration> webIdentityTokenCredentialsStaleTime = config.getWebIdentityTokenCredentialsStaleTime();
         Optional<String> staticIamRole = Optional.ofNullable(config.getIamRole());
         String staticRoleSessionName = config.getRoleSessionName();
         String externalId = config.getExternalId();
@@ -190,9 +193,10 @@ private static S3ClientFactory s3ClientFactory(SdkHttpClient httpClient, OpenTel
             s3.forcePathStyle(pathStyleAccess);
 
             if (useWebIdentityTokenCredentialsProvider) {
-                s3.credentialsProvider(WebIdentityTokenFileCredentialsProvider.builder()
-                        .asyncCredentialUpdateEnabled(true)
-                        .build());
+                WebIdentityTokenFileCredentialsProvider.Builder builder = WebIdentityTokenFileCredentialsProvider.builder().asyncCredentialUpdateEnabled(true);
+                webIdentityTokenCredentialsPrefetchTime.ifPresent(duration -> builder.prefetchTime(duration.toJavaTime()));
+                webIdentityTokenCredentialsStaleTime.ifPresent(duration -> builder.staleTime(duration.toJavaTime()));
+                s3.credentialsProvider(builder.build());
             }
             else if (iamRole.isPresent()) {
                 s3.credentialsProvider(StsAssumeRoleCredentialsProvider.builder()
@@ -219,6 +223,8 @@ private static S3Presigner s3PreSigner(SdkHttpClient httpClient, OpenTelemetry o
         Optional<String> staticEndpoint = Optional.ofNullable(config.getEndpoint());
         boolean pathStyleAccess = config.isPathStyleAccess();
         boolean useWebIdentityTokenCredentialsProvider = config.isUseWebIdentityTokenCredentialsProvider();
+        Optional<Duration> webIdentityTokenCredentialsPrefetchTime = config.getWebIdentityTokenCredentialsPrefetchTime();
+        Optional<Duration> webIdentityTokenCredentialsStaleTime = config.getWebIdentityTokenCredentialsStaleTime();
         Optional<String> staticIamRole = Optional.ofNullable(config.getIamRole());
         String staticRoleSessionName = config.getRoleSessionName();
         String externalId = config.getExternalId();
@@ -234,9 +240,10 @@ private static S3Presigner s3PreSigner(SdkHttpClient httpClient, OpenTelemetry o
                 .build());
 
         if (useWebIdentityTokenCredentialsProvider) {
-            s3.credentialsProvider(WebIdentityTokenFileCredentialsProvider.builder()
-                    .asyncCredentialUpdateEnabled(true)
-                    .build());
+            WebIdentityTokenFileCredentialsProvider.Builder builder = WebIdentityTokenFileCredentialsProvider.builder().asyncCredentialUpdateEnabled(true);
+            webIdentityTokenCredentialsPrefetchTime.ifPresent(duration -> builder.prefetchTime(duration.toJavaTime()));
+            webIdentityTokenCredentialsStaleTime.ifPresent(duration -> builder.staleTime(duration.toJavaTime()));
+            s3.credentialsProvider(builder.build());
         }
         else if (staticIamRole.isPresent()) {
             s3.credentialsProvider(StsAssumeRoleCredentialsProvider.builder()

lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemConfig.java

@@ -60,6 +60,8 @@ public void testDefaults()
                 .setMaxErrorRetries(20)
                 .setSseKmsKeyId(null)
                 .setUseWebIdentityTokenCredentialsProvider(false)
+                .setWebIdentityTokenCredentialsPrefetchTime(null)
+                .setWebIdentityTokenCredentialsStaleTime(null)
                 .setSseCustomerKey(null)
                 .setStreamingPartSize(DataSize.of(32, MEGABYTE))
                 .setRequesterPays(false)
@@ -102,6 +104,8 @@ public void testExplicitPropertyMappings()
                 .put("s3.sse.kms-key-id", "mykey")
                 .put("s3.sse.customer-key", "customerKey")
                 .put("s3.use-web-identity-token-credentials-provider", "true")
+                .put("s3.web-identity-token-credentials-prefetch-time", "10m")
+                .put("s3.web-identity-token-credentials-stale-time", "5m")
                 .put("s3.streaming.part-size", "42MB")
                 .put("s3.requester-pays", "true")
                 .put("s3.max-connections", "42")
@@ -140,6 +144,8 @@ public void testExplicitPropertyMappings()
                 .setSseType(S3SseType.KMS)
                 .setSseKmsKeyId("mykey")
                 .setUseWebIdentityTokenCredentialsProvider(true)
+                .setWebIdentityTokenCredentialsPrefetchTime(new Duration(10, MINUTES))
+                .setWebIdentityTokenCredentialsStaleTime(new Duration(5, MINUTES))
                 .setSseCustomerKey("customerKey")
                 .setRequesterPays(true)
                 .setMaxConnections(42)