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

[SongOf]

Motivation

ConsumerImpl's chunked-message reassembly state — the per-uuid ChunkedMessageCtx, its chunkedMsgBuffer,
pendingChunkedMessageCount and pendingChunkedMessageUuidQueue — is mutated from two different threads with no
synchronization:

• the receive/assembly path (processMessageChunk, and the last-chunk finalize in messageReceived) runs on
  the Netty IO event-loop thread (ClientCnx.handleMessage calls consumer.messageReceived(...) directly);
• the incomplete-chunk expiry path (removeExpireIncompleteChunkedMessages) runs on the client's
  internalPinnedExecutor — a separate single-thread pool (client.getInternalExecutorService()), not the
  eventLoopGroup.

When a late chunk for a uuid arrives while the expiry task is removing that same ctx, the expiry thread can
release() / recycle() a ChunkedMessageCtx and its buffer while the receive thread is still writing into it.
This races into:

• use-after-free / double-free of chunkedMsgBuffer (IllegalReferenceCountException, or worse — writing into
  memory the allocator already handed to someone else);
• double ChunkedMessageCtx.recycle(), which corrupts the Netty Recycler pool and can hand the same instance to
  two different chunked messages;
• pendingChunkedMessageCount drift (non-atomic int mutated from both threads).

Incomplete-chunk expiry is enabled by default (expireTimeOfIncompleteChunkedMessageMillis = 1 minute), so any
consumer of chunked messages is exposed.

The same bookkeeping also has several single-threaded defects:

• a redelivered first chunk (chunkId == 0 for a uuid that already has an in-progress ctx) replaced the old ctx
  without decrementing pendingChunkedMessageCount and re-enqueued the uuid, so the counter over-counted and the
  queue held duplicate / mis-ordered entries;
• a completed chunked message was removed from chunkedMessagesMap but its uuid stayed at the head of
  pendingChunkedMessageUuidQueue. removeExpireIncompleteChunkedMessages uses peek() and returned on that null
  ctx, so every incomplete message queued behind such a ghost head was never expired (buffers leaked, chunks
  never acked);
• the forward-gap discard path (a chunk that skips ids) removed the ctx from the map without decrementing the
  count or removing the queue entry.

Modifications

• Add a dedicated chunkedMessageLock. processMessageChunk, removeOldestPendingChunkedMessage and
  removeExpireIncompleteChunkedMessages run under it, and messageReceived holds it across the last-chunk
  assembly + finalize so the assemble→finalize window is closed. pendingChunkedMessageCount is mutated only
  under the lock. Non-chunked messages never touch it.
• No application code runs under the lock. doAcknowledge is not purely asynchronous: the persistent
  acknowledgments grouping tracker invokes the consumer's acknowledgment interceptors inline. Every path that decides
  an ack while holding the lock (expiry, oldest-eviction, corrupted-ctx replacement, duplicated chunks, expired
  unexpected chunks) now appends the id to a deferredAcks list, and the three entry points that take the lock issue
  those acks only after leaving the synchronized block (acknowledgeChunkedMessageIds). Decompression of a completed
  message happens outside the lock too: the last chunk is assembled and the ctx removed + recycled under the lock,
  the still-compressed buffer and chunk ids are captured into locals, and the codec runs after the lock is released.
• Final-chunk decompression failures dispose of everything. The detached assembled buffer is released in a
  finally, and uncompressPayloadIfNeeded now treats an unchecked codec exception (ZLib throws
  IllegalArgumentException when the decoded size differs from the advertised one) the same way as an
  IOException: log, discardCorruptedMessage, return null. The null branch then individually acks the earlier
  chunks' message ids, which the expiry sweep can no longer reach since the ctx is gone. This also changes the
  non-chunked path, where such an exception previously escaped messageReceived on the IO thread instead of
  discarding the message with DecompressionError.
• Keep count and queue in sync with the map on every removal: the redelivered-first-chunk replacement decrements the
  count and re-enqueues the uuid once at the tail; the final-chunk completion and the forward-gap discard remove the
  uuid from the queue (and the latter decrements the count); and removeExpireIncompleteChunkedMessages drains ghost
  heads instead of stopping at them, mirroring removeOldestPendingChunkedMessage.

No public API, wire protocol, schema, config defaults or metrics are changed.

Verifying this change

• Make sure that the change passes the CI checks.

This change added tests to ConsumerImplTest and can be verified as follows:

