diff --git a/.changes/next-release/bugfix-AmazonSQS-9b1a920.json b/.changes/next-release/bugfix-AmazonSQS-9b1a920.json new file mode 100644 index 000000000000..8c6cad398661 --- /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 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 d8d54f718a59..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 @@ -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 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 { 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..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 @@ -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) + .shutdownTimeout(builder.shutdownTimeout) .build(), scheduledExecutor, client @@ -61,18 +78,27 @@ private DefaultSqsAsyncBatchManager(DefaultBuilder builder) { 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> 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; 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..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 @@ -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,6 +89,10 @@ public int maxBatchBytesSize() { return maxBatchBytesSize; } + public Duration shutdownTimeout() { + return shutdownTimeout; + } + 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 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); } 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..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 @@ -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 shutdownTimeout; 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.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 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 closeAndDispatch() { + 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 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> - 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..7d558d41b5e2 --- /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#shutdownTimeout} + * test seam so a short timeout 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 shutdownTimeout) { + DefaultSqsAsyncBatchManager.DefaultBuilder builder = + (DefaultSqsAsyncBatchManager.DefaultBuilder) DefaultSqsAsyncBatchManager.builder(); + builder.client(client); + builder.scheduledExecutor(executor); + builder.overrideConfiguration(BatchOverrideConfiguration.builder() + .maxBatchSize(maxBatchSize) + .sendRequestFrequency(sendFrequency) + .build()); + builder.shutdownTimeout(shutdownTimeout); + return builder.build(); + } + + @Test + @Timeout(20) + void close_boundsShutdownToOneSharedTimeout_notPerManager() { + Duration timeout = 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), 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)); + + long start = System.nanoTime(); + batchManager.close(); + long closeMillis = (System.nanoTime() - start) / 1_000_000; + + assertThat(closeMillis).as("close() should wait about one timeout").isGreaterThanOrEqualTo(300); + assertThat(closeMillis) + .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 + @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_whenSendCompletesWithinTimeout() 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_thatDoesNotCompleteWithinTimeout() { + 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); + } +}