From 6fe0dca0d58894d89f950c63a331ceee5bdd48f8 Mon Sep 17 00:00:00 2001 From: Zixuan Liu Date: Thu, 3 Sep 2026 15:57:08 +0800 Subject: [PATCH] [fix][client] Defer op cmd release to the write event loop on send timeout Motivation The send-timeout path (failPendingMessages with cnx() == null) released op.cmd and recycled the op inline on the timer thread while the message frame may still have been queued for writing on a connection event loop: the write callback sits behind a reconnect/disconnect storm, or the buffers are already in the channel outbound buffer. Releasing the buffers from the wrong thread could return them to the pool while the in-flight write is still reading them, letting new batches overwrite the frame content and corrupting the wire bytes (broken checksum / lost frame sync on the broker). The existing cnx != null branch already deferred this cleanup to the event loop; the null (reconnect-window) branch did not. Modifications - OpSendMsg tracks the event loop each cmd was last handed to for writing (writeEventLoop), and cmd is now volatile so stale callbacks can never observe a released cmd. - releaseOpCmdAndRecycle defers the cmd release and op recycle to that event loop (serialized after the in-flight write), with an inline fallback when the loop is shutting down. releaseOpCmd clears op.cmd before releasing so a stale callback always fails its guard. - WriteInEventLoopCallback now skips stale writes (op re-sent on another connection or already disposed) and only drops the reference it took, instead of writing a released buffer or mutating a recycled op. - failPendingMessages fails inline when the connection event loop rejects the deferred task (shutting down during reconnect churn) instead of leaving the pending messages queued. - ackReceived / recoverChecksumError / recoverNotAllowedError now go through the same write-loop-aware release. --- .../pulsar/client/impl/ProducerImpl.java | 105 ++++- .../pulsar/client/impl/ProducerImplTest.java | 390 ++++++++++++++++++ 2 files changed, 474 insertions(+), 21 deletions(-) diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java index 052055b329a62..c914644bdae96 100644 --- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java +++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ProducerImpl.java @@ -37,6 +37,7 @@ import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelOutboundHandler; import io.netty.channel.ChannelPromise; +import io.netty.channel.EventLoop; import io.netty.util.AbstractReferenceCounted; import io.netty.util.Recycler; import io.netty.util.Recycler.Handle; @@ -59,6 +60,7 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -1176,8 +1178,18 @@ public void run() { .log("Sending message"); try { - cnx.ctx().writeAndFlush(cmd, cnx.ctx().voidPromise()); - op.updateSentTimestamp(); + // Stale guard: the op may have been re-sent on another connection or disposed while this callback + // was queued behind a disconnect storm. Only write if this loop is still the op's write loop and + // op.cmd still identifies our cmd; otherwise drop only the reference this callback took. + if (op.writeEventLoop == cnx.ctx().channel().eventLoop() && op.cmd == cmd) { + cnx.ctx().writeAndFlush(cmd, cnx.ctx().voidPromise()); + // Re-check before mutating: the op may have been disposed concurrently on another loop. + if (op.cmd == cmd) { + op.updateSentTimestamp(); + } + } else { + cmd.release(); + } } finally { recycle(); } @@ -1428,8 +1440,7 @@ protected void ackReceived(ClientCnx cnx, long sequenceId, long highestSequenceI .log("Got exception while completing the callback for msg"); } } - ReferenceCountUtil.safeRelease(op.cmd); - op.recycle(); + releaseOpCmdAndRecycle(op); } protected long getHighestSequenceId(OpSendMsg op) { @@ -1488,8 +1499,7 @@ protected synchronized void recoverChecksumError(ClientCnx cnx, long sequenceId) .exception(t) .log("Got exception while completing the callback for msg"); } - ReferenceCountUtil.safeRelease(op.cmd); - op.recycle(); + releaseOpCmdAndRecycle(op); return; } else { log.debug() @@ -1521,8 +1531,7 @@ protected synchronized void recoverNotAllowedError(long sequenceId, String error .exception(t) .log("Got exception while completing the callback for msg"); } - ReferenceCountUtil.safeRelease(op.cmd); - op.recycle(); + releaseOpCmdAndRecycle(op); } } @@ -1612,10 +1621,21 @@ protected static final class OpSendMsg { LatencyHistogram rpcLatencyHistogram; MessageImpl msg; List> msgs; - ByteBufPair cmd; + // Volatile so a stale WriteInEventLoopCallback on an older connection's event loop can never observe + // cmd == this after the disposal has released it; the disposal also clears this field before releasing. + volatile ByteBufPair cmd; SendCallback callback; Runnable rePopulate; ChunkedMessageCtx chunkedMessageCtx; + /** + * The connection event loop that this op's {@link #cmd} was last handed to for writing, or + * {@code null} if the {@link #cmd} has never been handed to a connection. + * + *

When a timeout or recovery path fails this op from another thread (e.g. the send-timeout timer + * while the producer is in the reconnect window), the cmd release and op recycle are deferred to this + * loop so they are serialized after any in-flight write. + */ + volatile EventLoop writeEventLoop; long uncompressedSize; long sequenceId; long createdAt; @@ -1635,6 +1655,7 @@ void initialize() { cmd = null; callback = null; rePopulate = null; + writeEventLoop = null; sequenceId = -1L; createdAt = -1L; firstSentAt = -1L; @@ -2406,8 +2427,7 @@ synchronized void failPendingMessages(ClientCnx cnx, PulsarClientException ex) { } client.getMemoryLimitController().releaseMemory(op.uncompressedSize); - ReferenceCountUtil.safeRelease(op.cmd); - op.recycle(); + releaseOpCmdAndRecycle(op); } semaphoreRelease(releaseCount.get()); @@ -2418,12 +2438,55 @@ synchronized void failPendingMessages(ClientCnx cnx, PulsarClientException ex) { } else { // If we have a connection, we schedule the callback and recycle on the event loop thread to avoid any // race condition since we also write the message on the socket from this thread - cnx.ctx().channel().eventLoop().execute(() -> { - synchronized (ProducerImpl.this) { - failPendingMessages(null, ex); - } - }); + try { + cnx.ctx().channel().eventLoop().execute(() -> { + synchronized (ProducerImpl.this) { + failPendingMessages(null, ex); + } + }); + } catch (RejectedExecutionException e) { + // The connection's event loop is shutting down; the deferred task would be dropped. Fail inline + // instead (per-op cmd release is still deferred to each op's write event loop). + log.warn() + .exception(e) + .log("Connection event loop is shutting down while failing pending messages; failing inline"); + failPendingMessages(null, ex); + } + } + } + + /** + * Releases the op's cmd and recycles the op, deferred to the event loop the cmd was last handed to for writing + * so the buffers are never released on another thread while a write may still be reading them (e.g. the + * send-timeout timer thread during the reconnect window). If the op was never handed to a connection, the + * cleanup runs inline. + */ + private void releaseOpCmdAndRecycle(OpSendMsg op) { + final EventLoop writeEventLoop = op.writeEventLoop; + if (writeEventLoop != null) { + try { + writeEventLoop.execute(() -> releaseOpCmd(op)); + return; + } catch (RejectedExecutionException e) { + // Event loop shutting down: release inline (the write either never ran, or the channel now owns + // the buffer reference it held). + log.warn() + .exception(e) + .log("Write event loop is shutting down while releasing an in-flight op cmd"); + } } + releaseOpCmd(op); + } + + /** + * Detaches {@code op.cmd} before releasing it so a stale write callback can never pass its guard and write a + * released (possibly re-pooled) buffer, then recycles the op. + */ + private void releaseOpCmd(OpSendMsg op) { + final ByteBufPair cmd = op.cmd; + op.cmd = null; + ReferenceCountUtil.safeRelease(cmd); + op.recycle(); } /** @@ -2556,8 +2619,7 @@ protected synchronized void processOpSendMsg(OpSendMsg op) { // releaseSemaphoreForSendOp() also gives the reserved memory back, so it must not be released again. releaseSemaphoreForSendOp(op); op.sendComplete(getTerminalException(state)); - ReferenceCountUtil.safeRelease(op.cmd); - op.recycle(); + releaseOpCmdAndRecycle(op); return; } pendingMessages.add(op); @@ -2577,6 +2639,7 @@ protected synchronized void processOpSendMsg(OpSendMsg op) { // If we do have a connection, the message is sent immediately, otherwise we'll try again once a new // connection is established op.cmd.retain(); + op.writeEventLoop = cnx.ctx().channel().eventLoop(); cnx.ctx().channel().eventLoop().execute(WriteInEventLoopCallback.create(this, cnx, op)); stats.updateNumMsgsSent(op.numMessagesInBatch, op.batchSizeByte); } else { @@ -2664,7 +2727,7 @@ private void recoverProcessOpSendMsgFrom(ClientCnx cnx, MessageImpl latestMsg + " but not, which is not expected."); releaseSemaphoreForSendOp(op); msgIterator.remove(); - op.recycle(); + releaseOpCmdAndRecycle(op); continue; } } else if (op.msg == latestMsgAttemptedRegisteredSchema && failedIncompatibleSchema @@ -2700,7 +2763,6 @@ private void recoverProcessOpSendMsgFrom(ClientCnx cnx, MessageImpl latestMsg op.rePopulate.run(); } msgIterator.remove(); - ReferenceCountUtil.safeRelease(op.cmd); try { // Need to protect ourselves from any exception being thrown in the future handler from the // application @@ -2711,7 +2773,7 @@ private void recoverProcessOpSendMsgFrom(ClientCnx cnx, MessageImpl latestMsg .log("Got exception while completing the failed publishing"); } releaseSemaphoreForSendOp(op); - op.recycle(); + releaseOpCmdAndRecycle(op); continue; } else if (op.msg.getSchemaState() == None) { // Event 1-1. @@ -2734,6 +2796,7 @@ private void recoverProcessOpSendMsgFrom(ClientCnx cnx, MessageImpl latestMsg stripChecksum(op); } op.cmd.retain(); + op.writeEventLoop = cnx.ctx().channel().eventLoop(); log.debug() .attr("sequenceId", op.sequenceId) .log("Re-Sending message"); diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ProducerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ProducerImplTest.java index c35f8a77b515b..37f63dbd45a1a 100644 --- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ProducerImplTest.java +++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ProducerImplTest.java @@ -18,19 +18,35 @@ */ package org.apache.pulsar.client.impl; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.fail; +import io.netty.buffer.Unpooled; +import io.netty.channel.Channel; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.EventLoop; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; +import java.util.concurrent.RejectedExecutionException; import org.apache.commons.lang3.reflect.FieldUtils; import org.apache.pulsar.client.api.PulsarClientException; import org.apache.pulsar.client.api.Schema; +import org.apache.pulsar.client.impl.ProducerImpl.OpSendMsg; +import org.apache.pulsar.client.impl.ProducerImpl.OpSendMsgQueue; import org.apache.pulsar.client.impl.metrics.LatencyHistogram; import org.apache.pulsar.common.api.proto.MessageMetadata; import org.apache.pulsar.common.protocol.ByteBufPair; @@ -193,4 +209,378 @@ public void testProcessOpSendMsgInTerminalStateReleasesMemoryOnce() throws Excep "The memory reserved for the message must be released exactly once in state " + state); } } + + /** + * Regression test for the send-timeout vs in-flight-write race (scenario B): + * + *

A batch frame is handed to a connection's event loop for writing, then the connection drops and the + * producer enters the reconnect window ({@code cnx() == null}). If the send timeout fires in that window, + * {@code failPendingMessages(null, ex)} runs on the timer thread. It must NOT release the op's cmd or recycle + * the op there: the write may still be in-flight on the (old) connection's event loop, and releasing the + * buffers on the timer thread would let new batches reuse (and overwrite) the memory that the in-flight write + * is still reading, corrupting the frame. The cmd release and op recycle must be deferred to the write's event + * loop so that they are serialized after the in-flight write. + */ + @Test + @SuppressWarnings("unchecked") + public void testSendTimeoutDuringInFlightWriteDefersCmdRelease() throws Exception { + ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doReturn(false).when(producer).isBatchMessagingEnabled(); + Mockito.doNothing().when(producer).semaphoreRelease(Mockito.anyInt()); + + PulsarClientImpl client = Mockito.mock(PulsarClientImpl.class); + Mockito.when(client.getMemoryLimitController()) + .thenReturn(Mockito.mock(MemoryLimitController.class)); + FieldUtils.writeField(producer, "client", client, true); + + OpSendMsgQueue pendingQueue = new OpSendMsgQueue(); + FieldUtils.writeField(producer, "pendingMessages", pendingQueue, true); + + // The connection event loop the write was handed to. It is still in-flight (the write callback is queued + // and has not run yet), so it must serialize the cmd release after the write. + List submittedTasks = new ArrayList<>(); + EventLoop writeEventLoop = Mockito.mock(EventLoop.class); + doAnswer(invocation -> { + submittedTasks.add(invocation.getArgument(0)); + return null; + }).when(writeEventLoop).execute(any(Runnable.class)); + + ByteBufPair cmd = ByteBufPair.get( + Unpooled.buffer().writeBytes("frame-header".getBytes(StandardCharsets.UTF_8)), + Unpooled.buffer().writeBytes("batch-payload".getBytes(StandardCharsets.UTF_8))); + MessageImpl msg = Mockito.mock(MessageImpl.class); + Mockito.when(msg.getUncompressedSize()).thenReturn(10); + OpSendMsg op = OpSendMsg.create( + Mockito.mock(LatencyHistogram.class), msg, cmd, 1L, Mockito.mock(SendCallback.class)); + op.totalChunks = 1; + op.chunkId = 0; + op.numMessagesInBatch = 1; + OpSendMsg opSpy = spy(op); + opSpy.writeEventLoop = writeEventLoop; + pendingQueue.add(opSpy); + + // The send timeout fires on the timer thread while cnx() == null and the write is still in-flight. + Thread timerThread = new Thread(() -> + producer.failPendingMessages(null, new PulsarClientException.TimeoutException("timeout"))); + timerThread.start(); + timerThread.join(); + + // The timer thread must not have released the cmd or recycled the op: the buffers must stay alive until + // the in-flight write completes, otherwise a new batch could reuse them and corrupt the frame. + assertEquals(cmd.refCnt(), 1, + "cmd must not be released on the timer thread while a write is in-flight"); + verify(opSpy, never()).recycle(); + assertEquals(submittedTasks.size(), 1, + "cmd release must be deferred to the in-flight write's event loop"); + assertEquals(pendingQueue.size(), 0, "timed-out op must be removed from the pending queue"); + + assertTrue(cmd.getFirst().readableBytes() > 0 && cmd.getSecond().readableBytes() > 0); + + submittedTasks.get(0).run(); + assertEquals(cmd.refCnt(), 0, "cmd must be released after the in-flight write completes"); + verify(opSpy).recycle(); + } + + /** + * When an op was never handed to a connection (no write can be in-flight), the timeout path must keep + * releasing the cmd and recycling the op inline to avoid holding them until the next event loop tick. + */ + @Test + @SuppressWarnings("unchecked") + public void testSendTimeoutReleasesCmdInlineWhenNoWriteInFlight() throws Exception { + ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doReturn(false).when(producer).isBatchMessagingEnabled(); + Mockito.doNothing().when(producer).semaphoreRelease(Mockito.anyInt()); + + PulsarClientImpl client = Mockito.mock(PulsarClientImpl.class); + Mockito.when(client.getMemoryLimitController()) + .thenReturn(Mockito.mock(MemoryLimitController.class)); + FieldUtils.writeField(producer, "client", client, true); + + OpSendMsgQueue pendingQueue = new OpSendMsgQueue(); + FieldUtils.writeField(producer, "pendingMessages", pendingQueue, true); + + EventLoop writeEventLoop = Mockito.mock(EventLoop.class); + doAnswer(invocation -> { + fail("No task must be submitted to an event loop when the op was never written"); + return null; + }).when(writeEventLoop).execute(any(Runnable.class)); + + ByteBufPair cmd = ByteBufPair.get( + Unpooled.buffer().writeBytes("frame-header".getBytes(StandardCharsets.UTF_8)), + Unpooled.buffer().writeBytes("batch-payload".getBytes(StandardCharsets.UTF_8))); + MessageImpl msg = Mockito.mock(MessageImpl.class); + Mockito.when(msg.getUncompressedSize()).thenReturn(10); + OpSendMsg op = OpSendMsg.create( + Mockito.mock(LatencyHistogram.class), msg, cmd, 1L, Mockito.mock(SendCallback.class)); + op.totalChunks = 1; + op.chunkId = 0; + op.numMessagesInBatch = 1; + // writeEventLoop stays null: this op was queued while disconnected and never handed to a connection. + pendingQueue.add(op); + + producer.failPendingMessages(null, new PulsarClientException.TimeoutException("timeout")); + + assertEquals(cmd.refCnt(), 0, "cmd must be released inline when no write is in-flight"); + assertEquals(pendingQueue.size(), 0); + } + + /** + * Verifies that {@link ProducerImpl#processOpSendMsg(OpSendMsg)} hands the op's cmd lifecycle to the + * connection's event loop when it schedules the write, so that a later timeout on another thread can defer + * the cmd release to the right event loop. + */ + @Test + @SuppressWarnings("unchecked") + public void testProcessOpSendMsgTracksWriteEventLoop() throws Exception { + ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doReturn(false).when(producer).isBatchMessagingEnabled(); + Mockito.doReturn(HandlerState.State.Ready).when(producer).getState(); + + List submittedTasks = new ArrayList<>(); + EventLoop eventLoop = Mockito.mock(EventLoop.class); + doAnswer(invocation -> { + submittedTasks.add(invocation.getArgument(0)); + return null; + }).when(eventLoop).execute(any(Runnable.class)); + Channel channel = Mockito.mock(Channel.class); + Mockito.when(channel.eventLoop()).thenReturn(eventLoop); + ChannelHandlerContext ctx = Mockito.mock(ChannelHandlerContext.class); + Mockito.when(ctx.channel()).thenReturn(channel); + ClientCnx cnx = Mockito.mock(ClientCnx.class); + Mockito.when(cnx.ctx()).thenReturn(ctx); + Mockito.doReturn(cnx).when(producer).getCnxIfReady(); + + ProducerStatsRecorder stats = Mockito.mock(ProducerStatsRecorder.class); + FieldUtils.writeField(producer, "stats", stats, true); + OpSendMsgQueue pendingQueue = new OpSendMsgQueue(); + FieldUtils.writeField(producer, "pendingMessages", pendingQueue, true); + + ByteBufPair cmd = ByteBufPair.get( + Unpooled.buffer().writeBytes("frame-header".getBytes(StandardCharsets.UTF_8)), + Unpooled.buffer().writeBytes("batch-payload".getBytes(StandardCharsets.UTF_8))); + MessageImpl msg = Mockito.mock(MessageImpl.class); + Mockito.when(msg.getUncompressedSize()).thenReturn(10); + OpSendMsg op = OpSendMsg.create( + Mockito.mock(LatencyHistogram.class), msg, cmd, 1L, Mockito.mock(SendCallback.class)); + op.totalChunks = 1; + op.chunkId = 0; + op.numMessagesInBatch = 1; + // Bypass the message-dependent paths (batch scheduling / schema registration / size checks) so the test + // focuses on the write hand-off. + op.msg = null; + + producer.processOpSendMsg(op); + + assertEquals(op.writeEventLoop, eventLoop, + "processOpSendMsg must record the event loop the cmd is handed to for writing"); + assertEquals(submittedTasks.size(), 1, "a write must have been scheduled on the connection event loop"); + assertEquals(cmd.refCnt(), 2, "the write path must hold an extra reference on the cmd"); + verify(stats).updateNumMsgsSent(1, 0); + } + + /** + * Regression test for the stale write-callback window: an op is handed to connection A's event loop, then + * re-sent on connection B after a reconnect (so {@code op.writeEventLoop} now points to B's loop), and the + * send timeout disposes it deferred on B's loop while the callback queued on A's loop has still not run. + * The stale callback must not write the released cmd or touch the recycled op; it must only drop the + * reference it took when it was scheduled. + */ + @Test + @SuppressWarnings("unchecked") + public void testStaleWriteCallbackSkipsDisposedOp() throws Exception { + ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doReturn(false).when(producer).isBatchMessagingEnabled(); + Mockito.doReturn(HandlerState.State.Ready).when(producer).getState(); + Mockito.doNothing().when(producer).semaphoreRelease(Mockito.anyInt()); + // The write callback logs through the instance logger, which a mock does not initialize. + FieldUtils.writeField(producer, "log", + Mockito.mock(io.github.merlimat.slog.Logger.class, Mockito.RETURNS_DEEP_STUBS), true); + PulsarClientImpl client = Mockito.mock(PulsarClientImpl.class); + Mockito.when(client.getMemoryLimitController()) + .thenReturn(Mockito.mock(MemoryLimitController.class)); + FieldUtils.writeField(producer, "client", client, true); + + List tasksOnLoopA = new ArrayList<>(); + EventLoop eventLoopA = Mockito.mock(EventLoop.class); + doAnswer(invocation -> { + tasksOnLoopA.add(invocation.getArgument(0)); + return null; + }).when(eventLoopA).execute(any(Runnable.class)); + Channel channelA = Mockito.mock(Channel.class); + Mockito.when(channelA.eventLoop()).thenReturn(eventLoopA); + ChannelHandlerContext ctxA = Mockito.mock(ChannelHandlerContext.class); + Mockito.when(ctxA.channel()).thenReturn(channelA); + ClientCnx cnxA = Mockito.mock(ClientCnx.class); + Mockito.when(cnxA.ctx()).thenReturn(ctxA); + Mockito.doReturn(cnxA).when(producer).getCnxIfReady(); + + ProducerStatsRecorder stats = Mockito.mock(ProducerStatsRecorder.class); + FieldUtils.writeField(producer, "stats", stats, true); + OpSendMsgQueue pendingQueue = new OpSendMsgQueue(); + FieldUtils.writeField(producer, "pendingMessages", pendingQueue, true); + + ByteBufPair cmd = ByteBufPair.get( + Unpooled.buffer().writeBytes("frame-header".getBytes(StandardCharsets.UTF_8)), + Unpooled.buffer().writeBytes("batch-payload".getBytes(StandardCharsets.UTF_8))); + MessageImpl msg = Mockito.mock(MessageImpl.class); + Mockito.when(msg.getUncompressedSize()).thenReturn(10); + OpSendMsg op = OpSendMsg.create( + Mockito.mock(LatencyHistogram.class), msg, cmd, 1L, Mockito.mock(SendCallback.class)); + op.totalChunks = 1; + op.chunkId = 0; + op.numMessagesInBatch = 1; + op.msg = null; + + // First write: the callback is queued on connection A's loop and has not run yet. + producer.processOpSendMsg(op); + Runnable staleWriteCallback = tasksOnLoopA.get(0); + assertEquals(cmd.refCnt(), 2, "base reference + the reference taken for the write"); + + // Reconnect: the op is re-sent on connection B, so the tracked write loop moves to B's loop. + List tasksOnLoopB = new ArrayList<>(); + EventLoop eventLoopB = Mockito.mock(EventLoop.class); + doAnswer(invocation -> { + tasksOnLoopB.add(invocation.getArgument(0)); + return null; + }).when(eventLoopB).execute(any(Runnable.class)); + op.writeEventLoop = eventLoopB; + + // The send timeout fires: the disposal is deferred to B's loop and recycles the op. + producer.failPendingMessages(null, new PulsarClientException.TimeoutException("timeout")); + assertEquals(tasksOnLoopB.size(), 1, "disposal must be deferred to the tracked write loop"); + tasksOnLoopB.get(0).run(); + assertEquals(cmd.refCnt(), 1, "only the stale callback's reference may remain"); + + // The stale callback from A's loop finally runs: it must skip the recycled op and only drop its + // own reference, instead of writing the released cmd or mutating the recycled op. + staleWriteCallback.run(); + assertEquals(cmd.refCnt(), 0, "the stale callback must release exactly its own reference"); + assertEquals(op.sequenceId, -1L, "the recycled op must not have been touched by the stale callback"); + assertEquals(op.retryCount, 0, "the recycled op must not have been touched by the stale callback"); + assertEquals(op.firstSentAt, -1L, "the recycled op must not have been touched by the stale callback"); + } + + /** + * The stale write callback must also skip the write when the op was re-sent on another connection but has + * NOT been disposed yet: the op's {@code writeEventLoop} now points at the new connection's loop, so the + * callback on the old connection must not write (that would double-send) and must not mutate the live op. It + * must only drop the reference it took when it was scheduled. + */ + @Test + @SuppressWarnings("unchecked") + public void testStaleWriteCallbackSkipsOpReSentOnAnotherLoop() throws Exception { + ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doReturn(false).when(producer).isBatchMessagingEnabled(); + Mockito.doReturn(HandlerState.State.Ready).when(producer).getState(); + Mockito.doNothing().when(producer).semaphoreRelease(Mockito.anyInt()); + FieldUtils.writeField(producer, "log", + Mockito.mock(io.github.merlimat.slog.Logger.class, Mockito.RETURNS_DEEP_STUBS), true); + PulsarClientImpl client = Mockito.mock(PulsarClientImpl.class); + Mockito.when(client.getMemoryLimitController()) + .thenReturn(Mockito.mock(MemoryLimitController.class)); + FieldUtils.writeField(producer, "client", client, true); + + List tasksOnLoopA = new ArrayList<>(); + EventLoop eventLoopA = Mockito.mock(EventLoop.class); + doAnswer(invocation -> { + tasksOnLoopA.add(invocation.getArgument(0)); + return null; + }).when(eventLoopA).execute(any(Runnable.class)); + Channel channelA = Mockito.mock(Channel.class); + Mockito.when(channelA.eventLoop()).thenReturn(eventLoopA); + ChannelHandlerContext ctxA = Mockito.mock(ChannelHandlerContext.class); + Mockito.when(ctxA.channel()).thenReturn(channelA); + ClientCnx cnxA = Mockito.mock(ClientCnx.class); + Mockito.when(cnxA.ctx()).thenReturn(ctxA); + Mockito.doReturn(cnxA).when(producer).getCnxIfReady(); + + ProducerStatsRecorder stats = Mockito.mock(ProducerStatsRecorder.class); + FieldUtils.writeField(producer, "stats", stats, true); + OpSendMsgQueue pendingQueue = new OpSendMsgQueue(); + FieldUtils.writeField(producer, "pendingMessages", pendingQueue, true); + + ByteBufPair cmd = ByteBufPair.get( + Unpooled.buffer().writeBytes("frame-header".getBytes(StandardCharsets.UTF_8)), + Unpooled.buffer().writeBytes("batch-payload".getBytes(StandardCharsets.UTF_8))); + MessageImpl msg = Mockito.mock(MessageImpl.class); + Mockito.when(msg.getUncompressedSize()).thenReturn(10); + OpSendMsg op = OpSendMsg.create( + Mockito.mock(LatencyHistogram.class), msg, cmd, 1L, Mockito.mock(SendCallback.class)); + op.totalChunks = 1; + op.chunkId = 0; + op.numMessagesInBatch = 1; + op.msg = null; + + // First write on connection A: the callback is queued but has not run yet. + producer.processOpSendMsg(op); + Runnable staleWriteCallback = tasksOnLoopA.get(0); + assertEquals(cmd.refCnt(), 2, "base reference + the reference taken for the write"); + + // The op is re-sent on connection B: the tracked write loop moves to B's loop, but the op is NOT disposed + // and op.cmd is unchanged. + EventLoop eventLoopB = Mockito.mock(EventLoop.class); + op.writeEventLoop = eventLoopB; + + // The stale callback from A finally runs. It must detect that this loop is no longer the op's write loop, + // skip the write and drop only its own reference, without touching the live op. + staleWriteCallback.run(); + assertEquals(cmd.refCnt(), 1, "the stale callback must release exactly its own reference"); + assertEquals(op.retryCount, 0, "the live op must not be mutated by the stale callback"); + assertEquals(op.firstSentAt, -1L, "the live op must not be mutated by the stale callback"); + // Skipping the stale write must not strand the message: the op stays pending, waiting for the + // response of the write on the new connection (receipt, or a later timeout as the last resort). + assertTrue(pendingQueue.peek() == op, "the live op must still be pending for the new write"); + assertEquals(pendingQueue.size(), 1); + } + + /** + * When the timeout fires while the current connection's event loop is shutting down (e.g. right after a + * reconnect churn), the deferred failPendingMessages task would be rejected and dropped. The producer must + * fall back to failing the pending messages inline instead of leaving them queued until the next timeout tick. + */ + @Test + @SuppressWarnings("unchecked") + public void testFailPendingMessagesHandlesShuttingDownEventLoop() throws Exception { + ProducerImpl producer = Mockito.mock(ProducerImpl.class, Mockito.CALLS_REAL_METHODS); + Mockito.doReturn(false).when(producer).isBatchMessagingEnabled(); + Mockito.doNothing().when(producer).semaphoreRelease(Mockito.anyInt()); + FieldUtils.writeField(producer, "log", + Mockito.mock(io.github.merlimat.slog.Logger.class, Mockito.RETURNS_DEEP_STUBS), true); + PulsarClientImpl client = Mockito.mock(PulsarClientImpl.class); + Mockito.when(client.getMemoryLimitController()) + .thenReturn(Mockito.mock(MemoryLimitController.class)); + FieldUtils.writeField(producer, "client", client, true); + + OpSendMsgQueue pendingQueue = new OpSendMsgQueue(); + FieldUtils.writeField(producer, "pendingMessages", pendingQueue, true); + + // The current connection's event loop rejects new tasks (it is shutting down). + EventLoop eventLoop = Mockito.mock(EventLoop.class); + doThrow(new RejectedExecutionException("shutting down")).when(eventLoop).execute(any(Runnable.class)); + Channel channel = Mockito.mock(Channel.class); + Mockito.when(channel.eventLoop()).thenReturn(eventLoop); + ChannelHandlerContext ctx = Mockito.mock(ChannelHandlerContext.class); + Mockito.when(ctx.channel()).thenReturn(channel); + ClientCnx cnx = Mockito.mock(ClientCnx.class); + Mockito.when(cnx.ctx()).thenReturn(ctx); + + ByteBufPair cmd = ByteBufPair.get( + Unpooled.buffer().writeBytes("frame-header".getBytes(StandardCharsets.UTF_8)), + Unpooled.buffer().writeBytes("batch-payload".getBytes(StandardCharsets.UTF_8))); + MessageImpl msg = Mockito.mock(MessageImpl.class); + Mockito.when(msg.getUncompressedSize()).thenReturn(10); + OpSendMsg op = OpSendMsg.create( + Mockito.mock(LatencyHistogram.class), msg, cmd, 1L, Mockito.mock(SendCallback.class)); + op.totalChunks = 1; + op.chunkId = 0; + op.numMessagesInBatch = 1; + // This op was never handed to a connection, so its disposal is safe inline. + pendingQueue.add(op); + + producer.failPendingMessages(cnx, new PulsarClientException.TimeoutException("timeout")); + + assertEquals(pendingQueue.size(), 0, "pending messages must still be failed when the event loop is down"); + assertEquals(cmd.refCnt(), 0, "the cmd must be released inline when the event loop rejects the task"); + } }