From 8c9533a0f81564b7c61ac0f9ce19020b453b3425 Mon Sep 17 00:00:00 2001
From: Viraj Jasani <vjasani@apache.org>
Date: Fri, 12 Aug 2022 03:59:15 -0700
Subject: [PATCH] HADOOP-18397. Shutdown AWSSecurityTokenService when its
 resources are no longer in use (#4722)

Contributed by Viraj Jasani.
---
 .../s3a/auth/MarshalledCredentialBinding.java |  8 +++---
 .../hadoop/fs/s3a/auth/STSClientFactory.java  | 13 +++------
 .../fs/s3a/ITestS3ATemporaryCredentials.java  | 27 ++++++++++---------
 3 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/auth/MarshalledCredentialBinding.java b/hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/auth/MarshalledCredentialBinding.java
index f9e840e776..29e815560a 100644
--- a/hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/auth/MarshalledCredentialBinding.java
+++ b/hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/auth/MarshalledCredentialBinding.java
@@ -207,9 +207,11 @@ public static MarshalledCredentials requestSessionCredentials(
               stsEndpoint.isEmpty() ? null : stsEndpoint,
               stsRegion)
               .build();
-      return fromSTSCredentials(
-          STSClientFactory.createClientConnection(tokenService, invoker)
-              .requestSessionCredentials(duration, TimeUnit.SECONDS));
+      try (STSClientFactory.STSClient stsClient = STSClientFactory.createClientConnection(
+          tokenService, invoker)) {
+        return fromSTSCredentials(stsClient.requestSessionCredentials(duration,
+            TimeUnit.SECONDS));
+      }
     } catch (SdkClientException e) {
       if (stsRegion.isEmpty()) {
         LOG.error("Region must be provided when requesting session credentials.",
diff --git a/hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/auth/STSClientFactory.java b/hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/auth/STSClientFactory.java
index 8c74afcc08..82d4fa5881 100644
--- a/hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/auth/STSClientFactory.java
+++ b/hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/auth/STSClientFactory.java
@@ -149,12 +149,10 @@ public static AWSSecurityTokenServiceClientBuilder builder(
    * @param tokenService STS instance
    * @param invoker invoker to use
    * @return an STS client bonded to that interface.
-   * @throws IOException on any failure
    */
   public static STSClient createClientConnection(
       final AWSSecurityTokenService tokenService,
-      final Invoker invoker)
-      throws IOException {
+      final Invoker invoker) {
     return new STSClient(tokenService, invoker);
   }
 
@@ -175,12 +173,9 @@ private STSClient(final AWSSecurityTokenService tokenService,
 
     @Override
     public void close() throws IOException {
-      try {
-        tokenService.shutdown();
-      } catch (UnsupportedOperationException ignored) {
-        // ignore this, as it is what the STS client currently
-        // does.
-      }
+      // Since we are not using AbstractAWSSecurityTokenService, we
+      // don't need to worry about catching UnsupportedOperationException.
+      tokenService.shutdown();
     }
 
     /**
diff --git a/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/ITestS3ATemporaryCredentials.java b/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/ITestS3ATemporaryCredentials.java
index 112d0fcb50..35525eb1fc 100644
--- a/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/ITestS3ATemporaryCredentials.java
+++ b/hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/ITestS3ATemporaryCredentials.java
@@ -125,13 +125,14 @@ public void testSTS() throws IOException {
         credentials,
         getStsEndpoint(conf),
         getStsRegion(conf));
-    STSClientFactory.STSClient clientConnection =
-        STSClientFactory.createClientConnection(
-            builder.build(),
-            new Invoker(new S3ARetryPolicy(conf), Invoker.LOG_EVENT));
-    Credentials sessionCreds = clientConnection
-        .requestSessionCredentials(TEST_SESSION_TOKEN_DURATION_SECONDS,
-            TimeUnit.SECONDS);
+    Credentials sessionCreds;
+    try (STSClientFactory.STSClient clientConnection =
+        STSClientFactory.createClientConnection(builder.build(),
+            new Invoker(new S3ARetryPolicy(conf), Invoker.LOG_EVENT))) {
+      sessionCreds = clientConnection
+          .requestSessionCredentials(
+              TEST_SESSION_TOKEN_DURATION_SECONDS, TimeUnit.SECONDS);
+    }
 
     // clone configuration so changes here do not affect the base FS.
     Configuration conf2 = new Configuration(conf);
@@ -379,11 +380,12 @@ public <E extends Exception> E expectedSessionRequestFailure(
             Invoker invoker = new Invoker(new S3ARetryPolicy(conf),
                 LOG_AT_ERROR);
 
-            STSClientFactory.STSClient stsClient
-                = STSClientFactory.createClientConnection(tokenService,
-                invoker);
-
-            return stsClient.requestSessionCredentials(30, TimeUnit.MINUTES);
+            try (STSClientFactory.STSClient stsClient =
+                STSClientFactory.createClientConnection(
+                    tokenService, invoker)) {
+              return stsClient.requestSessionCredentials(
+                  30, TimeUnit.MINUTES);
+            }
           });
     }
   }
@@ -413,6 +415,7 @@ public void testTemporaryCredentialValidationOnLoad() throws Throwable {
           return sc.toString();
         });
   }
+
   @Test
   public void testEmptyTemporaryCredentialValidation() throws Throwable {
     Configuration conf = new Configuration();