Fix message loss on stop() for PollableChannel reactive consumer

[arimu1]

Summary

• On cancel during an in-flight MessageSource.receive(), do not nack. Kafka/AMQP recover via their own redelivery when the delivery is left unacknowledged.
• PollableChannel sources adapted through messageChannelToFlux are best-effort re-queued (send(message, 0)). A direct messageSourceToFlux(() -> channel.receive(0)) lambda is not re-queued.
• autoAck runs on Flux.doOnNext of the returned flux (delivery to this flux), not on Mono.doOnSuccess.
• Discard rescue is doOnDiscard only (no private onNextDropped hook). Coverage is drops at or upstream of this flux.

Fixes #11262

Test plan

• ./gradlew :spring-integration-core:test --tests "*IntegrationReactiveUtilsTests" --tests "*ReactiveStreamsConsumerTests" --tests "*ReactiveMessageSourceProducerTests" (JDK 17)

[arimu1]

Thanks for the review — addressed all three points in 5ff484a:

1. Cancel/receive race — removed the synchronized cancelled[] flag. receive() no longer runs under any lock. Pending messages are tracked in an AtomicReference, cancel/reclaim is scheduled on the same Schedulers.single() instance as onRequest, and emission uses Sinks.One.tryEmitValue() so FAIL_CANCELLED / FAIL_TERMINATED reliably triggers reclaim instead of silently dropping after a discarded success().
2. No lock across receive/emit — onRequest only schedules work on the shared scheduler; receive() and tryEmitValue() run outside any monitor.
3. returnMessage send=false — PollableChannelMessageSource.returnMessage() now uses non-blocking send(message, 0) and nacks when re-queue fails.

Tests: ./gradlew :spring-integration-core:test --tests "org.springframework.integration.channel.reactive.ReactiveStreamsConsumerTests" --tests "org.springframework.integration.channel.IntegrationReactiveUtilsTests" (includes messageNotLostOnStopWithPollableChannel + new bounded QueueChannel(1) variant).

Please take another look when you have a moment.

[arimu1]

Addressed round-2 review on 5ff484ae (reactor-core 3.8.7 Sinks.One.tryEmitValue OK-with-zero-delivery after cancel).

Changes (c310c5f377)

• Clear pendingMessage only in inner sink.asMono().doOnSuccess (actual downstream delivery), not after tryEmitValue.
• After tryEmitValue, pendingMessage.getAndSet(null) and reclaim when still present — covers OK-but-undelivered without relying on FAIL_CANCELLED/FAIL_TERMINATED.
• Kept lock-free receive/emit and bounded send(message, 0) + nack fallback.

Test

• Added messageNotLostOnStopDuringFirstPoll — fails on the prior OK-clear path, passes with reclaim (1500-iteration race loop).

Verification

./gradlew :spring-integration-core:test \
  --tests '*ReactiveStreamsConsumerTests.messageNotLostOnStopDuringFirstPoll' \
  --tests '*ReactiveStreamsConsumerTests.messageNotLostOnStopWithPollableChannel' \
  --tests '*ReactiveStreamsConsumerTests.messageNotLostOnStopWithBoundedPollableChannel'

Ready for re-review.

[artembilan]

@arimu1 ,

can you agree with your Cursor locally and push into PR only after that?

I know, I did something similar before myself but just as experiment and I reviewed it anyway not bothering anyone else.

Here, this Cursor chat with itself sounds like not final solution, so a bit not fair to ask the team for review.

thanks for understanding , and for the help!
I planned to talk to Claude on the matter already after release next week 😅

[artembilan]

Thanks for the investigation in #11262 — that analysis is good and the race is real.
I have run this branch locally and instrumented it, and the conclusion is that the fix does not close the hole, while Reactor already offers a much smaller way to close it. Details below.

1. The new test does not pass on this branch

./gradlew :spring-integration-core:test --tests "*ReactiveStreamsConsumerTests" --tests "*IntegrationReactiveUtilsTests" at c310c5f377:

ReactiveStreamsConsumerTests > messageNotLostOnStopDuringFirstPoll() FAILED
    org.opentest4j.AssertionFailedError: expected: 0 but was: 41
    at ReactiveStreamsConsumerTests.java:289
13 tests completed, 1 failed

Build scan: https://ge.spring.io/s/h4q5mkpbj7cey

