Attempt to fix ThreadLeakControl issues in S3BlobContainerRetriesTests (#18201)

The Default aws client might create its own ScheduledExecutorService
which is closed via ExecutorService.shutdown() when the client is
closed. Calling shutdown might not ensure that all threads are gone
when ThreadLeakControl is ran.
Provide our own ScheduledExecutorService during this test so that it
runs thrown our tearDown method which calls ThreadPool#terminate waiting
for 5 seconds.

Signed-off-by: David Causse <dcausse@wikimedia.org>

Author: nomoa

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3AsyncService.java

@@ -52,6 +52,7 @@
 import java.nio.file.Path;
 import java.time.Duration;
 import java.util.Map;
+import java.util.concurrent.ScheduledExecutorService;
 
 import static java.util.Collections.emptyMap;
 
@@ -75,10 +76,20 @@ class S3AsyncService implements Closeable {
      */
     private volatile Map<Settings, S3ClientSettings> derivedClientSettings = emptyMap();
 
-    S3AsyncService(final Path configPath) {
+    /**
+     * Optional scheduled executor service to use for the client
+     */
+    private final @Nullable ScheduledExecutorService clientExecutorService;
+
+    S3AsyncService(final Path configPath, @Nullable ScheduledExecutorService clientExecutorService) {
         staticClientSettings = MapBuilder.<String, S3ClientSettings>newMapBuilder()
             .put("default", S3ClientSettings.getClientSettings(Settings.EMPTY, "default", configPath))
             .immutableMap();
+        this.clientExecutorService = clientExecutorService;
+    }
+
+    S3AsyncService(final Path configPath) {
+        this(configPath, null);
     }
 
     /**
@@ -173,7 +184,7 @@ synchronized AmazonAsyncS3WithCredentials buildClient(
     ) {
         setDefaultAwsProfilePath();
         final S3AsyncClientBuilder builder = S3AsyncClient.builder();
-        builder.overrideConfiguration(buildOverrideConfiguration(clientSettings));
+        builder.overrideConfiguration(buildOverrideConfiguration(clientSettings, clientExecutorService));
         final AwsCredentialsProvider credentials = buildCredentials(logger, clientSettings);
         builder.credentialsProvider(credentials);
 
@@ -234,7 +245,10 @@ synchronized AmazonAsyncS3WithCredentials buildClient(
         return AmazonAsyncS3WithCredentials.create(client, priorityClient, urgentClient, credentials);
     }
 
-    static ClientOverrideConfiguration buildOverrideConfiguration(final S3ClientSettings clientSettings) {
+    static ClientOverrideConfiguration buildOverrideConfiguration(
+        final S3ClientSettings clientSettings,
+        ScheduledExecutorService clientExecutorService
+    ) {
         RetryPolicy retryPolicy = SocketAccess.doPrivileged(
             () -> RetryPolicy.builder()
                 .numRetries(clientSettings.maxRetries)
@@ -243,11 +257,12 @@ static ClientOverrideConfiguration buildOverrideConfiguration(final S3ClientSett
                 )
                 .build()
         );
+        ClientOverrideConfiguration.Builder builder = ClientOverrideConfiguration.builder();
+        if (clientExecutorService != null) {
+            builder = builder.scheduledExecutorService(clientExecutorService);
+        }
 
-        return ClientOverrideConfiguration.builder()
-            .retryPolicy(retryPolicy)
-            .apiCallAttemptTimeout(Duration.ofMillis(clientSettings.requestTimeoutMillis))
-            .build();
+        return builder.retryPolicy(retryPolicy).apiCallAttemptTimeout(Duration.ofMillis(clientSettings.requestTimeoutMillis)).build();
     }
 
     // pkg private for tests

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3Service.java

@@ -92,6 +92,7 @@
 import java.time.Duration;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledExecutorService;
 
 import static java.util.Collections.emptyMap;
 
@@ -115,10 +116,20 @@ class S3Service implements Closeable {
      */
     private volatile Map<Settings, S3ClientSettings> derivedClientSettings = new ConcurrentHashMap<>();
 
-    S3Service(final Path configPath) {
+    /**
+     * Optional scheduled executor service to use for the client
+     */
+    private final @Nullable ScheduledExecutorService clientExecutorService;
+
+    S3Service(final Path configPath, @Nullable ScheduledExecutorService clientExecutorService) {
         staticClientSettings = MapBuilder.<String, S3ClientSettings>newMapBuilder()
             .put("default", S3ClientSettings.getClientSettings(Settings.EMPTY, "default", configPath))
             .immutableMap();
+        this.clientExecutorService = clientExecutorService;
+    }
+
+    S3Service(final Path configPath) {
+        this(configPath, null);
     }
 
     /**
@@ -204,7 +215,7 @@ AmazonS3WithCredentials buildClient(final S3ClientSettings clientSettings) {
         final AwsCredentialsProvider credentials = buildCredentials(logger, clientSettings);
         builder.credentialsProvider(credentials);
         builder.httpClientBuilder(buildHttpClient(clientSettings));
-        builder.overrideConfiguration(buildOverrideConfiguration(clientSettings));
+        builder.overrideConfiguration(buildOverrideConfiguration(clientSettings, clientExecutorService));
 
         String endpoint = Strings.hasLength(clientSettings.endpoint) ? clientSettings.endpoint : DEFAULT_S3_ENDPOINT;
         if ((endpoint.startsWith("http://") || endpoint.startsWith("https://")) == false) {
@@ -318,8 +329,14 @@ static ProxyConfiguration buildHttpProxyConfiguration(S3ClientSettings clientSet
         return proxyConfiguration.build();
     }
 
-    static ClientOverrideConfiguration buildOverrideConfiguration(final S3ClientSettings clientSettings) {
+    static ClientOverrideConfiguration buildOverrideConfiguration(
+        final S3ClientSettings clientSettings,
+        @Nullable ScheduledExecutorService clientExecutorService
+    ) {
         ClientOverrideConfiguration.Builder clientOverrideConfiguration = ClientOverrideConfiguration.builder();
+        if (clientExecutorService != null) {
+            clientOverrideConfiguration = clientOverrideConfiguration.scheduledExecutorService(clientExecutorService);
+        }
         if (Strings.hasLength(clientSettings.signerOverride)) {
             clientOverrideConfiguration = clientOverrideConfiguration.putAdvancedOption(
                 SdkAdvancedClientOption.SIGNER,

plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/AwsS3ServiceImplTests.java

@@ -358,7 +358,7 @@ private void launchAWSConfigurationTest(
             assertThat(proxyConfiguration.password(), is(expectedProxyPassword));
         }
 
-        final ClientOverrideConfiguration clientOverrideConfiguration = S3Service.buildOverrideConfiguration(clientSettings);
+        final ClientOverrideConfiguration clientOverrideConfiguration = S3Service.buildOverrideConfiguration(clientSettings, null);
 
         assertTrue(clientOverrideConfiguration.retryPolicy().isPresent());
         assertThat(clientOverrideConfiguration.retryPolicy().get().numRetries(), is(expectedMaxRetries));

plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerRetriesTests.java

@@ -130,15 +130,15 @@ public class S3BlobContainerRetriesTests extends AbstractBlobContainerRetriesTes
     @Before
     public void setUp() throws Exception {
         previousOpenSearchPathConf = SocketAccess.doPrivileged(() -> System.setProperty("opensearch.path.conf", configPath().toString()));
-        service = new S3Service(configPath());
-        asyncService = new S3AsyncService(configPath());
+        scheduler = new ScheduledThreadPoolExecutor(1);
+        service = new S3Service(configPath(), scheduler);
+        asyncService = new S3AsyncService(configPath(), scheduler);
 
         futureCompletionService = Executors.newSingleThreadExecutor();
         streamReaderService = Executors.newSingleThreadExecutor();
         transferNIOGroup = new AsyncTransferEventLoopGroup(1);
         remoteTransferRetry = Executors.newFixedThreadPool(20);
         transferQueueConsumerService = Executors.newFixedThreadPool(2);
-        scheduler = new ScheduledThreadPoolExecutor(1);
         GenericStatsMetricPublisher genericStatsMetricPublisher = new GenericStatsMetricPublisher(10000L, 10, 10000L, 10);
         normalPrioritySizeBasedBlockingQ = new SizeBasedBlockingQ(
             new ByteSizeValue(Runtime.getRuntime().availableProcessors() * 5L, ByteSizeUnit.GB),

plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3ClientSettingsTests.java

@@ -314,13 +314,13 @@ public void testSignerOverrideCanBeSet() {
         assertThat(settings.get("other").signerOverride, is(signerOverride));
 
         ClientOverrideConfiguration defaultConfiguration = SocketAccess.doPrivileged(
-            () -> S3Service.buildOverrideConfiguration(settings.get("default"))
+            () -> S3Service.buildOverrideConfiguration(settings.get("default"), null)
         );
         Optional<Signer> defaultSigner = defaultConfiguration.advancedOption(SdkAdvancedClientOption.SIGNER);
         assertFalse(defaultSigner.isPresent());
 
         ClientOverrideConfiguration configuration = SocketAccess.doPrivileged(
-            () -> S3Service.buildOverrideConfiguration(settings.get("other"))
+            () -> S3Service.buildOverrideConfiguration(settings.get("other"), null)
         );
         Optional<Signer> otherSigner = configuration.advancedOption(SdkAdvancedClientOption.SIGNER);
         assertTrue(otherSigner.isPresent());