• testChunkedMessageCountRaceBetweenReceiveAndExpiry — drives processMessageChunk on a "receiver" thread
  concurrently with the real removeExpireIncompleteChunkedMessages on an "expirer" thread for the "late chunk
  arrives for a uuid being expired" scenario. Fails before the fix (IllegalReferenceCountException / corrupted
  count).
• testDuplicateFirstChunkOvercountsPendingChunkedMessageCount — redelivered first chunk; count equals map size and
  the uuid appears once in the queue.
• testExpiryDrainsPastGhostQueueEntries — a completed uuid at the queue head must not stop an expired incomplete
  message behind it from being cleaned.
• testForwardGapDiscardKeepsCountAndQueueConsistent — the forward-gap discard keeps count and queue in sync.
• testFinalChunkThroughMessageReceivedDeliversAssembledMessage — through the real messageReceived entry point
  with CRC32C-framed ZLIB chunks: the message arrives with the original payload and a ChunkMessageIdImpl, the
  assembled buffer is released, chunk ids are registered with the unack tracker, bookkeeping is empty.
• testFinalChunkDecompressionFailureThroughMessageReceived (two variants: corrupt input → IOException;
  mismatched advertised size → IllegalArgumentException) — the detached buffer is released, chunk 0 is acked
  through the interceptor, chunk 1 is discarded on the connection, bookkeeping is empty, nothing escapes
  messageReceived. The size-mismatch variant fails before the fix with the escaping IllegalArgumentException.
• testChunkReceiveProceedsWhileExpiryAckInterceptorIsBlocked — an interceptor parks onAcknowledge on a latch
  while expiry is acking; a chunk for another uuid must still be processed on a second thread. Fails before the fix
  with the receiver blocked on the lock.
• testReplacedCorruptedChunkedMessageAcksItsEarlierChunks and testEvictingOldestPendingChunkedMessageAcksItsChunks
  — the replacement and max-pending eviction paths ack through the deferred-ack path and keep count/queue/map
  consistent.

ConsumerImplTest (31) passes, the full pulsar-client unit suite (952 tests) passes with no regressions, and
./gradlew quickCheck is clean.

Does this pull request potentially affect one of the following parts:

• Dependencies (add or upgrade a dependency)
• The public API
• The schema
• The default values of configurations
• The threading model
• The binary protocol
• The REST endpoints
• The admin CLI options
• The metrics
• Anything that affects deployment

Documentation

• doc-required
• doc-not-needed
  (internal bug fix; no user-facing behavior or config change)
• doc
• doc-complete

[congbobo184]

LGTM

[lhotari]

Review performed with AI assistance (Claude Code / Claude Fable 5 combined with a Codex gpt-5.6-sol review pass); findings below were verified against the code before posting.

Overall the change looks sound and worthwhile. The receive/assembly path (Netty IO thread) and the incomplete-chunk expiry path (internalPinnedExecutor) genuinely race on ChunkedMessageCtx, its buffer and pendingChunkedMessageCount, and the new chunkedMessageLock serializes them correctly. The lock scope is well chosen — decompression, newMessage and callback dispatch stay outside it, non-chunked messages never touch it — and I verified there are no lock-ordering cycles (nothing called under the lock — doAcknowledge, increaseAvailablePermits, trackMessage — re-enters chunk code) and that buffer refcounts stay balanced (uncompressPayloadIfNeeded does not consume its input). The PR also fixes several real pre-existing bugs beyond the headline race: the assemble→finalize NPE window, the duplicate-first-chunk overcount, the forward-gap discard desync, and — most impactful — the ghost-head bug where a completed uuid at the queue head permanently stalled expiry. All four new tests pass locally (:pulsar-client-original:test, race test ~4s).

One regression should be addressed before merge:

1. Decompression failure now permanently drops the earlier chunks' message IDs (final-chunk path in messageReceived)

The ctx is removed from chunkedMessagesMap and recycled inside the lock before decompression. If uncompressPayloadIfNeeded then fails, the code returns without acking or tracking the captured chunkedMessageIds; discardCorruptedMessage inside it only handles the final chunk's ID. Since the ctx is gone, the expiry sweep can never reach those IDs either — the first n−1 chunks stay unacked until the consumer reconnects (permanent ack hole / stuck backlog on that subscription). Previously the ctx stayed in the map and expiry eventually acked all recorded chunk IDs (before tripping the latent double-release this PR fixes — broken differently, but disposal did happen). Suggested fix: in the uncompressedPayload == null branch, individually doAcknowledge each non-null captured ID, mirroring the corrupted-chunk handling in doProcessMessageChunk / removeChunkMessage(..., autoAck=true) semantics.

