Fix cancel/receive race in messageSourceToFlux for PollableChannel

Replace synchronized cancellation flag with pending-message reclaim on a
shared scheduler and Sinks.One tryEmitValue to detect cancelled emission.
Handle failed PollableChannel re-queue via nack when send returns false.

Fixes #11262

Signed-off-by: arimu1 <19286898+arimu1@users.noreply.github.com>

Author: arimu1

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

@@ -17,6 +17,7 @@
 package org.springframework.integration.util;
 
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.LockSupport;
 
 import io.micrometer.context.ContextSnapshotFactory;
@@ -27,6 +28,7 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
+import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 import reactor.util.context.Context;
 import reactor.util.context.ContextView;
@@ -129,29 +131,18 @@ public static ContextView captureReactorContext() {
 	 */
 	@SuppressWarnings("NullAway")
 	public static <T> Flux<Message<T>> messageSourceToFlux(MessageSource<T> messageSource) {
-		return Mono.
-				<Message<T>>create(monoSink -> {
-					Object cancellationMonitor = new Object();
-					boolean[] cancelled = {false};
-					monoSink.onCancel(() -> {
-						synchronized (cancellationMonitor) {
-							cancelled[0] = true;
-						}
-					});
-					monoSink.onRequest(value -> {
-						synchronized (cancellationMonitor) {
-							if (cancelled[0]) {
-								return;
-							}
-							Message<T> message = messageSource.receive();
-							if (cancelled[0]) {
-								handleUndeliveredMessage(messageSource, message);
-								return;
-							}
-							monoSink.success(message);
-						}
-					});
-				})
+		Scheduler scheduler = Schedulers.single();
+		return Mono.defer(() -> {
+			AtomicReference<Message<?>> pendingMessage = new AtomicReference<>();
+			Sinks.One<Message<T>> sink = Sinks.one();
+
+			return sink.asMono()
+					.doOnCancel(() -> scheduler.schedule(() -> {
+						Message<?> undeliveredMessage = pendingMessage.getAndSet(null);
+						handleUndeliveredMessage(messageSource, undeliveredMessage);
+					}))
+					.doOnRequest(request -> scheduler.schedule(() -> emitReceivedMessage(messageSource, pendingMessage, sink)));
+		})
 				.doOnSuccess((message) -> {
 					if (message != null) {
 						AckUtils.autoAck(StaticMessageHeaderAccessor.getAcknowledgmentCallback(message));
@@ -165,7 +156,7 @@ public static <T> Flux<Message<T>> messageSourceToFlux(MessageSource<T> messageS
 							}
 							LOGGER.error("Error from Flux for : " + messageSource, ex);
 						})
-				.subscribeOn(Schedulers.single())
+				.subscribeOn(scheduler)
 				.repeatWhenEmpty((repeat) ->
 						repeat.flatMap((increment) ->
 								Mono.deferContextual(ctx ->
@@ -204,6 +195,26 @@ else if (messageChannel instanceof PollableChannel pollableChannel) {
 		}
 	}
 
+	@SuppressWarnings("NullAway")
+	private static <T> void emitReceivedMessage(MessageSource<T> messageSource,
+			AtomicReference<Message<?>> pendingMessage, Sinks.One<Message<T>> sink) {
+
+		try {
+			Message<T> message = messageSource.receive();
+			if (message != null) {
+				pendingMessage.set(message);
+			}
+			Sinks.EmitResult emitResult = sink.tryEmitValue(message);
+			if (emitResult == Sinks.EmitResult.FAIL_CANCELLED || emitResult == Sinks.EmitResult.FAIL_TERMINATED) {
+				handleUndeliveredMessage(messageSource, message);
+			}
+			pendingMessage.set(null);
+		}
+		catch (Exception ex) {
+			sink.tryEmitError(ex);
+		}
+	}
+
 	private static void handleUndeliveredMessage(MessageSource<?> messageSource, @Nullable Message<?> message) {
 		if (message == null) {
 			return;
@@ -231,7 +242,11 @@ private static final class PollableChannelMessageSource<T> implements MessageSou
 		}
 
 		void returnMessage(Message<?> message) {
-			this.channel.send(message);
+			if (!this.channel.send(message, 0)) {
+				LOGGER.warn("Failed to return undelivered message to pollable channel [" + this.channel
+						+ "]; nacking instead");
+				AckUtils.autoNack(StaticMessageHeaderAccessor.getAcknowledgmentCallback(message));
+			}
 		}
 
 	}

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

@@ -291,6 +291,49 @@ public void messageNotLostOnStopWithPollableChannel() throws InterruptedExceptio
 		assertThat(stealCount).isZero();
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	public void messageNotLostOnStopWithBoundedPollableChannel() throws InterruptedException {
+		int iterations = 500;
+		int stealCount = 0;
+		for (int iteration = 0; iteration < iterations; iteration++) {
+			QueueChannel testChannel = new QueueChannel(1);
+			Subscriber<Message<?>> testSubscriber = (Subscriber<Message<?>>) Mockito.mock(Subscriber.class);
+			BlockingQueue<Message<?>> messages = new LinkedBlockingQueue<>();
+
+			willAnswer(i -> {
+				messages.put((Message<?>) 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);
+
+			assertThat(messages.poll(10, TimeUnit.SECONDS)).isSameAs(testMessage);
+
+			reactiveConsumer.stop();
+			testChannel.send(new GenericMessage<>("bait-" + iteration));
+
+			for (int spin = 0; spin < 20000 && testChannel.getQueueSize() > 0; spin++) {
+				// busy-poll to hit the race window
+			}
+
+			if (testChannel.getQueueSize() == 0) {
+				stealCount++;
+			}
+		}
+		assertThat(stealCount).isZero();
+	}
+
 	@Test
 	public void testReactiveStreamsConsumerViaConsumerEndpointFactoryBean() throws Exception {
 		FluxMessageChannel testChannel = new FluxMessageChannel();