From 802c270811eea71febfc11242fd823280daa194c Mon Sep 17 00:00:00 2001 From: Zoe Wang <33073555+zoewangg@users.noreply.github.com> Date: Thu, 20 Aug 2026 13:58:44 -0700 Subject: [PATCH 1/2] fix(sqs): stop close() busy-loop duplicate sends SqsAsyncBatchManager.close() extracted one ready batch once and re-flushed that same map in a loop that only the async completion callback could clear, so on a real client the closing thread re-sent the identical SendMessageBatch thousands of times. It now drains the buffer by re-extracting and flushing each batch (including partials) exactly once, terminating when the buffer is empty. Shutdown is also graceful: DefaultSqsAsyncBatchManager.close() first dispatches every buffered batch across the three write managers (dispatchPending), then does a single bounded wait for the in-flight sends so callers receive real results, and finally cancels only the stragglers past the deadline (cancelPending). The wait is one shared grace period (internal, default 5s) across all three managers, not one each, and close() is idempotent and guarded against concurrent calls. Post-close submissions fail fast with IllegalStateException. Internal only; no public API change. --- .../bugfix-AmazonSQS-9b1a920.json | 6 + .../batchmanager/SqsAsyncBatchManager.java | 7 + .../DefaultSqsAsyncBatchManager.java | 65 +++++- .../RequestBatchConfiguration.java | 19 ++ .../batchmanager/RequestBatchManager.java | 50 ++++- .../batchmanager/RequestBatchManagerTest.java | 49 ----- .../DefaultSqsAsyncBatchManagerTest.java | 201 ++++++++++++++++++ 7 files changed, 334 insertions(+), 63 deletions(-) create mode 100644 .changes/next-release/bugfix-AmazonSQS-9b1a920.json create mode 100644 services/sqs/src/test/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManagerTest.java diff --git a/.changes/next-release/bugfix-AmazonSQS-9b1a920.json b/.changes/next-release/bugfix-AmazonSQS-9b1a920.json new file mode 100644 index 000000000000..334a45c6eaac --- /dev/null +++ b/.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 grace period for in-flight batch sends to complete so their callers receive the real result instead of a cancellation." +} diff --git a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java index d8d54f718a59..8c9cbf2cef9a 100644 --- a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java +++ b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java @@ -148,6 +148,13 @@ default CompletableFuture 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 grace period waiting for the in-flight batch sends to complete so their + * callers receive the real service result. Requests still outstanding when the grace period expires are cancelled. + */ + @Override + void close(); interface Builder { diff --git a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java index 08300d485249..76aec6776a21 100644 --- a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java +++ b/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> 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,6 +70,7 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) { new SendMessageBatchManager( RequestBatchConfiguration.builder(builder.overrideConfiguration) .maxBatchBytesSize(MAX_SEND_MESSAGE_PAYLOAD_SIZE_BYTES) + .shutdownGracePeriod(builder.shutdownGracePeriod) .build(), scheduledExecutor, client @@ -61,18 +78,27 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) { this.deleteMessageBatchManager = new DeleteMessageBatchManager( - RequestBatchConfiguration.builder(builder.overrideConfiguration).build(), + RequestBatchConfiguration.builder(builder.overrideConfiguration) + .shutdownGracePeriod(builder.shutdownGracePeriod) + .build(), scheduledExecutor, client ); this.changeMessageVisibilityBatchManager = new ChangeMessageVisibilityBatchManager( - RequestBatchConfiguration.builder(builder.overrideConfiguration).build(), + RequestBatchConfiguration.builder(builder.overrideConfiguration) + .shutdownGracePeriod(builder.shutdownGracePeriod) + .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> futures = + requestBatchManagers.stream().map(requestBatchManager -> requestBatchManager.dispatchPending()) + .collect(Collectors.toList()); + + awaitQuietly(CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])), + sendMessageBatchManager.shutdownGracePeriod()); + + requestBatchManagers.forEach(requestBatchManager -> requestBatchManager.cancelPending()); receiveMessageBatchManager.close(); } + private static void awaitQuietly(CompletableFuture pending, Duration grace) { + try { + pending.get(grace.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 shutdownGracePeriod; private DefaultBuilder() { } + @SdkTestInternalApi + DefaultBuilder shutdownGracePeriod(Duration shutdownGracePeriod) { + this.shutdownGracePeriod = shutdownGracePeriod; + return this; + } + @Override public SqsAsyncBatchManager.Builder overrideConfiguration(BatchOverrideConfiguration overrideConfiguration) { this.overrideConfiguration = overrideConfiguration; diff --git a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchConfiguration.java b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchConfiguration.java index 08cd1c0818ce..1f55e6e96bcf 100644 --- a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchConfiguration.java +++ b/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_GRACE_PERIOD = 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 shutdownGracePeriod; 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.shutdownGracePeriod = builder.shutdownGracePeriod != null ? + builder.shutdownGracePeriod : DEFAULT_SHUTDOWN_GRACE_PERIOD; } @@ -80,6 +89,10 @@ public int maxBatchBytesSize() { return maxBatchBytesSize; } + public Duration shutdownGracePeriod() { + return shutdownGracePeriod; + } + public static final class Builder { private Integer maxBatchItems; @@ -87,6 +100,7 @@ public static final class Builder { private Integer maxBufferSize; private Duration sendRequestFrequency; private Integer maxBatchBytesSize; + private Duration shutdownGracePeriod; private Builder() { } @@ -116,6 +130,11 @@ public Builder maxBatchBytesSize(Integer maxBatchBytesSize) { return this; } + public Builder shutdownGracePeriod(Duration shutdownGracePeriod) { + this.shutdownGracePeriod = shutdownGracePeriod; + return this; + } + public RequestBatchConfiguration build() { return new RequestBatchConfiguration(this); } diff --git a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchManager.java b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchManager.java index dd1de65cd2c9..779ea9abf4ea 100644 --- a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchManager.java +++ b/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 { - // abm stands for Automatic Batching Manager public static final Consumer USER_AGENT_APPLIER = b -> b.addApiName(ApiName.builder().version("abm").name("hll").build()); @@ -45,10 +45,12 @@ public abstract class RequestBatchManager { private final int maxBatchItems; private final Duration sendRequestFrequency; + private final Duration shutdownGracePeriod; private final BatchingMap requestsAndResponsesMaps; private final ScheduledExecutorService scheduledExecutor; private final Set> pendingBatchResponses ; private final Set> pendingResponses ; + private final AtomicBoolean closed = new AtomicBoolean(false); protected RequestBatchManager(RequestBatchConfiguration overrideConfiguration, @@ -56,6 +58,7 @@ protected RequestBatchManager(RequestBatchConfiguration overrideConfiguration, batchConfiguration = overrideConfiguration; this.maxBatchItems = batchConfiguration.maxBatchItems(); this.sendRequestFrequency = batchConfiguration.sendRequestFrequency(); + this.shutdownGracePeriod = batchConfiguration.shutdownGracePeriod(); this.scheduledExecutor = Validate.notNull(scheduledExecutor, "Null scheduledExecutor"); pendingBatchResponses = ConcurrentHashMap.newKeySet(); pendingResponses = ConcurrentHashMap.newKeySet(); @@ -65,6 +68,10 @@ protected RequestBatchManager(RequestBatchConfiguration overrideConfiguration, public CompletableFuture batchRequest(RequestT request) { CompletableFuture 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 dispatchPending() { + if (!closed.compareAndSet(false, true)) { + return CompletableFuture.completedFuture(null); + } + drainBuffers(); + List> snapshot = new ArrayList<>(pendingResponses); + return CompletableFuture.allOf(snapshot.toArray(new CompletableFuture[0])); + } + + /** + * Phase 2 of shutdown. Cancels any requests still pending after the grace wait, 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 shutdownGracePeriod() { + return shutdownGracePeriod; + } + + private void drainBuffers() { requestsAndResponsesMaps.forEach((batchKey, batchBuffer) -> { requestsAndResponsesMaps.cancelScheduledFlush(batchKey); - Map> - extractedEntries = requestsAndResponsesMaps.extractBatchIfReady(batchKey); - + Map> 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(); } } \ No newline at end of file diff --git a/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/RequestBatchManagerTest.java b/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/RequestBatchManagerTest.java index c82984bbecd1..b7d1e3f928f7 100644 --- a/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/batchmanager/RequestBatchManagerTest.java +++ b/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 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 response1 = batchManager.batchRequest(request1); - CompletableFuture 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 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 response = batchManager.batchRequest(request); - - batchManager.close(); - assertThrows(CancellationException.class, () -> response.join()); - - } - - @Test void batchRequest_MoreThanBufferSize_Fails() throws Exception { final int MAX_QUEUES_THRESHOLD = 10000; diff --git a/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManagerTest.java b/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManagerTest.java new file mode 100644 index 000000000000..00da91ea90cb --- /dev/null +++ b/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManagerTest.java @@ -0,0 +1,201 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.services.sqs.internal.batchmanager; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.time.Duration; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import software.amazon.awssdk.services.sqs.SqsAsyncClient; +import software.amazon.awssdk.services.sqs.batchmanager.BatchOverrideConfiguration; +import software.amazon.awssdk.services.sqs.batchmanager.SqsAsyncBatchManager; +import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchRequest; +import software.amazon.awssdk.services.sqs.model.ChangeMessageVisibilityBatchResponse; +import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchRequest; +import software.amazon.awssdk.services.sqs.model.DeleteMessageBatchResponse; +import software.amazon.awssdk.services.sqs.model.SendMessageBatchRequest; +import software.amazon.awssdk.services.sqs.model.SendMessageBatchResponse; +import software.amazon.awssdk.services.sqs.model.SendMessageBatchResultEntry; +import software.amazon.awssdk.services.sqs.model.SendMessageResponse; + +/** + * Close/shutdown behavior of the real {@link DefaultSqsAsyncBatchManager} (and, through it, the real write batch + * managers), driven over a mock {@link SqsAsyncClient} whose batch-send futures the test controls. Lives in the + * internal package to reach the package-private {@link DefaultSqsAsyncBatchManager.DefaultBuilder#shutdownGracePeriod} + * test seam so a short grace can be injected. + */ +class DefaultSqsAsyncBatchManagerTest { + + private static final String QUEUE_URL = "https://sqs.us-east-1.amazonaws.com/123456789012/q"; + + private ScheduledExecutorService executor; + + @BeforeEach + void setUp() { + executor = Executors.newScheduledThreadPool(2); + } + + @AfterEach + void tearDown() { + executor.shutdownNow(); + } + + private SqsAsyncBatchManager batchManager(SqsAsyncClient client, int maxBatchSize, Duration sendFrequency, + Duration shutdownGracePeriod) { + DefaultSqsAsyncBatchManager.DefaultBuilder builder = + (DefaultSqsAsyncBatchManager.DefaultBuilder) DefaultSqsAsyncBatchManager.builder(); + builder.client(client); + builder.scheduledExecutor(executor); + builder.overrideConfiguration(BatchOverrideConfiguration.builder() + .maxBatchSize(maxBatchSize) + .sendRequestFrequency(sendFrequency) + .build()); + builder.shutdownGracePeriod(shutdownGracePeriod); + return builder.build(); + } + + @Test + @Timeout(20) + void close_boundsShutdownToOneSharedGracePeriod_notPerManager() { + Duration grace = Duration.ofMillis(400); + SqsAsyncClient client = mock(SqsAsyncClient.class); + when(client.sendMessageBatch(any(SendMessageBatchRequest.class))) + .thenReturn(new CompletableFuture()); + when(client.deleteMessageBatch(any(DeleteMessageBatchRequest.class))) + .thenReturn(new CompletableFuture()); + when(client.changeMessageVisibilityBatch(any(ChangeMessageVisibilityBatchRequest.class))) + .thenReturn(new CompletableFuture()); + + SqsAsyncBatchManager batchManager = batchManager(client, 10, Duration.ofHours(1), grace); + batchManager.sendMessage(r -> r.queueUrl(QUEUE_URL).messageBody("m")); + batchManager.deleteMessage(r -> r.queueUrl(QUEUE_URL).receiptHandle("rh")); + batchManager.changeMessageVisibility(r -> r.queueUrl(QUEUE_URL).receiptHandle("rh").visibilityTimeout(30)); + + long start = System.nanoTime(); + batchManager.close(); + long closeMillis = (System.nanoTime() - start) / 1_000_000; + + assertThat(closeMillis).as("close() should wait about one grace period").isGreaterThanOrEqualTo(300); + assertThat(closeMillis) + .as("close() must be bounded by ONE shared grace period, not one per manager (3x would be ~%d ms)", + 3 * grace.toMillis()) + .isLessThan(2 * grace.toMillis()); + } + + @Test + @Timeout(20) + void close_flushesBufferedPartialBatchExactlyOnce() { + SqsAsyncClient client = mock(SqsAsyncClient.class); + when(client.sendMessageBatch(any(SendMessageBatchRequest.class))) + .thenReturn(new CompletableFuture()); + + SqsAsyncBatchManager batchManager = batchManager(client, 10, Duration.ofHours(1), Duration.ofMillis(300)); + for (int i = 0; i < 5; i++) { + int n = i; + batchManager.sendMessage(r -> r.queueUrl(QUEUE_URL).messageBody("m" + n)); + } + + batchManager.close(); + + verify(client, times(1)).sendMessageBatch(any(SendMessageBatchRequest.class)); + } + + @Test + @Timeout(20) + void close_flushesFullThenResidualPartialBatch_eachExactlyOnce() { + SqsAsyncClient client = mock(SqsAsyncClient.class); + when(client.sendMessageBatch(any(SendMessageBatchRequest.class))) + .thenReturn(new CompletableFuture()); + + SqsAsyncBatchManager batchManager = batchManager(client, 10, Duration.ofHours(1), Duration.ofMillis(300)); + for (int i = 0; i < 15; i++) { + int n = i; + batchManager.sendMessage(r -> r.queueUrl(QUEUE_URL).messageBody("m" + n)); + } + + batchManager.close(); + + // 10 auto-flush a full batch during sendMessage, the residual 5 are drained on close: two sends, each once. + verify(client, times(2)).sendMessageBatch(any(SendMessageBatchRequest.class)); + } + + @Test + @Timeout(20) + void close_completesCallerWithRealResult_whenSendCompletesWithinGracePeriod() throws Exception { + SqsAsyncClient client = mock(SqsAsyncClient.class); + CompletableFuture sendFuture = new CompletableFuture<>(); + when(client.sendMessageBatch(any(SendMessageBatchRequest.class))).thenReturn(sendFuture); + + SqsAsyncBatchManager batchManager = batchManager(client, 10, Duration.ofHours(1), Duration.ofSeconds(2)); + CompletableFuture response = batchManager.sendMessage(r -> r.queueUrl(QUEUE_URL).messageBody("m")); + executor.schedule(() -> sendFuture.complete( + SendMessageBatchResponse.builder() + .successful(SendMessageBatchResultEntry.builder() + .id("0") + .messageId("msg-0") + .md5OfMessageBody("d41d8cd98f00b204e9800998ecf8427e") + .build()) + .build()), + 100, TimeUnit.MILLISECONDS); + + batchManager.close(); + + assertThat(response.get(1, TimeUnit.SECONDS).messageId()).isEqualTo("msg-0"); + } + + @Test + @Timeout(20) + void close_cancelsStragglerSend_thatDoesNotCompleteWithinGracePeriod() { + SqsAsyncClient client = mock(SqsAsyncClient.class); + when(client.sendMessageBatch(any(SendMessageBatchRequest.class))) + .thenReturn(new CompletableFuture()); + + SqsAsyncBatchManager batchManager = batchManager(client, 10, Duration.ofHours(1), Duration.ofMillis(300)); + CompletableFuture response = batchManager.sendMessage(r -> r.queueUrl(QUEUE_URL).messageBody("m")); + + batchManager.close(); + + assertThatThrownBy(response::join).isInstanceOf(CancellationException.class); + } + + @Test + @Timeout(20) + void sendMessageAfterClose_completesExceptionallyWithIllegalStateException() { + SqsAsyncClient client = mock(SqsAsyncClient.class); + SqsAsyncBatchManager batchManager = batchManager(client, 10, Duration.ofHours(1), Duration.ofMillis(1)); + + batchManager.close(); + batchManager.close(); // idempotent: second close is a no-op + + CompletableFuture response = batchManager.sendMessage(r -> r.queueUrl(QUEUE_URL).messageBody("m")); + + assertThatThrownBy(() -> response.get(1, TimeUnit.SECONDS)).hasCauseInstanceOf(IllegalStateException.class); + } +} From 4623cd61f3ec9cd285dff1ff3ef013fc3efbb972 Mon Sep 17 00:00:00 2001 From: Zoe Wang <33073555+zoewangg@users.noreply.github.com> Date: Fri, 21 Aug 2026 15:07:00 -0700 Subject: [PATCH 2/2] refactor(sqs): rename grace period to timeout Renamed the internal shutdown-timeout terminology from "grace period" to "timeout": DEFAULT_SHUTDOWN_TIMEOUT, the shutdownTimeout field, accessor, and builder seam, and the SqsAsyncBatchManager.close() Javadoc. Also renamed RequestBatchManager.dispatchPending() to closeAndDispatch(). Both per PR review. Updated the changelog wording accordingly. Internal @SdkInternalApi only; no public API change. --- .../bugfix-AmazonSQS-9b1a920.json | 2 +- .../batchmanager/SqsAsyncBatchManager.java | 4 +-- .../DefaultSqsAsyncBatchManager.java | 20 +++++++------- .../RequestBatchConfiguration.java | 18 ++++++------- .../batchmanager/RequestBatchManager.java | 12 ++++----- .../DefaultSqsAsyncBatchManagerTest.java | 26 +++++++++---------- 6 files changed, 41 insertions(+), 41 deletions(-) diff --git a/.changes/next-release/bugfix-AmazonSQS-9b1a920.json b/.changes/next-release/bugfix-AmazonSQS-9b1a920.json index 334a45c6eaac..8c6cad398661 100644 --- a/.changes/next-release/bugfix-AmazonSQS-9b1a920.json +++ b/.changes/next-release/bugfix-AmazonSQS-9b1a920.json @@ -2,5 +2,5 @@ "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 grace period for in-flight batch sends to complete so their callers receive the real result instead of a cancellation." + "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." } diff --git a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java index 8c9cbf2cef9a..14f7756a7d8d 100644 --- a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java +++ b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/batchmanager/SqsAsyncBatchManager.java @@ -150,8 +150,8 @@ default CompletableFuture receiveMessage( /** * 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 grace period waiting for the in-flight batch sends to complete so their - * callers receive the real service result. Requests still outstanding when the grace period expires are cancelled. + * 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(); diff --git a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java index 76aec6776a21..f3b0fbb264d5 100644 --- a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java +++ b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManager.java @@ -70,7 +70,7 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) { new SendMessageBatchManager( RequestBatchConfiguration.builder(builder.overrideConfiguration) .maxBatchBytesSize(MAX_SEND_MESSAGE_PAYLOAD_SIZE_BYTES) - .shutdownGracePeriod(builder.shutdownGracePeriod) + .shutdownTimeout(builder.shutdownTimeout) .build(), scheduledExecutor, client @@ -79,7 +79,7 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) { this.deleteMessageBatchManager = new DeleteMessageBatchManager( RequestBatchConfiguration.builder(builder.overrideConfiguration) - .shutdownGracePeriod(builder.shutdownGracePeriod) + .shutdownTimeout(builder.shutdownTimeout) .build(), scheduledExecutor, client @@ -88,7 +88,7 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) { this.changeMessageVisibilityBatchManager = new ChangeMessageVisibilityBatchManager( RequestBatchConfiguration.builder(builder.overrideConfiguration) - .shutdownGracePeriod(builder.shutdownGracePeriod) + .shutdownTimeout(builder.shutdownTimeout) .build(), scheduledExecutor, client @@ -136,19 +136,19 @@ public void close() { } List> futures = - requestBatchManagers.stream().map(requestBatchManager -> requestBatchManager.dispatchPending()) + requestBatchManagers.stream().map(requestBatchManager -> requestBatchManager.closeAndDispatch()) .collect(Collectors.toList()); awaitQuietly(CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])), - sendMessageBatchManager.shutdownGracePeriod()); + sendMessageBatchManager.shutdownTimeout()); requestBatchManagers.forEach(requestBatchManager -> requestBatchManager.cancelPending()); receiveMessageBatchManager.close(); } - private static void awaitQuietly(CompletableFuture pending, Duration grace) { + private static void awaitQuietly(CompletableFuture pending, Duration timeout) { try { - pending.get(grace.toMillis(), TimeUnit.MILLISECONDS); + pending.get(timeout.toMillis(), TimeUnit.MILLISECONDS); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } catch (TimeoutException e) { @@ -163,14 +163,14 @@ public static final class DefaultBuilder implements SqsAsyncBatchManager.Builder private SqsAsyncClient client; private BatchOverrideConfiguration overrideConfiguration; private ScheduledExecutorService scheduledExecutor; - private Duration shutdownGracePeriod; + private Duration shutdownTimeout; private DefaultBuilder() { } @SdkTestInternalApi - DefaultBuilder shutdownGracePeriod(Duration shutdownGracePeriod) { - this.shutdownGracePeriod = shutdownGracePeriod; + DefaultBuilder shutdownTimeout(Duration shutdownTimeout) { + this.shutdownTimeout = shutdownTimeout; return this; } diff --git a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchConfiguration.java b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchConfiguration.java index 1f55e6e96bcf..d6d7583baad0 100644 --- a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchConfiguration.java +++ b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchConfiguration.java @@ -32,14 +32,14 @@ public final class RequestBatchConfiguration { * 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_GRACE_PERIOD = Duration.ofSeconds(5); + 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 shutdownGracePeriod; + private final Duration shutdownTimeout; private RequestBatchConfiguration(Builder builder) { @@ -50,8 +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.shutdownGracePeriod = builder.shutdownGracePeriod != null ? - builder.shutdownGracePeriod : DEFAULT_SHUTDOWN_GRACE_PERIOD; + this.shutdownTimeout = builder.shutdownTimeout != null ? + builder.shutdownTimeout : DEFAULT_SHUTDOWN_TIMEOUT; } @@ -89,8 +89,8 @@ public int maxBatchBytesSize() { return maxBatchBytesSize; } - public Duration shutdownGracePeriod() { - return shutdownGracePeriod; + public Duration shutdownTimeout() { + return shutdownTimeout; } public static final class Builder { @@ -100,7 +100,7 @@ public static final class Builder { private Integer maxBufferSize; private Duration sendRequestFrequency; private Integer maxBatchBytesSize; - private Duration shutdownGracePeriod; + private Duration shutdownTimeout; private Builder() { } @@ -130,8 +130,8 @@ public Builder maxBatchBytesSize(Integer maxBatchBytesSize) { return this; } - public Builder shutdownGracePeriod(Duration shutdownGracePeriod) { - this.shutdownGracePeriod = shutdownGracePeriod; + public Builder shutdownTimeout(Duration shutdownTimeout) { + this.shutdownTimeout = shutdownTimeout; return this; } diff --git a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchManager.java b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchManager.java index 779ea9abf4ea..0bbaaba1d9f7 100644 --- a/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchManager.java +++ b/services/sqs/src/main/java/software/amazon/awssdk/services/sqs/internal/batchmanager/RequestBatchManager.java @@ -45,7 +45,7 @@ public abstract class RequestBatchManager { private final int maxBatchItems; private final Duration sendRequestFrequency; - private final Duration shutdownGracePeriod; + private final Duration shutdownTimeout; private final BatchingMap requestsAndResponsesMaps; private final ScheduledExecutorService scheduledExecutor; private final Set> pendingBatchResponses ; @@ -58,7 +58,7 @@ protected RequestBatchManager(RequestBatchConfiguration overrideConfiguration, batchConfiguration = overrideConfiguration; this.maxBatchItems = batchConfiguration.maxBatchItems(); this.sendRequestFrequency = batchConfiguration.sendRequestFrequency(); - this.shutdownGracePeriod = batchConfiguration.shutdownGracePeriod(); + this.shutdownTimeout = batchConfiguration.shutdownTimeout(); this.scheduledExecutor = Validate.notNull(scheduledExecutor, "Null scheduledExecutor"); pendingBatchResponses = ConcurrentHashMap.newKeySet(); pendingResponses = ConcurrentHashMap.newKeySet(); @@ -173,7 +173,7 @@ private void performScheduledFlush(String batchKey) { * 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 dispatchPending() { + protected CompletableFuture closeAndDispatch() { if (!closed.compareAndSet(false, true)) { return CompletableFuture.completedFuture(null); } @@ -183,7 +183,7 @@ protected CompletableFuture dispatchPending() { } /** - * Phase 2 of shutdown. Cancels any requests still pending after the grace wait, surfacing + * 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() { @@ -192,8 +192,8 @@ protected void cancelPending() { requestsAndResponsesMaps.clear(); } - protected Duration shutdownGracePeriod() { - return shutdownGracePeriod; + protected Duration shutdownTimeout() { + return shutdownTimeout; } private void drainBuffers() { diff --git a/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManagerTest.java b/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManagerTest.java index 00da91ea90cb..7d558d41b5e2 100644 --- a/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManagerTest.java +++ b/services/sqs/src/test/java/software/amazon/awssdk/services/sqs/internal/batchmanager/DefaultSqsAsyncBatchManagerTest.java @@ -48,8 +48,8 @@ /** * Close/shutdown behavior of the real {@link DefaultSqsAsyncBatchManager} (and, through it, the real write batch * managers), driven over a mock {@link SqsAsyncClient} whose batch-send futures the test controls. Lives in the - * internal package to reach the package-private {@link DefaultSqsAsyncBatchManager.DefaultBuilder#shutdownGracePeriod} - * test seam so a short grace can be injected. + * internal package to reach the package-private {@link DefaultSqsAsyncBatchManager.DefaultBuilder#shutdownTimeout} + * test seam so a short timeout can be injected. */ class DefaultSqsAsyncBatchManagerTest { @@ -68,7 +68,7 @@ void tearDown() { } private SqsAsyncBatchManager batchManager(SqsAsyncClient client, int maxBatchSize, Duration sendFrequency, - Duration shutdownGracePeriod) { + Duration shutdownTimeout) { DefaultSqsAsyncBatchManager.DefaultBuilder builder = (DefaultSqsAsyncBatchManager.DefaultBuilder) DefaultSqsAsyncBatchManager.builder(); builder.client(client); @@ -77,14 +77,14 @@ private SqsAsyncBatchManager batchManager(SqsAsyncClient client, int maxBatchSiz .maxBatchSize(maxBatchSize) .sendRequestFrequency(sendFrequency) .build()); - builder.shutdownGracePeriod(shutdownGracePeriod); + builder.shutdownTimeout(shutdownTimeout); return builder.build(); } @Test @Timeout(20) - void close_boundsShutdownToOneSharedGracePeriod_notPerManager() { - Duration grace = Duration.ofMillis(400); + void close_boundsShutdownToOneSharedTimeout_notPerManager() { + Duration timeout = Duration.ofMillis(400); SqsAsyncClient client = mock(SqsAsyncClient.class); when(client.sendMessageBatch(any(SendMessageBatchRequest.class))) .thenReturn(new CompletableFuture()); @@ -93,7 +93,7 @@ void close_boundsShutdownToOneSharedGracePeriod_notPerManager() { when(client.changeMessageVisibilityBatch(any(ChangeMessageVisibilityBatchRequest.class))) .thenReturn(new CompletableFuture()); - SqsAsyncBatchManager batchManager = batchManager(client, 10, Duration.ofHours(1), grace); + SqsAsyncBatchManager batchManager = batchManager(client, 10, Duration.ofHours(1), timeout); batchManager.sendMessage(r -> r.queueUrl(QUEUE_URL).messageBody("m")); batchManager.deleteMessage(r -> r.queueUrl(QUEUE_URL).receiptHandle("rh")); batchManager.changeMessageVisibility(r -> r.queueUrl(QUEUE_URL).receiptHandle("rh").visibilityTimeout(30)); @@ -102,11 +102,11 @@ void close_boundsShutdownToOneSharedGracePeriod_notPerManager() { batchManager.close(); long closeMillis = (System.nanoTime() - start) / 1_000_000; - assertThat(closeMillis).as("close() should wait about one grace period").isGreaterThanOrEqualTo(300); + assertThat(closeMillis).as("close() should wait about one timeout").isGreaterThanOrEqualTo(300); assertThat(closeMillis) - .as("close() must be bounded by ONE shared grace period, not one per manager (3x would be ~%d ms)", - 3 * grace.toMillis()) - .isLessThan(2 * grace.toMillis()); + .as("close() must be bounded by ONE shared timeout, not one per manager (3x would be ~%d ms)", + 3 * timeout.toMillis()) + .isLessThan(2 * timeout.toMillis()); } @Test @@ -148,7 +148,7 @@ void close_flushesFullThenResidualPartialBatch_eachExactlyOnce() { @Test @Timeout(20) - void close_completesCallerWithRealResult_whenSendCompletesWithinGracePeriod() throws Exception { + void close_completesCallerWithRealResult_whenSendCompletesWithinTimeout() throws Exception { SqsAsyncClient client = mock(SqsAsyncClient.class); CompletableFuture sendFuture = new CompletableFuture<>(); when(client.sendMessageBatch(any(SendMessageBatchRequest.class))).thenReturn(sendFuture); @@ -172,7 +172,7 @@ void close_completesCallerWithRealResult_whenSendCompletesWithinGracePeriod() th @Test @Timeout(20) - void close_cancelsStragglerSend_thatDoesNotCompleteWithinGracePeriod() { + void close_cancelsStragglerSend_thatDoesNotCompleteWithinTimeout() { SqsAsyncClient client = mock(SqsAsyncClient.class); when(client.sendMessageBatch(any(SendMessageBatchRequest.class))) .thenReturn(new CompletableFuture());