[fix][broker] Shared subscription stalls forever on a chunked message

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SharedConsumerAssignor.java

@@ -123,16 +123,44 @@ private Consumer getConsumer(final int numConsumers) {
         return null;
     }
 
+    /**
+     * Remove every uuid mapping that points to the given consumer. Must be called when a consumer is removed from the
+     * dispatcher: a mapping to a disconnected consumer has no permits left, so the remaining chunks of that uuid stay
+     * unassignable and stall the subscription. Also bounds the map when the last chunk is never published.
+     */
+    public void removeConsumer(final Consumer consumer) {
+        uuidToConsumer.values().removeIf(cachedConsumer -> cachedConsumer == consumer);
+    }
+
+    /** Drop all uuid mappings, e.g. after all consumers have been removed from the dispatcher. */
+    public void clear() {
+        uuidToConsumer.clear();
+    }
+
     private Consumer getConsumerForUuid(final MessageMetadata metadata, final Consumer defaultConsumer) {
         final String uuid = metadata.getUuid();
         Consumer consumer = uuidToConsumer.get(uuid);
         if (consumer == null) {
-            if (metadata.getChunkId() != 0) {
-                // Not the first chunk, skip it
-                return null;
-            }
             consumer = defaultConsumer;
-            uuidToConsumer.put(uuid, consumer);
+            if (metadata.getChunkId() == 0) {
+                uuidToConsumer.put(uuid, consumer);
+            } else {
+                // Orphan chunk: only chunk 0 creates the mapping, and the cursor always re-reads chunk 0 before
+                // this one while it is still unacked, so a missing mapping means chunk 0 is acknowledged and gone.
+                // Never replay it: the dispatcher drains the replay queue before every normal read, so an entry that
+                // can never be assigned stalls the whole subscription. Dispatch it and let the client discard it.
+                if (subscription != null) {
+                    log.warn()
+                            .attr("topic", subscription.getTopic().getName())
+                            .attr("subscription", subscription.getName())
+                            .attr("uuid", uuid)
+                            .attr("chunkId", metadata.getChunkId())
+                            .attr("numChunks", metadata.getNumChunksFromMsg())
+                            .attr("consumer", defaultConsumer)
+                            .log("Dispatching orphan chunk whose first chunk is gone. The client should discard and"
+                                    + " acknowledge it");
+                }
+            }
         }
         final int permits = consumerToPermits.computeIfAbsent(consumer, Consumer::getAvailablePermits);
         if (permits <= 0) {

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java

@@ -241,6 +241,8 @@ protected boolean isConsumersExceededOnSubscription() {
     public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceException {
         // decrement unack-message count for removed consumer
         addUnAckedMessages(-consumer.getUnackedMessages());
+        // Drop this consumer's chunk uuid mappings, otherwise its pending chunks stay unassignable forever.
+        assignor.removeConsumer(consumer);
         if (consumerSet.removeAll(consumer) == 1) {
             consumerList.remove(consumer);
             log.info()
@@ -291,6 +293,7 @@ protected synchronized void clearComponentsAfterRemovedAllConsumers() {
 
         redeliveryMessages.clear();
         redeliveryTracker.clear();
+        assignor.clear();
         if (closeFuture != null) {
             log.info("All consumers removed. Subscription is disconnected");
             closeFuture.complete(null);

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java

@@ -231,6 +231,8 @@ protected boolean isConsumersExceededOnSubscription() {
     public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceException {
         // decrement unack-message count for removed consumer
         addUnAckedMessages(-consumer.getUnackedMessages());
+        // Drop this consumer's chunk uuid mappings, otherwise its pending chunks stay unassignable forever.
+        assignor.removeConsumer(consumer);
         if (consumerSet.removeAll(consumer) == 1) {
             consumerList.remove(consumer);
             log.info()
@@ -271,6 +273,7 @@ private synchronized void clearComponentsAfterRemovedAllConsumers() {
 
         redeliveryMessages.clear();
         redeliveryTracker.clear();
+        assignor.clear();
         if (closeFuture != null) {
             log.info("All consumers removed. Subscription is disconnected");
             closeFuture.complete(null);

pulsar-broker/src/test/java/org/apache/pulsar/broker/service/SharedConsumerAssignorTest.java

@@ -153,6 +153,88 @@ public void testMultiConsumerWithSmallPermits() {
         assertTrue(assignor.getUuidToConsumer().isEmpty());
     }
 
+    /**
+     * An orphan chunk is a chunk whose uuid is not tracked and whose chunk id is not 0, e.g. entry 5 (A-1-2-3) when
+     * the earlier chunks of "A-1" were acknowledged before the mapping existed. That happens after a topic unload, a
+     * broker restart, or when the consumer that owned the uuid disconnected.
+     * <p>
+     * Such an entry can never become assignable, so handing it to the unassigned message processor (the replay queue)
+     * stalls the subscription forever: {@code PersistentDispatcherMultipleConsumers#readMoreEntries} drains the replay
+     * queue before every normal read, so the dispatcher keeps replaying the same entry, never reads anything new and
+     * the mark-delete position never moves. It must be dispatched instead, so that the client can discard and
+     * acknowledge the incomplete chunked message.
+     */
+    @Test
+    public void testOrphanChunkIsDispatchedInsteadOfReplayedForever() {
+        final Consumer consumer = new Consumer("A", 100);
+        roundRobinConsumerSelector.addConsumers(consumer);
+        assertTrue(assignor.getUuidToConsumer().isEmpty());
+
+        // The last chunk of "A-1" is read while the uuid was never cached
+        Map<Consumer, List<EntryAndMetadata>> result = assignor.assign(entryAndMetadataList.subList(5, 6), 1);
+        assertEquals(toString(result.getOrDefault(consumer, Collections.emptyList())),
+                Collections.singletonList("0:5@A-1-2-3"));
+        // The entry must not go back to the replay queue, otherwise the dispatcher can never make progress
+        assertTrue(replayQueue.isEmpty());
+        // An orphan chunk must not create a mapping either, there is no first chunk to pin to a consumer
+        assertTrue(assignor.getUuidToConsumer().isEmpty());
+
+        // The same holds for a middle chunk, not only for the last one
+        result = assignor.assign(entryAndMetadataList.subList(2, 3), 1);
+        assertEquals(toString(result.getOrDefault(consumer, Collections.emptyList())),
+                Collections.singletonList("0:2@A-1-1-3"));
+        assertTrue(replayQueue.isEmpty());
+        assertTrue(assignor.getUuidToConsumer().isEmpty());
+    }
+
+    /**
+     * When the consumer that owns a uuid is removed from the dispatcher, the mapping must be dropped. Otherwise the
+     * remaining chunks stay pinned to a disconnected consumer that has no available permits, so they are replayed
+     * forever and the subscription stalls.
+     */
+    @Test
+    public void testUuidMappingIsRemovedWhenConsumerIsRemoved() {
+        final Consumer consumerA = new Consumer("A", 3);
+        final Consumer consumerB = new Consumer("B", 100);
+        roundRobinConsumerSelector.addConsumers(consumerA);
+
+        // consumerA receives A-0, A-1-0-3 and A-1-1-3, so the uuid "A-1" gets pinned to consumerA
+        assignor.assign(entryAndMetadataList.subList(0, 3), 1);
+        assertTrue(replayQueue.isEmpty());
+        assertEquals(assignor.getUuidToConsumer().keySet(), Sets.newHashSet("A-1"));
+        assertSame(assignor.getUuidToConsumer().get("A-1"), consumerA);
+
+        // consumerA disconnects while the chunked message "A-1" is still incomplete
+        assignor.removeConsumer(consumerA);
+        assertTrue(assignor.getUuidToConsumer().isEmpty());
+
+        roundRobinConsumerSelector.clear();
+        roundRobinConsumerSelector.addConsumers(consumerB);
+
+        // The last chunk of "A-1" must be dispatched to the remaining consumer instead of getting stuck
+        Map<Consumer, List<EntryAndMetadata>> result = assignor.assign(entryAndMetadataList.subList(5, 6), 1);
+        assertNull(result.get(consumerA));
+        assertEquals(toString(result.getOrDefault(consumerB, Collections.emptyList())),
+                Collections.singletonList("0:5@A-1-2-3"));
+        assertTrue(replayQueue.isEmpty());
+    }
+
+    /**
+     * The uuid mappings must not outlive the consumers of the dispatcher, otherwise they leak for every chunked
+     * message whose last chunk is never published.
+     */
+    @Test
+    public void testClearRemovesAllUuidMappings() {
+        final Consumer consumer = new Consumer("A", 100);
+        roundRobinConsumerSelector.addConsumers(consumer);
+
+        assignor.assign(entryAndMetadataList.subList(0, 5), 1);
+        assertEquals(assignor.getUuidToConsumer().keySet(), Sets.newHashSet("A-1", "B-1"));
+
+        assignor.clear();
+        assertTrue(assignor.getUuidToConsumer().isEmpty());
+    }
+
     @RequiredArgsConstructor
     static class ConsumerSelector implements Supplier<Consumer> {