Skip to content

[fix][client] Serialize chunked-message bookkeeping to fix use-after-… #26084

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SongOf wants to merge 1 commit into
base: master
Choose a base branch
from
+354 −27
Open

[fix][client] Serialize chunked-message bookkeeping to fix use-after-… #26084

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 102 additions & 27 deletions pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,11 @@ public class ConsumerImpl<T> extends ConsumerBase<T> implements ConnectionHandle

protected Map<String, ChunkedMessageCtx> chunkedMessagesMap = new ConcurrentHashMap<>();
private 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).
private final Object chunkedMessageLock = new Object();
protected long expireTimeOfIncompleteChunkedMessageMillis = 0;
private final AtomicBoolean expireChunkMessageTaskScheduled = new AtomicBoolean(false);
private final int maxPendingChunkedMessage;
Expand Down Expand Up @@ -1507,28 +1512,54 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien
if (isMessageUndecryptable || (numMessages == 1 && !msgMetadata.hasNumMessagesInBatch())) {

if (isChunkedMessage) {
uncompressedPayload = processMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, cnx);
if (uncompressedPayload == null) {
return;
ByteBuf compressedAssembledPayload;
MessageIdImpl[] chunkedMessageIds;
// 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) {

@Denovo1998 Denovo1998 Jun 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The lock is described as guarding only chunk bookkeeping, but this path currently holds chunkedMessageLock through final-chunk decompression in doProcessMessageChunk() (uncompressPayloadIfNeeded). That can be a relatively expensive codec operation and, on decompression failure, can call discardCorruptedMessage() while the lock is held. This serializes unrelated chunked-message completions and expiry behind decompression work.

Could we narrow the critical section to only append/remove the ctx and capture the assembled buffer / chunk ids, then decompress and build the final message outside the lock? That would still close the assemble -> finalize race without turning the lock into a codec/decompression gate.

compressedAssembledPayload =
processMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, cnx);
if (compressedAssembledPayload == null) {
return;
}

// last chunk received: so, stitch chunked-messages and clear up chunkedMsgBuffer
log.debug().attr("chunkid", msgMetadata.getChunkId())
.attr("totalChunks", msgMetadata.getNumChunksFromMsg())
.attr("msgid", msgId)
.attr("sequenceid", msgMetadata.getSequenceId())
.log("Chunked message completed chunkId, total-chunks, msgId sequenceId");

// 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());

@Denovo1998 Denovo1998 Jun 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this still leaves a stale-queue-head case. The PR now states that pendingChunkedMessageUuidQueue has exactly one entry per in-progress uuid, but not every map removal keeps that invariant. The final-chunk path removes the ctx from chunkedMessagesMap and decrements the count without removing the uuid from pendingChunkedMessageUuidQueue; the unexpected/out-of-order branch also removes/recycles the ctx without removing the queue entry or decrementing pendingChunkedMessageCount.

That matters because doRemoveExpireIncompleteChunkedMessages() uses peek(). If the queue head points to a completed/removed uuid, chunkedMsgCtx is null and the method returns, so later expired incomplete chunked messages behind that stale head are never expired. I reproduced this on the PR head by sending chunk 0 then out-of-order chunk 2 for one uuid, then an expired chunk 0 for another uuid; removeExpireIncompleteChunkedMessages() returns at the stale head and leaves the second ctx in the map.

Could we either remove the uuid from pendingChunkedMessageUuidQueue on every path that removes from chunkedMessagesMap (final completion and unexpected/out-of-order cleanup included), or make the expiry loop drain stale null heads before deciding that the oldest live ctx is not expired?

@SongOf SongOf Jun 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed in the final-chunk path. and fixed in out-of-order chunk path. At the same time, modify doRemoveExpireIncompleteChunkedMessages to handle null elements at the head of the queue, aligning its logic with the doRemoveOldestPendingChunkedMessage method.

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();
}

// last chunk received: so, stitch chunked-messages and clear up chunkedMsgBuffer
log.debug().attr("chunkid", msgMetadata.getChunkId())
.attr("totalChunks", msgMetadata.getNumChunksFromMsg())
.attr("msgid", msgId)
.attr("sequenceid", msgMetadata.getSequenceId())
.log("Chunked message completed chunkId, total-chunks, msgId sequenceId");

