fix(sqs): stop close() busy-loop duplicate sends

.changes/next-release/bugfix-AmazonSQS-9b1a920.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "Amazon SQS",
+    "contributor": "",
+    "description": "Fixed `SqsAsyncBatchManager.close()` re-sending the same buffered batch in a busy loop. On close, each buffered batch (including partial batches) is now flushed exactly once, and close waits a bounded timeout (approximately 5 seconds) for in-flight batch sends to complete so their callers receive the real result instead of a cancellation."
+}

services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java

@@ -148,6 +148,13 @@ default CompletableFuture<ReceiveMessageResponse> receiveMessage(
         return receiveMessage(ReceiveMessageRequest.builder().applyMutation(request).build());
     }
 
+    /**
+     * Closes the batch manager and releases its resources. Buffered requests are first flushed to the service, then
+     * this method blocks for up to a bounded timeout waiting for the in-flight batch sends to complete so their
+     * callers receive the real service result. Requests still outstanding when the timeout expires are cancelled.
+     */
+    @Override
+    void close();
 
     interface Builder {
 

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

@@ -18,9 +18,18 @@
 
 import static software.amazon.awssdk.services.sqs.internal.batchmanager.ResponseBatchConfiguration.MAX_SEND_MESSAGE_PAYLOAD_SIZE_BYTES;
 
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.services.sqs.SqsAsyncClient;
 import software.amazon.awssdk.services.sqs.batchmanager.BatchOverrideConfiguration;
 import software.amazon.awssdk.services.sqs.batchmanager.SqsAsyncBatchManager;
@@ -32,10 +41,13 @@
 import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse;
 import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
 import software.amazon.awssdk.services.sqs.model.SendMessageResponse;
+import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
 
 @SdkInternalApi
 public final class DefaultSqsAsyncBatchManager implements SqsAsyncBatchManager {
+    private static final Logger log = Logger.loggerFor(DefaultSqsAsyncBatchManager.class);
+
     private final SqsAsyncClient client;
 
     private final SendMessageBatchManager sendMessageBatchManager;
@@ -46,6 +58,10 @@ public final class DefaultSqsAsyncBatchManager implements SqsAsyncBatchManager {
 
     private final ReceiveMessageBatchManager receiveMessageBatchManager;
 
+    private final List<RequestBatchManager<?, ?, ?>> requestBatchManagers;
+
+    private final AtomicBoolean closed = new AtomicBoolean(false);
+
     private DefaultSqsAsyncBatchManager(DefaultBuilder builder) {
         this.client = Validate.notNull(builder.client, "client cannot be null");
         ScheduledExecutorService scheduledExecutor  = Validate.notNull(builder.scheduledExecutor,
@@ -54,25 +70,35 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) {
             new SendMessageBatchManager(
                 RequestBatchConfiguration.builder(builder.overrideConfiguration)
                                          .maxBatchBytesSize(MAX_SEND_MESSAGE_PAYLOAD_SIZE_BYTES)
+                                         .shutdownTimeout(builder.shutdownTimeout)
                                          .build(),
                 scheduledExecutor,
                 client
             );
 
         this.deleteMessageBatchManager =
             new DeleteMessageBatchManager(
-                RequestBatchConfiguration.builder(builder.overrideConfiguration).build(),
+                RequestBatchConfiguration.builder(builder.overrideConfiguration)
+                                         .shutdownTimeout(builder.shutdownTimeout)
+                                         .build(),
                 scheduledExecutor,
                 client
             );
 
         this.changeMessageVisibilityBatchManager =
             new ChangeMessageVisibilityBatchManager(
-                RequestBatchConfiguration.builder(builder.overrideConfiguration).build(),
+                RequestBatchConfiguration.builder(builder.overrideConfiguration)
+                                         .shutdownTimeout(builder.shutdownTimeout)
+                                         .build(),
                 scheduledExecutor,
                 client
             );
 
+        requestBatchManagers = new ArrayList<>(3);
+        requestBatchManagers.add(sendMessageBatchManager);
+        requestBatchManagers.add(deleteMessageBatchManager);
+        requestBatchManagers.add(changeMessageVisibilityBatchManager);
+
         this.receiveMessageBatchManager =
             new ReceiveMessageBatchManager(client,
                                            scheduledExecutor,
@@ -105,20 +131,49 @@ public static SqsAsyncBatchManager.Builder builder() {
 
     @Override
     public void close() {
-        sendMessageBatchManager.close();
-        deleteMessageBatchManager.close();
-        changeMessageVisibilityBatchManager.close();
+        if (!closed.compareAndSet(false, true)) {
+            return;
+        }
+
+        List<CompletableFuture<Void>> futures =
+            requestBatchManagers.stream().map(requestBatchManager -> requestBatchManager.closeAndDispatch())
+                                                                    .collect(Collectors.toList());
+
+        awaitQuietly(CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])),
+                     sendMessageBatchManager.shutdownTimeout());
+
+        requestBatchManagers.forEach(requestBatchManager -> requestBatchManager.cancelPending());
         receiveMessageBatchManager.close();
     }
 
+    private static void awaitQuietly(CompletableFuture<?> pending, Duration timeout) {
+        try {
+            pending.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+        } catch (TimeoutException e) {
+            log.debug(() -> "Timed out waiting for in-flight batch sends to complete during close; "
+                            + "cancelling outstanding requests.");
+        } catch (ExecutionException e) {
+            // A send failed; its caller future is already settled with the failure. Nothing to do here.
+        }
+    }
+
     public static final class DefaultBuilder implements SqsAsyncBatchManager.Builder {
         private SqsAsyncClient client;
         private BatchOverrideConfiguration overrideConfiguration;
         private ScheduledExecutorService scheduledExecutor;
+        private Duration shutdownTimeout;
 
         private DefaultBuilder() {
         }
 
+        @SdkTestInternalApi
+        DefaultBuilder shutdownTimeout(Duration shutdownTimeout) {
+            this.shutdownTimeout = shutdownTimeout;
+            return this;
+        }
+
         @Override
         public SqsAsyncBatchManager.Builder overrideConfiguration(BatchOverrideConfiguration overrideConfiguration) {
             this.overrideConfiguration = overrideConfiguration;

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

@@ -28,11 +28,18 @@ public final class RequestBatchConfiguration {
     public static final int DEFAULT_MAX_BUFFER_SIZE = 500;
     public static final Duration DEFAULT_MAX_BATCH_OPEN_IN_MS = Duration.ofMillis(200);
 
+    /**
+     * Time {@code close()} waits for in-flight batch sends to complete so their callers receive real results before
+     * outstanding requests are cancelled.
+     */
+    public static final Duration DEFAULT_SHUTDOWN_TIMEOUT = Duration.ofSeconds(5);
+
     private final Integer maxBatchItems;
     private final Integer maxBatchKeys;
     private final Integer maxBufferSize;
     private final Duration sendRequestFrequency;
     private final Integer maxBatchBytesSize;
+    private final Duration shutdownTimeout;
 
     private RequestBatchConfiguration(Builder builder) {
 
@@ -43,6 +50,8 @@ private RequestBatchConfiguration(Builder builder) {
                                                   builder.sendRequestFrequency :
                                          DEFAULT_MAX_BATCH_OPEN_IN_MS;
         this.maxBatchBytesSize = builder.maxBatchBytesSize != null ? builder.maxBatchBytesSize : DEFAULT_MAX_BATCH_BYTES_SIZE;
+        this.shutdownTimeout = builder.shutdownTimeout != null ?
+                                   builder.shutdownTimeout : DEFAULT_SHUTDOWN_TIMEOUT;
 
     }
 
@@ -80,13 +89,18 @@ public int maxBatchBytesSize() {
         return maxBatchBytesSize;
     }
 
+    public Duration shutdownTimeout() {
+        return shutdownTimeout;
+    }
+
     public static final class Builder {
 
         private Integer maxBatchItems;
         private Integer maxBatchKeys;
         private Integer maxBufferSize;
         private Duration sendRequestFrequency;
         private Integer maxBatchBytesSize;
+        private Duration shutdownTimeout;
 
         private Builder() {
         }
@@ -116,6 +130,11 @@ public Builder maxBatchBytesSize(Integer maxBatchBytesSize) {
             return this;
         }
 
+        public Builder shutdownTimeout(Duration shutdownTimeout) {
+            this.shutdownTimeout = shutdownTimeout;
+            return this;
+        }
+
         public RequestBatchConfiguration build() {
             return new RequestBatchConfiguration(this);
         }

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

@@ -26,6 +26,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
@@ -36,7 +37,6 @@
 @SdkInternalApi
 public abstract class RequestBatchManager<RequestT, ResponseT, BatchResponseT> {
 
-
     // abm stands for Automatic Batching Manager
     public static final Consumer<AwsRequestOverrideConfiguration.Builder> USER_AGENT_APPLIER =
         b -> b.addApiName(ApiName.builder().version("abm").name("hll").build());
@@ -45,17 +45,20 @@ public abstract class RequestBatchManager<RequestT, ResponseT, BatchResponseT> {
 
     private final int maxBatchItems;
     private final Duration sendRequestFrequency;
+    private final Duration shutdownTimeout;
     private final BatchingMap<RequestT, ResponseT> requestsAndResponsesMaps;
     private final ScheduledExecutorService scheduledExecutor;
     private final Set<CompletableFuture<BatchResponseT>> pendingBatchResponses ;
     private final Set<CompletableFuture<ResponseT>> pendingResponses ;
+    private final AtomicBoolean closed = new AtomicBoolean(false);
 
 
     protected RequestBatchManager(RequestBatchConfiguration overrideConfiguration,
                                   ScheduledExecutorService scheduledExecutor) {
         batchConfiguration = overrideConfiguration;
         this.maxBatchItems = batchConfiguration.maxBatchItems();
         this.sendRequestFrequency = batchConfiguration.sendRequestFrequency();
+        this.shutdownTimeout = batchConfiguration.shutdownTimeout();
         this.scheduledExecutor = Validate.notNull(scheduledExecutor, "Null scheduledExecutor");
         pendingBatchResponses = ConcurrentHashMap.newKeySet();
         pendingResponses = ConcurrentHashMap.newKeySet();
@@ -65,6 +68,10 @@ protected RequestBatchManager(RequestBatchConfiguration overrideConfiguration,
 
     public CompletableFuture<ResponseT> batchRequest(RequestT request) {
         CompletableFuture<ResponseT> response = new CompletableFuture<>();
+        if (closed.get()) {
+            response.completeExceptionally(new IllegalStateException("The client has been shut down."));
+            return response;
+        }
         pendingResponses.add(response);
         response.whenComplete((r, t) -> pendingResponses.remove(response));
 
@@ -160,20 +167,45 @@ private void performScheduledFlush(String batchKey) {
         }
     }
 
-    public void close() {
+    /**
+     * Phase 1 of shutdown, non-blocking. Marks the manager closed, dispatches every buffered batch (each exactly
+     * once), and returns a future that completes when all in-flight sends dispatched here have completed, so the
+     * caller can wait on it. Does not block and does not cancel; a second call is a no-op that returns an
+     * already-completed future.
+     */
+    protected CompletableFuture<Void> closeAndDispatch() {
+        if (!closed.compareAndSet(false, true)) {
+            return CompletableFuture.completedFuture(null);
+        }
+        drainBuffers();
+        List<CompletableFuture<ResponseT>> snapshot = new ArrayList<>(pendingResponses);
+        return CompletableFuture.allOf(snapshot.toArray(new CompletableFuture[0]));
+    }
+
+    /**
+     * Phase 2 of shutdown. Cancels any requests still pending after the timeout, surfacing
+     * {@link java.util.concurrent.CancellationException} to their callers, and releases the buffers. Idempotent.
+     */
+    protected void cancelPending() {
+        pendingResponses.forEach(future -> future.cancel(true));
+        pendingBatchResponses.forEach(future -> future.cancel(true));
+        requestsAndResponsesMaps.clear();
+    }
+
+    protected Duration shutdownTimeout() {
+        return shutdownTimeout;
+    }
+
+    private void drainBuffers() {
         requestsAndResponsesMaps.forEach((batchKey, batchBuffer) -> {
             requestsAndResponsesMaps.cancelScheduledFlush(batchKey);
-            Map<String, BatchingExecutionContext<RequestT, ResponseT>>
-                extractedEntries = requestsAndResponsesMaps.extractBatchIfReady(batchKey);
-
+            Map<String, BatchingExecutionContext<RequestT, ResponseT>> extractedEntries =
+                requestsAndResponsesMaps.extractEntriesForScheduledFlush(batchKey, maxBatchItems);
             while (!extractedEntries.isEmpty()) {
                 flushBuffer(batchKey, extractedEntries);
+                extractedEntries = requestsAndResponsesMaps.extractEntriesForScheduledFlush(batchKey, maxBatchItems);
             }
-
         });
-        pendingBatchResponses.forEach(future -> future.cancel(true));
-        pendingResponses.forEach(future -> future.cancel(true));
-        requestsAndResponsesMaps.clear();
     }
 
 }

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

@@ -21,7 +21,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
@@ -152,54 +151,6 @@ void batchRequest_WithNetworkError_throwsException() throws Exception {
         assertThrows(ExecutionException.class, () -> response.get(1, TimeUnit.SECONDS));
     }
 
-    @Test
-    void close_FlushesAllBatches() throws Exception {
-        String request1 = "testRequest:0";
-        String batchKey = "testRequest";
-        String request2 = "testRequest:1";
-        CompletableFuture<BatchResponse> batchResponseFuture = CompletableFuture.completedFuture(batchedResponse(2,
-                                                                                                                 "testResponse"));
-
-        when(mockClient.sendBatchAsync(any(), eq(batchKey))).thenReturn(batchResponseFuture);
-
-        SampleBatchManager batchManager=
-            new SampleBatchManager(BatchOverrideConfiguration.builder().maxBatchSize(2).sendRequestFrequency(Duration.ofHours(1)).build(), scheduledExecutor, mockClient);
-
-        CompletableFuture<String> response1 = batchManager.batchRequest(request1);
-        CompletableFuture<String> response2 = batchManager.batchRequest(request2);
-        // Even though the mock returns results immediately, since this is asynchronous execution, the test environment may take
-        // additional time due to the Scheduled Executors execution on that machine.
-        Thread.sleep(200);
-        batchManager.close();
-
-        assertEquals("testResponse0", response1.get(1, TimeUnit.SECONDS));
-
-        assertEquals("testResponse1", response2.get(1, TimeUnit.SECONDS));
-    }
-
-
-    @Test
-    void batchRequest_ClosedWhenWaitingForResponse() throws Exception {
-        String request = "testRequest:1";
-        String batchKey = "testRequest";
-        CompletableFuture<BatchResponse> batchResponseFuture = new CompletableFuture<>();
-
-        // Simulate successful response with delay
-        scheduledExecutor.schedule(() -> batchResponseFuture.complete(batchedResponse(1, "testResponse")),
-                                   10, TimeUnit.HOURS);
-
-        when(mockClient.sendBatchAsync(any(), eq(batchKey))).thenReturn(batchResponseFuture);
-
-        SampleBatchManager batchManager =
-            new SampleBatchManager(BatchOverrideConfiguration.builder().maxBatchSize(1).build(), scheduledExecutor, mockClient);
-        CompletableFuture<String> response = batchManager.batchRequest(request);
-
-        batchManager.close();
-        assertThrows(CancellationException.class, () -> response.join());
-
-    }
-
-
     @Test
     void batchRequest_MoreThanBufferSize_Fails() throws Exception {
         final int MAX_QUEUES_THRESHOLD = 10000;