You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pulsar.apache.org by lh...@apache.org on 2026/10/16 22:31:43 UTC
(pulsar) 20/24: [fix][client] Serialize chunked-message bookkeeping to fix use-after-free and count/queue drift (#26084)
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a commit to branch branch-4.0-backports-2026-10-17
in repository https://gitbox.apache.org/repos/asf/pulsar.git
commit 5f89be58e6ed533dac0409aa15d88f7ac98fa488
Author: SongOf <46...@users.noreply.github.com>
AuthorDate: Fri Oct 16 08:27:23 2026 +0800
[fix][client] Serialize chunked-message bookkeeping to fix use-after-free and count/queue drift (#26084)
Co-authored-by: maxlisongsong <ma...@didiglobal.com>
Co-authored-by: Lari Hotari <lh...@users.noreply.github.com>
Co-authored-by: Penghui Li <pe...@apache.org>
(cherry picked from commit 779380666628b5fe62cfd0b2a0b4531669905659)
---
.../pulsar/client/impl/MessageChunkingTest.java | 30 +-
.../apache/pulsar/client/impl/ConsumerImpl.java | 241 +++++--
.../pulsar/client/impl/ConsumerImplTest.java | 693 ++++++++++++++++++++-
3 files changed, 900 insertions(+), 64 deletions(-)
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/MessageChunkingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/MessageChunkingTest.java
index ecd5a9c46ceb..0aa2768a2f02 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/MessageChunkingTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/impl/MessageChunkingTest.java
@@ -20,6 +20,7 @@ package org.apache.pulsar.client.impl;
import static org.assertj.core.api.Assertions.assertThat;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
@@ -416,15 +417,15 @@ public class MessageChunkingTest extends ProducerConsumerBase {
/**
* Chunks are processed on the connection's IO thread while the expiry task runs on another thread.
- * With queue [A, B], A's completion removes A's context and then A's uuid. If expiry inspects A in
- * between (no context, so it skips A as a stale head) and only removes it after completion already
- * did, it must remove A itself, not whatever is now at the head: dropping B would leave B's context
- * in chunkedMessagesMap with nothing in the queue referencing it, so it could never expire or be
- * evicted.
+ * With queue [A, B], A's completion removes A's context and then A's uuid. If expiry could inspect A in
+ * between (no context, so it skips A as a stale head) and only remove it after completion already did,
+ * it might drop whatever is now at the head: dropping B would leave B's context in chunkedMessagesMap
+ * with nothing in the queue referencing it, so it could never expire or be evicted.
*
- * The interleaving is forced deterministically by wrapping chunkedMessagesMap: A's completion
- * (map.remove("A") on the IO thread) starts expiry and waits until expiry has inspected A, and
- * expiry (map.get("A")) then waits until completion has removed A's uuid before it continues.
+ * Completion and expiry both hold chunkedMessageLock, so expiry must not observe that intermediate
+ * state at all. The test wraps chunkedMessagesMap: A's completion (map.remove("A") on the IO thread,
+ * under the lock) starts expiry and waits until expiry either blocks or inspects A (map.get("A")),
+ * and inspecting A is a failure. Either way B must stay pending and reachable for expiry.
*/
@Test
public void testExpiryStaleHeadSkipDoesNotDropConcurrentLiveEntry() throws Exception {
@@ -457,6 +458,7 @@ public class MessageChunkingTest extends ProducerConsumerBase {
final CountDownLatch expiryInspectedA = new CountDownLatch(1);
final AtomicBoolean raceTriggered = new AtomicBoolean();
+ final AtomicBoolean expiryInspectedPartialRemoval = new AtomicBoolean();
final AtomicReference<Thread> expiryThread = new AtomicReference<>();
final AtomicReference<Throwable> expiryError = new AtomicReference<>();
Map<String, ConsumerImpl.ChunkedMessageCtx> hookedMap =
@@ -476,11 +478,11 @@ public class MessageChunkingTest extends ProducerConsumerBase {
}, "test-chunk-expiry");
expiryThread.set(t);
t.start();
- try {
- expiryInspectedA.await(10, TimeUnit.SECONDS);
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- }
+ // Completion holds chunkedMessageLock here, so expiry has to block on it instead of
+ // inspecting A between A's map removal and its uuid removal.
+ Awaitility.await().atMost(10, TimeUnit.SECONDS).until(() ->
+ t.getState() == Thread.State.BLOCKED || expiryInspectedA.getCount() == 0);
+ expiryInspectedPartialRemoval.set(expiryInspectedA.getCount() == 0);
}
return removed;
}
@@ -509,6 +511,8 @@ public class MessageChunkingTest extends ProducerConsumerBase {
assertEquals(a.getValue(), "chunk-A-0|chunk-A-1|");
assertTrue(raceTriggered.get(), "A's completion did not trigger the interleaving");
+ assertFalse(expiryInspectedPartialRemoval.get(),
+ "expiry inspected A after its context was removed but before its uuid was");
expiryThread.get().join(TimeUnit.SECONDS.toMillis(10));
assertNull(expiryError.get());
diff --git a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java
index 2b9f0373fea3..f413451d51c7 100644
--- a/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java
+++ b/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java
@@ -216,15 +216,25 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
protected volatile boolean paused;
protected Map<String, ChunkedMessageCtx> chunkedMessagesMap = new ConcurrentHashMap<>();
- private int pendingChunkedMessageCount = 0;
+ @VisibleForTesting
+ int pendingChunkedMessageCount = 0;
+ // Serializes all chunked-message bookkeeping. The receive/assembly path (processMessageChunk and
+ // removeOldestPendingChunkedMessage, on the Netty IO thread) and the eviction/expiry path
+ // (removeExpireIncompleteChunkedMessages, on the internalPinnedExecutor) both mutate the same ChunkedMessageCtx,
+ // its buffer and pendingChunkedMessageCount; without this lock they can race (use-after-free / double-recycle).
+ // No application code runs under it: acks decided inside are collected and issued after it is released (see
+ // acknowledgeChunkedMessageIds), and decompression of a completed message happens outside it as well.
+ private final Object chunkedMessageLock = new Object();
protected long expireTimeOfIncompleteChunkedMessageMillis = 0;
- private final AtomicBoolean expireChunkMessageTaskScheduled = new AtomicBoolean(false);
+ @VisibleForTesting
+ final AtomicBoolean expireChunkMessageTaskScheduled = new AtomicBoolean(false);
private final int maxPendingChunkedMessage;
// if queue size is reasonable (most of the time equal to number of producers try to publish messages concurrently
// on the topic) then it guards against broken chunked message which was not fully published
private final boolean autoAckOldestChunkedMessageOnQueueFull;
// it will be used to manage N outstanding chunked message buffers
- private final BlockingQueue<String> pendingChunkedMessageUuidQueue;
+ @VisibleForTesting
+ final BlockingQueue<String> pendingChunkedMessageUuidQueue;
@VisibleForTesting
int getPendingChunkedMessageUuidQueueSizeForTest() {
@@ -1504,33 +1514,81 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
MessageIdImpl[] chunkedMessageIds = null;
if (isChunkedMessage) {
- uncompressedPayload = processMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, cnx);
- if (uncompressedPayload == null) {
+ ByteBuf compressedAssembledPayload;
+ // Acks decided under the lock are issued only after it is released, see acknowledgeChunkedMessageIds.
+ // They are issued in a finally: once a path has queued them it has already removed the ctx they belong
+ // to, so if chunk processing then throws (a malformed chunk overflowing the buffer, say) nothing else
+ // could ever ack those ids.
+ List<MessageIdImpl> deferredAcks = new ArrayList<>();
+ try {
+ // Critical section: append the final chunk, then atomically remove + recycle the ctx, capturing
+ // the assembled (still-compressed) buffer and the chunk ids into locals. Keeping assembly and
+ // finalize in one locked region closes the assemble->finalize race with the expiry/eviction path.
+ // Once the ctx is out of chunkedMessagesMap and recycled, the captured buffer is reachable only
+ // through this thread, so decompression and message building are done below, OUTSIDE the lock, to
+ // avoid serializing other completions and expiry behind codec work (and discardCorruptedMessage
+ // on failure).
+ synchronized (chunkedMessageLock) {
+ compressedAssembledPayload =
+ doProcessMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, cnx,
+ deferredAcks);
+ if (compressedAssembledPayload != null) {
+ // last chunk received: so, stitch chunked-messages and clear up chunkedMsgBuffer
+ if (log.isDebugEnabled()) {
+ log.debug("Chunked message completed chunkId {}, total-chunks {}, msgId {}"
+ + " sequenceId {}", msgMetadata.getChunkId(),
+ msgMetadata.getNumChunksFromMsg(), msgId, msgMetadata.getSequenceId());
+ }
+
+ // remove the ctx from the map and recycle it, capturing the chunk ids before recycle lets
+ // it be reused. The returned buffer is this ctx's chunkedMsgBuffer; recycle only nulls the
+ // field, the buffer object stays alive and is now owned solely by
+ // compressedAssembledPayload.
+ ChunkedMessageCtx chunkedMsgCtx = chunkedMessagesMap.remove(msgMetadata.getUuid());
+ chunkedMessageIds = chunkedMsgCtx.chunkedMessageIds;
+ // Drop the completed uuid from the queue too, mirroring the map removal, so it doesn't
+ // linger as a ghost entry.
+ pendingChunkedMessageUuidQueue.remove(msgMetadata.getUuid());
+ pendingChunkedMessageCount--;
+ chunkedMsgCtx.recycle();
+ }
+ }
+ } finally {
+ acknowledgeChunkedMessageIds(deferredAcks);
+ }
+ if (compressedAssembledPayload == null) {
return;
}
- // last chunk received: so, stitch chunked-messages and clear up chunkedMsgBuffer
- if (log.isDebugEnabled()) {
- log.debug("Chunked message completed chunkId {}, total-chunks {}, msgId {} sequenceId {}",
- msgMetadata.getChunkId(), msgMetadata.getNumChunksFromMsg(), msgId,
- msgMetadata.getSequenceId());
+ // Outside the lock: set the chunk message id and decompress the assembled payload. The chunk ids are
+ // registered with the unack tracker below, only if the message also passes the startMessageId filter.
+ if (chunkedMessageIds.length > 0) {
+ msgId = new ChunkMessageIdImpl(chunkedMessageIds[0],
+ chunkedMessageIds[chunkedMessageIds.length - 1]);
}
-
- // remove buffer from the map, set the chunk message id
- ChunkedMessageCtx chunkedMsgCtx = chunkedMessagesMap.remove(msgMetadata.getUuid());
- if (chunkedMsgCtx.chunkedMessageIds.length > 0) {
- msgId = new ChunkMessageIdImpl(chunkedMsgCtx.chunkedMessageIds[0],
- chunkedMsgCtx.chunkedMessageIds[chunkedMsgCtx.chunkedMessageIds.length - 1]);
+ try {
+ uncompressedPayload = uncompressPayloadIfNeeded(messageId, msgMetadata, compressedAssembledPayload,
+ cnx, false);
+ } finally {
+ // The assembled buffer is detached from any ctx, so nothing else can release it: do so here even
+ // if decompression throws, or it would leak together with the message ids acked below.
+ compressedAssembledPayload.release();
+ }
+ if (uncompressedPayload == null) {
+ // The ctx is already out of chunkedMessagesMap and the expiry queue, so the earlier chunks'
+ // message ids can no longer be reached by the expiry sweep. Ack them here, mirroring
+ // removeChunkMessage's autoAck semantics, so a decompression failure doesn't leave them as a
+ // permanent ack hole. The last entry is skipped: it's the same (ledgerId, entryId) as the
+ // `messageId` passed to uncompressPayloadIfNeeded above, which its failure path
+ // (discardCorruptedMessage) already acked directly.
+ for (int i = 0; i < chunkedMessageIds.length - 1; i++) {
+ MessageIdImpl chunkedMessageId = chunkedMessageIds[i];
+ if (chunkedMessageId != null) {
+ doAcknowledge(chunkedMessageId, AckType.Individual, Collections.emptyMap(), null);
+ }
+ }
+ return;
}
- chunkedMessageIds = chunkedMsgCtx.chunkedMessageIds;
- pendingChunkedMessageCount--;
- // The completed message's uuid was added to pendingChunkedMessageUuidQueue when its
- // first chunk arrived, but is only ever removed by the eviction/expiry paths. On the
- // normal completion path it was never removed, so the queue accumulated one entry per
- // completed chunked message unboundedly (a memory leak for long-running consumers).
- // Remove it here to keep the queue in sync with chunkedMessagesMap.
- pendingChunkedMessageUuidQueue.remove(msgMetadata.getUuid());
- chunkedMsgCtx.recycle();
}
// If the topic is non-persistent, we should not ignore any messages.
@@ -1588,8 +1646,45 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
}
- private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata msgMetadata, MessageIdImpl msgId,
+ /**
+ * Processes one chunk under {@code chunkedMessageLock} and issues any acks it decided on after the lock is
+ * released. {@code messageReceived} calls {@link #doProcessMessageChunk} directly instead, because it has to keep
+ * the final-chunk assembly and the finalize in one locked region.
+ */
+ @VisibleForTesting
+ ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata msgMetadata, MessageIdImpl msgId,
MessageIdData messageId, ClientCnx cnx) {
+ List<MessageIdImpl> deferredAcks = new ArrayList<>();
+ try {
+ synchronized (chunkedMessageLock) {
+ return doProcessMessageChunk(compressedPayload, msgMetadata, msgId, messageId, cnx, deferredAcks);
+ }
+ } finally {
+ // Also on an exception: see the finally in messageReceived.
+ acknowledgeChunkedMessageIds(deferredAcks);
+ }
+ }
+
+ /**
+ * Issues the individual acks that were decided while {@code chunkedMessageLock} was held. {@code doAcknowledge} is
+ * not purely asynchronous: the persistent acknowledgments grouping tracker invokes the consumer's acknowledgment
+ * interceptors inline, and application code must never run under the lock. Otherwise a slow interceptor on the
+ * expiry thread stalls chunk reception on the IO thread, and one that waits for that IO thread deadlocks it.
+ * Callers invoke this from a {@code finally}: the ids were queued by a path that already removed their ctx, so if
+ * processing throws afterwards they would otherwise be lost as a permanent ack hole.
+ */
+ private void acknowledgeChunkedMessageIds(List<MessageIdImpl> messageIds) {
+ for (MessageIdImpl messageId : messageIds) {
+ doAcknowledge(messageId, AckType.Individual, Collections.emptyMap(), null);
+ }
+ }
+
+ /**
+ * Must be called with {@code chunkedMessageLock} held. Acks are never issued here: the caller collects them in
+ * {@code deferredAcks} and issues them once the lock is released, see {@link #acknowledgeChunkedMessageIds}.
+ */
+ private ByteBuf doProcessMessageChunk(ByteBuf compressedPayload, MessageMetadata msgMetadata, MessageIdImpl msgId,
+ MessageIdData messageId, ClientCnx cnx, List<MessageIdImpl> deferredAcks) {
if (msgMetadata.getChunkId() != (msgMetadata.getNumChunksFromMsg() - 1)) {
increaseAvailablePermits(cnx);
}
@@ -1633,11 +1728,11 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
.noneMatch(messageId1 -> messageId1 != null && messageId1.ledgerId == messageId.getLedgerId()
&& messageId1.entryId == messageId.getEntryId());
if (isCorruptedChunkMessageDetected) {
- Arrays.stream(chunkedMsgCtx.chunkedMessageIds).forEach(messageId1 -> {
+ for (MessageIdImpl messageId1 : chunkedMsgCtx.chunkedMessageIds) {
if (messageId1 != null) {
- doAcknowledge(messageId1, AckType.Individual, Collections.emptyMap(), null);
+ deferredAcks.add(messageId1);
}
- });
+ }
}
// The first chunk of a new chunked-message received before receiving other chunks of previous
// chunked-message
@@ -1647,16 +1742,23 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
}
chunkedMsgCtx.recycle();
chunkedMessagesMap.remove(msgMetadata.getUuid());
+ // Drop the stale queue entry; it is re-added at the tail below, keeping the queue ordered by
+ // receivedTime with exactly one entry per in-progress uuid. The replaced ctx was already counted,
+ // so this '--' balances the unconditional '++' below.
+ pendingChunkedMessageUuidQueue.remove(msgMetadata.getUuid());
+ pendingChunkedMessageCount--;
}
pendingChunkedMessageCount++;
if (maxPendingChunkedMessage > 0 && pendingChunkedMessageCount > maxPendingChunkedMessage) {
- removeOldestPendingChunkedMessage();
+ removeOldestPendingChunkedMessage(deferredAcks);
}
int totalChunks = msgMetadata.getNumChunksFromMsg();
ByteBuf chunkedMsgBuffer = PulsarByteBufAllocator.DEFAULT.buffer(msgMetadata.getTotalChunkMsgSize(),
msgMetadata.getTotalChunkMsgSize());
chunkedMsgCtx = chunkedMessagesMap.computeIfAbsent(msgMetadata.getUuid(),
(key) -> ChunkedMessageCtx.get(totalChunks, chunkedMsgBuffer));
+ // Enqueue at the tail so the queue stays ordered oldest-first by receivedTime. For a redelivered first
+ // chunk the stale entry was removed above, so this keeps exactly one entry per in-progress uuid.
pendingChunkedMessageUuidQueue.add(msgMetadata.getUuid());
}
@@ -1684,7 +1786,7 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
.noneMatch(messageId1 -> messageId1 != null && messageId1.ledgerId == messageId.getLedgerId()
&& messageId1.entryId == messageId.getEntryId());
if (isDuplicatedChunk) {
- doAcknowledge(msgId, AckType.Individual, Collections.emptyMap(), null);
+ deferredAcks.add(msgId);
}
return null;
}
@@ -1698,6 +1800,11 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
ReferenceCountUtil.safeRelease(chunkedMsgCtx.chunkedMsgBuffer);
}
chunkedMsgCtx.recycle();
+ // The ctx was counted and enqueued when its first chunk created it, so keep the rest of the
+ // bookkeeping in sync with the map removal: drop the queue entry (no ghost) and decrement the count
+ // (otherwise it drifts upward and prematurely triggers removeOldestPendingChunkedMessage).
+ pendingChunkedMessageUuidQueue.remove(msgMetadata.getUuid());
+ pendingChunkedMessageCount--;
}
chunkedMessagesMap.remove(msgMetadata.getUuid());
compressedPayload.release();
@@ -1712,7 +1819,7 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
if (expireTimeOfIncompleteChunkedMessageMillis > 0
&& System.currentTimeMillis() > (msgMetadata.getPublishTime()
+ expireTimeOfIncompleteChunkedMessageMillis)) {
- doAcknowledge(msgId, AckType.Individual, Collections.emptyMap(), null);
+ deferredAcks.add(msgId);
} else {
trackMessage(msgId);
}
@@ -1721,7 +1828,13 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
chunkedMsgCtx.chunkedMessageIds[msgMetadata.getChunkId()] = msgId;
// append the chunked payload and update lastChunkedMessage-id
- chunkedMsgCtx.chunkedMsgBuffer.writeBytes(compressedPayload);
+ try {
+ chunkedMsgCtx.chunkedMsgBuffer.writeBytes(compressedPayload);
+ } catch (RuntimeException e) {
+ // A malformed chunk that overflows the advertised total size: the incoming chunk is still owned here.
+ compressedPayload.release();
+ throw e;
+ }
chunkedMsgCtx.lastChunkedMessageId = msgMetadata.getChunkId();
// if final chunk is not received yet then release payload and return
@@ -1730,11 +1843,11 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
return null;
}
+ // Last chunk assembled. Release the incoming chunk and hand the assembled (still-compressed) buffer back to the
+ // caller. Decompression is intentionally left to the caller, outside chunkedMessageLock, so the codec work (and
+ // any discardCorruptedMessage on failure) does not serialize other chunked-message completions or expiry.
compressedPayload.release();
- compressedPayload = chunkedMsgCtx.chunkedMsgBuffer;
- ByteBuf uncompressedPayload = uncompressPayloadIfNeeded(messageId, msgMetadata, compressedPayload, cnx, false);
- compressedPayload.release();
- return uncompressedPayload;
+ return chunkedMsgCtx.chunkedMsgBuffer;
}
/**
@@ -2155,7 +2268,10 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
try {
ByteBuf uncompressedPayload = codec.decode(payload, uncompressedSize);
return uncompressedPayload;
- } catch (IOException e) {
+ } catch (IOException | RuntimeException e) {
+ // A codec reports a corrupt payload either as an IOException or, like ZLib on a decoded size that differs
+ // from the advertised one, as an unchecked exception. Both mean the message cannot be decoded and must be
+ // discarded the same way rather than escaping the IO thread.
log.error("[{}][{}] Failed to decompress message with {} at {}: {}", topic, subscription, compressionType,
messageId, e.getMessage(), e);
discardCorruptedMessage(messageId, currentCnx, ValidationError.DecompressionError);
@@ -3162,7 +3278,10 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
}
}
- private void removeOldestPendingChunkedMessage() {
+ /**
+ * Must be called with {@code chunkedMessageLock} held; acks go to {@code deferredAcks}.
+ */
+ private void removeOldestPendingChunkedMessage(List<MessageIdImpl> deferredAcks) {
ChunkedMessageCtx chunkedMsgCtx = null;
String firstPendingMsgUuid = null;
while (chunkedMsgCtx == null && !pendingChunkedMessageUuidQueue.isEmpty()) {
@@ -3171,37 +3290,59 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
chunkedMsgCtx = StringUtils.isNotBlank(firstPendingMsgUuid) ? chunkedMessagesMap.get(firstPendingMsgUuid)
: null;
}
- removeChunkMessage(firstPendingMsgUuid, chunkedMsgCtx, this.autoAckOldestChunkedMessageOnQueueFull);
+ removeChunkMessage(firstPendingMsgUuid, chunkedMsgCtx, this.autoAckOldestChunkedMessageOnQueueFull,
+ deferredAcks);
}
protected void removeExpireIncompleteChunkedMessages() {
if (expireTimeOfIncompleteChunkedMessageMillis <= 0) {
return;
}
- ChunkedMessageCtx chunkedMsgCtx;
+ // Acks decided under the lock are issued only after it is released, see acknowledgeChunkedMessageIds.
+ List<MessageIdImpl> deferredAcks = new ArrayList<>();
+ try {
+ synchronized (chunkedMessageLock) {
+ doRemoveExpireIncompleteChunkedMessages(deferredAcks);
+ }
+ } finally {
+ // Also on an exception: see the finally in messageReceived.
+ acknowledgeChunkedMessageIds(deferredAcks);
+ }
+ }
+
+ /**
+ * Must be called with {@code chunkedMessageLock} held; acks go to {@code deferredAcks}.
+ */
+ private void doRemoveExpireIncompleteChunkedMessages(List<MessageIdImpl> deferredAcks) {
String messageUUID;
while ((messageUUID = pendingChunkedMessageUuidQueue.peek()) != null) {
- chunkedMsgCtx = StringUtils.isNotBlank(messageUUID) ? chunkedMessagesMap.get(messageUUID) : null;
+ ChunkedMessageCtx chunkedMsgCtx =
+ StringUtils.isNotBlank(messageUUID) ? chunkedMessagesMap.get(messageUUID) : null;
if (chunkedMsgCtx == null) {
- // Stale queue head with no backing context (e.g. the message already completed, or a
- // discard path left the uuid behind). Skip it and continue, so a ghost head cannot block
- // expiry of live entries queued behind it.
+ // Ghost head: the chunked message already completed/was removed from chunkedMessagesMap but its uuid
+ // lingers in the queue. Drop it and keep scanning (mirroring removeOldestPendingChunkedMessage) so
+ // genuinely-expired incomplete chunks queued behind it are still cleaned.
pendingChunkedMessageUuidQueue.remove(messageUUID);
continue;
}
if (System.currentTimeMillis()
> (chunkedMsgCtx.receivedTime + expireTimeOfIncompleteChunkedMessageMillis)) {
pendingChunkedMessageUuidQueue.remove(messageUUID);
- removeChunkMessage(messageUUID, chunkedMsgCtx, true);
+ removeChunkMessage(messageUUID, chunkedMsgCtx, true, deferredAcks);
} else {
- // Head is a live, not-yet-expired context. Entries are queued in arrival order, so
- // nothing behind it can be older; stop here.
+ // The queue is ordered oldest-first, so the first live not-yet-expired head means nothing behind it
+ // is expired either.
return;
}
}
}
- private void removeChunkMessage(String msgUUID, ChunkedMessageCtx chunkedMsgCtx, boolean autoAck) {
+ /**
+ * Must be called with {@code chunkedMessageLock} held. With {@code autoAck} the chunk ids are appended to
+ * {@code deferredAcks} for the caller to ack after releasing the lock, see {@link #acknowledgeChunkedMessageIds}.
+ */
+ private void removeChunkMessage(String msgUUID, ChunkedMessageCtx chunkedMsgCtx, boolean autoAck,
+ List<MessageIdImpl> deferredAcks) {
if (chunkedMsgCtx == null) {
return;
}
@@ -3214,7 +3355,7 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle
}
if (autoAck) {
log.info("Removing chunk message-id {}", msgId);
- doAcknowledge(msgId, AckType.Individual, Collections.emptyMap(), null);
+ deferredAcks.add(msgId);
} else {
trackMessage(msgId);
}
diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerImplTest.java
index d2405c241e87..5421f47a1856 100644
--- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerImplTest.java
+++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/ConsumerImplTest.java
@@ -20,6 +20,7 @@ package org.apache.pulsar.client.impl;
import static org.apache.pulsar.common.protocol.Commands.DEFAULT_CONSUMER_EPOCH;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.atLeast;
@@ -33,20 +34,30 @@ import static org.mockito.Mockito.when;
import static org.testng.Assert.assertTrue;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
+import java.util.Collections;
import java.util.List;
+import java.util.Random;
+import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import lombok.Cleanup;
import org.apache.pulsar.client.api.Consumer;
+import org.apache.pulsar.client.api.ConsumerInterceptor;
import org.apache.pulsar.client.api.Message;
+import org.apache.pulsar.client.api.MessageId;
import org.apache.pulsar.client.api.MessageIdAdv;
import org.apache.pulsar.client.api.MessageListener;
import org.apache.pulsar.client.api.MessagePayload;
@@ -59,12 +70,18 @@ import org.apache.pulsar.client.impl.conf.ConsumerConfigurationData;
import org.apache.pulsar.client.impl.conf.TopicConsumerConfigurationData;
import org.apache.pulsar.client.util.ExecutorProvider;
import org.apache.pulsar.client.util.ScheduledExecutorProvider;
+import org.apache.pulsar.common.api.proto.CommandMessage;
+import org.apache.pulsar.common.api.proto.CompressionType;
+import org.apache.pulsar.common.api.proto.MessageIdData;
import org.apache.pulsar.common.api.proto.MessageMetadata;
+import org.apache.pulsar.common.compression.CompressionCodecProvider;
+import org.apache.pulsar.common.protocol.Commands;
import org.apache.pulsar.common.util.Backoff;
import org.awaitility.Awaitility;
import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
public class ConsumerImplTest {
@@ -86,6 +103,13 @@ public class ConsumerImplTest {
}
private void createConsumer(ConsumerConfigurationData consumerConf, String topicName) {
+ createConsumer(consumerConf, topicName, null);
+ }
+
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ private void createConsumer(ConsumerConfigurationData consumerConf, String topicName,
+ ConsumerInterceptor<byte[]> interceptor) {
+ cleanup();
executorProvider = new ExecutorProvider(1, "ConsumerImplTest");
internalExecutor = Executors.newSingleThreadScheduledExecutor();
@@ -96,8 +120,10 @@ public class ConsumerImplTest {
CompletableFuture<Consumer<byte[]>> subscribeFuture = new CompletableFuture<>();
consumerConf.setSubscriptionName("test-sub");
+ ConsumerInterceptors<byte[]> interceptors =
+ interceptor == null ? null : new ConsumerInterceptors<>(Collections.singletonList(interceptor));
consumer = ConsumerImpl.newConsumerImpl(client, topicName, consumerConf,
- executorProvider, -1, false, subscribeFuture, null, null, null,
+ executorProvider, -1, false, subscribeFuture, null, null, interceptors,
true);
consumer.setState(HandlerState.State.Ready);
}
@@ -114,6 +140,296 @@ public class ConsumerImplTest {
}
}
+ /**
+ * Reproduces the chunked-message bookkeeping data race in {@link ConsumerImpl}, using the real
+ * production cross-thread pair:
+ * <ul>
+ * <li>{@code processMessageChunk(...)} on the Netty IO event-loop thread (here: the "receiver" thread);</li>
+ * <li>{@code removeExpireIncompleteChunkedMessages()} on the internalPinnedExecutor (here: the "expirer"
+ * thread).</li>
+ * </ul>
+ *
+ * <p>The faithful scenario is "a late chunk arrives for a uuid that is concurrently being expired": for each uuid
+ * the receiver delivers chunk 0, marks the ctx as already expired (receivedTime = 0), then delivers a late chunk 1
+ * that writes into the same {@code chunkedMsgBuffer} — exactly when the expiry task may release/recycle that ctx.
+ * Without serialization this races (use-after-free / double-recycle). With the fix, both paths take
+ * chunkedMessageLock, so it passes and pendingChunkedMessageCount stays equal to chunkedMessagesMap.size().
+ */
+ @Test(timeOut = 60000)
+ public void testChunkedMessageCountRaceBetweenReceiveAndExpiry() throws Exception {
+ // No max-pending eviction on the receive thread, so only the expiry path removes entries
+ consumerConf.setMaxPendingChunkedMessage(0);
+ createConsumer(consumerConf);
+ // Enable the expiry path (so removeExpireIncompleteChunkedMessages actually works) but skip its lazy
+ // self-scheduling so the test doesn't need a scheduled-executor mock.
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 1L;
+ consumer.expireChunkMessageTaskScheduled.set(true);
+
+ final int rounds = 5;
+ final int iterations = 50000;
+ final AtomicReference<Throwable> error = new AtomicReference<>();
+ // The workers use this consumer rather than the field, which the next test's setup reassigns, and they stop
+ // on this flag: if the test is interrupted (TestNG's timeout) the finally below stops them before teardown.
+ final ConsumerImpl<byte[]> target = consumer;
+ final AtomicBoolean stop = new AtomicBoolean();
+ final List<Thread> workers = new ArrayList<>();
+
+ try {
+ for (int round = 0; round < rounds && error.get() == null; round++) {
+ final int currentRound = round;
+ final CountDownLatch start = new CountDownLatch(1);
+ Thread receiver = new Thread(() -> {
+ try {
+ start.await();
+ for (int i = 0; i < iterations && error.get() == null && !stop.get(); i++) {
+ String uuid = "uuid-" + currentRound + "-" + i;
+ sendChunk(target, uuid, 0, 3);
+ ConsumerImpl.ChunkedMessageCtx ctx = target.chunkedMessagesMap.get(uuid);
+ if (ctx != null) {
+ // Make it eligible for removeExpireIncompleteChunkedMessages right away.
+ ctx.receivedTime = 0;
+ }
+ // Late chunk for the now-expiring uuid -> writes into the same ctx buffer.
+ sendChunk(target, uuid, 1, 3);
+ }
+ } catch (Throwable t) {
+ error.compareAndSet(null, t);
+ }
+ }, "receiver-netty-sim");
+
+ Thread expirer = new Thread(() -> {
+ try {
+ start.await();
+ // also stop on a failure: corrupted state may never drain the map
+ while (!stop.get() && error.get() == null
+ && (receiver.isAlive() || !target.chunkedMessagesMap.isEmpty())) {
+ target.removeExpireIncompleteChunkedMessages();
+ }
+ } catch (Throwable t) {
+ error.compareAndSet(null, t);
+ }
+ }, "expirer-pinned");
+
+ workers.add(receiver);
+ workers.add(expirer);
+ receiver.start();
+ expirer.start();
+ start.countDown();
+ receiver.join();
+ expirer.join();
+ }
+
+ if (error.get() != null) {
+ throw new AssertionError(
+ "chunked-message bookkeeping race: concurrent ctx/buffer access corrupted state", error.get());
+ }
+ int count = target.pendingChunkedMessageCount;
+ int mapSize = target.chunkedMessagesMap.size();
+ Assert.assertEquals(count, mapSize, "pendingChunkedMessageCount (" + count
+ + ") drifted from chunkedMessagesMap.size() (" + mapSize + ")");
+ } finally {
+ // Partial buffers can be released only once no worker can touch them any more; if one is still running
+ // after the bounded wait, leaking them is preferable to a use-after-free.
+ if (stopWorkers(stop, workers, 10_000)) {
+ releasePendingChunkedMessages(target);
+ }
+ }
+ }
+
+ /**
+ * Asks the workers to stop and waits for them, in total at most {@code timeoutMillis}; returns whether all of them
+ * have exited. An interrupt ends the wait early and is preserved for the caller.
+ */
+ private static boolean stopWorkers(AtomicBoolean stop, List<Thread> workers, long timeoutMillis) {
+ stop.set(true);
+ long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(timeoutMillis);
+ for (Thread worker : workers) {
+ long remainingMillis = TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime());
+ if (remainingMillis <= 0) {
+ break;
+ }
+ try {
+ worker.join(remainingMillis);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ break;
+ }
+ }
+ return workers.stream().noneMatch(Thread::isAlive);
+ }
+
+ private void sendChunk(String uuid, int chunkId, int numChunks) {
+ sendChunk(consumer, uuid, chunkId, numChunks);
+ }
+
+ private static void sendChunk(ConsumerImpl<byte[]> target, String uuid, int chunkId, int numChunks) {
+ sendChunk(target, uuid, chunkId, numChunks, chunkId, new byte[] {1, 2, 3, 4});
+ }
+
+ /**
+ * Delivers one chunk through {@code processMessageChunk} with an explicit entry id, so a test can redeliver a
+ * chunk id under a different message id.
+ */
+ private void sendChunk(String uuid, int chunkId, int numChunks, long entryId) {
+ sendChunk(uuid, chunkId, numChunks, entryId, new byte[] {1, 2, 3, 4});
+ }
+
+ /** Every chunk advertises a total message size of 64 bytes, so a larger payload is a malformed chunk. */
+ private static final int CHUNK_TOTAL_MSG_SIZE = 64;
+
+ private void sendChunk(String uuid, int chunkId, int numChunks, long entryId, byte[] payload) {
+ sendChunk(consumer, uuid, chunkId, numChunks, entryId, payload);
+ }
+
+ private static void sendChunk(ConsumerImpl<byte[]> target, String uuid, int chunkId, int numChunks, long entryId,
+ byte[] payload) {
+ MessageMetadata md = new MessageMetadata()
+ .setProducerName("p").setSequenceId(0).setPublishTime(System.currentTimeMillis())
+ .setUuid(uuid).setChunkId(chunkId).setNumChunksFromMsg(numChunks)
+ .setTotalChunkMsgSize(CHUNK_TOTAL_MSG_SIZE);
+ md.setCompression(CompressionType.NONE);
+ MessageIdData idData = new MessageIdData().setLedgerId(1L).setEntryId(entryId);
+ MessageIdImpl msgId = new MessageIdImpl(1L, entryId, -1);
+ ByteBuf chunk = Unpooled.wrappedBuffer(payload);
+ target.processMessageChunk(chunk, md, msgId, idData, null);
+ }
+
+ /**
+ * When a duplicated first chunk (chunkId == 0, redelivered) arrives for a uuid that already has an
+ * in-progress ctx, the old ctx is replaced: removing/recycling it must be paired with decrementing
+ * pendingChunkedMessageCount before the unconditional increment for the new ctx, so the counter stays
+ * equal to the real number of in-progress chunked messages (chunkedMessagesMap.size()).
+ */
+ @Test
+ public void testDuplicateFirstChunkOvercountsPendingChunkedMessageCount() throws Exception {
+ // Disable the lazy expiry self-scheduling so this test doesn't need a scheduled-executor mock.
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 0L;
+
+ final String uuid = "uuid-dup";
+ try {
+ // Deliver the first chunk (chunkId == 0) twice for the SAME uuid and SAME (ledgerId, entryId): a
+ // redelivered first chunk of a 2-chunk message. The message never completes, so only the ++/--
+ // bookkeeping is exercised.
+ for (int call = 0; call < 2; call++) {
+ sendChunk(uuid, 0, 2);
+ }
+
+ int count = consumer.pendingChunkedMessageCount;
+ int mapSize = consumer.chunkedMessagesMap.size();
+ Assert.assertEquals(count, mapSize,
+ "a redelivered first chunk over-counted pendingChunkedMessageCount (" + count
+ + ") vs chunkedMessagesMap.size() (" + mapSize + ")");
+
+ // Only this single uuid is ever enqueued, so the queue size equals its occurrence count. After a
+ // redelivered first chunk it must be exactly 1 (stale entry removed, re-added once) — not 2 as the old
+ // double-enqueue did. (GrowableArrayBlockingQueue intentionally doesn't support iteration, so assert on
+ // size().)
+ Assert.assertEquals(consumer.pendingChunkedMessageUuidQueue.size(), 1,
+ "uuid should appear exactly once in pendingChunkedMessageUuidQueue but queue size was "
+ + consumer.pendingChunkedMessageUuidQueue.size());
+ } finally {
+ // The message is never completed, so its partial buffer is still live: release it even if an assertion
+ // above failed.
+ releasePendingChunkedMessages();
+ }
+ }
+
+ /**
+ * Releases every partial chunked message still held by the consumer, for tests that intentionally leave one.
+ */
+ private void releasePendingChunkedMessages() {
+ releasePendingChunkedMessages(consumer);
+ }
+
+ private static void releasePendingChunkedMessages(ConsumerImpl<byte[]> target) {
+ for (ConsumerImpl.ChunkedMessageCtx ctx : target.chunkedMessagesMap.values()) {
+ if (ctx.chunkedMsgBuffer != null) {
+ ctx.chunkedMsgBuffer.release();
+ }
+ }
+ target.chunkedMessagesMap.clear();
+ }
+
+ /**
+ * A uuid whose ctx is no longer in {@code chunkedMessagesMap} but still sits at the head of
+ * {@code pendingChunkedMessageUuidQueue} is a "ghost" head. Every removal path now drops the queue entry along
+ * with the map entry, so this is a defensive invariant of {@code removeExpireIncompleteChunkedMessages}: it must
+ * drop such null-ctx heads and keep scanning (as {@code removeOldestPendingChunkedMessage} does), because a ghost
+ * that stopped the scan would leave every genuinely-expired incomplete chunk behind it uncleaned (its
+ * {@code chunkedMsgBuffer} leaked and its chunks never acked), which is what happened before this invariant was
+ * enforced.
+ */
+ @Test
+ public void testExpiryDrainsPastGhostQueueEntries() throws Exception {
+ // No max-pending eviction, so only the expiry path consumes the queue.
+ consumerConf.setMaxPendingChunkedMessage(0);
+ createConsumer(consumerConf);
+ // Enable expiry but skip the lazy self-scheduling so the test doesn't need a scheduled-executor mock.
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 1L;
+ consumer.expireChunkMessageTaskScheduled.set(true);
+
+ // uuid-A: first chunk of a 3-chunk message -> in-progress. Enqueued first (queue head).
+ sendChunk("uuid-A", 0, 3);
+ // uuid-B: first chunk of a 2-chunk message -> in-progress, enqueued after A. Never completed.
+ sendChunk("uuid-B", 0, 2);
+
+ // Plant a ghost head by hand: take uuid-A's ctx out of the map and recycle it while leaving its uuid in
+ // pendingChunkedMessageUuidQueue. (Releasing the buffer mirrors the assembled-payload release that a real
+ // completion performs.)
+ ConsumerImpl.ChunkedMessageCtx ctxA = consumer.chunkedMessagesMap.remove("uuid-A");
+ assertThat(ctxA).isNotNull();
+ if (ctxA.chunkedMsgBuffer != null) {
+ ctxA.chunkedMsgBuffer.release();
+ }
+ ctxA.recycle();
+
+ // Make uuid-B eligible for expiry right away.
+ ConsumerImpl.ChunkedMessageCtx ctxB = consumer.chunkedMessagesMap.get("uuid-B");
+ assertThat(ctxB).isNotNull();
+ ctxB.receivedTime = 0;
+ ByteBuf bufB = ctxB.chunkedMsgBuffer;
+
+ consumer.removeExpireIncompleteChunkedMessages();
+
+ // With the bug, the ghost head (uuid-A) makes expiry return before reaching the expired uuid-B.
+ assertThat(consumer.chunkedMessagesMap.containsKey("uuid-B"))
+ .as("expired incomplete chunked message behind a ghost queue head must be cleaned")
+ .isFalse();
+ assertThat(bufB.refCnt()).as("expired chunk buffer must be released").isZero();
+ assertThat(consumer.pendingChunkedMessageUuidQueue)
+ .as("ghost and expired entries must both be drained from the queue").isEmpty();
+ }
+
+ /**
+ * The "lost first chunk" / forward-gap discard path in {@code processMessageChunk} removes an in-progress ctx from
+ * {@code chunkedMessagesMap}, but must also keep the other bookkeeping in sync: decrement
+ * {@code pendingChunkedMessageCount} (the ctx was counted when its first chunk created it) and drop the uuid from
+ * {@code pendingChunkedMessageUuidQueue} (otherwise a ghost lingers). Without the decrement the count drifts upward
+ * and prematurely triggers {@code removeOldestPendingChunkedMessage}.
+ */
+ @Test
+ public void testForwardGapDiscardKeepsCountAndQueueConsistent() throws Exception {
+ // Disable the lazy expiry self-scheduling so this test doesn't need a scheduled-executor mock.
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 0L;
+
+ final String uuid = "uuid-gap";
+ // First chunk (chunkId 0) of a 3-chunk message -> in-progress: count=1, map={uuid}, queue=[uuid].
+ sendChunk(uuid, 0, 3);
+ // A chunk that skips chunkId 1 (forward gap): chunkId 2 != lastChunkedMessageId(0)+1 and > lastChunkedMessageId
+ // -> hits the "lost first chunk" discard path that removes the ctx from chunkedMessagesMap.
+ sendChunk(uuid, 2, 3);
+
+ int count = consumer.pendingChunkedMessageCount;
+ int mapSize = consumer.chunkedMessagesMap.size();
+ Assert.assertEquals(count, mapSize,
+ "forward-gap discard over-counted pendingChunkedMessageCount (" + count
+ + ") vs chunkedMessagesMap.size() (" + mapSize + ")");
+
+ Assert.assertEquals(consumer.pendingChunkedMessageUuidQueue.size(), mapSize,
+ "discarded uuid must be removed from pendingChunkedMessageUuidQueue (no ghost), queue size was "
+ + consumer.pendingChunkedMessageUuidQueue.size());
+ }
+
@Test(invocationTimeOut = 1000)
public void testNotifyPendingReceivedCallback_EmptyQueueNotThrowsException() {
consumer.notifyPendingReceivedCallback(null, null);
@@ -475,4 +791,379 @@ public class ConsumerImplTest {
"ack timeout tracking was disabled on a persistent topic");
}
+ /**
+ * A chunk-ack interceptor that records every individually acked message id and can be paused on a latch, so a test
+ * can hold an acknowledgment callback open and observe what the consumer does meanwhile.
+ */
+ private static class RecordingAckInterceptor implements ConsumerInterceptor<byte[]> {
+ final List<MessageId> acked = new CopyOnWriteArrayList<>();
+ final CountDownLatch ackStarted = new CountDownLatch(1);
+ final CountDownLatch ackRelease = new CountDownLatch(1);
+ volatile boolean pauseAcks;
+
+ @Override
+ public void close() {
+ }
+
+ @Override
+ public Message<byte[]> beforeConsume(Consumer<byte[]> consumer, Message<byte[]> message) {
+ return message;
+ }
+
+ @Override
+ public void onAcknowledge(Consumer<byte[]> consumer, MessageId messageId, Throwable exception) {
+ acked.add(messageId);
+ if (pauseAcks) {
+ ackStarted.countDown();
+ try {
+ ackRelease.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ @Override
+ public void onAcknowledgeCumulative(Consumer<byte[]> consumer, MessageId messageId, Throwable exception) {
+ }
+
+ @Override
+ public void onNegativeAcksSend(Consumer<byte[]> consumer, Set<MessageId> messageIds) {
+ }
+
+ @Override
+ public void onAckTimeoutSend(Consumer<byte[]> consumer, Set<MessageId> messageIds) {
+ }
+ }
+
+ private static final String PERSISTENT_TOPIC = "persistent://tenant/ns1/my-topic";
+
+ /**
+ * Delivers one chunk through the real {@code messageReceived} entry point, framed as the broker would send it
+ * (checksum, metadata, payload): chunk {@code chunkId} of a {@code numChunks}-chunk message carrying the given
+ * still-compressed bytes. The entry id doubles as the chunk id.
+ */
+ private void deliverChunk(ClientCnx cnx, String uuid, int chunkId, int numChunks, int totalChunkMsgSize,
+ CompressionType compressionType, int uncompressedSize, byte[] chunkBytes) {
+ ByteBuf headersAndPayload = frameChunk(uuid, chunkId, numChunks, totalChunkMsgSize, compressionType,
+ uncompressedSize, chunkBytes);
+ try {
+ consumer.messageReceived(chunkCommand(chunkId), headersAndPayload, cnx);
+ } finally {
+ headersAndPayload.release();
+ }
+ }
+
+ /** Frames one chunk as the broker would send it: checksum, metadata and the given payload bytes. */
+ private static ByteBuf frameChunk(String uuid, int chunkId, int numChunks, int totalChunkMsgSize,
+ CompressionType compressionType, int uncompressedSize, byte[] chunkBytes) {
+ MessageMetadata md = new MessageMetadata()
+ .setProducerName("p").setSequenceId(0).setPublishTime(System.currentTimeMillis())
+ .setUuid(uuid).setChunkId(chunkId).setNumChunksFromMsg(numChunks)
+ .setTotalChunkMsgSize(totalChunkMsgSize)
+ .setCompression(compressionType).setUncompressedSize(uncompressedSize);
+ return Commands.serializeMetadataAndPayload(Commands.ChecksumType.Crc32c, md,
+ Unpooled.wrappedBuffer(chunkBytes));
+ }
+
+ private CommandMessage chunkCommand(long entryId) {
+ CommandMessage cmd = new CommandMessage().setConsumerId(consumer.consumerId);
+ cmd.setMessageId().setLedgerId(1L).setEntryId(entryId);
+ return cmd;
+ }
+
+ private static byte[] zlibCompress(byte[] data) {
+ ByteBuf source = Unpooled.wrappedBuffer(data);
+ ByteBuf compressed = CompressionCodecProvider.getCompressionCodec(CompressionType.ZLIB).encode(source);
+ try {
+ byte[] bytes = new byte[compressed.readableBytes()];
+ compressed.readBytes(bytes);
+ return bytes;
+ } finally {
+ compressed.release();
+ source.release();
+ }
+ }
+
+ private static void assertNoPendingChunkedMessages(ConsumerImpl<?> consumer) {
+ assertThat(consumer.chunkedMessagesMap).as("chunkedMessagesMap").isEmpty();
+ assertThat(consumer.pendingChunkedMessageUuidQueue).as("pendingChunkedMessageUuidQueue").isEmpty();
+ assertThat(consumer.pendingChunkedMessageCount).as("pendingChunkedMessageCount").isZero();
+ }
+
+ /**
+ * The final chunk, delivered through {@code messageReceived}, transfers the assembled buffer out of the context
+ * (removed and recycled under the lock) and decompresses it outside the lock. On success the message must reach
+ * the receive queue with the original payload and a chunk message id, all bookkeeping must be empty, the assembled
+ * buffer released, and the chunk ids registered with the unack tracker.
+ */
+ @Test(timeOut = 30000)
+ public void testFinalChunkThroughMessageReceivedDeliversAssembledMessage() throws Exception {
+ RecordingAckInterceptor interceptor = new RecordingAckInterceptor();
+ createConsumer(consumerConf, PERSISTENT_TOPIC, interceptor);
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 0L;
+ ClientCnx cnx = ClientTestFixtures.mockClientCnx();
+
+ byte[] original = new byte[4096];
+ new Random(1).nextBytes(original);
+ byte[] compressed = zlibCompress(original);
+ int split = compressed.length / 2;
+ byte[] chunk0 = Arrays.copyOfRange(compressed, 0, split);
+ byte[] chunk1 = Arrays.copyOfRange(compressed, split, compressed.length);
+
+ deliverChunk(cnx, "uuid-ok", 0, 2, compressed.length, CompressionType.ZLIB, original.length, chunk0);
+ ByteBuf assembled = consumer.chunkedMessagesMap.get("uuid-ok").chunkedMsgBuffer;
+ assertThat(assembled).isNotNull();
+ deliverChunk(cnx, "uuid-ok", 1, 2, compressed.length, CompressionType.ZLIB, original.length, chunk1);
+
+ Awaitility.await().untilAsserted(() -> assertThat(consumer.numMessagesInQueue()).isEqualTo(1));
+ Message<byte[]> message = consumer.incomingMessages.poll();
+ assertThat(message).isNotNull();
+ try {
+ assertThat(message.getData()).isEqualTo(original);
+ assertThat(message.getMessageId()).isInstanceOf(ChunkMessageIdImpl.class);
+ assertThat(consumer.unAckedChunkedMessageIdSequenceMap.containsKey(message.getMessageId()))
+ .as("chunk ids registered with the unack tracker").isTrue();
+ assertThat(assembled.refCnt()).as("assembled buffer released after decompression").isZero();
+ assertNoPendingChunkedMessages(consumer);
+ assertThat(interceptor.acked).as("a successful completion acks nothing").isEmpty();
+ } finally {
+ message.release();
+ }
+ }
+
+ @DataProvider
+ public Object[][] decompressionFailures() {
+ return new Object[][] {
+ // ZLib reports corrupt input as an IOException
+ {"corrupt-input"},
+ // ZLib reports a decoded size that differs from the advertised one as an IllegalArgumentException
+ {"size-mismatch"},
+ };
+ }
+
+ /**
+ * When the final chunk's decompression fails, the assembled buffer has already been detached from its context, so
+ * nothing but this path can release it or dispose of the earlier chunks' message ids. Whether the codec fails with
+ * a checked or an unchecked exception, the buffer must be released, the earlier chunk ids individually acked, the
+ * final chunk discarded, all bookkeeping left empty, and nothing must escape {@code messageReceived}.
+ */
+ @Test(dataProvider = "decompressionFailures", timeOut = 30000)
+ public void testFinalChunkDecompressionFailureThroughMessageReceived(String failure) throws Exception {
+ RecordingAckInterceptor interceptor = new RecordingAckInterceptor();
+ createConsumer(consumerConf, PERSISTENT_TOPIC, interceptor);
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 0L;
+ ClientCnx cnx = ClientTestFixtures.mockClientCnx();
+
+ byte[] original = new byte[4096];
+ new Random(2).nextBytes(original);
+ byte[] compressed = zlibCompress(original);
+ int uncompressedSize = original.length;
+ if ("corrupt-input".equals(failure)) {
+ // keep the stream header, corrupt the deflate body
+ for (int i = 4; i < compressed.length; i++) {
+ compressed[i] = (byte) ~compressed[i];
+ }
+ } else {
+ uncompressedSize = original.length + 1;
+ }
+ int split = compressed.length / 2;
+ byte[] chunk0 = Arrays.copyOfRange(compressed, 0, split);
+ byte[] chunk1 = Arrays.copyOfRange(compressed, split, compressed.length);
+
+ deliverChunk(cnx, "uuid-bad", 0, 2, compressed.length, CompressionType.ZLIB, uncompressedSize, chunk0);
+ ByteBuf assembled = consumer.chunkedMessagesMap.get("uuid-bad").chunkedMsgBuffer;
+ assertThat(assembled).isNotNull();
+ deliverChunk(cnx, "uuid-bad", 1, 2, compressed.length, CompressionType.ZLIB, uncompressedSize, chunk1);
+
+ assertThat(assembled.refCnt()).as("detached assembled buffer released on decompression failure").isZero();
+ assertNoPendingChunkedMessages(consumer);
+ assertThat(consumer.numMessagesInQueue()).as("no message delivered").isZero();
+ assertThat(consumer.unAckedChunkedMessageIdSequenceMap).as("nothing registered with the unack tracker")
+ .isEmpty();
+ // The earlier chunk (entry 0) is acked individually; the final chunk (entry 1) is discarded with a validation
+ // error straight on the connection, so it is not routed through the interceptor. Two chunks are far below the
+ // flow-control threshold, so that discard is the only write on the connection.
+ assertThat(interceptor.acked).containsExactly(new MessageIdImpl(1L, 0L, -1));
+ verify(cnx.ctx(), times(1)).writeAndFlush(any(), any());
+ }
+
+ /**
+ * Acks decided under {@code chunkedMessageLock} are issued after it is released. The persistent acknowledgments
+ * grouping tracker runs the consumer's acknowledgment interceptors inline, so while the expiry thread is held in
+ * an application interceptor, a chunk arriving on the IO thread must still be processed rather than block on the
+ * lock behind that callback.
+ */
+ @Test(timeOut = 30000)
+ public void testChunkReceiveProceedsWhileExpiryAckInterceptorIsBlocked() throws Exception {
+ RecordingAckInterceptor interceptor = new RecordingAckInterceptor();
+ consumerConf.setMaxPendingChunkedMessage(0);
+ createConsumer(consumerConf, PERSISTENT_TOPIC, interceptor);
+ // Enable expiry but skip the lazy self-scheduling so the test doesn't need a scheduled-executor mock.
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 1L;
+ consumer.expireChunkMessageTaskScheduled.set(true);
+
+ sendChunk("uuid-expired", 0, 3);
+ consumer.chunkedMessagesMap.get("uuid-expired").receivedTime = 0;
+ interceptor.pauseAcks = true;
+
+ Thread expirer = new Thread(consumer::removeExpireIncompleteChunkedMessages, "expirer-pinned");
+ expirer.start();
+ try {
+ assertThat(interceptor.ackStarted.await(10, TimeUnit.SECONDS))
+ .as("expiry reached the acknowledgment interceptor").isTrue();
+
+ // The expiry thread is now parked inside the application callback. A chunk for another uuid must still
+ // be processed on the "IO thread".
+ CountDownLatch received = new CountDownLatch(1);
+ Thread receiver = new Thread(() -> {
+ sendChunk("uuid-live", 0, 3);
+ received.countDown();
+ }, "receiver-netty-sim");
+ receiver.start();
+ try {
+ assertThat(received.await(5, TimeUnit.SECONDS))
+ .as("chunk reception must not wait for the acknowledgment interceptor to return").isTrue();
+ } finally {
+ // never leave the receiver parked on the lock if the assertion fails
+ interceptor.ackRelease.countDown();
+ receiver.join();
+ }
+ assertThat(consumer.chunkedMessagesMap).containsKey("uuid-live");
+ assertThat(consumer.chunkedMessagesMap).doesNotContainKey("uuid-expired");
+ } finally {
+ interceptor.ackRelease.countDown();
+ expirer.join();
+ releasePendingChunkedMessages();
+ }
+ assertThat(interceptor.acked).containsExactly(new MessageIdImpl(1L, 0L, -1));
+ }
+
+ /**
+ * A first chunk arriving under a new message id for a uuid that still has an in-progress context means the earlier
+ * chunks belong to a corrupted chunked message that will never be delivered. Replacing the context must ack those
+ * earlier chunk ids, through the deferred-ack path, and leave the count and queue consistent with the map.
+ */
+ @Test(timeOut = 30000)
+ public void testReplacedCorruptedChunkedMessageAcksItsEarlierChunks() throws Exception {
+ RecordingAckInterceptor interceptor = new RecordingAckInterceptor();
+ createConsumer(consumerConf, PERSISTENT_TOPIC, interceptor);
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 0L;
+
+ try {
+ // entries 10 and 11: chunks 0 and 1 of a corrupted 3-chunk message that never completes
+ sendChunk("uuid-corrupt", 0, 3, 10L);
+ sendChunk("uuid-corrupt", 1, 3, 11L);
+ ByteBuf corruptedBuffer = consumer.chunkedMessagesMap.get("uuid-corrupt").chunkedMsgBuffer;
+ // Hold an extra reference, so that the pooled buffer cannot be recycled for the next chunked message
+ corruptedBuffer.retain();
+ // entry 20: chunk 0 of the re-published message under the same uuid
+ sendChunk("uuid-corrupt", 0, 3, 20L);
+
+ assertThat(interceptor.acked).as("earlier chunks of the replaced message are acked")
+ .containsExactly(new MessageIdImpl(1L, 10L, -1), new MessageIdImpl(1L, 11L, -1));
+ assertThat(corruptedBuffer.refCnt()).as("replaced buffer released").isEqualTo(1);
+ corruptedBuffer.release();
+ assertThat(consumer.chunkedMessagesMap).containsOnlyKeys("uuid-corrupt");
+ assertThat(consumer.pendingChunkedMessageCount).isEqualTo(1);
+ assertThat(consumer.pendingChunkedMessageUuidQueue.size()).isEqualTo(1);
+ } finally {
+ releasePendingChunkedMessages();
+ }
+ }
+
+ /**
+ * Exceeding {@code maxPendingChunkedMessage} evicts the oldest in-progress chunked message on the receive path.
+ * With {@code autoAckOldestChunkedMessageOnQueueFull} its chunk ids are acked, through the deferred-ack path, and
+ * the eviction keeps the count and queue consistent with the map.
+ */
+ @Test(timeOut = 30000)
+ public void testEvictingOldestPendingChunkedMessageAcksItsChunks() throws Exception {
+ RecordingAckInterceptor interceptor = new RecordingAckInterceptor();
+ consumerConf.setMaxPendingChunkedMessage(1);
+ consumerConf.setAutoAckOldestChunkedMessageOnQueueFull(true);
+ createConsumer(consumerConf, PERSISTENT_TOPIC, interceptor);
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 0L;
+
+ try {
+ sendChunk("uuid-old", 0, 3, 10L);
+ ByteBuf oldBuffer = consumer.chunkedMessagesMap.get("uuid-old").chunkedMsgBuffer;
+ // Hold an extra reference, so that the pooled buffer cannot be recycled for the next chunked message
+ oldBuffer.retain();
+ // a second in-progress message exceeds maxPendingChunkedMessage = 1 and evicts the oldest
+ sendChunk("uuid-new", 0, 3, 20L);
+
+ assertThat(interceptor.acked).as("evicted chunks are acked")
+ .containsExactly(new MessageIdImpl(1L, 10L, -1));
+ assertThat(oldBuffer.refCnt()).as("evicted buffer released").isEqualTo(1);
+ oldBuffer.release();
+ assertThat(consumer.chunkedMessagesMap).containsOnlyKeys("uuid-new");
+ assertThat(consumer.pendingChunkedMessageCount).isEqualTo(1);
+ assertThat(consumer.pendingChunkedMessageUuidQueue.size()).isEqualTo(1);
+ } finally {
+ releasePendingChunkedMessages();
+ }
+ }
+
+ /**
+ * Acks queued under {@code chunkedMessageLock} belong to a context that has already been removed, so they must be
+ * issued even when chunk processing then throws: here a replacement first chunk queues the old context's ids and
+ * is itself malformed (its payload overflows the advertised total size), so appending it fails with
+ * {@link IndexOutOfBoundsException}. The old ids must still be acked, or they are lost as a permanent ack hole.
+ */
+ @Test(timeOut = 30000)
+ public void testAcksQueuedUnderLockAreIssuedWhenChunkProcessingThrows() throws Exception {
+ RecordingAckInterceptor interceptor = new RecordingAckInterceptor();
+ createConsumer(consumerConf, PERSISTENT_TOPIC, interceptor);
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 0L;
+
+ try {
+ sendChunk("uuid-throw", 0, 3, 10L);
+ sendChunk("uuid-throw", 1, 3, 11L);
+ // replacement chunk 0 under a new message id, larger than the buffer it allocates for itself
+ byte[] oversized = new byte[CHUNK_TOTAL_MSG_SIZE + 1];
+ assertThatThrownBy(() -> sendChunk("uuid-throw", 0, 3, 20L, oversized))
+ .isInstanceOf(IndexOutOfBoundsException.class);
+
+ assertThat(interceptor.acked).as("ids of the replaced context are acked despite the failure")
+ .containsExactly(new MessageIdImpl(1L, 10L, -1), new MessageIdImpl(1L, 11L, -1));
+ } finally {
+ releasePendingChunkedMessages();
+ }
+ }
+
+ /**
+ * The same failure as {@link #testAcksQueuedUnderLockAreIssuedWhenChunkProcessingThrows}, delivered through the
+ * real {@code messageReceived} entry point so that its own acknowledgment {@code finally} is the one under test,
+ * rather than the one in the {@code processMessageChunk} wrapper. The failed chunk must also give up its reference
+ * to the frame.
+ */
+ @Test(timeOut = 30000)
+ public void testAcksQueuedUnderLockAreIssuedWhenFramedChunkProcessingThrows() throws Exception {
+ RecordingAckInterceptor interceptor = new RecordingAckInterceptor();
+ createConsumer(consumerConf, PERSISTENT_TOPIC, interceptor);
+ consumer.expireTimeOfIncompleteChunkedMessageMillis = 0L;
+ ClientCnx cnx = ClientTestFixtures.mockClientCnx();
+
+ byte[] chunk = new byte[] {1, 2, 3, 4};
+ try {
+ deliverChunk(cnx, "uuid-framed-throw", 0, 3, CHUNK_TOTAL_MSG_SIZE, CompressionType.NONE, 0, chunk);
+ deliverChunk(cnx, "uuid-framed-throw", 1, 3, CHUNK_TOTAL_MSG_SIZE, CompressionType.NONE, 0, chunk);
+ // replacement chunk 0 under a new entry id, larger than the buffer it allocates for itself
+ ByteBuf oversized = frameChunk("uuid-framed-throw", 0, 3, CHUNK_TOTAL_MSG_SIZE, CompressionType.NONE, 0,
+ new byte[CHUNK_TOTAL_MSG_SIZE + 1]);
+ try {
+ assertThatThrownBy(() -> consumer.messageReceived(chunkCommand(20L), oversized, cnx))
+ .isInstanceOf(IndexOutOfBoundsException.class);
+ assertThat(oversized.refCnt()).as("the failed chunk keeps no reference to the frame").isEqualTo(1);
+ } finally {
+ oversized.release();
+ }
+
+ assertThat(interceptor.acked).as("ids of the replaced context are acked despite the failure")
+ .containsExactly(new MessageIdImpl(1L, 0L, -1), new MessageIdImpl(1L, 1L, -1));
+ } finally {
+ releasePendingChunkedMessages();
+ }
+ }
}