Test / hygiene items:

2. Reflection into private state in the new tests. Project convention is no reflection into private state — use @VisibleForTesting package-private accessors instead. ConsumerImplTest is already in org.apache.pulsar.client.impl, so making processMessageChunk, pendingChunkedMessageCount, pendingChunkedMessageUuidQueue and expireChunkMessageTaskScheduled package-private removes all setAccessible/FieldUtils use. Some reflection is unnecessary even now: expireTimeOfIncompleteChunkedMessageMillis is protected and can be assigned directly, as the test already does with chunkedMessagesMap.

3. Stale javadoc on testDuplicateFirstChunkOvercountsPendingChunkedMessageCount: it still says "Currently FAILS deterministically … Enable once the duplicate-first-chunk path decrements the counter", which is pre-fix wording — the fix is included in this PR and the test passes. It also hardcodes source line references ("ConsumerImpl.java:1628-1634") that will rot. Please rewrite it to describe the guarded invariant.

4. The race test has no timeout. testChunkedMessageCountRaceBetweenReceiveAndExpiry uses unbounded receiver.join()/expirer.join() and no @Test(timeOut = …). The scenario it guards is exactly the kind whose future regression could be a deadlock between the two paths — which would hang the suite instead of failing. Please add a generous timeOut.

5. The PR description understates the change (in a good way): the ghost-head expiry-stall fix in doRemoveExpireIncompleteChunkedMessages (previously a completed uuid at the queue head blocked all expiry behind it indefinitely) and the forward-gap discard count/queue sync fix aren't mentioned, and "Verifying this change" lists 2 tests while 4 were added. Please update Motivation/Modifications so the full behavior change is captured for reviewers and release notes.

