Flush buffered acknowledgements from the shutdown wait

spring-cloud-aws-sqs/src/main/java/io/awspring/cloud/sqs/listener/acknowledgement/AbstractOrderingAcknowledgementProcessor.java

@@ -64,7 +64,8 @@ public abstract class AbstractOrderingAcknowledgementProcessor<T>
 	private AsyncAcknowledgementResultCallback<T> acknowledgementResultCallback = new AsyncAcknowledgementResultCallback<T>() {
 	};
 
-	private boolean running;
+	// Written under lifecycleMonitor, but read without synchronization from the polling and scheduler threads
+	private volatile boolean running;
 
 	private String id;
 

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();

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<Void> sendToExecutor(Collection<Message<String>> mes
 		assertThat(ackLatch.await(1, TimeUnit.SECONDS)).isEqualTo(shouldWaitAllAcks);
 	}
 
+	@Test
+	void givenScheduledAcknowledgement_whenStoppedWithRemainderBelowThreshold_shouldAcknowledgeRemainingMessages()
+			throws Exception {
+		List<Message<String>> messages = buildMessages(ACK_THRESHOLD_TEN + 5);
+		List<Message<String>> thresholdBatch = messages.subList(0, ACK_THRESHOLD_TEN);
+		List<Message<String>> remainder = messages.subList(ACK_THRESHOLD_TEN, messages.size());
+		Collection<Message<String>> 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<String> 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<Message<String>> messages = buildMessages(ACK_THRESHOLD_TEN + 5);
+		Collection<Message<String>> 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<String> 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<Message<String>> buildMessages(int count) {
+		return IntStream.range(0, count).mapToObj(index -> MessageBuilder.withPayload(String.valueOf(index)).build())
+				.collect(Collectors.toList());
+	}
+
+	private static BatchingAcknowledgementProcessor<String> createProcessor(Duration acknowledgementInterval,
+			AcknowledgementExecutor<String> acknowledgementExecutor) {
+		BatchingAcknowledgementProcessor<String> 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);