This is an automated email from the ASF dual-hosted git repository.
gaborgsomogyi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 42e8aed66f0 [FLINK-40855][filesystems] Shut down STS clients after S3
delegation-token acquisition
42e8aed66f0 is described below
commit 42e8aed66f067564f733f36adef43d2e6482746c
Author: Aleksandr Savonin <[email protected]>
AuthorDate: Fri Oct 2 11:53:53 2026 +0200
[FLINK-40855][filesystems] Shut down STS clients after S3 delegation-token
acquisition
---
.../token/AbstractS3DelegationTokenProvider.java | 47 ++++++---
.../AbstractS3DelegationTokenProviderTest.java | 117 +++++++++++++++++++++
2 files changed, 148 insertions(+), 16 deletions(-)
diff --git
a/flink-filesystems/flink-s3-fs-base/src/main/java/org/apache/flink/fs/s3/common/token/AbstractS3DelegationTokenProvider.java
b/flink-filesystems/flink-s3-fs-base/src/main/java/org/apache/flink/fs/s3/common/token/AbstractS3DelegationTokenProvider.java
index 82e99dcde89..a311a6aad45 100644
---
a/flink-filesystems/flink-s3-fs-base/src/main/java/org/apache/flink/fs/s3/common/token/AbstractS3DelegationTokenProvider.java
+++
b/flink-filesystems/flink-s3-fs-base/src/main/java/org/apache/flink/fs/s3/common/token/AbstractS3DelegationTokenProvider.java
@@ -19,6 +19,7 @@
package org.apache.flink.fs.s3.common.token;
import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.GlobalConfiguration;
import org.apache.flink.core.security.token.DelegationTokenProvider;
@@ -87,22 +88,36 @@ public abstract class AbstractS3DelegationTokenProvider
implements DelegationTok
public ObtainedDelegationTokens obtainDelegationTokens() throws Exception {
LOG.info("Obtaining session credentials token with access key: {}",
accessKey);
- AWSSecurityTokenService stsClient =
- AWSSecurityTokenServiceClientBuilder.standard()
- .withRegion(region)
- .withCredentials(
- new AWSStaticCredentialsProvider(
- new BasicAWSCredentials(accessKey,
secretKey)))
- .build();
- GetSessionTokenResult sessionTokenResult = stsClient.getSessionToken();
- Credentials credentials = sessionTokenResult.getCredentials();
- LOG.info(
- "Session credentials obtained successfully with access key: {}
expiration: {}",
- credentials.getAccessKeyId(),
- credentials.getExpiration());
+ final AWSSecurityTokenService stsClient = createStsClient();
+ try {
+ final GetSessionTokenResult sessionTokenResult =
stsClient.getSessionToken();
+ final Credentials credentials =
sessionTokenResult.getCredentials();
+ LOG.info(
+ "Session credentials obtained successfully with access
key: {} expiration: {}",
+ credentials.getAccessKeyId(),
+ credentials.getExpiration());
- return new ObtainedDelegationTokens(
- InstantiationUtil.serializeObject(credentials),
- Optional.of(credentials.getExpiration().getTime()));
+ return new ObtainedDelegationTokens(
+ InstantiationUtil.serializeObject(credentials),
+ Optional.of(credentials.getExpiration().getTime()));
+ } finally {
+ // Shutdown is best effort, so a failure must not discard the
obtained tokens or
+ // mask an acquisition failure.
+ try {
+ stsClient.shutdown();
+ } catch (RuntimeException e) {
+ LOG.warn("Failed to shut down STS client for {}",
serviceName(), e);
+ }
+ }
+ }
+
+ @VisibleForTesting
+ AWSSecurityTokenService createStsClient() {
+ return AWSSecurityTokenServiceClientBuilder.standard()
+ .withRegion(region)
+ .withCredentials(
+ new AWSStaticCredentialsProvider(
+ new BasicAWSCredentials(accessKey, secretKey)))
+ .build();
}
}
diff --git
a/flink-filesystems/flink-s3-fs-base/src/test/java/org/apache/flink/fs/s3/common/token/AbstractS3DelegationTokenProviderTest.java
b/flink-filesystems/flink-s3-fs-base/src/test/java/org/apache/flink/fs/s3/common/token/AbstractS3DelegationTokenProviderTest.java
index 43f3eb528f5..d4d2f697ef2 100644
---
a/flink-filesystems/flink-s3-fs-base/src/test/java/org/apache/flink/fs/s3/common/token/AbstractS3DelegationTokenProviderTest.java
+++
b/flink-filesystems/flink-s3-fs-base/src/test/java/org/apache/flink/fs/s3/common/token/AbstractS3DelegationTokenProviderTest.java
@@ -19,12 +19,21 @@
package org.apache.flink.fs.s3.common.token;
import org.apache.flink.configuration.Configuration;
+import
org.apache.flink.core.security.token.DelegationTokenProvider.ObtainedDelegationTokens;
+import org.apache.flink.util.InstantiationUtil;
+import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
+import com.amazonaws.services.securitytoken.AbstractAWSSecurityTokenService;
+import com.amazonaws.services.securitytoken.model.Credentials;
+import com.amazonaws.services.securitytoken.model.GetSessionTokenResult;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import java.util.Date;
+
import static
org.apache.flink.core.security.token.DelegationTokenProvider.CONFIG_PREFIX;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for {@link AbstractS3DelegationTokenProvider}. */
class AbstractS3DelegationTokenProviderTest {
@@ -33,6 +42,8 @@ class AbstractS3DelegationTokenProviderTest {
private static final String ACCESS_KEY_ID = "testAccessKeyId";
private static final String SECRET_ACCESS_KEY = "testSecretAccessKey";
+ private final TestingStsClient stsClient = new TestingStsClient();
+
private AbstractS3DelegationTokenProvider provider;
@BeforeEach
@@ -43,6 +54,11 @@ class AbstractS3DelegationTokenProviderTest {
public String serviceName() {
return "s3";
}
+
+ @Override
+ AWSSecurityTokenService createStsClient() {
+ return stsClient;
+ }
};
}
@@ -62,4 +78,105 @@ class AbstractS3DelegationTokenProviderTest {
assertThat(provider.delegationTokensRequired()).isTrue();
}
+
+ @Test
+ void obtainDelegationTokensShouldShutdownClientAfterSuccess() throws
Exception {
+ final Credentials credentials = createCredentials();
+ stsClient.credentials = credentials;
+
+ final ObtainedDelegationTokens tokens =
provider.obtainDelegationTokens();
+
+ final Credentials deserializedCredentials =
+ InstantiationUtil.deserializeObject(
+ tokens.getTokens(), getClass().getClassLoader());
+ assertThat(deserializedCredentials).isEqualTo(credentials);
+
assertThat(tokens.getValidUntil()).contains(credentials.getExpiration().getTime());
+ assertThat(stsClient.requestCount).isEqualTo(1);
+ assertThat(stsClient.shutdownCount).isEqualTo(1);
+ }
+
+ @Test
+ void obtainDelegationTokensShouldShutdownClientAfterRequestFailure() {
+ final RuntimeException requestFailure = new RuntimeException("STS
request failed");
+ stsClient.requestFailure = requestFailure;
+
+
assertThatThrownBy(provider::obtainDelegationTokens).isSameAs(requestFailure);
+
+ assertThat(stsClient.shutdownCount).isEqualTo(1);
+ }
+
+ @Test
+ void
obtainDelegationTokensShouldShutdownClientAfterPostprocessingFailure() {
+ stsClient.credentials = createCredentials().withExpiration(null);
+
+ assertThatThrownBy(provider::obtainDelegationTokens)
+ .isInstanceOf(NullPointerException.class);
+
+ assertThat(stsClient.shutdownCount).isEqualTo(1);
+ }
+
+ @Test
+ void obtainDelegationTokensShouldPreserveRequestFailureWhenShutdownFails()
{
+ final RuntimeException requestFailure = new RuntimeException("STS
request failed");
+ final RuntimeException shutdownFailure = new RuntimeException("STS
shutdown failed");
+ stsClient.requestFailure = requestFailure;
+ stsClient.shutdownFailure = shutdownFailure;
+
+
assertThatThrownBy(provider::obtainDelegationTokens).isSameAs(requestFailure);
+ assertThat(stsClient.shutdownCount).isEqualTo(1);
+ }
+
+ @Test
+ void obtainDelegationTokensShouldReturnTokensWhenShutdownFails() throws
Exception {
+ final Credentials credentials = createCredentials();
+ stsClient.credentials = credentials;
+ final RuntimeException shutdownFailure = new RuntimeException("STS
shutdown failed");
+ stsClient.shutdownFailure = shutdownFailure;
+
+ final ObtainedDelegationTokens tokens =
provider.obtainDelegationTokens();
+
+ final Credentials deserializedCredentials =
+ InstantiationUtil.deserializeObject(
+ tokens.getTokens(), getClass().getClassLoader());
+ assertThat(deserializedCredentials).isEqualTo(credentials);
+
assertThat(tokens.getValidUntil()).contains(credentials.getExpiration().getTime());
+ assertThat(stsClient.requestCount).isEqualTo(1);
+ assertThat(stsClient.shutdownCount).isEqualTo(1);
+ }
+
+ private static Credentials createCredentials() {
+ return new Credentials()
+ .withAccessKeyId(ACCESS_KEY_ID)
+ .withSecretAccessKey(SECRET_ACCESS_KEY)
+ .withSessionToken("testSessionToken")
+ .withExpiration(new Date(123456789L));
+ }
+
+ private static class TestingStsClient extends
AbstractAWSSecurityTokenService {
+ private Credentials credentials = createCredentials();
+ private RuntimeException requestFailure;
+ private RuntimeException shutdownFailure;
+ private int requestCount;
+ private int shutdownCount;
+
+ @Override
+ public GetSessionTokenResult getSessionToken() {
+ if (shutdownCount > 0) {
+ throw new IllegalStateException("STS client already shut
down");
+ }
+ requestCount++;
+ if (requestFailure != null) {
+ throw requestFailure;
+ }
+ return new GetSessionTokenResult().withCredentials(credentials);
+ }
+
+ @Override
+ public void shutdown() {
+ shutdownCount++;
+ if (shutdownFailure != null) {
+ throw shutdownFailure;
+ }
+ }
+ }
}