6. The PR title is truncated — it literally ends with a Unicode ellipsis ("…to fix use-after-…", GitHub's auto-fill from a long commit subject). Please complete it, e.g. [fix][client] Serialize chunked-message bookkeeping to fix use-after-free and count drift.

[SongOf]

@lhotari
Fixed 1–4; Edited 5–6.

---

Thanks for the thorough review — fixed 1–4:

1. Fixed: the ack loop in messageReceived now skips the last entry (chunkedMessageIds[0..length-2]), since it's already acked by discardCorruptedMessage inside uncompressPayloadIfNeeded's failure path. The earlier n−1 chunks
   are now acked individually instead of being dropped.
2. Fixed: pendingChunkedMessageCount, pendingChunkedMessageUuidQueue, expireChunkMessageTaskScheduled, and processMessageChunk are now package-private + @VisibleForTesting; all setAccessible/FieldUtils reflection removed from
   the test file.
3. Fixed: rewrote the javadoc on testDuplicateFirstChunkOvercountsPendingChunkedMessageCount to describe the guarded invariant instead of the pre-fix failure state; dropped the hardcoded line references.
4. Fixed: added @test(timeOut = 60000) to testChunkedMessageCountRaceBetweenReceiveAndExpiry.

[lhotari]

Thanks for tightening the chunk bookkeeping. The map/count/queue transitions and normal final-chunk ownership transfer look consistent. Two remaining failure paths need attention: an unchecked decompression exception can leave the detached assembled buffer unreleased, and expiry can invoke application acknowledgment interceptors while holding the new lock, blocking the receive path. One added test also leaves its final partial-message buffer allocated.

Please add coverage through messageReceived for final-chunk success and decompression failures, including an unchecked codec exception. The new race test exercises only non-final chunks, so it does not cover the changed ownership-transfer path.

[lhotari]

[BUG] Release the detached assembled buffer when decoding throws unchecked

The context has already been removed and recycled here, but compressedAssembledPayload.release() runs only if decoding returns normally. ConsumerImpl.java:2199-2224 catches only IOException; ZLib can throw IllegalArgumentException when the decoded size differs from the advertised size:

pulsar/pulsar-common/src/main/java/org/apache/pulsar/common/compression/CompressionCodecZLib.java

Lines 119 to 128 in 81f8ff7

	try {
	resultLength = inflater.inflate(uncompressed.array(), uncompressed.arrayOffset(), uncompressedLength);
	} catch (DataFormatException e) {
	throw new IOException(e);
	}
	checkArgument(resultLength == uncompressedLength);
	uncompressed.writerIndex(uncompressedLength);
	return uncompressed;

That exception skips both the release and the null-result cleanup. Previously the context remained available to expiry; now neither the assembled buffer nor its chunk IDs remain registered there. Please release the detached buffer in finally and handle this failure so the chunk IDs are disposed of consistently. Add a regression test through the actual final-chunk messageReceived path with a mismatched decoded size; the new non-final-chunk tests do not exercise this transfer.

[SongOf]

Fixed in 3a30fe4. The release is now in a finally, and uncompressPayloadIfNeeded treats an unchecked codec exception the same way as an IOException: log, discardCorruptedMessage, return null. The existing null branch then acks the earlier chunk ids, so disposal is the same for both failure kinds. This also changes the non-chunked path, where such an exception previously escaped messageReceived on the IO thread instead of discarding the message with DecompressionError; the PR description now calls that out.

Added testFinalChunkDecompressionFailureThroughMessageReceived, driven through messageReceived with real CRC32C-framed ZLIB chunks, in two variants: corrupt input (IOException) and a mismatched advertised size (IllegalArgumentException). Both assert that the detached buffer is released, chunk 0 is acked through the interceptor, chunk 1 is discarded on the connection, and map/queue/count are empty. The size-mismatch variant fails without the fix with the escaping IllegalArgumentException. Added testFinalChunkThroughMessageReceivedDeliversAssembledMessage for the success path of the same ownership transfer.

Separately, CompressionCodecZLib.decode leaks the uncompressed buffer it allocated when that size check throws. That is a pre-existing codec issue independent of this PR; I can open a follow-up for it.

[lhotari]

[BUG] Run acknowledgment callbacks after releasing the chunk-state lock

This lock is held through removeChunkMessage → doAcknowledge, which is not entirely asynchronous: the persistent acknowledgment tracker invokes the consumer callback inline:

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

Lines 200 to 209 in 81f8ff7

	private CompletableFuture<Void> addIndividualAcknowledgment(
	MessageIdAdv msgId,
	@Nullable MessageIdAdv batchMessageId,
	Map<String, Long> properties,
	boolean groupedByCaller) {
	if (batchMessageId != null) {
	consumer.onAcknowledge(batchMessageId, null);
	} else {
	consumer.onAcknowledge(msgId, null);
	}

ConsumerBase.java:953-956 forwards that to application interceptors. If an expiry-thread interceptor waits for a producer send receipt on the same IO event loop, a concurrently arriving chunk can block that loop acquiring this monitor. The receipt then cannot be processed until the interceptor returns, causing a timeout or a deadlock when that wait has no timeout. Even a slow callback now stalls chunk reception.

Please capture the acknowledgment work under the lock and invoke it after releasing the monitor, including the replacement/discard paths. A latch-controlled interceptor test should verify that a concurrent receive can progress while that callback is paused.

[SongOf]

Fixed in 3a30fe4. Every path that decides an ack under chunkedMessageLock (expiry, oldest-eviction, the corrupted-first-chunk replacement, duplicated chunks, and the expired unexpected-chunk case) now appends the id to a deferredAcks list, and the three entry points that take the lock (processMessageChunk, removeExpireIncompleteChunkedMessages, messageReceived) call doAcknowledge for the list only after leaving the synchronized block. trackMessage stays under the lock since it does not reach application code.

Added testChunkReceiveProceedsWhileExpiryAckInterceptorIsBlocked: a persistent-topic consumer with an interceptor that parks onAcknowledge on a latch. The expiry thread is held inside the interceptor while a second thread delivers a chunk for another uuid, which must complete within 5 seconds. Without the change the receiver blocks on the monitor and the test fails. testReplacedCorruptedChunkedMessageAcksItsEarlierChunks and testEvictingOldestPendingChunkedMessageAcksItsChunks cover the replacement and eviction paths through the same deferred-ack mechanism.

[lhotari]

[QUALITY] Release the partial-message buffer left by the duplicate-first-chunk test

Both sends are chunk 0 of a two-chunk message, and expiry is disabled. Replacing the first context releases its buffer, but the second context deliberately remains in chunkedMessagesMap with a live chunkedMsgBuffer. ConsumerImplTest.java:110-119 only shuts down executors, so the surviving buffer is never released. Please drain/release the remaining partial context in a finally block so cleanup also runs when an assertion fails.

[SongOf]

Fixed in 3a30fe4. The test body is wrapped in try/finally and the surviving partial context is released by a releasePendingChunkedMessages() helper, which the new tests that intentionally leave a partial message use as well.