diff --git a/spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/AbstractOrderingAcknowledgementProcessor.java b/spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/AbstractOrderingAcknowledgementProcessor.java index d61257345..e179756c9 100644 --- a/spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/AbstractOrderingAcknowledgementProcessor.java +++ b/spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/AbstractOrderingAcknowledgementProcessor.java @@ -64,7 +64,8 @@ public abstract class AbstractOrderingAcknowledgementProcessor private AsyncAcknowledgementResultCallback acknowledgementResultCallback = new AsyncAcknowledgementResultCallback() { }; - private boolean running; + // Written under lifecycleMonitor, but read without synchronization from the polling and scheduler threads + private volatile boolean running; private String id; diff --git a/spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/BatchingAcknowledgementProcessor.java b/spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/BatchingAcknowledgementProcessor.java index 5f9ff0ff9..44ecc9726 100644 --- a/spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/BatchingAcknowledgementProcessor.java +++ b/spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/BatchingAcknowledgementProcessor.java @@ -250,6 +250,7 @@ private void waitOnAcknowledgementsIfTimeoutSet() { if (LocalDateTime.now().isAfter(endTime)) { throw new TimeoutException(); } + flushRemainingAcks(); Thread.sleep(200); } logger.debug("All acknowledgements completed."); @@ -269,6 +270,28 @@ private void waitOnAcknowledgementsIfTimeoutSet() { } } + /** + * Flushes whatever is left in the buffer, so that messages below the acknowledgement threshold are acknowledged + * on shutdown instead of being discarded when the buffer is cleared. Waits for the ack queue to drain first so + * that batches are not needlessly split. + * + * This has to be retried on every iteration of the shutdown wait loop rather than executed once: the polling + * thread may be holding a message it has already polled but not yet added to the buffer, in which case the + * queue looks empty while the message is in neither the queue nor the buffer. + */ + private void flushRemainingAcks() { + if (!this.acks.isEmpty()) { + return; + } + this.context.lock(); + try { + this.context.executeAllAcks(); + } + finally { + this.context.unlock(); + } + } + private boolean hasUnfinishedAcks() { var unfinishedAcks = this.context.runningAcks.stream().filter(Predicate.not(CompletableFuture::isDone)) .toList().size(); diff --git a/spring-cloud-aws-sqs/src/test/java/io/awspring/cloud/sqs/listener/acknowledgement/BatchingAcknowledgementProcessorTests.java b/spring-cloud-aws-sqs/src/test/java/io/awspring/cloud/sqs/listener/acknowledgement/BatchingAcknowledgementProcessorTests.java index 1d9fbc351..1362ab73b 100644 --- a/spring-cloud-aws-sqs/src/test/java/io/awspring/cloud/sqs/listener/acknowledgement/BatchingAcknowledgementProcessorTests.java +++ b/spring-cloud-aws-sqs/src/test/java/io/awspring/cloud/sqs/listener/acknowledgement/BatchingAcknowledgementProcessorTests.java @@ -68,6 +68,8 @@ class BatchingAcknowledgementProcessorTests { private static final Duration ACK_INTERVAL_ZERO = Duration.ZERO; + private static final Duration ACK_INTERVAL_THIRTY_SECONDS = Duration.ofSeconds(30); + private static final int MAX_ACKNOWLEDGEMENTS_PER_BATCH_TEN = 10; private static final Integer ACK_THRESHOLD_TEN = 10; @@ -189,6 +191,71 @@ protected CompletableFuture sendToExecutor(Collection> mes assertThat(ackLatch.await(1, TimeUnit.SECONDS)).isEqualTo(shouldWaitAllAcks); } + @Test + void givenScheduledAcknowledgement_whenStoppedWithRemainderBelowThreshold_shouldAcknowledgeRemainingMessages() + throws Exception { + List> messages = buildMessages(ACK_THRESHOLD_TEN + 5); + List> thresholdBatch = messages.subList(0, ACK_THRESHOLD_TEN); + List> remainder = messages.subList(ACK_THRESHOLD_TEN, messages.size()); + Collection> acknowledgedMessages = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch thresholdFlushLatch = new CountDownLatch(1); + + // The interval outlasts the test, so the remainder can only be flushed by the shutdown drain and never by a + // scheduled execution. A scheduled execution does exist here, unlike in the interval ZERO scenario below. + BatchingAcknowledgementProcessor processor = createProcessor(ACK_INTERVAL_THIRTY_SECONDS, + messagesToAck -> { + acknowledgedMessages.addAll(messagesToAck); + thresholdFlushLatch.countDown(); + return CompletableFuture.completedFuture(null); + }); + processor.start(); + + processor.doOnAcknowledge(thresholdBatch); + assertThat(thresholdFlushLatch.await(10, TimeUnit.SECONDS)).isTrue(); + processor.doOnAcknowledge(remainder); + processor.stop(); + + assertThat(acknowledgedMessages).containsExactlyInAnyOrderElementsOf(messages); + } + + @Test + void givenZeroIntervalAndThreshold_whenStoppedWithRemainderBelowThreshold_shouldAcknowledgeRemainingMessages() + throws Exception { + List> messages = buildMessages(ACK_THRESHOLD_TEN + 5); + Collection> acknowledgedMessages = Collections.synchronizedList(new ArrayList<>()); + + // No scheduled execution is ever created with interval ZERO, so the shutdown drain is the only thing that can + // flush the remainder. + BatchingAcknowledgementProcessor processor = createProcessor(ACK_INTERVAL_ZERO, messagesToAck -> { + acknowledgedMessages.addAll(messagesToAck); + return CompletableFuture.completedFuture(null); + }); + processor.start(); + + processor.doOnAcknowledge(messages); + processor.stop(); + + assertThat(acknowledgedMessages).containsExactlyInAnyOrderElementsOf(messages); + } + + private static List> buildMessages(int count) { + return IntStream.range(0, count).mapToObj(index -> MessageBuilder.withPayload(String.valueOf(index)).build()) + .collect(Collectors.toList()); + } + + private static BatchingAcknowledgementProcessor createProcessor(Duration acknowledgementInterval, + AcknowledgementExecutor acknowledgementExecutor) { + BatchingAcknowledgementProcessor processor = new BatchingAcknowledgementProcessor<>(); + processor.configure(SqsContainerOptions.builder().acknowledgementInterval(acknowledgementInterval) + .acknowledgementThreshold(ACK_THRESHOLD_TEN).acknowledgementOrdering(AcknowledgementOrdering.PARALLEL) + .acknowledgementShutdownTimeout(Duration.ofSeconds(10)).build()); + processor.setTaskExecutor(new SimpleAsyncTaskExecutor()); + processor.setAcknowledgementExecutor(acknowledgementExecutor); + processor.setMaxAcknowledgementsPerBatch(MAX_ACKNOWLEDGEMENTS_PER_BATCH_TEN); + processor.setId(ID); + return processor; + } + @Test void shouldAckAfterTime() throws Exception { given(message.getHeaders()).willReturn(messageHeaders);