[fix][client] Serialize chunked-message bookkeeping to fix use-after-free and count/queue drift

Author: maxlisongsong

pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java

@@ -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;
@@ -1507,28 +1512,35 @@ 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;
-                }
+                // Keep assembly of the final chunk and the finalize (map removal + recycle) in one critical section so
+                // the expiry/eviction path can't release/recycle this ctx in between.
+                synchronized (chunkedMessageLock) {
+                    uncompressedPayload = processMessageChunk(uncompressedPayload, msgMetadata, msgId, messageId, cnx);
+                    if (uncompressedPayload == 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 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]);
+                    // 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]);
+                    }
+                    // add chunked messageId to unack-message tracker, and reduce pending-chunked-message count
+                    unAckedChunkedMessageIdSequenceMap.put(msgId, 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();
                 }
-                // add chunked messageId to unack-message tracker, and reduce pending-chunked-message count
-                unAckedChunkedMessageIdSequenceMap.put(msgId, chunkedMsgCtx.chunkedMessageIds);
-                pendingChunkedMessageCount--;
-                chunkedMsgCtx.recycle();
             }
 
             // If the topic is non-persistent, we should not ignore any messages.
@@ -1573,6 +1585,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);
         }
@@ -1630,6 +1651,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) {
@@ -1640,6 +1666,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());
         }
 
@@ -1686,6 +1714,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();
@@ -3127,6 +3160,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()) {
@@ -3142,15 +3183,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;
             }
         }