HADOOP-18397. Shutdown AWSSecurityTokenService when its resources are no longer in use (#4722)

Contributed by Viraj Jasani.
This commit is contained in:
Viraj Jasani 2022-08-12 03:59:15 -07:00 committed by Steve Loughran
parent 93c4704b33
commit 6b7c1329b2
No known key found for this signature in database
GPG Key ID: D22CF846DBB162A0
3 changed files with 24 additions and 24 deletions

View File

@ -207,9 +207,11 @@ public static MarshalledCredentials requestSessionCredentials(
stsEndpoint.isEmpty() ? null : stsEndpoint, stsEndpoint.isEmpty() ? null : stsEndpoint,
stsRegion) stsRegion)
.build(); .build();
return fromSTSCredentials( try (STSClientFactory.STSClient stsClient = STSClientFactory.createClientConnection(
STSClientFactory.createClientConnection(tokenService, invoker) tokenService, invoker)) {
.requestSessionCredentials(duration, TimeUnit.SECONDS)); return fromSTSCredentials(stsClient.requestSessionCredentials(duration,
TimeUnit.SECONDS));
}
} catch (SdkClientException e) { } catch (SdkClientException e) {
if (stsRegion.isEmpty()) { if (stsRegion.isEmpty()) {
LOG.error("Region must be provided when requesting session credentials.", LOG.error("Region must be provided when requesting session credentials.",

View File

@ -149,12 +149,10 @@ public static AWSSecurityTokenServiceClientBuilder builder(
* @param tokenService STS instance * @param tokenService STS instance
* @param invoker invoker to use * @param invoker invoker to use
* @return an STS client bonded to that interface. * @return an STS client bonded to that interface.
* @throws IOException on any failure
*/ */
public static STSClient createClientConnection( public static STSClient createClientConnection(
final AWSSecurityTokenService tokenService, final AWSSecurityTokenService tokenService,
final Invoker invoker) final Invoker invoker) {
throws IOException {
return new STSClient(tokenService, invoker); return new STSClient(tokenService, invoker);
} }
@ -175,12 +173,9 @@ private STSClient(final AWSSecurityTokenService tokenService,
@Override @Override
public void close() throws IOException { public void close() throws IOException {
try { // Since we are not using AbstractAWSSecurityTokenService, we
// don't need to worry about catching UnsupportedOperationException.
tokenService.shutdown(); tokenService.shutdown();
} catch (UnsupportedOperationException ignored) {
// ignore this, as it is what the STS client currently
// does.
}
} }
/** /**

View File

@ -125,13 +125,14 @@ public void testSTS() throws IOException {
credentials, credentials,
getStsEndpoint(conf), getStsEndpoint(conf),
getStsRegion(conf)); getStsRegion(conf));
STSClientFactory.STSClient clientConnection = Credentials sessionCreds;
STSClientFactory.createClientConnection( try (STSClientFactory.STSClient clientConnection =
builder.build(), STSClientFactory.createClientConnection(builder.build(),
new Invoker(new S3ARetryPolicy(conf), Invoker.LOG_EVENT)); new Invoker(new S3ARetryPolicy(conf), Invoker.LOG_EVENT))) {
Credentials sessionCreds = clientConnection sessionCreds = clientConnection
.requestSessionCredentials(TEST_SESSION_TOKEN_DURATION_SECONDS, .requestSessionCredentials(
TimeUnit.SECONDS); TEST_SESSION_TOKEN_DURATION_SECONDS, TimeUnit.SECONDS);
}
// clone configuration so changes here do not affect the base FS. // clone configuration so changes here do not affect the base FS.
Configuration conf2 = new Configuration(conf); Configuration conf2 = new Configuration(conf);
@ -379,11 +380,12 @@ public <E extends Exception> E expectedSessionRequestFailure(
Invoker invoker = new Invoker(new S3ARetryPolicy(conf), Invoker invoker = new Invoker(new S3ARetryPolicy(conf),
LOG_AT_ERROR); LOG_AT_ERROR);
STSClientFactory.STSClient stsClient try (STSClientFactory.STSClient stsClient =
= STSClientFactory.createClientConnection(tokenService, STSClientFactory.createClientConnection(
invoker); tokenService, invoker)) {
return stsClient.requestSessionCredentials(
return stsClient.requestSessionCredentials(30, TimeUnit.MINUTES); 30, TimeUnit.MINUTES);
}
}); });
} }
} }
@ -413,6 +415,7 @@ public void testTemporaryCredentialValidationOnLoad() throws Throwable {
return sc.toString(); return sc.toString();
}); });
} }
@Test @Test
public void testEmptyTemporaryCredentialValidation() throws Throwable { public void testEmptyTemporaryCredentialValidation() throws Throwable {
Configuration conf = new Configuration(); Configuration conf = new Configuration();