// 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]);
// Outside the lock: set the chunk message id, decompress the assembled payload, and (only on success)
// register it with the unack tracker.
if (chunkedMessageIds.length > 0) {
msgId = new ChunkMessageIdImpl(chunkedMessageIds[0],
chunkedMessageIds[chunkedMessageIds.length - 1]);
}
// add chunked messageId to unack-message tracker, and reduce pending-chunked-message count
unAckedChunkedMessageIdSequenceMap.put(msgId, chunkedMsgCtx.chunkedMessageIds);
pendingChunkedMessageCount--;
chunkedMsgCtx.recycle();
uncompressedPayload =
uncompressPayloadIfNeeded(messageId, msgMetadata, compressedAssembledPayload, cnx, false);
compressedAssembledPayload.release();
if (uncompressedPayload == null) {
return;
}
// add chunked messageId to unack-message tracker
unAckedChunkedMessageIdSequenceMap.put(msgId, chunkedMessageIds);
}

// If the topic is non-persistent, we should not ignore any messages.
Expand Down Expand Up @@ -1573,6 +1604,15 @@ void messageReceived(CommandMessage cmdMessage, ByteBuf headersAndPayload, Clien

private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata msgMetadata, MessageIdImpl msgId,
MessageIdData messageId, ClientCnx cnx) {
// Self-guard: this method does not rely on the caller already holding chunkedMessageLock. messageReceived does
// hold it (to keep the final-chunk assembly + finalize atomic), so that path re-enters this lock harmlessly.
synchronized (chunkedMessageLock) {
return doProcessMessageChunk(compressedPayload, msgMetadata, msgId, messageId, cnx);
}
}

private ByteBuf doProcessMessageChunk(ByteBuf compressedPayload, MessageMetadata msgMetadata, MessageIdImpl msgId,
MessageIdData messageId, ClientCnx cnx) {
if (msgMetadata.getChunkId() != (msgMetadata.getNumChunksFromMsg() - 1)) {
increaseAvailablePermits(cnx);
}
Expand Down Expand Up @@ -1630,6 +1670,11 @@ private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata m
}
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) {
Expand All @@ -1640,6 +1685,8 @@ private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata m
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());
}

Expand Down Expand Up @@ -1686,6 +1733,11 @@ private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata m
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();
Expand All @@ -1710,11 +1762,11 @@ private ByteBuf processMessageChunk(ByteBuf compressedPayload, MessageMetadata m
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;
}

/**
Expand Down Expand Up @@ -3127,6 +3179,14 @@ public void recycle() {
}

private void removeOldestPendingChunkedMessage() {
// Self-guard, same as processMessageChunk: today this only runs inside doProcessMessageChunk (already under the
// lock, so this re-enters), but locking here keeps the method safe for any direct caller.
synchronized (chunkedMessageLock) {
doRemoveOldestPendingChunkedMessage();
}
}

private void doRemoveOldestPendingChunkedMessage() {
ChunkedMessageCtx chunkedMsgCtx = null;
String firstPendingMsgUuid = null;
while (chunkedMsgCtx == null && !pendingChunkedMessageUuidQueue.isEmpty()) {
Expand All @@ -3142,15 +3202,30 @@ protected void removeExpireIncompleteChunkedMessages() {
if (expireTimeOfIncompleteChunkedMessageMillis <= 0) {
return;
}
ChunkedMessageCtx chunkedMsgCtx = null;
synchronized (chunkedMessageLock) {
doRemoveExpireIncompleteChunkedMessages();
}
}

private void doRemoveExpireIncompleteChunkedMessages() {
String messageUUID;
while ((messageUUID = pendingChunkedMessageUuidQueue.peek()) != null) {
chunkedMsgCtx = StringUtils.isNotBlank(messageUUID) ? chunkedMessagesMap.get(messageUUID) : null;
if (chunkedMsgCtx != null && System
.currentTimeMillis() > (chunkedMsgCtx.receivedTime + expireTimeOfIncompleteChunkedMessageMillis)) {
ChunkedMessageCtx chunkedMsgCtx =
StringUtils.isNotBlank(messageUUID) ? chunkedMessagesMap.get(messageUUID) : null;
if (chunkedMsgCtx == null) {
// Ghost head: the chunked message already completed/was removed from chunkedMessagesMap but its uuid
// lingers in the queue. Drop it and keep scanning (mirroring doRemoveOldestPendingChunkedMessage) 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);
} else {
// The queue is ordered oldest-first, so the first live not-yet-expired head means nothing behind it
// is expired either.
return;
}
}
Expand Down
Loading