Address code review comments

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
opensearch-project · Jun 2, 2022 · b88d2a5 · b88d2a5
1 parent 2e330af
commit b88d2a5
Show file tree

Hide file tree

Showing 6 changed files with 107 additions and 25 deletions.
diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/AmazonS3Reference.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/AmazonS3Reference.java
@@ -32,17 +32,39 @@
 
 package org.opensearch.repositories.s3;
 
+import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.AmazonS3Client;
+
+import org.opensearch.common.Nullable;
 import org.opensearch.common.concurrent.RefCountedReleasable;
 
+import java.io.Closeable;
+import java.io.IOException;
+
 /**
  * Handles the shutdown of the wrapped {@link AmazonS3Client} using reference
  * counting.
  */
 public class AmazonS3Reference extends RefCountedReleasable<AmazonS3> {
-
     AmazonS3Reference(AmazonS3 client) {
-        super("AWS_S3_CLIENT", client, client::shutdown);
+        this(client, null);
+    }
+
+    AmazonS3Reference(AmazonS3WithCredentials client) {
+        this(client.client(), client.credentials());
+    }
+
+    AmazonS3Reference(AmazonS3 client, @Nullable AWSCredentialsProvider credentials) {
+        super("AWS_S3_CLIENT", client, () -> {
+            client.shutdown();
+            if (credentials instanceof Closeable) {
+                try {
+                    ((Closeable) credentials).close();
+                } catch (IOException e) {
+                    /* Do nothing here */
+                }
+            }
+        });
     }
 }
diff --git a/...s/repository-s3/src/main/java/org/opensearch/repositories/s3/AmazonS3WithCredentials.java b/...s/repository-s3/src/main/java/org/opensearch/repositories/s3/AmazonS3WithCredentials.java
@@ -0,0 +1,47 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.repositories.s3;
+
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.services.s3.AmazonS3;
+
+import org.opensearch.common.Nullable;
+
+/**
+ * The holder of the AmazonS3 and AWSCredentialsProvider
+ */
+final class AmazonS3WithCredentials {
+    private final AmazonS3 client;
+    private final AWSCredentialsProvider credentials;
+
+    private AmazonS3WithCredentials(final AmazonS3 client, @Nullable final AWSCredentialsProvider credentials) {
+        this.client = client;
+        this.credentials = credentials;
+    }
+
+    AmazonS3 getClient() {
+        return client;
+    }
+
+    AWSCredentialsProvider getCredentials() {
+        return credentials;
+    }
+
+    AmazonS3 client() {
+        return client;
+    }
+
+    AWSCredentialsProvider credentials() {
+        return credentials;
+    }
+
+    static AmazonS3WithCredentials create(final AmazonS3 client, @Nullable final AWSCredentialsProvider credentials) {
+        return new AmazonS3WithCredentials(client, credentials);
+    }
+}
diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3Service.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3Service.java
@@ -73,8 +73,6 @@
 import java.net.Socket;
 import java.security.SecureRandom;
 import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 
 import static com.amazonaws.SDKGlobalConfiguration.AWS_ROLE_ARN_ENV_VAR;
 import static com.amazonaws.SDKGlobalConfiguration.AWS_ROLE_SESSION_NAME_ENV_VAR;
@@ -85,7 +83,6 @@ class S3Service implements Closeable {
     private static final Logger logger = LogManager.getLogger(S3Service.class);
 
     private volatile Map<S3ClientSettings, AmazonS3Reference> clientsCache = emptyMap();
-    private Set<Closeable> credentialsCache = ConcurrentHashMap.newKeySet();
 
     /**
      * Client settings calculated from static configuration and settings in the keystore.
@@ -176,14 +173,10 @@ S3ClientSettings settings(RepositoryMetadata repositoryMetadata) {
     }
 
     // proxy for testing
-    AmazonS3 buildClient(final S3ClientSettings clientSettings) {
+    AmazonS3WithCredentials buildClient(final S3ClientSettings clientSettings) {
         final AmazonS3ClientBuilder builder = AmazonS3ClientBuilder.standard();
 
         final AWSCredentialsProvider credentials = buildCredentials(logger, clientSettings);
-        if (credentials instanceof Closeable) {
-            credentialsCache.add((Closeable) credentials);
-        }
-
         builder.withCredentials(credentials);
         builder.withClientConfiguration(buildConfiguration(clientSettings));
 
@@ -211,7 +204,8 @@ AmazonS3 buildClient(final S3ClientSettings clientSettings) {
         if (clientSettings.disableChunkedEncoding) {
             builder.disableChunkedEncoding();
         }
-        return SocketAccess.doPrivileged(builder::build);
+        final AmazonS3 client = SocketAccess.doPrivileged(builder::build);
+        return AmazonS3WithCredentials.create(client, credentials);
     }
 
     // pkg private for tests
@@ -350,18 +344,9 @@ private synchronized void releaseCachedClients() {
             clientReference.decRef();
         }
 
-        for (final Closeable closeable : credentialsCache) {
-            try {
-                closeable.close();
-            } catch (IOException e) {
-                /* Ignoring */
-            }
-        }
-
         // clear previously cached clients, they will be build lazily
         clientsCache = emptyMap();
         derivedClientSettings = emptyMap();
