From 81e356dc6c8eadb0b40086cc11d4c30fd58e0e26 Mon Sep 17 00:00:00 2001
From: Ashish Singh <ssashish@amazon.com>
Date: Wed, 18 Dec 2024 17:12:24 +0530
Subject: [PATCH] Fix failures in S3RepositoryThirdPartyTests

Signed-off-by: Ashish Singh <ssashish@amazon.com>
---
 .../repositories/s3/AmazonAsyncS3Reference.java      |  3 +--
 .../opensearch/repositories/s3/S3AsyncService.java   |  7 +------
 .../repositories/s3/S3RepositoryPlugin.java          |  3 +++
 .../s3/S3BlobContainerMockClientTests.java           |  2 +-
 .../repositories/s3/S3BlobStoreContainerTests.java   | 12 ++++--------
 5 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/AmazonAsyncS3Reference.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/AmazonAsyncS3Reference.java
index 3fdf8b5b27a97..45170ea1ad209 100644
--- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/AmazonAsyncS3Reference.java
+++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/AmazonAsyncS3Reference.java
@@ -25,7 +25,7 @@ public class AmazonAsyncS3Reference extends RefCountedReleasable<AmazonAsyncS3Wi
 
     private static final Logger logger = LogManager.getLogger(AmazonAsyncS3Reference.class);
 
-    AmazonAsyncS3Reference(AmazonAsyncS3WithCredentials client, Runnable onClose) {
+    AmazonAsyncS3Reference(AmazonAsyncS3WithCredentials client) {
         super("AWS_S3_CLIENT", client, () -> {
             client.client().close();
             client.priorityClient().close();
@@ -38,7 +38,6 @@ public class AmazonAsyncS3Reference extends RefCountedReleasable<AmazonAsyncS3Wi
                     logger.error("Exception while closing AwsCredentialsProvider", e);
                 }
             }
-            onClose.run();
         });
     }
 }
diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3AsyncService.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3AsyncService.java
index 0e41fb6d16eb3..7397c3132c17c 100644
--- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3AsyncService.java
+++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3AsyncService.java
@@ -121,12 +121,7 @@ public AmazonAsyncS3Reference client(
             }
 
             final AmazonAsyncS3Reference clientReference = new AmazonAsyncS3Reference(
-                buildClient(clientSettings, urgentExecutorBuilder, priorityExecutorBuilder, normalExecutorBuilder),
-                () -> {
-                    urgentExecutorBuilder.getAsyncTransferEventLoopGroup().close();
-                    priorityExecutorBuilder.getAsyncTransferEventLoopGroup().close();
-                    normalExecutorBuilder.getAsyncTransferEventLoopGroup().close();
-                }
+                buildClient(clientSettings, urgentExecutorBuilder, priorityExecutorBuilder, normalExecutorBuilder)
             );
             clientReference.incRef();
             clientsCache = MapBuilder.newMapBuilder(clientsCache).put(clientSettings, clientReference).immutableMap();
diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3RepositoryPlugin.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3RepositoryPlugin.java
index c199b5249df20..72a812339e387 100644
--- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3RepositoryPlugin.java
+++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3RepositoryPlugin.java
@@ -387,5 +387,8 @@ public void reload(Settings settings) {
     public void close() throws IOException {
         service.close();
         s3AsyncService.close();
+        urgentExecutorBuilder.getAsyncTransferEventLoopGroup().close();
+        priorityExecutorBuilder.getAsyncTransferEventLoopGroup().close();
+        normalExecutorBuilder.getAsyncTransferEventLoopGroup().close();
     }
 }
diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java
index 0240a76c87f3e..9b413ac81d766 100644
--- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java
+++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java
@@ -294,7 +294,7 @@ public AmazonAsyncS3Reference client(
             AsyncExecutorContainer priorityExecutorBuilder,
             AsyncExecutorContainer normalExecutorBuilder
         ) {
-            return new AmazonAsyncS3Reference(AmazonAsyncS3WithCredentials.create(asyncClient, asyncClient, asyncClient, null), () -> {});
+            return new AmazonAsyncS3Reference(AmazonAsyncS3WithCredentials.create(asyncClient, asyncClient, asyncClient, null));
         }
     }
 
diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java
index 0aa86aa905fad..53371cd1529ce 100644
--- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java
+++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java
@@ -1136,8 +1136,7 @@ public void testReadBlobAsyncMultiPart() throws Exception {
 
         final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class);
         final AmazonAsyncS3Reference amazonAsyncS3Reference = new AmazonAsyncS3Reference(
-            AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null),
-            () -> {}
+            AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null)
         );
 
         final S3BlobStore blobStore = mock(S3BlobStore.class);
@@ -1195,8 +1194,7 @@ public void testReadBlobAsyncSinglePart() throws Exception {
 
         final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class);
         final AmazonAsyncS3Reference amazonAsyncS3Reference = new AmazonAsyncS3Reference(
-            AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null),
-            () -> {}
+            AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null)
         );
         final S3BlobStore blobStore = mock(S3BlobStore.class);
         final BlobPath blobPath = new BlobPath();
@@ -1251,8 +1249,7 @@ public void testReadBlobAsyncFailure() throws Exception {
 
         final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class);
         final AmazonAsyncS3Reference amazonAsyncS3Reference = new AmazonAsyncS3Reference(
-            AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null),
-            () -> {}
+            AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null)
         );
 
         final S3BlobStore blobStore = mock(S3BlobStore.class);
@@ -1295,8 +1292,7 @@ public void testReadBlobAsyncOnCompleteFailureMissingData() throws Exception {
 
         final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class);
         final AmazonAsyncS3Reference amazonAsyncS3Reference = new AmazonAsyncS3Reference(
-            AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null),
-            () -> {}
+            AmazonAsyncS3WithCredentials.create(s3AsyncClient, s3AsyncClient, s3AsyncClient, null)
         );
 
         final S3BlobStore blobStore = mock(S3BlobStore.class);