Fix message loss on stop() for PollableChannel reactive consumer

spring-integration-core/src/main/java/org/springframework/integration/util/IntegrationReactiveUtils.java

@@ -49,6 +49,7 @@
  * Utilities for adapting integration components to/from reactive types.
  *
  * @author Artem Bilan
+ * @author Fardan An
  *
  * @since 5.3
  */
@@ -122,21 +123,29 @@ public static ContextView captureReactorContext() {
 	 * or falls back to 1-second duration.
 	 * If a produced message has an
 	 * {@link org.springframework.integration.IntegrationMessageHeaderAccessor#ACKNOWLEDGMENT_CALLBACK} header
-	 * it is ack'ed in the {@link Mono#doOnSuccess} and nack'ed in the {@link Mono#doOnError}.
+	 * it is ack'ed in {@link Flux#doOnNext} when this flux emits to its subscriber, and nack'ed in
+	 * the {@link Mono#doOnError}.
+	 * <p>
+	 * Cancellation while {@link MessageSource#receive()} is in flight (the sink is already
+	 * {@code CANCELLED}) is handled by Reactor's discard hook on this flux. Only a
+	 * {@link PollableChannel} adapted through {@link #messageChannelToFlux(MessageChannel)}
+	 * is best-effort re-queued via non-blocking {@link PollableChannel#send(Message, long)}
+	 * (interceptor chain is re-entered; FIFO is not preserved if the queue is not empty; a
+	 * full bounded channel is logged and the message is not nack'd). A direct
+	 * {@code messageSourceToFlux(() -> channel.receive(0))} lambda is not re-queued.
+	 * Other sources are left unacknowledged so the source's own redelivery
+	 * (broker requeue, uncommitted offset) can recover them. The discard hook only sees drops
+	 * at or upstream of this flux; operators a caller adds afterward
+	 * ({@code ReactiveStreamsConsumer.setReactiveCustomizer}, a {@code flatMap} on the
+	 * {@code ReactiveMessageHandler} path) use a different context and are not rescued here.
 	 * @param messageSource the {@link MessageSource} to adapt.
 	 * @param <T> the expected payload type.
 	 * @return a {@link Flux} which pulls messages from the {@link MessageSource} on demand.
 	 */
 	@SuppressWarnings("NullAway")
 	public static <T> Flux<Message<T>> messageSourceToFlux(MessageSource<T> messageSource) {
-		return Mono.
-				<Message<T>>create(monoSink ->
-				monoSink.onRequest(value -> monoSink.success(messageSource.receive())))
-				.doOnSuccess((message) -> {
-					if (message != null) {
-						AckUtils.autoAck(StaticMessageHeaderAccessor.getAcknowledgmentCallback(message));
-					}
-				})
+		return Mono.<Message<T>>create(monoSink ->
+						monoSink.onRequest((value) -> monoSink.success(messageSource.receive())))
 				.doOnError(MessagingException.class,
 						(ex) -> {
 							Message<?> failedMessage = ex.getFailedMessage();
@@ -152,7 +161,10 @@ public static <T> Flux<Message<T>> messageSourceToFlux(MessageSource<T> messageS
 										Mono.delay(ctx.getOrDefault(DELAY_WHEN_EMPTY_KEY,
 												DEFAULT_DELAY_WHEN_EMPTY)))))
 				.repeat()
-				.retryWhen(Retry.indefinitely().filter(MessagingException.class::isInstance));
+				.doOnDiscard(Message.class, (message) -> handleUndeliveredMessage(messageSource, message))
+				.retryWhen(Retry.indefinitely().filter(MessagingException.class::isInstance))
+				.doOnNext((message) ->
+						AckUtils.autoAck(StaticMessageHeaderAccessor.getAcknowledgmentCallback(message)));
 	}
 
 	/**
@@ -161,7 +173,8 @@ public static <T> Flux<Message<T>> messageSourceToFlux(MessageSource<T> messageS
 	 * is returned as is because it is already a {@link Publisher};
 	 * - a {@link SubscribableChannel} is subscribed with a {@link MessageHandler}
 	 * for the {@link Sinks.Many#tryEmitNext(Object)} which is returned from this method;
-	 * - a {@link PollableChannel} is wrapped into a {@link MessageSource} lambda and reuses
+	 * - a {@link PollableChannel} is wrapped into a {@link PollableChannelMessageSource} so
+	 * cancelled in-flight receives can be re-queued, and reuses
 	 * {@link #messageSourceToFlux(MessageSource)}.
 	 * @param messageChannel the {@link MessageChannel} to adapt.
 	 * @param <T> the expected payload type.
@@ -175,15 +188,46 @@ public static <T> Flux<Message<T>> messageChannelToFlux(MessageChannel messageCh
 		else if (messageChannel instanceof SubscribableChannel) {
 			return adaptSubscribableChannelToPublisher((SubscribableChannel) messageChannel);
 		}
-		else if (messageChannel instanceof PollableChannel) {
-			return messageSourceToFlux(() -> (Message<T>) ((PollableChannel) messageChannel).receive(0));
+		else if (messageChannel instanceof PollableChannel pollableChannel) {
+			return messageSourceToFlux(new PollableChannelMessageSource<>(pollableChannel));
 		}
 		else {
 			throw new IllegalArgumentException("The 'messageChannel' must be an instance of Publisher, " +
 					"SubscribableChannel or PollableChannel, not: " + messageChannel);
 		}
 	}
 
+	private static void handleUndeliveredMessage(MessageSource<?> messageSource, @Nullable Message<?> message) {
+		if (message == null) {
+			return;
+		}
+		if (messageSource instanceof PollableChannelMessageSource<?> pollableChannelMessageSource) {
+			pollableChannelMessageSource.returnMessage(message);
+		}
+	}
+
+	private static final class PollableChannelMessageSource<T> implements MessageSource<T> {
+
+		private final PollableChannel channel;
+
+		PollableChannelMessageSource(PollableChannel channel) {
+			this.channel = channel;
+		}
+
+		@Override
+		@SuppressWarnings("unchecked")
+		public @Nullable Message<T> receive() {
+			return (Message<T>) this.channel.receive(0);
+		}
+
+		void returnMessage(Message<?> message) {
+			if (!this.channel.send(message, 0)) {
+				LOGGER.warn("Failed to return undelivered message to pollable channel [" + this.channel + "]");
+			}
+		}
+
+	}
+
 	@SuppressWarnings("unchecked")
 	private static <T> Flux<Message<T>> adaptSubscribableChannelToPublisher(SubscribableChannel inputChannel) {
 		return Flux.defer(() -> {

spring-integration-core/src/test/java/org/springframework/integration/channel/IntegrationReactiveUtilsTests.java

@@ -19,8 +19,10 @@
 import java.time.Duration;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.LockSupport;
 
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Test;
@@ -32,16 +34,21 @@
 import reactor.test.StepVerifier;
 import reactor.util.concurrent.Queues;
 
+import org.springframework.integration.IntegrationMessageHeaderAccessor;
+import org.springframework.integration.acks.AcknowledgmentCallback;
 import org.springframework.integration.core.MessageSource;
 import org.springframework.integration.util.IntegrationReactiveUtils;
+import org.springframework.messaging.Message;
 import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.support.GenericMessage;
+import org.springframework.messaging.support.MessageBuilder;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
  * @author Sergei Egorov
  * @author Artem Bilan
+ * @author Fardan An
  *
  * @since 5.1.9
  */
@@ -150,4 +157,51 @@ void testRetryOnMessagingExceptionOnly() {
 		assertThat(finalException.get()).hasMessage("non-retryable RuntimeException");
 	}
 
+	@Test
+	void undeliveredMessageNotAcknowledgedWhenCancelledDuringBlockingReceive() throws InterruptedException {
+		CountDownLatch receiveBlocked = new CountDownLatch(1);
+		CountDownLatch receiveCompleted = new CountDownLatch(1);
+		AtomicBoolean releaseReceive = new AtomicBoolean();
+		AtomicReference<AcknowledgmentCallback.Status> ackStatus = new AtomicReference<>();
+		AtomicReference<Message<?>> delivered = new AtomicReference<>();
+
+		Message<?> testMessage = MessageBuilder.withPayload("test")
+				.setHeader(IntegrationMessageHeaderAccessor.ACKNOWLEDGMENT_CALLBACK,
+						(AcknowledgmentCallback) ackStatus::set)
+				.build();
+
+		MessageSource<Object> blockingSource = () -> {
+			receiveBlocked.countDown();
+			try {
+				while (!releaseReceive.get()) {
+					LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(10)); // NOSONAR busy wait
+				}
+				@SuppressWarnings("unchecked")
+				Message<Object> message = (Message<Object>) testMessage;
+				return message;
+			}
+			finally {
+				receiveCompleted.countDown();
+			}
+		};
+
+		Disposable subscription = IntegrationReactiveUtils.messageSourceToFlux(blockingSource)
+				.subscribe(delivered::set);
+
+		try {
+			assertThat(receiveBlocked.await(10, TimeUnit.SECONDS)).isTrue();
+			subscription.dispose();
+		}
+		finally {
+			releaseReceive.set(true);
+		}
+
+		assertThat(receiveCompleted.await(10, TimeUnit.SECONDS)).isTrue();
+		CountDownLatch singleIdle = new CountDownLatch(1);
+		Schedulers.single().schedule(singleIdle::countDown);
+		assertThat(singleIdle.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(delivered.get()).isNull();
+		assertThat(ackStatus.get()).isNull();
+	}
+
 }

spring-integration-core/src/test/java/org/springframework/integration/channel/reactive/ReactiveStreamsConsumerTests.java

@@ -23,7 +23,9 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.LockSupport;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -65,6 +67,7 @@
 
 /**
  * @author Artem Bilan
+ * @author Fardan An
  *
  * @since 5.0
  */
@@ -248,6 +251,66 @@ public void testReactiveStreamsConsumerPollableChannel() throws InterruptedExcep
 		reactiveConsumer.stop();
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	public void messageNotLostWhenStopDuringBlockingReceive() throws InterruptedException {
+		CountDownLatch receiveBlocked = new CountDownLatch(1);
+		AtomicBoolean releaseReceive = new AtomicBoolean();
+
+		QueueChannel testChannel = new QueueChannel(1) {
+
+			@Override
+			public Message<?> receive(long timeout) {
+				Message<?> message = super.receive(timeout);
+				if (message != null) {
+					receiveBlocked.countDown();
+					while (!releaseReceive.get()) {
+						LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(10)); // NOSONAR busy wait
+					}
+				}
+				return message;
+			}
+
+		};
+
+		Subscriber<Message<?>> testSubscriber = (Subscriber<Message<?>>) Mockito.mock(Subscriber.class);
+		AtomicReference<Message<?>> delivered = new AtomicReference<>();
+
+		willAnswer(i -> {
+			delivered.set(i.getArgument(0));
+			return null;
+		}).given(testSubscriber).onNext(any(Message.class));
+
+		ReactiveStreamsConsumer reactiveConsumer = new ReactiveStreamsConsumer(testChannel, testSubscriber);
+		reactiveConsumer.setBeanFactory(TEST_INTEGRATION_CONTEXT);
+		reactiveConsumer.afterPropertiesSet();
+		reactiveConsumer.start();
+
+		Message<?> testMessage = new GenericMessage<>("test");
+		testChannel.send(testMessage);
+
+		ArgumentCaptor<Subscription> captor = ArgumentCaptor.forClass(Subscription.class);
+		verify(testSubscriber).onSubscribe(captor.capture());
+		captor.getValue().request(1);
+
+		try {
+			assertThat(receiveBlocked.await(10, TimeUnit.SECONDS)).isTrue();
+			assertThat(testChannel.getQueueSize()).isZero();
+			reactiveConsumer.stop();
+		}
+		finally {
+			releaseReceive.set(true);
+		}
+
+		await().atMost(Duration.ofSeconds(5))
+				.until(() -> testChannel.getQueueSize() > 0);
+
+		assertThat(delivered.get()).isNull();
+		assertThat(testChannel.receive(0)).isSameAs(testMessage);
+
+		reactiveConsumer.stop();
+	}
+
 	@Test
 	public void testReactiveStreamsConsumerViaConsumerEndpointFactoryBean() throws Exception {
 		FluxMessageChannel testChannel = new FluxMessageChannel();