-        credentialsCache.clear();
 
         // shutdown IdleConnectionReaper background thread
         // it will be restarted on new client usage

diff --git a/...epository-s3/src/test/java/org/opensearch/repositories/s3/RepositoryCredentialsTests.java b/...epository-s3/src/test/java/org/opensearch/repositories/s3/RepositoryCredentialsTests.java
@@ -317,9 +317,10 @@ public static final class ProxyS3Service extends S3Service {
             private static final Logger logger = LogManager.getLogger(ProxyS3Service.class);
 
             @Override
-            AmazonS3 buildClient(final S3ClientSettings clientSettings) {
-                final AmazonS3 client = super.buildClient(clientSettings);
-                return new ClientAndCredentials(client, buildCredentials(logger, clientSettings));
+            AmazonS3WithCredentials buildClient(final S3ClientSettings clientSettings) {
+                final AmazonS3WithCredentials client = super.buildClient(clientSettings);
+                final AWSCredentialsProvider credentials = buildCredentials(logger, clientSettings);
+                return AmazonS3WithCredentials.create(new ClientAndCredentials(client.client(), credentials), credentials);
             }
 
         }

diff --git a/...ins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3ClientSettingsTests.java b/...ins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3ClientSettingsTests.java
@@ -246,7 +246,7 @@ public void testRegionCanBeSet() {
         assertThat(settings.get("default").region, is(""));
         assertThat(settings.get("other").region, is(region));
         try (S3Service s3Service = new S3Service()) {
-            AmazonS3Client other = (AmazonS3Client) s3Service.buildClient(settings.get("other"));
+            AmazonS3Client other = (AmazonS3Client) s3Service.buildClient(settings.get("other")).client();
             assertThat(other.getSignerRegionOverride(), is(region));
         }
     }

diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3ServiceTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3ServiceTests.java
@@ -32,10 +32,12 @@
 package org.opensearch.repositories.s3;
 
 import org.opensearch.cluster.metadata.RepositoryMetadata;
-
+import org.opensearch.common.settings.MockSecureSettings;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.util.Map;
+
 public class S3ServiceTests extends OpenSearchTestCase {
 
     public void testCachedClientsAreReleased() {
@@ -56,4 +58,29 @@ public void testCachedClientsAreReleased() {
         final S3ClientSettings clientSettingsReloaded = s3Service.settings(metadata1);
         assertNotSame(clientSettings, clientSettingsReloaded);
     }
+
+    public void testCachedClientsWithCredentialsAreReleased() {
+        final MockSecureSettings secureSettings = new MockSecureSettings();
+        secureSettings.setString("s3.client.default.role_arn", "role");
+        final Map<String, S3ClientSettings> defaults = S3ClientSettings.load(
+            Settings.builder().setSecureSettings(secureSettings).put("s3.client.default.identity_token_file", "file").build()
+        );
+        final S3Service s3Service = new S3Service();
+        s3Service.refreshAndClearCache(defaults);
+        final Settings settings = Settings.builder().put("endpoint", "http://first").put("region", "us-east-2").build();
+        final RepositoryMetadata metadata1 = new RepositoryMetadata("first", "s3", settings);
+        final RepositoryMetadata metadata2 = new RepositoryMetadata("second", "s3", settings);
+        final S3ClientSettings clientSettings = s3Service.settings(metadata2);
+        final S3ClientSettings otherClientSettings = s3Service.settings(metadata2);
+        assertSame(clientSettings, otherClientSettings);
+        final AmazonS3Reference reference = s3Service.client(metadata1);
+        reference.close();
+        s3Service.close();
+        final AmazonS3Reference referenceReloaded = s3Service.client(metadata1);
+        assertNotSame(referenceReloaded, reference);
+        referenceReloaded.close();
+        s3Service.close();
+        final S3ClientSettings clientSettingsReloaded = s3Service.settings(metadata1);
+        assertNotSame(clientSettings, clientSettingsReloaded);
+    }
 }