Only the DCO check ran on this PR, so there is no CI signal — that local run is all we have, and 41 losses out of 1500 is an order of magnitude worse than the ~0.2% you measured on main.

2. Where those messages actually go: Operators.onDiscard

The issue says the dropped value "doesn't even go through Reactor's Operators.onNextDropped hook". True of onNextDropped — but the value is not lost to Reactor. MonoCreate.DefaultMonoSink (reactor-core 3.8.7):

public void success(@Nullable T value) {
    if (value == null) { success(); return; }
    Disposable d = this.disposable;
    if (d == CANCELLED) {
        Operators.onDiscard(value, actual.currentContext());   // <-- our case
        return;
    }
    else if (d == TERMINATED) {
        Operators.onNextDropped(value, actual.currentContext());
        return;
    }
    ...

cancel() sets disposable to CANCELLED before success(messageSource.receive()) runs, so the message we just drained is handed to the discard handler of the current subscriber context.

I measured this. Harness: QueueChannel(1), subscribe with a BaseSubscriber that does not request in hookOnSubscribe (same as SubscriberDecorator), send, request(1), dispose(), then spin until the message is either delivered or back in the channel; a .doOnDiscard(Message.class, ...) counter on the flux; 3000 iterations. Every message is accounted for exactly once:

	delivered	to onDiscard	back in channel	unaccounted
main, Mono.create	3	25	2971	1
this PR, Sinks.One	2	16	2982	0

Two things follow.

The reclaim isn't catching them. On this branch 16 of 3000 messages are still handed to Reactor's discard hook and are not re-queued (the totals add to exactly 3000, so none of those 16 also landed in the channel). In production there is no discard handler installed, so those 16 are silently lost — the same defect the PR is meant to fix, consistent with messageNotLostOnStopDuringFirstPoll failing above.

A discard handler catches them on unmodified main. 25 of 3000, with no production change at all. So the rescue can be a handler on the existing chain — no Sinks.One, no AtomicReference, no extra scheduler task, no restructuring of the poll loop:

.doOnDiscard(Message.class, (message) -> /* return to the channel, or nack */)

Note the single remaining unaccounted on main: that one goes through the TERMINATED branch above, i.e. Operators.onNextDropped, which also honours a per-sequence context key (Hooks.KEY_ON_NEXT_DROPPED). A complete fix wants both. Still a handful of lines.

That route also covers what this PR structurally cannot. The reclaim relies on the value reaching the subscriber synchronously inside tryEmitValue(): the inner doOnSuccess clears pendingMessage, and only then does emitReceivedMessage do its getAndSet reclaim. Any async boundary downstream breaks that, and the framework has two of its own — ReactiveStreamsConsumer.setReactiveCustomizer(...) (a user publishOn, onBackpressureBuffer, ...) and the ReactiveMessageHandler branch of doStart(), which is a flatMap. There, tryEmitValue returns OK, doOnSuccess clears pendingMessage, the message sits in the operator queue, cancellation drops it, and the reclaim finds null. Reactor routes those queued-but-cancelled values through the same discard hook.

3. This branch also introduces a spurious ack

Follows from the numbers above rather than from reading alone. Those 16 messages were not re-queued, so pendingMessage had been cleared, so the inner doOnSuccess ran — and the outer .doOnSuccess(... AckUtils.autoAck ...) is the next operator in the same synchronous chain. So they were acknowledged and then discarded downstream.

Mono.create cannot do that: it checks CANCELLED before actual.onNext(), so nothing is acked when the value is dropped. For ack-capable sources going through messageSourceToFlux (which is public API — ReactiveMessageSourceProducer and every MessageSource-backed reactive adapter), "acked but never delivered" is a worse outcome than the bug being fixed.

4. Notes on the current implementation

Whichever route we take:

1. The doOnCancel reclaim is unreachable. pendingMessage is only assigned inside emitReceivedMessage, which always ends with getAndSet(null) plus reclaim, and doOnSuccess clears it on delivery. By the time that task runs on the single scheduler, the reference can only be null.

2. The scheme silently depends on Schedulers.single() being single-threaded. MonoSubscribeOn.trySchedule() trampolines the downstream request onto the scheduler worker, so the doOnRequest callback queues emitReceivedMessage behind the running request task and the sink's demand is registered first. Invert that order and SinkOneMulticast parks the value (tryEmitValue returns OK with an empty subscriber array and nothing delivered), the reclaim pushes the message back into the channel, and the later request drains the parked value as well — a silent duplicate. Nothing in the code or the tests states that invariant.

3. channel.send(message, 0) is not a neutral rescue. It re-enters the channel's interceptor chain (wire taps fire twice, metrics double-count, message-store-backed channels re-persist) and appends at the tail, so FIFO is not preserved when the queue is not empty.

4. On a bounded channel the fix degrades to today's behaviour. send(..., 0) on a full queue fails, we autoNack, and a plain QueueChannel message has no AcknowledgmentCallback — still lost, now with a WARN. Same for the generic MessageSource branch: file, JDBC and friends have no ack callback, so autoNack is a no-op. Best-effort is acceptable, but messageSourceToFlux is public API and the Javadoc must state what happens to an in-flight message on cancellation.

5. sink.tryEmitError(ex) on a cancelled sink swallows the exception outright, where Mono.create's error() goes to Operators.onErrorDropped — logged and hookable. Let's keep that.

6. An extra scheduler dispatch per poll, plus a Sinks.One and an AtomicReference per poll, on the hot path of every pollable-channel-backed reactive flow — to close a race that only happens on stop().

5. Tests

3500 iterations across three tests, each building an endpoint and busy-spinning up to 20000 Thread.yield()s, asserting on probabilistic counters, in the standard test task. Two of them (...WithPollableChannel, ...WithBoundedPollableChannel) are near-identical.

They are also weak guards: assertThat(stealCount).isZero() only proves the bait did not vanish. It passes when the message was stolen and re-queued, and it would pass on a duplicate delivery. And as my table shows, the settle window matters — gating the spin on the wrong condition changes the measurement, which is exactly what you do not want in a regression test.

Please replace them with one deterministic test: a MessageSource that blocks inside receive() on a latch the test releases only after stop() has cancelled. Exact window, one iteration, and it can assert the real invariant — delivered or back in the channel, exactly once, and not acked when it is not delivered. If you want a soak loop too, it belongs behind @LongRunningIntegrationTest.

Also please confirm testReactiveStreamsConsumerPollableChannel (the test that actually flakes in CI) is stable afterwards.

6. Small things

• messageChannelToFlux Javadoc still says a PollableChannel "is wrapped into a MessageSource lambda" — it is a PollableChannelMessageSource now.
• Please add an @author tag with your real first and last name to the classes you touch (see CONTRIBUTING).
• The last commit message attributes the OK-but-not-delivered result to a "reactor-core 3.8.7 cancel path". It is not version-specific: Sinks.one() is SinkOneMulticast, and tryEmitValue returns OK whenever the subscriber array is empty but not terminated. I would rather not carry a fix in core whose rationale reads as pinned to a snapshot.

So: not this shape, please. The discard-hook route looks much more promising, and it needs a decision on re-queue-vs-nack semantics for a PollableChannel first — let's settle that on #11262 before the next push.

[artembilan]

Correction to §3 of my review, plus one more data point. I re-ran the probe with a real AcknowledgmentCallback on the message so the ack is measured rather than inferred — same harness, messageSourceToFlux with a () -> channel.receive(0) source (i.e. the public-API path, not the PollableChannelMessageSource one), QueueChannel(1), request(1) then immediate dispose(), 3000 iterations:

	delivered	to onDiscard	onNextDropped	still in channel	ACCEPT without delivery	REJECT
main, Mono.create	19	16	3	2962	8	0
this PR, Sinks.One	9	14	0	2974	14	3

So I was wrong to call the spurious ack something this PR introduces — main does it too, 8 times in 3000. The CANCELLED check in MonoCreate.DefaultMonoSink.success() only guards the case where cancellation has reached the sink itself; when a downstream operator is already cancelled, actual.onNext(value) still runs, the outer .doOnSuccess(... autoAck ...) fires, and the value is dropped after that. My §3 reasoning was right about the mechanism and wrong about it being new.

What the numbers do say:

• This PR roughly doubles it — 8 → 14 acked-but-never-delivered per 3000 — and rescues only 3 of the 17 drained-and-not-delivered messages (the 3 REJECTs, which are also the 3 "unaccounted" in my earlier table since a lambda source gets nacked rather than re-queued).
• onNextDropped=3 on main confirms what I guessed at in the review: the remainder that doOnDiscard does not see go through the TERMINATED branch. So a discard handler alone is not the whole answer — Hooks.KEY_ON_NEXT_DROPPED needs covering too.

And it exposes something both routes have to deal with, which I had not appreciated: autoAck runs too early. The current .doOnSuccess(... autoAck ...) acknowledges as soon as the value leaves the source Mono, before we know any subscriber took it. Once a message is acked, a later discard handler cannot rescue it — AckUtils.autoNack is a no-op on an already-acknowledged callback, so those 8 (or 14) are unrecoverable no matter what hook we install. Any real fix has to move the acknowledgement to the point of actual delivery, otherwise the discard/dropped handlers only ever rescue the subset that was dropped before the ack.

That does not change my conclusion on this PR — the reclaim mechanism still misses the majority of cases and the tests still need to be deterministic — but it does mean the design discussion on #11262 should cover ack placement, not just re-queue-vs-nack.

[arimu1]

Replaced the Sinks.One reclaim with Reactor discard hooks on the existing Mono.create chain (doOnDiscard(Message.class, …) + sequence-local onNextDropped). Probabilistic soak tests are gone; added a deterministic blocking-receive test instead. Semantics proposal is on #11262.

Tip: e39ca9691f

[artembilan]

Thanks for reworking this onto the discard hooks — that shape is right, and I confirmed the branch is green locally (:spring-integration-core:test --tests "*IntegrationReactiveUtilsTests" --tests "*ReactiveStreamsConsumerTests" → 12 tests, 0 failures, both new tests passing).

Unfortunately the else branch of handleUndeliveredMessage() is a data-loss regression, so I can't merge this.

1. autoNack on cancellation destroys the message for ack-capable sources

Every non-PollableChannel source goes down AckUtils.autoNack(...), which is acknowledge(Status.REJECT). REJECT is not "put it back" — its javadoc is "Mark the message as rejected", as distinct from REQUEUE, "Reject the message and requeue so that it will be redelivered".

For the same in-flight message with the sink CANCELLED:

	ack call	KafkaMessageSource.KafkaAckCallback	AmqpMessageSource
main	none — MonoCreate.DefaultMonoSink.success() calls Operators.onDiscard(value, ctx) and no hook is installed	position advanced, offset not committed	delivery tag left unacked
this PR	autoNack → Status.REJECT	case ACCEPT, REJECT -> commitIfPossible(record)	basicReject(deliveryTag, false)

KafkaAckCallback.autoAckEnabled defaults to true (KafkaMessageSource.java:810), so the isAutoAck() gate in autoNack does not save us.

To be precise about the delta, since it is not "main is safe": on main, a stop() during an in-flight poll does lose the message for the life of that source instance — but it is recovered, because Kafka has not committed the offset (redelivered on restart/rebalance) and AMQP has not settled the delivery (requeued when the channel closes). With this PR the Kafka offset is committed and the AMQP message is rejected without requeue, so it is gone across restarts too. A PR titled "Fix message loss" cannot make loss permanent where it was previously recoverable.

Your own new test pins this: undeliveredMessageNackedWhenCancelledDuringBlockingReceive asserts Status.REJECT.

2. ...and AckUtils.requeue is not the fix either

Please don't just swap autoNack → requeue — that's a patch on the wrong shape, and it would still be wrong.

The real problem is the one I raised on #11262: autoAck runs in .doOnSuccess(...), i.e. as soon as the value leaves the source Mono, before any subscriber has taken it. Move the acknowledgement to the point of actual delivery and this whole branch disappears: an in-flight message that was never delivered is simply never acknowledged, and each source's own redelivery semantics (broker redelivery, offset rollback, filter rollback) handle it. No explicit nack is needed for ack-capable sources at all. The discard hook then only has to exist for the PollableChannel case, where there is no ack protocol to fall back on.

So "Ack placement ... is a separate concern worth addressing later" doesn't hold — ack placement is precisely the constraint that makes the else branch unsound. It has to be part of this change, not after it.

3. The onNextDropped branch needs a test, or it should go

REACTOR_ON_NEXT_DROPPED_CONTEXT_KEY = "reactor.onNextDropped.local" is hardcoded because Hooks.KEY_ON_NEXT_DROPPED is package-private in reactor-core (3.8.6, Hooks.java:649). There's no public contract for that key, and if Reactor changes it we get a silent regression — nothing fails to compile and no test goes red.

I'd accept that coupling if it demonstrably bought us something, but neither new test reaches the branch: both block inside receive() until after cancel(), which is the d == CANCELLED → Operators.onDiscard path that doOnDiscard already covers. So either:

• add a test that provably lands in Operators.onNextDropped — MonoCreate.DefaultMonoSink.request() has no once-guard on requestConsumer, so a second request() re-invokes it and the second success(value) hits d == TERMINATED; that looks reachable, but I've only read it, not measured it — or
• drop the context hook and keep doOnDiscard alone.

Note also that unlike Operators.discardLocalAdapter, which composes (safeConsumer.andThen(consumer)), a plain ctx.put(...) replaces any local onNextDropped hook set downstream.

4. Both new tests can wedge Schedulers.single() for the rest of the JVM

messageSourceToFlux does subscribeOn(Schedulers.single()) — one cached, JVM-wide worker — and the blocking receive() occupies it. In messageNotLostWhenStopDuringBlockingReceive there are two assertions between the latch and the release:

assertThat(receiveBlocked.await(10, TimeUnit.SECONDS)).isTrue();
assertThat(testChannel.getQueueSize()).isZero();

reactiveConsumer.stop();
releaseReceive.set(true);

If either assertion fails, the busy-wait never exits, that shared worker parks forever, and every later test in the same JVM using messageSourceToFlux stalls — one failure cascades into unrelated timeouts. Please wrap the gate in try { ... } finally { releaseReceive.set(true); } in both tests.

Smaller things

• assertThat(wasDelivered ^ wasRequeued).isTrue() — you replaced the soak loops with a deterministic test, so assert the deterministic outcome. As written it also passes if the message is delivered after stop(), which is a different bug, not a pass.
• The try/catch → monoSink.error(ex) around messageSource.receive() is unrelated to this fix. MonoCreate.DefaultMonoSink.request() calls requestConsumer.accept(n) with no try/catch, so this genuinely changes where a receive() exception surfaces, and now feeds it to retryWhen. It may well be an improvement — please raise it separately with its own test rather than folding it in.
• The new javadoc — "Reactor routes the value through discard hooks which rescue it here" — is only true for drops at or upstream of the contextWrite. Anything a caller adds downstream (ReactiveStreamsConsumer.setReactiveCustomizer, the flatMap on the ReactiveMessageHandler path, the final subscriber) discards with a context that does not carry our hook, and those messages were already ACCEPTed by doOnSuccess. Please state the limits rather than implying full coverage.
• Commit message: the headline needs the GH-11262: prefix and the trailer needs the full URL, Fixes: https://github.com/spring-projects/spring-integration/issues/11262 (see CONTRIBUTING.md).
• The git author is arimu1 but the javadoc @author says Fardan An. Please set user.name to your real first and last name so the two match.

[arimu1]

Thanks @artembilan — addressed on dff9f151d5.

Ack / nack

• Discard no longer nacks. Non-PollableChannel sources are left unacknowledged so Kafka/AMQP keep their own redelivery.
• autoAck moved from Mono.doOnSuccess to Flux.doOnNext on the returned flux (after retryWhen).
• Failed send(..., 0) on a full bounded channel is logged only.

onNextDropped

• Removed the private reactor.onNextDropped.local hook. Rescue is doOnDiscard only.

Tests

• Latch release is in finally in both tests so Schedulers.single() cannot stick.
• messageNotLostWhenStopDuringBlockingReceive now asserts re-queue (not XOR).
• The source-callback test asserts the status stays unset after cancel (undeliveredMessageNotAcknowledgedWhenCancelledDuringBlockingReceive).

Smaller

• Restored receive() without a wrapping try/catch.
• Javadoc states the discard-hook limit (drops at/upstream of this flux only).
• Commit uses GH-11262: / Fixes: URL; author is Fardan An.

:spring-integration-core:test --tests "*IntegrationReactiveUtilsTests" --tests "*ReactiveStreamsConsumerTests" --tests "*ReactiveMessageSourceProducerTests" is green locally (JDK 17).