SdkClosable implemented

Author: joviegas

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/QueueAttributesManager.java

@@ -113,22 +113,19 @@ private CompletableFuture<Map<QueueAttributeName, String>> getAttributeMap() {
             CompletableFuture<Map<QueueAttributeName, String>> newFuture = new CompletableFuture<>();
 
             if (queueAttributeMap.compareAndSet(future, newFuture)) {
-                // Only one thread will execute this block and fetch the attributes.
                 fetchQueueAttributes().whenComplete((r, t) -> {
                     if (t != null) {
-                        newFuture.completeExceptionally(t); // Complete the future exceptionally
+                        newFuture.completeExceptionally(t);
                     } else {
-                        newFuture.complete(r); // Complete the future with the result
+                        newFuture.complete(r);
                     }
                 });
                 return newFuture;
             } else {
-                // If another thread already set the future, cancel this one and use the existing one.
                 newFuture.cancel(true);
-                return queueAttributeMap.get(); // Return the future set by the winning thread
+                return queueAttributeMap.get();
             }
         }
-
         return future;
     }
 

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/ReceiveBatchManager.java

@@ -24,9 +24,10 @@
 import software.amazon.awssdk.services.sqs.SqsAsyncClient;
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest;
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
+import software.amazon.awssdk.utils.SdkAutoCloseable;
 
 @SdkInternalApi
-public class ReceiveBatchManager {
+public class ReceiveBatchManager implements SdkAutoCloseable {
 
     private final SqsAsyncClient sqsClient;
     private final ScheduledExecutorService executor;
@@ -62,8 +63,8 @@ public CompletableFuture<ReceiveMessageResponse> processRequest(ReceiveMessageRe
         });
     }
 
-    public void shutdown() {
-        receiveQueueBuffer.shutdown();
-        executor.shutdown();
+    @Override
+    public void close() {
+        receiveQueueBuffer.close();
     }
 }

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/ReceiveQueueBuffer.java

@@ -28,9 +28,10 @@
 import software.amazon.awssdk.services.sqs.SqsAsyncClient;
 import software.amazon.awssdk.services.sqs.model.Message;
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
+import software.amazon.awssdk.utils.SdkAutoCloseable;
 
 @SdkInternalApi
-public class ReceiveQueueBuffer {
+public class ReceiveQueueBuffer implements SdkAutoCloseable {
 
     private final ScheduledExecutorService executor;
     private final SqsAsyncClient sqsClient;
@@ -66,9 +67,9 @@ public boolean isShutDown() {
         return shutDown.get();
     }
 
-    public void shutdown() {
+    @Override
+    public void close() {
         if (this.shutDown.compareAndSet(false, true)) {
-            // Clear all finished tasks
             while (!finishedTasks.isEmpty()) {
                 ReceiveSqsMessageHelper batch = finishedTasks.poll();
                 if (inflightReceiveMessageBatches.get() > 0) {
@@ -78,8 +79,6 @@ public void shutdown() {
                     batch.clear();
                 }
             }
-
-            // Clear futures
             futures.forEach(futureWrapper -> {
                 if (!futureWrapper.getFuture().isDone()) {
                     futureWrapper.getFuture().completeExceptionally(new CancellationException("Shutdown in progress"));
@@ -190,7 +189,6 @@ private void reportBatchFinished(ReceiveSqsMessageHelper batch) {
         spawnMoreReceiveTasks();
     }
 
-
     private static class FutureRequestWrapper {
         private final CompletableFuture<ReceiveMessageResponse> future;
         private final int requestedSize;

services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/ReceiveBatchManagerTest.java

@@ -165,7 +165,7 @@ public void testShutdown() throws Exception {
         ReceiveMessageRequest request = ReceiveMessageRequest.builder().maxNumberOfMessages(10).build();
         CompletableFuture<ReceiveMessageResponse> futureResponse = receiveBatchManager.processRequest(request);
 
-        receiveBatchManager.shutdown();
+        receiveBatchManager.close();
 
         assertThrows(IllegalStateException.class, () -> receiveBatchManager.processRequest(request));
 

services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/ReceiveQueueBufferTest.java

@@ -298,7 +298,7 @@ public void receiveMessageShutDown() throws Exception {
         receiveQueueBuffer.receiveMessage(future, 10);
 
         // Shutdown receiveQueueBuffer
-        receiveQueueBuffer.shutdown();
+        receiveQueueBuffer.close();
 
         // Verify that the future is completed exceptionally
         assertTrue(future.isCompletedExceptionally());
@@ -441,7 +441,7 @@ public void testShutdownExceptionallyCompletesAllIncompleteFutures() throws Exce
     }
 
     // Shutdown the queue buffer and assert no exceptions are thrown
-    assertDoesNotThrow(() -> receiveQueueBuffer.shutdown());
+    assertDoesNotThrow(() -> receiveQueueBuffer.close());
 
     // Verify that each future completes exceptionally with CancellationException
     for (CompletableFuture<ReceiveMessageResponse> future : futures) {
@@ -477,7 +477,7 @@ public void visibilityTimeOutErrorsAreLogged() throws Exception {
 
         try (LogCaptor logCaptor = LogCaptor.create(Level.DEBUG)) {
             // Shutdown the receiveQueueBuffer to trigger the visibility timeout errors
-            assertDoesNotThrow(() -> receiveQueueBuffer.shutdown());
+            assertDoesNotThrow(receiveQueueBuffer::close);
 
             // Verify that an error was logged for failing to change visibility timeout
             assertThat(logCaptor.loggedEvents()).anySatisfy(logEvent -> {