From fdc69f6e456753d9df5d2c2e398b21e4fa9e9f60 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sat, 8 Aug 2026 08:40:16 +0800 Subject: [PATCH 1/7] Fix Shared dispatcher permit accounting during consumer churn --- .../pulsar/broker/service/Consumer.java | 30 +++- ...PersistentDispatcherMultipleConsumers.java | 6 +- ...entDispatcherMultipleConsumersClassic.java | 6 +- .../SharedDispatcherPermitAccountingTest.java | 145 ++++++++++++++++++ 4 files changed, 181 insertions(+), 6 deletions(-) create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java index a7f296536881e..7952dafb39f89 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java @@ -117,6 +117,10 @@ public class Consumer { private static final AtomicIntegerFieldUpdater MESSAGE_PERMITS_UPDATER = AtomicIntegerFieldUpdater.newUpdater(Consumer.class, "messagePermits"); private volatile int messagePermits = 0; + // Flow commands update the consumer before the dispatcher processes the update asynchronously. Track permits in + // between those steps so consumer removal only subtracts permits already included in the dispatcher total. + private final Object flowPermitAccountingLock = new Object(); + private int pendingDispatcherFlowPermits = 0; // It starts keep tracking of messagePermits once consumer gets blocked, as consumer needs two separate counts: // messagePermits (1) before and (2) after being blocked: to dispatch only blockedPermit number of messages at the // time of redelivery @@ -916,7 +920,7 @@ public void flowPermits(int additionalNumberOfMessages) { } int oldPermits; if (!blockedConsumerOnUnackedMsgs) { - oldPermits = MESSAGE_PERMITS_UPDATER.getAndAdd(this, additionalNumberOfMessages); + oldPermits = addPermitsPendingDispatcherUpdate(additionalNumberOfMessages); log.debug() .attr("additionalNumberOfMessages", additionalNumberOfMessages) .log("Added message permits before updating dispatcher"); @@ -943,7 +947,7 @@ public void flowPermits(int additionalNumberOfMessages) { void flowConsumerBlockedPermits(Consumer consumer) { int additionalNumberOfPermits = PERMITS_RECEIVED_WHILE_CONSUMER_BLOCKED_UPDATER.getAndSet(consumer, 0); // add newly flow permits to actual consumer.messagePermits - MESSAGE_PERMITS_UPDATER.getAndAdd(consumer, additionalNumberOfPermits); + consumer.addPermitsPendingDispatcherUpdate(additionalNumberOfPermits); log.debug() .attr("additionalNumberOfPermits", additionalNumberOfPermits) .log("Added blocked permits"); @@ -955,6 +959,28 @@ public int getAvailablePermits() { return MESSAGE_PERMITS_UPDATER.get(this); } + private int addPermitsPendingDispatcherUpdate(int additionalNumberOfPermits) { + synchronized (flowPermitAccountingLock) { + pendingDispatcherFlowPermits += additionalNumberOfPermits; + return MESSAGE_PERMITS_UPDATER.getAndAdd(this, additionalNumberOfPermits); + } + } + + /** Marks an asynchronous Flow update as processed by the dispatcher. */ + public void completePendingDispatcherFlow(int additionalNumberOfPermits) { + synchronized (flowPermitAccountingLock) { + pendingDispatcherFlowPermits = Math.max(0, + pendingDispatcherFlowPermits - additionalNumberOfPermits); + } + } + + /** Returns the permits already included in the dispatcher total when this consumer is removed. */ + public int getAvailablePermitsForDispatcherRemoval() { + synchronized (flowPermitAccountingLock) { + return MESSAGE_PERMITS_UPDATER.get(this) - pendingDispatcherFlowPermits; + } + } + /** * return 0 if there is no entry dispatched yet. */ diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java index b659f6e2200d8..937d22a1f6d2b 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java @@ -258,9 +258,10 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE notifyAddedToReplay.setTrue(); } }); - totalAvailablePermits -= consumer.getAvailablePermits(); + int availablePermits = consumer.getAvailablePermitsForDispatcherRemoval(); + totalAvailablePermits -= availablePermits; log.debug() - .attr("diffAvailablePermits", consumer.getAvailablePermits()) + .attr("availablePermits", availablePermits) .attr("totalAvailablePermits", totalAvailablePermits) .log("Decreased totalAvailablePermits"); if (notifyAddedToReplay.booleanValue()) { @@ -306,6 +307,7 @@ public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { } private synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { + consumer.completePendingDispatcherFlow(additionalNumberOfMessages); if (!consumerSet.contains(consumer)) { log.debug() .attr("consumer", consumer) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java index 3de50042b592d..c1b73076aefab 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java @@ -244,9 +244,10 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE consumer.getPendingAcks().forEach((ledgerId, entryId, batchSize, stickyKeyHash) -> { addMessageToReplay(ledgerId, entryId, stickyKeyHash); }); - totalAvailablePermits -= consumer.getAvailablePermits(); + int availablePermits = consumer.getAvailablePermitsForDispatcherRemoval(); + totalAvailablePermits -= availablePermits; log.debug() - .attr("availablePermits", consumer.getAvailablePermits()) + .attr("availablePermits", availablePermits) .attr("totalAvailablePermits", totalAvailablePermits) .log("Decreased totalAvailablePermits by in PersistentDispatcherMultipleConsumers. " + "New dispatcher permit count is"); @@ -286,6 +287,7 @@ public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { } private synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { + consumer.completePendingDispatcherFlow(additionalNumberOfMessages); if (!consumerSet.contains(consumer)) { log.debug() .attr("consumer", consumer) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java new file mode 100644 index 0000000000000..e19a3f4d55399 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java @@ -0,0 +1,145 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pulsar.broker.service.persistent; + +import static java.util.Collections.emptyMap; +import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Shared; +import static org.apache.pulsar.common.api.proto.KeySharedMode.AUTO_SPLIT; +import static org.apache.pulsar.common.protocol.Commands.DEFAULT_CONSUMER_EPOCH; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import io.netty.util.concurrent.EventExecutor; +import org.apache.bookkeeper.mledger.ManagedCursor; +import org.apache.bookkeeper.mledger.impl.ManagedCursorImpl; +import org.apache.pulsar.broker.service.Consumer; +import org.apache.pulsar.broker.service.Dispatcher; +import org.apache.pulsar.broker.service.SharedPulsarBaseTest; +import org.apache.pulsar.broker.service.Subscription; +import org.apache.pulsar.broker.service.TransportCnx; +import org.apache.pulsar.client.api.MessageId; +import org.apache.pulsar.common.api.proto.KeySharedMeta; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +@Test(groups = "broker-api") +public class SharedDispatcherPermitAccountingTest extends SharedPulsarBaseTest { + + @DataProvider(name = "dispatcherImplementations") + public Object[][] dispatcherImplementations() { + return new Object[][] {{false}, {true}}; + } + + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) + public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean classic) throws Exception { + String topicName = newTopicName(); + String subscriptionName = "shared-sub"; + admin.topics().createNonPartitionedTopic(topicName); + + PersistentTopic topic = (PersistentTopic) getTopic(topicName, false).join().get(); + ManagedCursor cursor = mock(ManagedCursorImpl.class); + when(cursor.getName()).thenReturn(subscriptionName); + Subscription subscription = mock(PersistentSubscription.class); + when(subscription.getName()).thenReturn(subscriptionName); + when(subscription.getTopic()).thenReturn(topic); + + Dispatcher dispatcher = classic + ? new NoopReadClassicDispatcher(topic, cursor, subscription) + : new NoopReadDispatcher(topic, cursor, subscription); + doAnswer(invocation -> { + dispatcher.consumerFlow(invocation.getArgument(0), invocation.getArgument(1)); + return null; + }).when(subscription).consumerFlow(any(), anyInt()); + + Consumer remainingConsumer = createConsumer(subscription, topicName, 1); + Consumer removedConsumer = createConsumer(subscription, topicName, 2); + dispatcher.addConsumer(remainingConsumer).join(); + dispatcher.addConsumer(removedConsumer).join(); + + remainingConsumer.flowPermits(10); + drainBrokerWorkerGroup(topic); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(10); + + synchronized (dispatcher) { + removedConsumer.flowPermits(1_000); + dispatcher.removeConsumer(removedConsumer); + } + drainBrokerWorkerGroup(topic); + + assertThat(totalAvailablePermits(dispatcher)) + .isEqualTo(remainingConsumer.getAvailablePermits()); + } + + private Consumer createConsumer(Subscription subscription, String topicName, long consumerId) { + TransportCnx cnx = mock(TransportCnx.class); + when(cnx.isActive()).thenReturn(true); + when(cnx.isWritable()).thenReturn(true); + return new Consumer(subscription, Shared, topicName, consumerId, 0, "consumer-" + consumerId, + false, cnx, "role", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), + MessageId.latest, DEFAULT_CONSUMER_EPOCH); + } + + private static void drainBrokerWorkerGroup(PersistentTopic topic) throws Exception { + for (EventExecutor eventExecutor : topic.getBrokerService().executor()) { + eventExecutor.submit(() -> { }).sync(); + } + } + + private static int totalAvailablePermits(Dispatcher dispatcher) { + if (dispatcher instanceof PersistentDispatcherMultipleConsumers pip379Dispatcher) { + return pip379Dispatcher.totalAvailablePermits; + } + return ((PersistentDispatcherMultipleConsumersClassic) dispatcher).totalAvailablePermits; + } + + private static class NoopReadDispatcher extends PersistentDispatcherMultipleConsumers { + NoopReadDispatcher(PersistentTopic topic, ManagedCursor cursor, Subscription subscription) { + super(topic, cursor, subscription); + } + + @Override + public void readMoreEntriesAsync() { + // No-op for permit accounting test. + } + + @Override + public synchronized void readMoreEntries() { + // No-op for permit accounting test. + } + } + + private static class NoopReadClassicDispatcher extends PersistentDispatcherMultipleConsumersClassic { + NoopReadClassicDispatcher(PersistentTopic topic, ManagedCursor cursor, Subscription subscription) { + super(topic, cursor, subscription); + } + + @Override + public void readMoreEntriesAsync() { + // No-op for permit accounting test. + } + + @Override + public synchronized void readMoreEntries() { + // No-op for permit accounting test. + } + } +} From 045ab145c8cac32e3d4f876b9b29aa8724c3748a Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sat, 8 Aug 2026 08:58:33 +0800 Subject: [PATCH 2/7] Add permit accounting coverage and documentation --- .../pulsar/broker/service/Consumer.java | 28 ++++++-- ...PersistentDispatcherMultipleConsumers.java | 2 + ...entDispatcherMultipleConsumersClassic.java | 2 + .../SharedDispatcherPermitAccountingTest.java | 70 +++++++++++++++---- 4 files changed, 85 insertions(+), 17 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java index 7952dafb39f89..486049bda60e5 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java @@ -117,8 +117,16 @@ public class Consumer { private static final AtomicIntegerFieldUpdater MESSAGE_PERMITS_UPDATER = AtomicIntegerFieldUpdater.newUpdater(Consumer.class, "messagePermits"); private volatile int messagePermits = 0; - // Flow commands update the consumer before the dispatcher processes the update asynchronously. Track permits in - // between those steps so consumer removal only subtracts permits already included in the dispatcher total. + /** + * Guards the Flow-side compound update of {@link #messagePermits} and + * {@link #pendingDispatcherFlowPermits}. A Flow command increases the consumer permits before the dispatcher + * processes the corresponding update asynchronously. Consumer removal can happen between those two operations, + * so both values must be observed consistently when calculating how many permits are already included in the + * dispatcher total. + * + *

The dispatcher callback is invoked only after this lock is released. This avoids holding the lock while + * calling into the subscription and preserves the lock order used by dispatcher flow processing and removal. + */ private final Object flowPermitAccountingLock = new Object(); private int pendingDispatcherFlowPermits = 0; // It starts keep tracking of messagePermits once consumer gets blocked, as consumer needs two separate counts: @@ -959,6 +967,10 @@ public int getAvailablePermits() { return MESSAGE_PERMITS_UPDATER.get(this); } + /** + * Adds permits after a Flow command is accepted and immediately before notifying the dispatcher. The pending + * count covers the interval until the dispatcher's asynchronous Flow task starts processing the same permits. + */ private int addPermitsPendingDispatcherUpdate(int additionalNumberOfPermits) { synchronized (flowPermitAccountingLock) { pendingDispatcherFlowPermits += additionalNumberOfPermits; @@ -966,7 +978,11 @@ private int addPermitsPendingDispatcherUpdate(int additionalNumberOfPermits) { } } - /** Marks an asynchronous Flow update as processed by the dispatcher. */ + /** + * Called at the start of the dispatcher's asynchronous Flow task, before checking whether this consumer is still + * connected. At this point the Flow update is no longer pending: the dispatcher will either add the permits to + * its total or ignore them because the consumer has already been removed. + */ public void completePendingDispatcherFlow(int additionalNumberOfPermits) { synchronized (flowPermitAccountingLock) { pendingDispatcherFlowPermits = Math.max(0, @@ -974,7 +990,11 @@ public void completePendingDispatcherFlow(int additionalNumberOfPermits) { } } - /** Returns the permits already included in the dispatcher total when this consumer is removed. */ + /** + * Called while the dispatcher removes this consumer. Permits belonging to Flow tasks that have not started yet + * are excluded because those permits have not been added to the dispatcher total and must not be subtracted from + * it during removal. + */ public int getAvailablePermitsForDispatcherRemoval() { synchronized (flowPermitAccountingLock) { return MESSAGE_PERMITS_UPDATER.get(this) - pendingDispatcherFlowPermits; diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java index 937d22a1f6d2b..df5805694ce26 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java @@ -258,6 +258,7 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE notifyAddedToReplay.setTrue(); } }); + // Exclude permits from Flow tasks that have not updated the dispatcher total yet. int availablePermits = consumer.getAvailablePermitsForDispatcherRemoval(); totalAvailablePermits -= availablePermits; log.debug() @@ -307,6 +308,7 @@ public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { } private synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { + // The queued Flow task is no longer pending, even if the consumer was removed while the task was waiting. consumer.completePendingDispatcherFlow(additionalNumberOfMessages); if (!consumerSet.contains(consumer)) { log.debug() diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java index c1b73076aefab..9227fa9545fe6 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java @@ -244,6 +244,7 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE consumer.getPendingAcks().forEach((ledgerId, entryId, batchSize, stickyKeyHash) -> { addMessageToReplay(ledgerId, entryId, stickyKeyHash); }); + // Exclude permits from Flow tasks that have not updated the dispatcher total yet. int availablePermits = consumer.getAvailablePermitsForDispatcherRemoval(); totalAvailablePermits -= availablePermits; log.debug() @@ -287,6 +288,7 @@ public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { } private synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { + // The queued Flow task is no longer pending, even if the consumer was removed while the task was waiting. consumer.completePendingDispatcherFlow(additionalNumberOfMessages); if (!consumerSet.contains(consumer)) { log.debug() diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java index e19a3f4d55399..ad5f9691dbd29 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java @@ -38,6 +38,7 @@ import org.apache.pulsar.broker.service.TransportCnx; import org.apache.pulsar.client.api.MessageId; import org.apache.pulsar.common.api.proto.KeySharedMeta; +import org.apache.pulsar.common.policies.data.stats.ConsumerStatsImpl; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -51,6 +52,56 @@ public Object[][] dispatcherImplementations() { @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean classic) throws Exception { + TestContext context = createTestContext(classic); + Consumer remainingConsumer = context.remainingConsumer(); + Consumer removedConsumer = context.removedConsumer(); + + remainingConsumer.flowPermits(10); + removedConsumer.flowPermits(20); + drainBrokerWorkerGroup(context.topic()); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); + + synchronized (context.dispatcher()) { + removedConsumer.flowPermits(400); + removedConsumer.flowPermits(600); + context.dispatcher().removeConsumer(removedConsumer); + } + drainBrokerWorkerGroup(context.topic()); + + assertThat(totalAvailablePermits(context.dispatcher())) + .isEqualTo(remainingConsumer.getAvailablePermits()); + } + + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) + public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean classic) throws Exception { + TestContext context = createTestContext(classic); + Consumer remainingConsumer = context.remainingConsumer(); + Consumer removedConsumer = context.removedConsumer(); + + remainingConsumer.flowPermits(10); + drainBrokerWorkerGroup(context.topic()); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(10); + + ConsumerStatsImpl blockedStats = new ConsumerStatsImpl(); + blockedStats.blockedConsumerOnUnackedMsgs = true; + removedConsumer.updateStats(blockedStats); + assertThat(removedConsumer.isBlocked()).isTrue(); + assertThat(removedConsumer.getMaxUnackedMessages()).isPositive(); + + removedConsumer.flowPermits(1_000); + assertThat(removedConsumer.getAvailablePermits()).isZero(); + + synchronized (context.dispatcher()) { + removedConsumer.updateBlockedConsumerOnUnackedMsgs(removedConsumer); + context.dispatcher().removeConsumer(removedConsumer); + } + drainBrokerWorkerGroup(context.topic()); + + assertThat(totalAvailablePermits(context.dispatcher())) + .isEqualTo(remainingConsumer.getAvailablePermits()); + } + + private TestContext createTestContext(boolean classic) throws Exception { String topicName = newTopicName(); String subscriptionName = "shared-sub"; admin.topics().createNonPartitionedTopic(topicName); @@ -75,18 +126,7 @@ public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean cla dispatcher.addConsumer(remainingConsumer).join(); dispatcher.addConsumer(removedConsumer).join(); - remainingConsumer.flowPermits(10); - drainBrokerWorkerGroup(topic); - assertThat(totalAvailablePermits(dispatcher)).isEqualTo(10); - - synchronized (dispatcher) { - removedConsumer.flowPermits(1_000); - dispatcher.removeConsumer(removedConsumer); - } - drainBrokerWorkerGroup(topic); - - assertThat(totalAvailablePermits(dispatcher)) - .isEqualTo(remainingConsumer.getAvailablePermits()); + return new TestContext(topic, dispatcher, remainingConsumer, removedConsumer); } private Consumer createConsumer(Subscription subscription, String topicName, long consumerId) { @@ -94,7 +134,7 @@ private Consumer createConsumer(Subscription subscription, String topicName, lon when(cnx.isActive()).thenReturn(true); when(cnx.isWritable()).thenReturn(true); return new Consumer(subscription, Shared, topicName, consumerId, 0, "consumer-" + consumerId, - false, cnx, "role", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), + true, cnx, "role", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), MessageId.latest, DEFAULT_CONSUMER_EPOCH); } @@ -111,6 +151,10 @@ private static int totalAvailablePermits(Dispatcher dispatcher) { return ((PersistentDispatcherMultipleConsumersClassic) dispatcher).totalAvailablePermits; } + private record TestContext(PersistentTopic topic, Dispatcher dispatcher, + Consumer remainingConsumer, Consumer removedConsumer) { + } + private static class NoopReadDispatcher extends PersistentDispatcherMultipleConsumers { NoopReadDispatcher(PersistentTopic topic, ManagedCursor cursor, Subscription subscription) { super(topic, cursor, subscription); From f102ddfabb4f741eba60181d1b91f6ee91cbd5e4 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sat, 8 Aug 2026 09:23:07 +0800 Subject: [PATCH 3/7] Add Shared permit accounting behavior test --- .../SharedDispatcherPermitAccountingTest.java | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java index ad5f9691dbd29..b66d8293c1f5c 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java @@ -29,6 +29,8 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import io.netty.util.concurrent.EventExecutor; +import java.time.Duration; +import java.util.concurrent.TimeUnit; import org.apache.bookkeeper.mledger.ManagedCursor; import org.apache.bookkeeper.mledger.impl.ManagedCursorImpl; import org.apache.pulsar.broker.service.Consumer; @@ -36,9 +38,13 @@ import org.apache.pulsar.broker.service.SharedPulsarBaseTest; import org.apache.pulsar.broker.service.Subscription; import org.apache.pulsar.broker.service.TransportCnx; +import org.apache.pulsar.client.api.Message; import org.apache.pulsar.client.api.MessageId; +import org.apache.pulsar.client.api.Producer; +import org.apache.pulsar.client.api.SubscriptionType; import org.apache.pulsar.common.api.proto.KeySharedMeta; import org.apache.pulsar.common.policies.data.stats.ConsumerStatsImpl; +import org.awaitility.Awaitility; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -101,6 +107,71 @@ public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(bool .isEqualTo(remainingConsumer.getAvailablePermits()); } + @Test(timeOut = 30_000) + public void testRemainingConsumerCanContinueAfterFlowAndCloseRace() throws Exception { + int receiverQueueSize = 10; + int messagesToConsume = receiverQueueSize * 2; + int pendingFlowPermits = 1_000; + String topicName = newTopicName(); + String subscriptionName = "shared-sub"; + admin.topics().createNonPartitionedTopic(topicName); + + try (Producer producer = pulsarClient.newProducer().topic(topicName).create(); + org.apache.pulsar.client.api.Consumer remainingClient = pulsarClient.newConsumer() + .topic(topicName) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .consumerName("remaining-consumer") + .receiverQueueSize(receiverQueueSize) + .subscribe(); + org.apache.pulsar.client.api.Consumer removedClient = pulsarClient.newConsumer() + .topic(topicName) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .consumerName("removed-consumer") + .receiverQueueSize(receiverQueueSize) + .subscribe()) { + PersistentTopic topic = (PersistentTopic) getTopic(topicName, false).join().orElseThrow(); + PersistentSubscription subscription = + (PersistentSubscription) topic.getSubscription(subscriptionName); + PersistentDispatcherMultipleConsumers dispatcher = + (PersistentDispatcherMultipleConsumers) subscription.getDispatcher(); + Consumer remainingBrokerConsumer = findConsumer(dispatcher, "remaining-consumer"); + Consumer removedBrokerConsumer = findConsumer(dispatcher, "removed-consumer"); + + Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(remainingBrokerConsumer.getAvailablePermits()).isEqualTo(receiverQueueSize); + assertThat(removedBrokerConsumer.getAvailablePermits()).isEqualTo(receiverQueueSize); + }); + drainBrokerWorkerGroup(topic); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); + + // Hold the dispatcher monitor so the asynchronous Flow task cannot run before the production + // Consumer.close -> PersistentSubscription.removeConsumer -> dispatcher.removeConsumer lifecycle. + synchronized (dispatcher) { + removedBrokerConsumer.flowPermits(pendingFlowPermits); + assertThat(removedBrokerConsumer.getAvailablePermits()) + .isEqualTo(receiverQueueSize + pendingFlowPermits); + removedBrokerConsumer.close(); + } + removedClient.close(); + drainBrokerWorkerGroup(topic); + + assertThat(dispatcher.getConsumers()).containsExactly(remainingBrokerConsumer); + assertThat(totalAvailablePermits(dispatcher)) + .isEqualTo(remainingBrokerConsumer.getAvailablePermits()); + + for (int i = 0; i < messagesToConsume; i++) { + producer.send(new byte[] {(byte) i}); + } + for (int i = 0; i < messagesToConsume; i++) { + Message message = remainingClient.receive(1, TimeUnit.SECONDS); + assertThat(message).isNotNull(); + remainingClient.acknowledge(message); + } + } + } + private TestContext createTestContext(boolean classic) throws Exception { String topicName = newTopicName(); String subscriptionName = "shared-sub"; @@ -151,6 +222,13 @@ private static int totalAvailablePermits(Dispatcher dispatcher) { return ((PersistentDispatcherMultipleConsumersClassic) dispatcher).totalAvailablePermits; } + private static Consumer findConsumer(Dispatcher dispatcher, String consumerName) { + return dispatcher.getConsumers().stream() + .filter(consumer -> consumerName.equals(consumer.consumerName())) + .findFirst() + .orElseThrow(); + } + private record TestContext(PersistentTopic topic, Dispatcher dispatcher, Consumer remainingConsumer, Consumer removedConsumer) { } From e2467d03493d57f271db7cfb7deee7ecd9dc2597 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Tue, 18 Aug 2026 12:03:57 +0800 Subject: [PATCH 4/7] Address Shared dispatcher permit accounting review feedback --- .../pulsar/broker/service/Consumer.java | 21 +- ...PersistentDispatcherMultipleConsumers.java | 3 +- ...entDispatcherMultipleConsumersClassic.java | 3 +- .../SharedDispatcherPermitAccountingTest.java | 249 ++++++++++++++---- 4 files changed, 223 insertions(+), 53 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java index 486049bda60e5..b45947b76e833 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java @@ -972,12 +972,19 @@ public int getAvailablePermits() { * count covers the interval until the dispatcher's asynchronous Flow task starts processing the same permits. */ private int addPermitsPendingDispatcherUpdate(int additionalNumberOfPermits) { + if (!shouldTrackPendingDispatcherFlowPermits()) { + return MESSAGE_PERMITS_UPDATER.getAndAdd(this, additionalNumberOfPermits); + } synchronized (flowPermitAccountingLock) { pendingDispatcherFlowPermits += additionalNumberOfPermits; return MESSAGE_PERMITS_UPDATER.getAndAdd(this, additionalNumberOfPermits); } } + private boolean shouldTrackPendingDispatcherFlowPermits() { + return isPersistentTopic && Subscription.isIndividualAckMode(subType); + } + /** * Called at the start of the dispatcher's asynchronous Flow task, before checking whether this consumer is still * connected. At this point the Flow update is no longer pending: the dispatcher will either add the permits to @@ -985,8 +992,8 @@ private int addPermitsPendingDispatcherUpdate(int additionalNumberOfPermits) { */ public void completePendingDispatcherFlow(int additionalNumberOfPermits) { synchronized (flowPermitAccountingLock) { - pendingDispatcherFlowPermits = Math.max(0, - pendingDispatcherFlowPermits - additionalNumberOfPermits); + // Preserve the accounting delta across signed int wrap, matching the other permit counters. + pendingDispatcherFlowPermits -= additionalNumberOfPermits; } } @@ -994,6 +1001,16 @@ public void completePendingDispatcherFlow(int additionalNumberOfPermits) { * Called while the dispatcher removes this consumer. Permits belonging to Flow tasks that have not started yet * are excluded because those permits have not been added to the dispatcher total and must not be subtracted from * it during removal. + * + *

This accounting is enabled for persistent Shared and Key_Shared dispatchers. It relies on every dispatcher + * Flow task calling {@link #completePendingDispatcherFlow(int)} before applying or ignoring the update. For these + * dispatchers, when observed under the dispatcher monitor, the total available permits equal the sum of this + * balance over all connected consumers. + * + *

The returned balance can be negative. A pending Flow makes permits visible on the consumer before its + * asynchronous dispatcher update runs, so the dispatcher can consume those permits while they are still counted + * as pending. Subtracting the negative balance during removal is required to restore the dispatcher total; callers + * must not clamp it to zero. */ public int getAvailablePermitsForDispatcherRemoval() { synchronized (flowPermitAccountingLock) { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java index df5805694ce26..e7a6d586e0aee 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java @@ -258,7 +258,8 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE notifyAddedToReplay.setTrue(); } }); - // Exclude permits from Flow tasks that have not updated the dispatcher total yet. + // Restore the invariant that the dispatcher total equals the sum of the removal balances of the + // remaining consumers. Exclude Flow permits that have not updated the dispatcher total yet. int availablePermits = consumer.getAvailablePermitsForDispatcherRemoval(); totalAvailablePermits -= availablePermits; log.debug() diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java index 9227fa9545fe6..3016d0fa8dceb 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java @@ -244,7 +244,8 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE consumer.getPendingAcks().forEach((ledgerId, entryId, batchSize, stickyKeyHash) -> { addMessageToReplay(ledgerId, entryId, stickyKeyHash); }); - // Exclude permits from Flow tasks that have not updated the dispatcher total yet. + // Restore the invariant that the dispatcher total equals the sum of the removal balances of the + // remaining consumers. Exclude Flow permits that have not updated the dispatcher total yet. int availablePermits = consumer.getAvailablePermitsForDispatcherRemoval(); totalAvailablePermits -= availablePermits; log.debug() diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java index b66d8293c1f5c..667c2d8e19d55 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java @@ -19,29 +19,46 @@ package org.apache.pulsar.broker.service.persistent; import static java.util.Collections.emptyMap; +import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Exclusive; +import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Failover; +import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Key_Shared; import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Shared; import static org.apache.pulsar.common.api.proto.KeySharedMode.AUTO_SPLIT; import static org.apache.pulsar.common.protocol.Commands.DEFAULT_CONSUMER_EPOCH; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import io.netty.util.concurrent.EventExecutor; +import io.netty.util.concurrent.ImmediateEventExecutor; import java.time.Duration; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.TimeUnit; +import org.apache.bookkeeper.mledger.Entry; import org.apache.bookkeeper.mledger.ManagedCursor; import org.apache.bookkeeper.mledger.impl.ManagedCursorImpl; import org.apache.pulsar.broker.service.Consumer; import org.apache.pulsar.broker.service.Dispatcher; +import org.apache.pulsar.broker.service.EntryBatchIndexesAcks; +import org.apache.pulsar.broker.service.EntryBatchSizes; +import org.apache.pulsar.broker.service.PulsarCommandSender; +import org.apache.pulsar.broker.service.RedeliveryTracker; +import org.apache.pulsar.broker.service.ServerCnx; import org.apache.pulsar.broker.service.SharedPulsarBaseTest; import org.apache.pulsar.broker.service.Subscription; +import org.apache.pulsar.broker.service.Topic; import org.apache.pulsar.broker.service.TransportCnx; import org.apache.pulsar.client.api.Message; import org.apache.pulsar.client.api.MessageId; import org.apache.pulsar.client.api.Producer; import org.apache.pulsar.client.api.SubscriptionType; +import org.apache.pulsar.common.api.proto.CommandSubscribe.SubType; import org.apache.pulsar.common.api.proto.KeySharedMeta; import org.apache.pulsar.common.policies.data.stats.ConsumerStatsImpl; import org.awaitility.Awaitility; @@ -56,9 +73,54 @@ public Object[][] dispatcherImplementations() { return new Object[][] {{false}, {true}}; } - @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) - public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean classic) throws Exception { - TestContext context = createTestContext(classic); + @DataProvider(name = "flowRaceDispatcherVariants") + public Object[][] flowRaceDispatcherVariants() { + return new Object[][] { + {false, Shared}, + {true, Shared}, + {false, Key_Shared}, + {true, Key_Shared} + }; + } + + @DataProvider(name = "untrackedSubscriptionVariants") + public Object[][] untrackedSubscriptionVariants() { + return new Object[][] { + {true, Exclusive}, + {true, Failover}, + {false, Shared}, + {false, Key_Shared} + }; + } + + @Test(dataProvider = "untrackedSubscriptionVariants", timeOut = 30_000) + public void testUnrelatedSubscriptionDoesNotAccumulatePendingDispatcherPermits( + boolean persistent, SubType subType) throws Exception { + String topicName = newTopicName(); + String subscriptionName = "sub"; + admin.topics().createNonPartitionedTopic(topicName); + PersistentTopic persistentTopic = (PersistentTopic) getTopic(topicName, false).join().orElseThrow(); + Topic consumerTopic = persistent ? persistentTopic : mock(Topic.class); + if (!persistent) { + when(consumerTopic.getBrokerService()).thenReturn(persistentTopic.getBrokerService()); + when(consumerTopic.getHierarchyTopicPolicies()) + .thenReturn(persistentTopic.getHierarchyTopicPolicies()); + } + Subscription subscription = mock(Subscription.class); + when(subscription.getName()).thenReturn(subscriptionName); + when(subscription.getTopic()).thenReturn(consumerTopic); + Consumer consumer = createConsumer(subscription, subType, topicName, 1); + + consumer.flowPermits(100); + + assertThat(consumer.getAvailablePermits()).isEqualTo(100); + assertThat(consumer.getAvailablePermitsForDispatcherRemoval()).isEqualTo(100); + } + + @Test(dataProvider = "flowRaceDispatcherVariants", timeOut = 30_000) + public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean classic, SubType subType) + throws Exception { + TestContext context = createTestContext(classic, subType); Consumer remainingConsumer = context.remainingConsumer(); Consumer removedConsumer = context.removedConsumer(); @@ -68,8 +130,10 @@ public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean cla assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); synchronized (context.dispatcher()) { + // Keep the asynchronous Flow tasks queued until removal completes. removedConsumer.flowPermits(400); removedConsumer.flowPermits(600); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); context.dispatcher().removeConsumer(removedConsumer); } drainBrokerWorkerGroup(context.topic()); @@ -78,6 +142,70 @@ public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean cla .isEqualTo(remainingConsumer.getAvailablePermits()); } + @Test(timeOut = 30_000) + public void testPendingPermitAccountingSurvivesSignedIntegerWrap() throws Exception { + String topicName = newTopicName(); + String subscriptionName = "shared-sub"; + admin.topics().createNonPartitionedTopic(topicName); + PersistentTopic topic = (PersistentTopic) getTopic(topicName, false).join().orElseThrow(); + Subscription subscription = mock(PersistentSubscription.class); + when(subscription.getName()).thenReturn(subscriptionName); + when(subscription.getTopic()).thenReturn(topic); + Consumer consumer = createConsumer(subscription, Shared, topicName, 1); + + consumer.flowPermits(Integer.MAX_VALUE); + consumer.flowPermits(Integer.MAX_VALUE); + consumer.flowPermits(Integer.MAX_VALUE); + + assertThat(consumer.getAvailablePermits()).isEqualTo(Integer.MAX_VALUE - 2); + assertThat(consumer.getAvailablePermitsForDispatcherRemoval()).isZero(); + + consumer.completePendingDispatcherFlow(Integer.MAX_VALUE); + assertThat(consumer.getAvailablePermitsForDispatcherRemoval()).isEqualTo(Integer.MAX_VALUE); + consumer.completePendingDispatcherFlow(Integer.MAX_VALUE); + assertThat(consumer.getAvailablePermitsForDispatcherRemoval()).isEqualTo(-2); + consumer.completePendingDispatcherFlow(Integer.MAX_VALUE); + assertThat(consumer.getAvailablePermitsForDispatcherRemoval()).isEqualTo(Integer.MAX_VALUE - 2); + } + + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) + public void testRemovalAppliesNegativeAccountedPermitBalance(boolean classic) throws Exception { + TestContext context = createTestContext(classic); + Consumer remainingConsumer = context.remainingConsumer(); + Consumer removedConsumer = context.removedConsumer(); + Consumer secondRemainingConsumer = createConsumer( + remainingConsumer.getSubscription(), Shared, context.topic().getName(), 3); + context.dispatcher().addConsumer(secondRemainingConsumer).join(); + + remainingConsumer.flowPermits(10); + secondRemainingConsumer.flowPermits(15); + removedConsumer.flowPermits(20); + drainBrokerWorkerGroup(context.topic()); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(45); + + synchronized (context.dispatcher()) { + // Keep the Flow task queued, then dispatch more than the removed consumer's 20 accounted permits. + // Its removal balance becomes 70 available - 100 pending = -30 and must be applied as-is. + removedConsumer.flowPermits(100); + simulateDispatch(context.dispatcher(), removedConsumer, 50); + + assertThat(removedConsumer.getAvailablePermits()).isEqualTo(70); + assertThat(removedConsumer.getAvailablePermitsForDispatcherRemoval()).isEqualTo(-30); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(-5); + + context.dispatcher().removeConsumer(removedConsumer); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(25); + } + drainBrokerWorkerGroup(context.topic()); + + // Keep two consumers after removal so a single-consumer read floor cannot mask dispatcher permit drift. + assertThat(context.dispatcher().getConsumers()) + .containsExactlyInAnyOrder(remainingConsumer, secondRemainingConsumer); + assertThat(totalAvailablePermits(context.dispatcher())) + .isEqualTo(remainingConsumer.getAvailablePermits() + + secondRemainingConsumer.getAvailablePermits()); + } + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean classic) throws Exception { TestContext context = createTestContext(classic); @@ -98,7 +226,9 @@ public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(bool assertThat(removedConsumer.getAvailablePermits()).isZero(); synchronized (context.dispatcher()) { + // Keep the asynchronous Flow task queued until removal completes. removedConsumer.updateBlockedConsumerOnUnackedMsgs(removedConsumer); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(10); context.dispatcher().removeConsumer(removedConsumer); } drainBrokerWorkerGroup(context.topic()); @@ -138,6 +268,7 @@ public void testRemainingConsumerCanContinueAfterFlowAndCloseRace() throws Excep (PersistentDispatcherMultipleConsumers) subscription.getDispatcher(); Consumer remainingBrokerConsumer = findConsumer(dispatcher, "remaining-consumer"); Consumer removedBrokerConsumer = findConsumer(dispatcher, "removed-consumer"); + ServerCnx removedConsumerCnx = (ServerCnx) removedBrokerConsumer.cnx(); Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { assertThat(remainingBrokerConsumer.getAvailablePermits()).isEqualTo(receiverQueueSize); @@ -146,14 +277,22 @@ public void testRemainingConsumerCanContinueAfterFlowAndCloseRace() throws Excep drainBrokerWorkerGroup(topic); assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); - // Hold the dispatcher monitor so the asynchronous Flow task cannot run before the production - // Consumer.close -> PersistentSubscription.removeConsumer -> dispatcher.removeConsumer lifecycle. - synchronized (dispatcher) { - removedBrokerConsumer.flowPermits(pendingFlowPermits); - assertThat(removedBrokerConsumer.getAvailablePermits()) - .isEqualTo(receiverQueueSize + pendingFlowPermits); - removedBrokerConsumer.close(); + // Follow the production subscription -> dispatcher lock order and hold the dispatcher monitor so the + // asynchronous Flow task cannot run before Consumer.close removes the consumer. + synchronized (subscription) { + synchronized (dispatcher) { + removedBrokerConsumer.flowPermits(pendingFlowPermits); + assertThat(removedBrokerConsumer.getAvailablePermits()) + .isEqualTo(receiverQueueSize + pendingFlowPermits); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); + removedBrokerConsumer.close(); + } } + // Wait for the broker-side close to remove the consumer from the connection map. The client close is then + // handled idempotently instead of attempting a second dispatcher removal. + Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> + assertThat(removedConsumerCnx.getConsumers() + .containsKey(removedBrokerConsumer.consumerId())).isFalse()); removedClient.close(); drainBrokerWorkerGroup(topic); @@ -173,6 +312,10 @@ public void testRemainingConsumerCanContinueAfterFlowAndCloseRace() throws Excep } private TestContext createTestContext(boolean classic) throws Exception { + return createTestContext(classic, Shared); + } + + private TestContext createTestContext(boolean classic, SubType subType) throws Exception { String topicName = newTopicName(); String subscriptionName = "shared-sub"; admin.topics().createNonPartitionedTopic(topicName); @@ -180,35 +323,75 @@ private TestContext createTestContext(boolean classic) throws Exception { PersistentTopic topic = (PersistentTopic) getTopic(topicName, false).join().get(); ManagedCursor cursor = mock(ManagedCursorImpl.class); when(cursor.getName()).thenReturn(subscriptionName); + when(cursor.isClosed()).thenReturn(true); Subscription subscription = mock(PersistentSubscription.class); when(subscription.getName()).thenReturn(subscriptionName); when(subscription.getTopic()).thenReturn(topic); - Dispatcher dispatcher = classic - ? new NoopReadClassicDispatcher(topic, cursor, subscription) - : new NoopReadDispatcher(topic, cursor, subscription); + Dispatcher dispatcher; + if (subType == Key_Shared) { + dispatcher = classic + ? new PersistentStickyKeyDispatcherMultipleConsumersClassic( + topic, cursor, subscription, getConfig(), new KeySharedMeta().setKeySharedMode(AUTO_SPLIT)) + : new PersistentStickyKeyDispatcherMultipleConsumers( + topic, cursor, subscription, getConfig(), new KeySharedMeta().setKeySharedMode(AUTO_SPLIT)); + } else { + dispatcher = classic + ? new PersistentDispatcherMultipleConsumersClassic(topic, cursor, subscription) + : new PersistentDispatcherMultipleConsumers(topic, cursor, subscription); + } doAnswer(invocation -> { dispatcher.consumerFlow(invocation.getArgument(0), invocation.getArgument(1)); return null; }).when(subscription).consumerFlow(any(), anyInt()); - Consumer remainingConsumer = createConsumer(subscription, topicName, 1); - Consumer removedConsumer = createConsumer(subscription, topicName, 2); + Consumer remainingConsumer = createConsumer(subscription, subType, topicName, 1); + Consumer removedConsumer = createConsumer(subscription, subType, topicName, 2); dispatcher.addConsumer(remainingConsumer).join(); dispatcher.addConsumer(removedConsumer).join(); return new TestContext(topic, dispatcher, remainingConsumer, removedConsumer); } - private Consumer createConsumer(Subscription subscription, String topicName, long consumerId) { + private Consumer createConsumer(Subscription subscription, SubType subType, String topicName, long consumerId) { TransportCnx cnx = mock(TransportCnx.class); + PulsarCommandSender commandSender = mock(PulsarCommandSender.class); when(cnx.isActive()).thenReturn(true); when(cnx.isWritable()).thenReturn(true); - return new Consumer(subscription, Shared, topicName, consumerId, 0, "consumer-" + consumerId, + when(cnx.getCommandSender()).thenReturn(commandSender); + when(commandSender.sendMessagesToConsumer(anyLong(), anyString(), any(), anyInt(), anyList(), + any(EntryBatchSizes.class), any(EntryBatchIndexesAcks.class), any(RedeliveryTracker.class), anyLong())) + .thenReturn(ImmediateEventExecutor.INSTANCE.newSucceededFuture(null)); + return new Consumer(subscription, subType, topicName, consumerId, 0, "consumer-" + consumerId, true, cnx, "role", emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), MessageId.latest, DEFAULT_CONSUMER_EPOCH); } + private static void simulateDispatch(Dispatcher dispatcher, Consumer consumer, int permits) { + Entry entry = mock(Entry.class); + when(entry.getLedgerId()).thenReturn(1L); + when(entry.getEntryId()).thenReturn(1L); + EntryBatchSizes batchSizes = EntryBatchSizes.get(1); + batchSizes.setBatchSize(0, permits); + EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(1); + RedeliveryTracker redeliveryTracker = mock(RedeliveryTracker.class); + + consumer.sendMessages(new ArrayList<>(List.of(entry)), batchSizes, batchIndexesAcks, + permits, 0, 0, redeliveryTracker).syncUninterruptibly(); + decrementTotalAvailablePermits(dispatcher, permits); + + batchSizes.recyle(); + batchIndexesAcks.recycle(); + } + + private static void decrementTotalAvailablePermits(Dispatcher dispatcher, int permits) { + if (dispatcher instanceof PersistentDispatcherMultipleConsumers pip379Dispatcher) { + pip379Dispatcher.totalAvailablePermits -= permits; + } else { + ((PersistentDispatcherMultipleConsumersClassic) dispatcher).totalAvailablePermits -= permits; + } + } + private static void drainBrokerWorkerGroup(PersistentTopic topic) throws Exception { for (EventExecutor eventExecutor : topic.getBrokerService().executor()) { eventExecutor.submit(() -> { }).sync(); @@ -232,36 +415,4 @@ private static Consumer findConsumer(Dispatcher dispatcher, String consumerName) private record TestContext(PersistentTopic topic, Dispatcher dispatcher, Consumer remainingConsumer, Consumer removedConsumer) { } - - private static class NoopReadDispatcher extends PersistentDispatcherMultipleConsumers { - NoopReadDispatcher(PersistentTopic topic, ManagedCursor cursor, Subscription subscription) { - super(topic, cursor, subscription); - } - - @Override - public void readMoreEntriesAsync() { - // No-op for permit accounting test. - } - - @Override - public synchronized void readMoreEntries() { - // No-op for permit accounting test. - } - } - - private static class NoopReadClassicDispatcher extends PersistentDispatcherMultipleConsumersClassic { - NoopReadClassicDispatcher(PersistentTopic topic, ManagedCursor cursor, Subscription subscription) { - super(topic, cursor, subscription); - } - - @Override - public void readMoreEntriesAsync() { - // No-op for permit accounting test. - } - - @Override - public synchronized void readMoreEntries() { - // No-op for permit accounting test. - } - } } From cf57407fe13f171a48dd787d760b1ae9e37391a0 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Fri, 28 Aug 2026 00:06:24 +0800 Subject: [PATCH 5/7] [fix][broker] Inline Shared dispatcher flow accounting --- .../pulsar/broker/service/Consumer.java | 33 +- ...PersistentDispatcherMultipleConsumers.java | 23 +- ...entDispatcherMultipleConsumersClassic.java | 26 +- .../SharedDispatcherPermitAccountingTest.java | 337 ++++++++++++------ 4 files changed, 278 insertions(+), 141 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java index b45947b76e833..b6738570f9a3a 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java @@ -119,10 +119,10 @@ public class Consumer { private volatile int messagePermits = 0; /** * Guards the Flow-side compound update of {@link #messagePermits} and - * {@link #pendingDispatcherFlowPermits}. A Flow command increases the consumer permits before the dispatcher - * processes the corresponding update asynchronously. Consumer removal can happen between those two operations, - * so both values must be observed consistently when calculating how many permits are already included in the - * dispatcher total. + * {@link #pendingDispatcherFlowPermits}. A Flow command increases the consumer permits before entering the + * dispatcher. Consumer removal can happen while the Flow caller is waiting for the dispatcher monitor, so both + * values must be observed consistently when calculating how many permits are already included in the dispatcher + * total. * *

The dispatcher callback is invoked only after this lock is released. This avoids holding the lock while * calling into the subscription and preserves the lock order used by dispatcher flow processing and removal. @@ -969,7 +969,8 @@ public int getAvailablePermits() { /** * Adds permits after a Flow command is accepted and immediately before notifying the dispatcher. The pending - * count covers the interval until the dispatcher's asynchronous Flow task starts processing the same permits. + * count covers the interval until the dispatcher starts processing the same permits, including time spent waiting + * for the dispatcher monitor. */ private int addPermitsPendingDispatcherUpdate(int additionalNumberOfPermits) { if (!shouldTrackPendingDispatcherFlowPermits()) { @@ -986,9 +987,9 @@ private boolean shouldTrackPendingDispatcherFlowPermits() { } /** - * Called at the start of the dispatcher's asynchronous Flow task, before checking whether this consumer is still - * connected. At this point the Flow update is no longer pending: the dispatcher will either add the permits to - * its total or ignore them because the consumer has already been removed. + * Called at the start of dispatcher Flow processing, before checking whether this consumer is still connected. At + * this point the Flow update is no longer pending: the dispatcher will either add the permits to its total or + * ignore them because the consumer has already been removed. */ public void completePendingDispatcherFlow(int additionalNumberOfPermits) { synchronized (flowPermitAccountingLock) { @@ -998,19 +999,19 @@ public void completePendingDispatcherFlow(int additionalNumberOfPermits) { } /** - * Called while the dispatcher removes this consumer. Permits belonging to Flow tasks that have not started yet - * are excluded because those permits have not been added to the dispatcher total and must not be subtracted from - * it during removal. + * Called while the dispatcher removes this consumer. Permits belonging to Flow updates that the dispatcher has + * not started processing are excluded because those permits have not been added to the dispatcher total and must + * not be subtracted from it during removal. * - *

This accounting is enabled for persistent Shared and Key_Shared dispatchers. It relies on every dispatcher - * Flow task calling {@link #completePendingDispatcherFlow(int)} before applying or ignoring the update. For these + *

This accounting is enabled for persistent Shared and Key_Shared dispatchers. It requires every dispatcher + * Flow update to call {@link #completePendingDispatcherFlow(int)} before applying or ignoring it. For these * dispatchers, when observed under the dispatcher monitor, the total available permits equal the sum of this * balance over all connected consumers. * *

The returned balance can be negative. A pending Flow makes permits visible on the consumer before its - * asynchronous dispatcher update runs, so the dispatcher can consume those permits while they are still counted - * as pending. Subtracting the negative balance during removal is required to restore the dispatcher total; callers - * must not clamp it to zero. + * dispatcher update runs, so the dispatcher can consume those permits while they are still counted as pending. + * Subtracting the negative balance during removal is required to restore the dispatcher total; callers must not + * clamp it to zero. */ public int getAvailablePermitsForDispatcherRemoval() { synchronized (flowPermitAccountingLock) { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java index e7a6d586e0aee..1a5cab940e22a 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java @@ -303,26 +303,27 @@ protected synchronized void clearComponentsAfterRemovedAllConsumers() { @Override public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { - topic.getBrokerService().executor().execute(() -> { - internalConsumerFlow(consumer, additionalNumberOfMessages); - }); - } + boolean connected; + int updatedTotalAvailablePermits = 0; + synchronized (this) { + consumer.completePendingDispatcherFlow(additionalNumberOfMessages); + connected = consumerSet.contains(consumer); + if (connected) { + totalAvailablePermits += additionalNumberOfMessages; + updatedTotalAvailablePermits = totalAvailablePermits; + } + } - private synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { - // The queued Flow task is no longer pending, even if the consumer was removed while the task was waiting. - consumer.completePendingDispatcherFlow(additionalNumberOfMessages); - if (!consumerSet.contains(consumer)) { + if (!connected) { log.debug() .attr("consumer", consumer) .log("Ignoring flow control from disconnected consumer"); return; } - totalAvailablePermits += additionalNumberOfMessages; - log.debug() .attr("consumer", consumer) - .attr("totalAvailablePermits", totalAvailablePermits) + .attr("totalAvailablePermits", updatedTotalAvailablePermits) .attr("additionalNumberOfMessages", additionalNumberOfMessages) .log("Trigger new read after receiving flow control message"); readMoreEntriesAsync(); diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java index 3016d0fa8dceb..8cf7f051f86d4 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java @@ -283,35 +283,37 @@ private synchronized void clearComponentsAfterRemovedAllConsumers() { @Override public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { - topic.getBrokerService().executor().execute(() -> { - internalConsumerFlow(consumer, additionalNumberOfMessages); - }); - } + boolean connected; + int updatedTotalAvailablePermits = 0; + synchronized (this) { + consumer.completePendingDispatcherFlow(additionalNumberOfMessages); + connected = consumerSet.contains(consumer); + if (connected) { + totalAvailablePermits += additionalNumberOfMessages; + updatedTotalAvailablePermits = totalAvailablePermits; + } + } - private synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { - // The queued Flow task is no longer pending, even if the consumer was removed while the task was waiting. - consumer.completePendingDispatcherFlow(additionalNumberOfMessages); - if (!consumerSet.contains(consumer)) { + if (!connected) { log.debug() .attr("consumer", consumer) .log("Ignoring flow control from disconnected consumer"); return; } - totalAvailablePermits += additionalNumberOfMessages; - log.debug() .attr("consumer", consumer) - .attr("totalAvailablePermits", totalAvailablePermits) + .attr("totalAvailablePermits", updatedTotalAvailablePermits) .attr("additionalNumberOfMessages", additionalNumberOfMessages) .log("- Trigger new read after receiving flow control message with permits " + "after adding permits"); - readMoreEntries(); + readMoreEntriesAsync(); } /** * We should not call readMoreEntries() recursively in the same thread as there is a risk of StackOverflowError. * */ + @Override public void readMoreEntriesAsync() { topic.getBrokerService().executor().execute(this::readMoreEntries); } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java index 667c2d8e19d55..5d18b59c76b33 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java @@ -33,13 +33,20 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; import io.netty.util.concurrent.EventExecutor; import io.netty.util.concurrent.ImmediateEventExecutor; import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.bookkeeper.mledger.Entry; import org.apache.bookkeeper.mledger.ManagedCursor; import org.apache.bookkeeper.mledger.impl.ManagedCursorImpl; @@ -51,12 +58,14 @@ import org.apache.pulsar.broker.service.RedeliveryTracker; import org.apache.pulsar.broker.service.ServerCnx; import org.apache.pulsar.broker.service.SharedPulsarBaseTest; +import org.apache.pulsar.broker.service.StickyKeyDispatcher; import org.apache.pulsar.broker.service.Subscription; import org.apache.pulsar.broker.service.Topic; import org.apache.pulsar.broker.service.TransportCnx; import org.apache.pulsar.client.api.Message; import org.apache.pulsar.client.api.MessageId; import org.apache.pulsar.client.api.Producer; +import org.apache.pulsar.client.api.Range; import org.apache.pulsar.client.api.SubscriptionType; import org.apache.pulsar.common.api.proto.CommandSubscribe.SubType; import org.apache.pulsar.common.api.proto.KeySharedMeta; @@ -129,13 +138,11 @@ public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean cla drainBrokerWorkerGroup(context.topic()); assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); - synchronized (context.dispatcher()) { - // Keep the asynchronous Flow tasks queued until removal completes. - removedConsumer.flowPermits(400); - removedConsumer.flowPermits(600); - assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); - context.dispatcher().removeConsumer(removedConsumer); - } + runWithFlowWaitingForDispatcher(context.dispatcher(), removedConsumer, 1_020, + () -> removedConsumer.flowPermits(1_000), () -> { + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); + context.dispatcher().removeConsumer(removedConsumer); + }); drainBrokerWorkerGroup(context.topic()); assertThat(totalAvailablePermits(context.dispatcher())) @@ -168,13 +175,13 @@ public void testPendingPermitAccountingSurvivesSignedIntegerWrap() throws Except assertThat(consumer.getAvailablePermitsForDispatcherRemoval()).isEqualTo(Integer.MAX_VALUE - 2); } - @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) - public void testRemovalAppliesNegativeAccountedPermitBalance(boolean classic) throws Exception { - TestContext context = createTestContext(classic); + @Test(dataProvider = "flowRaceDispatcherVariants", timeOut = 30_000) + public void testRemovalAppliesNegativeAccountedPermitBalance(boolean classic, SubType subType) throws Exception { + TestContext context = createTestContext(classic, subType); Consumer remainingConsumer = context.remainingConsumer(); Consumer removedConsumer = context.removedConsumer(); Consumer secondRemainingConsumer = createConsumer( - remainingConsumer.getSubscription(), Shared, context.topic().getName(), 3); + remainingConsumer.getSubscription(), subType, context.topic().getName(), 3); context.dispatcher().addConsumer(secondRemainingConsumer).join(); remainingConsumer.flowPermits(10); @@ -183,19 +190,19 @@ public void testRemovalAppliesNegativeAccountedPermitBalance(boolean classic) th drainBrokerWorkerGroup(context.topic()); assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(45); - synchronized (context.dispatcher()) { - // Keep the Flow task queued, then dispatch more than the removed consumer's 20 accounted permits. - // Its removal balance becomes 70 available - 100 pending = -30 and must be applied as-is. - removedConsumer.flowPermits(100); - simulateDispatch(context.dispatcher(), removedConsumer, 50); + runWithFlowWaitingForDispatcher(context.dispatcher(), removedConsumer, 120, + () -> removedConsumer.flowPermits(100), () -> { + // Dispatch more than the removed consumer's 20 accounted permits. Its removal balance becomes + // 70 available - 100 pending = -30 and must be applied as-is. + simulateDispatch(context.dispatcher(), removedConsumer, 50); - assertThat(removedConsumer.getAvailablePermits()).isEqualTo(70); - assertThat(removedConsumer.getAvailablePermitsForDispatcherRemoval()).isEqualTo(-30); - assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(-5); + assertThat(removedConsumer.getAvailablePermits()).isEqualTo(70); + assertThat(removedConsumer.getAvailablePermitsForDispatcherRemoval()).isEqualTo(-30); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(-5); - context.dispatcher().removeConsumer(removedConsumer); - assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(25); - } + context.dispatcher().removeConsumer(removedConsumer); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(25); + }); drainBrokerWorkerGroup(context.topic()); // Keep two consumers after removal so a single-consumer read floor cannot mask dispatcher permit drift. @@ -206,9 +213,10 @@ public void testRemovalAppliesNegativeAccountedPermitBalance(boolean classic) th + secondRemainingConsumer.getAvailablePermits()); } - @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) - public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean classic) throws Exception { - TestContext context = createTestContext(classic); + @Test(dataProvider = "flowRaceDispatcherVariants", timeOut = 30_000) + public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits( + boolean classic, SubType subType) throws Exception { + TestContext context = createTestContext(classic, subType); Consumer remainingConsumer = context.remainingConsumer(); Consumer removedConsumer = context.removedConsumer(); @@ -225,89 +233,172 @@ public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(bool removedConsumer.flowPermits(1_000); assertThat(removedConsumer.getAvailablePermits()).isZero(); - synchronized (context.dispatcher()) { - // Keep the asynchronous Flow task queued until removal completes. - removedConsumer.updateBlockedConsumerOnUnackedMsgs(removedConsumer); - assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(10); - context.dispatcher().removeConsumer(removedConsumer); - } + runWithFlowWaitingForDispatcher(context.dispatcher(), removedConsumer, 1_000, + () -> removedConsumer.updateBlockedConsumerOnUnackedMsgs(removedConsumer), () -> { + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(10); + context.dispatcher().removeConsumer(removedConsumer); + }); drainBrokerWorkerGroup(context.topic()); assertThat(totalAvailablePermits(context.dispatcher())) .isEqualTo(remainingConsumer.getAvailablePermits()); } - @Test(timeOut = 30_000) - public void testRemainingConsumerCanContinueAfterFlowAndCloseRace() throws Exception { + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) + public void testReadTriggerRunsOutsideDispatcherMonitorAfterConsumerRemoval(boolean classic) throws Exception { + TestContext context = createTestContext(classic, Shared, true); + Consumer remainingConsumer = context.remainingConsumer(); + Consumer removedConsumer = context.removedConsumer(); + AbstractPersistentDispatcherMultipleConsumers dispatcher = + (AbstractPersistentDispatcherMultipleConsumers) context.dispatcher(); + + remainingConsumer.flowPermits(10); + removedConsumer.flowPermits(20); + drainBrokerWorkerGroup(context.topic()); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(30); + + CountDownLatch readTriggerEntered = new CountDownLatch(1); + CountDownLatch allowReadTrigger = new CountDownLatch(1); + AtomicBoolean readTriggerHeldDispatcherMonitor = new AtomicBoolean(); + AtomicInteger readMoreEntriesCalls = new AtomicInteger(); + if (classic) { + doAnswer(invocation -> { + readMoreEntriesCalls.incrementAndGet(); + return invocation.callRealMethod(); + }).when((PersistentDispatcherMultipleConsumersClassic) dispatcher).readMoreEntries(); + } else { + doAnswer(invocation -> { + readMoreEntriesCalls.incrementAndGet(); + return invocation.callRealMethod(); + }).when((PersistentDispatcherMultipleConsumers) dispatcher).readMoreEntries(); + } + doAnswer(invocation -> { + readTriggerHeldDispatcherMonitor.set(Thread.holdsLock(dispatcher)); + readTriggerEntered.countDown(); + assertThat(allowReadTrigger.await(5, TimeUnit.SECONDS)).isTrue(); + return invocation.callRealMethod(); + }).when(dispatcher).readMoreEntriesAsync(); + + ExecutorService flowExecutor = Executors.newSingleThreadExecutor(); + Future flowFuture = null; + try { + flowFuture = flowExecutor.submit(() -> { + removedConsumer.flowPermits(1_000); + return null; + }); + assertThat(readTriggerEntered.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(readTriggerHeldDispatcherMonitor).isFalse(); + assertThat(flowFuture.isDone()).isFalse(); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(1_030); + + dispatcher.removeConsumer(removedConsumer); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(remainingConsumer.getAvailablePermits()); + // Classic removal invokes readMoreEntries synchronously. Count only the task submitted by the Flow + // trigger after the removal has completed. + readMoreEntriesCalls.set(0); + } finally { + allowReadTrigger.countDown(); + try { + if (flowFuture != null) { + flowFuture.get(5, TimeUnit.SECONDS); + } + } finally { + flowExecutor.shutdownNow(); + assertThat(flowExecutor.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); + } + } + drainBrokerWorkerGroup(context.topic()); + + assertThat(dispatcher.getConsumers()).containsExactly(remainingConsumer); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(remainingConsumer.getAvailablePermits()); + assertThat(readMoreEntriesCalls).hasValue(1); + } + + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) + public void testRemainingConsumerCanContinueAfterFlowAndCloseRace(boolean classic) throws Exception { int receiverQueueSize = 10; int messagesToConsume = receiverQueueSize * 2; - int pendingFlowPermits = 1_000; + int additionalFlowPermits = 1_000; String topicName = newTopicName(); String subscriptionName = "shared-sub"; - admin.topics().createNonPartitionedTopic(topicName); - - try (Producer producer = pulsarClient.newProducer().topic(topicName).create(); - org.apache.pulsar.client.api.Consumer remainingClient = pulsarClient.newConsumer() - .topic(topicName) - .subscriptionName(subscriptionName) - .subscriptionType(SubscriptionType.Shared) - .consumerName("remaining-consumer") - .receiverQueueSize(receiverQueueSize) - .subscribe(); - org.apache.pulsar.client.api.Consumer removedClient = pulsarClient.newConsumer() - .topic(topicName) - .subscriptionName(subscriptionName) - .subscriptionType(SubscriptionType.Shared) - .consumerName("removed-consumer") - .receiverQueueSize(receiverQueueSize) - .subscribe()) { - PersistentTopic topic = (PersistentTopic) getTopic(topicName, false).join().orElseThrow(); - PersistentSubscription subscription = - (PersistentSubscription) topic.getSubscription(subscriptionName); - PersistentDispatcherMultipleConsumers dispatcher = - (PersistentDispatcherMultipleConsumers) subscription.getDispatcher(); - Consumer remainingBrokerConsumer = findConsumer(dispatcher, "remaining-consumer"); - Consumer removedBrokerConsumer = findConsumer(dispatcher, "removed-consumer"); - ServerCnx removedConsumerCnx = (ServerCnx) removedBrokerConsumer.cnx(); - - Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(remainingBrokerConsumer.getAvailablePermits()).isEqualTo(receiverQueueSize); - assertThat(removedBrokerConsumer.getAvailablePermits()).isEqualTo(receiverQueueSize); - }); - drainBrokerWorkerGroup(topic); - assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); - - // Follow the production subscription -> dispatcher lock order and hold the dispatcher monitor so the - // asynchronous Flow task cannot run before Consumer.close removes the consumer. - synchronized (subscription) { - synchronized (dispatcher) { - removedBrokerConsumer.flowPermits(pendingFlowPermits); - assertThat(removedBrokerConsumer.getAvailablePermits()) - .isEqualTo(receiverQueueSize + pendingFlowPermits); - assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); - removedBrokerConsumer.close(); + boolean previousClassicSetting = getConfig().isSubscriptionSharedUseClassicPersistentImplementation(); + getConfig().setSubscriptionSharedUseClassicPersistentImplementation(classic); + try { + admin.topics().createNonPartitionedTopic(topicName); + + try (Producer producer = pulsarClient.newProducer() + .topic(topicName) + .enableBatching(false) + .create(); + org.apache.pulsar.client.api.Consumer remainingClient = pulsarClient.newConsumer() + .topic(topicName) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .consumerName("remaining-consumer") + .receiverQueueSize(receiverQueueSize) + .subscribe(); + org.apache.pulsar.client.api.Consumer removedClient = pulsarClient.newConsumer() + .topic(topicName) + .subscriptionName(subscriptionName) + .subscriptionType(SubscriptionType.Shared) + .consumerName("removed-consumer") + .receiverQueueSize(receiverQueueSize) + .subscribe()) { + PersistentTopic topic = (PersistentTopic) getTopic(topicName, false).join().orElseThrow(); + PersistentSubscription subscription = + (PersistentSubscription) topic.getSubscription(subscriptionName); + Dispatcher dispatcher = subscription.getDispatcher(); + assertThat(dispatcher).isInstanceOf(classic + ? PersistentDispatcherMultipleConsumersClassic.class + : PersistentDispatcherMultipleConsumers.class); + Consumer remainingBrokerConsumer = findConsumer(dispatcher, "remaining-consumer"); + Consumer removedBrokerConsumer = findConsumer(dispatcher, "removed-consumer"); + ServerCnx removedConsumerCnx = (ServerCnx) removedBrokerConsumer.cnx(); + + Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(remainingBrokerConsumer.getAvailablePermits()).isEqualTo(receiverQueueSize); + assertThat(removedBrokerConsumer.getAvailablePermits()).isEqualTo(receiverQueueSize); + }); + drainBrokerWorkerGroup(topic); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); + + // Follow the production subscription -> dispatcher lock order. The Flow thread first updates the + // consumer permits and then waits for the dispatcher monitor. Closing from this thread removes the + // consumer before the Flow update reaches the dispatcher. + synchronized (subscription) { + runWithFlowWaitingForDispatcher(dispatcher, removedBrokerConsumer, + receiverQueueSize + additionalFlowPermits, + () -> removedBrokerConsumer.flowPermits(additionalFlowPermits), () -> { + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); + removedBrokerConsumer.close(); + assertThat(dispatcher.getConsumers()).containsExactly(remainingBrokerConsumer); + assertThat(totalAvailablePermits(dispatcher)) + .isEqualTo(remainingBrokerConsumer.getAvailablePermits()); + }); + } + // Wait for the broker-side close to remove the consumer from the connection map. The client close is + // then handled idempotently instead of attempting a second dispatcher removal. + Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> + assertThat(removedConsumerCnx.getConsumers() + .containsKey(removedBrokerConsumer.consumerId())).isFalse()); + removedClient.close(); + drainBrokerWorkerGroup(topic); + + assertThat(dispatcher.getConsumers()).containsExactly(remainingBrokerConsumer); + assertThat(totalAvailablePermits(dispatcher)) + .isEqualTo(remainingBrokerConsumer.getAvailablePermits()); + + for (int i = 0; i < messagesToConsume; i++) { + producer.send(new byte[] {(byte) i}); + } + for (int i = 0; i < messagesToConsume; i++) { + Message message = remainingClient.receive(5, TimeUnit.SECONDS); + assertThat(message).isNotNull(); + remainingClient.acknowledge(message); } } - // Wait for the broker-side close to remove the consumer from the connection map. The client close is then - // handled idempotently instead of attempting a second dispatcher removal. - Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> - assertThat(removedConsumerCnx.getConsumers() - .containsKey(removedBrokerConsumer.consumerId())).isFalse()); - removedClient.close(); - drainBrokerWorkerGroup(topic); - - assertThat(dispatcher.getConsumers()).containsExactly(remainingBrokerConsumer); - assertThat(totalAvailablePermits(dispatcher)) - .isEqualTo(remainingBrokerConsumer.getAvailablePermits()); - - for (int i = 0; i < messagesToConsume; i++) { - producer.send(new byte[] {(byte) i}); - } - for (int i = 0; i < messagesToConsume; i++) { - Message message = remainingClient.receive(1, TimeUnit.SECONDS); - assertThat(message).isNotNull(); - remainingClient.acknowledge(message); - } + } finally { + getConfig().setSubscriptionSharedUseClassicPersistentImplementation(previousClassicSetting); } } @@ -316,6 +407,10 @@ private TestContext createTestContext(boolean classic) throws Exception { } private TestContext createTestContext(boolean classic, SubType subType) throws Exception { + return createTestContext(classic, subType, false); + } + + private TestContext createTestContext(boolean classic, SubType subType, boolean spyDispatcher) throws Exception { String topicName = newTopicName(); String subscriptionName = "shared-sub"; admin.topics().createNonPartitionedTopic(topicName); @@ -340,17 +435,21 @@ topic, cursor, subscription, getConfig(), new KeySharedMeta().setKeySharedMode(A ? new PersistentDispatcherMultipleConsumersClassic(topic, cursor, subscription) : new PersistentDispatcherMultipleConsumers(topic, cursor, subscription); } + if (spyDispatcher) { + dispatcher = spy(dispatcher); + } + Dispatcher configuredDispatcher = dispatcher; doAnswer(invocation -> { - dispatcher.consumerFlow(invocation.getArgument(0), invocation.getArgument(1)); + configuredDispatcher.consumerFlow(invocation.getArgument(0), invocation.getArgument(1)); return null; }).when(subscription).consumerFlow(any(), anyInt()); Consumer remainingConsumer = createConsumer(subscription, subType, topicName, 1); Consumer removedConsumer = createConsumer(subscription, subType, topicName, 2); - dispatcher.addConsumer(remainingConsumer).join(); - dispatcher.addConsumer(removedConsumer).join(); + configuredDispatcher.addConsumer(remainingConsumer).join(); + configuredDispatcher.addConsumer(removedConsumer).join(); - return new TestContext(topic, dispatcher, remainingConsumer, removedConsumer); + return new TestContext(topic, configuredDispatcher, remainingConsumer, removedConsumer); } private Consumer createConsumer(Subscription subscription, SubType subType, String topicName, long consumerId) { @@ -375,9 +474,15 @@ private static void simulateDispatch(Dispatcher dispatcher, Consumer consumer, i batchSizes.setBatchSize(0, permits); EntryBatchIndexesAcks batchIndexesAcks = EntryBatchIndexesAcks.get(1); RedeliveryTracker redeliveryTracker = mock(RedeliveryTracker.class); + List stickyKeyHashes = null; + if (dispatcher instanceof StickyKeyDispatcher stickyKeyDispatcher) { + List ranges = stickyKeyDispatcher.getConsumerKeyHashRanges().get(consumer); + assertThat(ranges).isNotEmpty(); + stickyKeyHashes = List.of(ranges.get(0).getStart()); + } - consumer.sendMessages(new ArrayList<>(List.of(entry)), batchSizes, batchIndexesAcks, - permits, 0, 0, redeliveryTracker).syncUninterruptibly(); + consumer.sendMessages(new ArrayList<>(List.of(entry)), stickyKeyHashes, batchSizes, batchIndexesAcks, + permits, 0, 0, redeliveryTracker, DEFAULT_CONSUMER_EPOCH).syncUninterruptibly(); decrementTotalAvailablePermits(dispatcher, permits); batchSizes.recyle(); @@ -392,6 +497,29 @@ private static void decrementTotalAvailablePermits(Dispatcher dispatcher, int pe } } + private static void runWithFlowWaitingForDispatcher(Dispatcher dispatcher, Consumer consumer, + int expectedConsumerPermits, CheckedRunnable flow, + CheckedRunnable whileFlowWaits) throws Exception { + ExecutorService flowExecutor = Executors.newSingleThreadExecutor(); + try { + Future flowFuture; + synchronized (dispatcher) { + flowFuture = flowExecutor.submit(() -> { + flow.run(); + return null; + }); + Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> + assertThat(consumer.getAvailablePermits()).isEqualTo(expectedConsumerPermits)); + assertThat(flowFuture.isDone()).isFalse(); + whileFlowWaits.run(); + } + flowFuture.get(5, TimeUnit.SECONDS); + } finally { + flowExecutor.shutdownNow(); + assertThat(flowExecutor.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); + } + } + private static void drainBrokerWorkerGroup(PersistentTopic topic) throws Exception { for (EventExecutor eventExecutor : topic.getBrokerService().executor()) { eventExecutor.submit(() -> { }).sync(); @@ -415,4 +543,9 @@ private static Consumer findConsumer(Dispatcher dispatcher, String consumerName) private record TestContext(PersistentTopic topic, Dispatcher dispatcher, Consumer remainingConsumer, Consumer removedConsumer) { } + + @FunctionalInterface + private interface CheckedRunnable { + void run() throws Exception; + } } From ecdb4b52fed3804694f2ab10a945cfcf1eeedcc3 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sat, 29 Aug 2026 22:09:19 +0800 Subject: [PATCH 6/7] [fix][broker] Keep Shared Flow processing off I/O threads The inline Flow path can wait for the persistent Shared dispatcher monitor on the connection EventLoop while dispatch or filter work owns the monitor. Route Flow accounting through the dispatcher's selected dispatchMessagesThread in both modern and classic implementations. Complete pending accounting and update total permits under the monitor, then trigger the read from the same lane. Leave rejected tasks pending so removal still excludes unapplied permits. Update the race tests for queued Flow processing and add deterministic coverage for EventLoop progress, lane affinity, removal, and shutdown rejection across Shared and Key_Shared. --- .../pulsar/broker/service/Consumer.java | 31 +- ...PersistentDispatcherMultipleConsumers.java | 18 +- ...entDispatcherMultipleConsumersClassic.java | 18 +- .../SharedDispatcherFlowThreadingTest.java | 318 ++++++++++++++++++ .../SharedDispatcherPermitAccountingTest.java | 166 +++------ 5 files changed, 419 insertions(+), 132 deletions(-) create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java index b6738570f9a3a..bd274a56d9069 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java @@ -119,10 +119,10 @@ public class Consumer { private volatile int messagePermits = 0; /** * Guards the Flow-side compound update of {@link #messagePermits} and - * {@link #pendingDispatcherFlowPermits}. A Flow command increases the consumer permits before entering the - * dispatcher. Consumer removal can happen while the Flow caller is waiting for the dispatcher monitor, so both - * values must be observed consistently when calculating how many permits are already included in the dispatcher - * total. + * {@link #pendingDispatcherFlowPermits}. A Flow command increases the consumer permits before the corresponding + * task runs on the dispatcher's message thread. Consumer removal can happen while that task is queued or waiting + * for the dispatcher monitor, so both values must be observed consistently when calculating how many permits are + * already included in the dispatcher total. * *

The dispatcher callback is invoked only after this lock is released. This avoids holding the lock while * calling into the subscription and preserves the lock order used by dispatcher flow processing and removal. @@ -969,8 +969,7 @@ public int getAvailablePermits() { /** * Adds permits after a Flow command is accepted and immediately before notifying the dispatcher. The pending - * count covers the interval until the dispatcher starts processing the same permits, including time spent waiting - * for the dispatcher monitor. + * count covers the interval until the dispatcher's asynchronous Flow task starts processing the same permits. */ private int addPermitsPendingDispatcherUpdate(int additionalNumberOfPermits) { if (!shouldTrackPendingDispatcherFlowPermits()) { @@ -987,9 +986,9 @@ private boolean shouldTrackPendingDispatcherFlowPermits() { } /** - * Called at the start of dispatcher Flow processing, before checking whether this consumer is still connected. At - * this point the Flow update is no longer pending: the dispatcher will either add the permits to its total or - * ignore them because the consumer has already been removed. + * Called at the start of the dispatcher's asynchronous Flow task, before checking whether this consumer is still + * connected. At this point the Flow update is no longer pending: the dispatcher will either add the permits to + * its total or ignore them because the consumer has already been removed. */ public void completePendingDispatcherFlow(int additionalNumberOfPermits) { synchronized (flowPermitAccountingLock) { @@ -999,19 +998,19 @@ public void completePendingDispatcherFlow(int additionalNumberOfPermits) { } /** - * Called while the dispatcher removes this consumer. Permits belonging to Flow updates that the dispatcher has - * not started processing are excluded because those permits have not been added to the dispatcher total and must - * not be subtracted from it during removal. + * Called while the dispatcher removes this consumer. Permits belonging to Flow tasks that have not started yet + * are excluded because those permits have not been added to the dispatcher total and must not be subtracted from + * it during removal. * *

This accounting is enabled for persistent Shared and Key_Shared dispatchers. It requires every dispatcher - * Flow update to call {@link #completePendingDispatcherFlow(int)} before applying or ignoring it. For these + * Flow task to call {@link #completePendingDispatcherFlow(int)} before applying or ignoring it. For these * dispatchers, when observed under the dispatcher monitor, the total available permits equal the sum of this * balance over all connected consumers. * *

The returned balance can be negative. A pending Flow makes permits visible on the consumer before its - * dispatcher update runs, so the dispatcher can consume those permits while they are still counted as pending. - * Subtracting the negative balance during removal is required to restore the dispatcher total; callers must not - * clamp it to zero. + * asynchronous dispatcher update runs, so the dispatcher can consume those permits while they are still counted + * as pending. Subtracting the negative balance during removal is required to restore the dispatcher total; callers + * must not clamp it to zero. */ public int getAvailablePermitsForDispatcherRemoval() { synchronized (flowPermitAccountingLock) { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java index 1a5cab940e22a..5147e2b3b3096 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java @@ -36,6 +36,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.ExecutorService; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -303,6 +304,21 @@ protected synchronized void clearComponentsAfterRemovedAllConsumers() { @Override public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { + Runnable flowTask = () -> internalConsumerFlow(consumer, additionalNumberOfMessages); + try { + dispatchMessagesThread.execute(flowTask); + } catch (RejectedExecutionException e) { + // Leave the permits pending so removal excludes this unapplied Flow. Never wait for the dispatcher + // monitor on the connection EventLoop, including while the broker is shutting down. + log.debug() + .attr("consumer", consumer) + .attr("executorShutdown", dispatchMessagesThread.isShutdown()) + .exception(e) + .log("Unable to schedule flow control update"); + } + } + + private void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { boolean connected; int updatedTotalAvailablePermits = 0; synchronized (this) { @@ -326,7 +342,7 @@ public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { .attr("totalAvailablePermits", updatedTotalAvailablePermits) .attr("additionalNumberOfMessages", additionalNumberOfMessages) .log("Trigger new read after receiving flow control message"); - readMoreEntriesAsync(); + readMoreEntries(); } /** diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java index 8cf7f051f86d4..31dae2cd8c175 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java @@ -35,6 +35,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.ExecutorService; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -283,6 +284,21 @@ private synchronized void clearComponentsAfterRemovedAllConsumers() { @Override public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { + Runnable flowTask = () -> internalConsumerFlow(consumer, additionalNumberOfMessages); + try { + dispatchMessagesThread.execute(flowTask); + } catch (RejectedExecutionException e) { + // Leave the permits pending so removal excludes this unapplied Flow. Never wait for the dispatcher + // monitor on the connection EventLoop, including while the broker is shutting down. + log.debug() + .attr("consumer", consumer) + .attr("executorShutdown", dispatchMessagesThread.isShutdown()) + .exception(e) + .log("Unable to schedule flow control update"); + } + } + + private void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { boolean connected; int updatedTotalAvailablePermits = 0; synchronized (this) { @@ -306,7 +322,7 @@ public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { .attr("totalAvailablePermits", updatedTotalAvailablePermits) .attr("additionalNumberOfMessages", additionalNumberOfMessages) .log("- Trigger new read after receiving flow control message with permits " + "after adding permits"); - readMoreEntriesAsync(); + readMoreEntries(); } /** diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java new file mode 100644 index 0000000000000..34f13b117e6b2 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java @@ -0,0 +1,318 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pulsar.broker.service.persistent; + +import static java.util.Collections.emptyMap; +import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Key_Shared; +import static org.apache.pulsar.common.api.proto.CommandSubscribe.SubType.Shared; +import static org.apache.pulsar.common.api.proto.KeySharedMode.AUTO_SPLIT; +import static org.apache.pulsar.common.protocol.Commands.DEFAULT_CONSUMER_EPOCH; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import io.netty.util.concurrent.EventExecutor; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.bookkeeper.common.util.OrderedExecutor; +import org.apache.bookkeeper.mledger.ManagedCursor; +import org.apache.bookkeeper.mledger.impl.ManagedCursorImpl; +import org.apache.pulsar.broker.service.BrokerService; +import org.apache.pulsar.broker.service.Consumer; +import org.apache.pulsar.broker.service.Dispatcher; +import org.apache.pulsar.broker.service.SharedPulsarBaseTest; +import org.apache.pulsar.broker.service.Subscription; +import org.apache.pulsar.broker.service.TransportCnx; +import org.apache.pulsar.client.api.MessageId; +import org.apache.pulsar.common.api.proto.CommandSubscribe.SubType; +import org.apache.pulsar.common.api.proto.KeySharedMeta; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +@Test(groups = "broker-api") +public class SharedDispatcherFlowThreadingTest extends SharedPulsarBaseTest { + + @DataProvider(name = "dispatcherVariants") + public Object[][] dispatcherVariants() { + return new Object[][] { + {false, Shared}, + {true, Shared}, + {false, Key_Shared}, + {true, Key_Shared} + }; + } + + @DataProvider(name = "dispatcherImplementations") + public Object[][] dispatcherImplementations() { + return new Object[][] {{false}, {true}}; + } + + @Test(dataProvider = "dispatcherVariants", timeOut = 30_000) + public void testFlowDoesNotBlockIoEventLoopWhileDispatchThreadHoldsMonitor( + boolean classic, SubType subType) throws Exception { + TestContext context = createTestContext(classic, subType); + CountDownLatch dispatchThreadHasMonitor = new CountDownLatch(1); + CountDownLatch releaseDispatchThread = new CountDownLatch(1); + Future dispatchTask = dispatchMessagesThread(context.dispatcher()).submit(() -> { + synchronized (context.dispatcher()) { + dispatchThreadHasMonitor.countDown(); + assertThat(releaseDispatchThread.await(5, TimeUnit.SECONDS)).isTrue(); + } + return null; + }); + + Future flowFuture = null; + Future heartbeatFuture = null; + try { + assertThat(dispatchThreadHasMonitor.await(5, TimeUnit.SECONDS)).isTrue(); + EventExecutor ioEventLoop = context.topic().getBrokerService().executor().next(); + flowFuture = ioEventLoop.submit(() -> context.consumer().flowPermits(100)); + heartbeatFuture = ioEventLoop.submit(() -> { }); + + // The Flow callback must only enqueue dispatcher work, allowing the next I/O event to run immediately. + heartbeatFuture.get(2, TimeUnit.SECONDS); + assertThat(context.consumer().getAvailablePermits()).isEqualTo(100); + assertThat(totalAvailablePermits(context.dispatcher())).isZero(); + } finally { + releaseDispatchThread.countDown(); + dispatchTask.get(5, TimeUnit.SECONDS); + if (flowFuture != null) { + flowFuture.get(5, TimeUnit.SECONDS); + } + if (heartbeatFuture != null) { + heartbeatFuture.get(5, TimeUnit.SECONDS); + } + } + + drainDispatchMessagesThread(context.dispatcher()); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(100); + assertThat(context.consumer().getAvailablePermitsForDispatcherRemoval()).isEqualTo(100); + } + + @Test(dataProvider = "dispatcherVariants", timeOut = 30_000) + public void testFlowAccountingAndReadRunOnDispatchMessagesThread( + boolean classic, SubType subType) throws Exception { + TestContext context = createTestContext(classic, subType); + CountDownLatch dispatchThreadBlocked = new CountDownLatch(1); + CountDownLatch releaseDispatchThread = new CountDownLatch(1); + AtomicReference dispatchThread = new AtomicReference<>(); + AtomicReference readThread = new AtomicReference<>(); + stubReadMoreEntries(context.dispatcher(), () -> readThread.set(Thread.currentThread())); + Future dispatchTask = dispatchMessagesThread(context.dispatcher()).submit(() -> { + dispatchThread.set(Thread.currentThread()); + dispatchThreadBlocked.countDown(); + assertThat(releaseDispatchThread.await(5, TimeUnit.SECONDS)).isTrue(); + return null; + }); + + try { + assertThat(dispatchThreadBlocked.await(5, TimeUnit.SECONDS)).isTrue(); + EventExecutor ioEventLoop = context.topic().getBrokerService().executor().next(); + ioEventLoop.submit(() -> context.consumer().flowPermits(100)).get(5, TimeUnit.SECONDS); + + assertThat(context.consumer().getAvailablePermits()).isEqualTo(100); + assertThat(totalAvailablePermits(context.dispatcher())).isZero(); + } finally { + releaseDispatchThread.countDown(); + dispatchTask.get(5, TimeUnit.SECONDS); + } + + drainDispatchMessagesThread(context.dispatcher()); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(100); + assertThat(context.consumer().getAvailablePermitsForDispatcherRemoval()).isEqualTo(100); + assertThat(readThread.get()).isSameAs(dispatchThread.get()); + } + + @Test(dataProvider = "dispatcherVariants", timeOut = 30_000) + public void testQueuedFlowCompletesPendingAfterConsumerRemoval(boolean classic, SubType subType) + throws Exception { + TestContext context = createTestContext(classic, subType); + CountDownLatch dispatchThreadBlocked = new CountDownLatch(1); + CountDownLatch releaseDispatchThread = new CountDownLatch(1); + AtomicInteger readTriggers = new AtomicInteger(); + stubReadMoreEntries(context.dispatcher(), readTriggers::incrementAndGet); + Future dispatchTask = dispatchMessagesThread(context.dispatcher()).submit(() -> { + dispatchThreadBlocked.countDown(); + assertThat(releaseDispatchThread.await(5, TimeUnit.SECONDS)).isTrue(); + return null; + }); + + try { + assertThat(dispatchThreadBlocked.await(5, TimeUnit.SECONDS)).isTrue(); + EventExecutor ioEventLoop = context.topic().getBrokerService().executor().next(); + ioEventLoop.submit(() -> context.consumer().flowPermits(100)).get(5, TimeUnit.SECONDS); + + assertThat(context.consumer().getAvailablePermits()).isEqualTo(100); + assertThat(context.consumer().getAvailablePermitsForDispatcherRemoval()).isZero(); + context.dispatcher().removeConsumer(context.consumer()); + assertThat(context.dispatcher().getConsumers()).isEmpty(); + assertThat(totalAvailablePermits(context.dispatcher())).isZero(); + } finally { + releaseDispatchThread.countDown(); + dispatchTask.get(5, TimeUnit.SECONDS); + } + + drainDispatchMessagesThread(context.dispatcher()); + assertThat(context.consumer().getAvailablePermitsForDispatcherRemoval()).isEqualTo(100); + assertThat(totalAvailablePermits(context.dispatcher())).isZero(); + assertThat(readTriggers).hasValue(0); + } + + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) + public void testRejectedFlowStaysPendingAndIsExcludedFromRemoval(boolean classic) throws Exception { + String topicName = newTopicName(); + String subscriptionName = "shared-sub"; + admin.topics().createNonPartitionedTopic(topicName); + PersistentTopic realTopic = (PersistentTopic) getTopic(topicName, false).join().orElseThrow(); + BrokerService isolatedBroker = spy(realTopic.getBrokerService()); + PersistentTopic isolatedTopic = spy(realTopic); + doReturn(isolatedBroker).when(isolatedTopic).getBrokerService(); + + ExecutorService rejectingLane = mock(ExecutorService.class); + doThrow(new RejectedExecutionException("test rejection")).when(rejectingLane).execute(any()); + when(rejectingLane.isShutdown()).thenReturn(true); + OrderedExecutor isolatedOrderedExecutor = mock(OrderedExecutor.class); + doReturn(rejectingLane).when(isolatedOrderedExecutor).chooseThread(); + doReturn(isolatedOrderedExecutor).when(isolatedBroker).getTopicOrderedExecutor(); + + ManagedCursor cursor = mock(ManagedCursorImpl.class); + when(cursor.getName()).thenReturn(subscriptionName); + when(cursor.isClosed()).thenReturn(true); + Subscription subscription = mock(PersistentSubscription.class); + when(subscription.getName()).thenReturn(subscriptionName); + when(subscription.getTopic()).thenReturn(isolatedTopic); + Dispatcher dispatcher = classic + ? new PersistentDispatcherMultipleConsumersClassic(isolatedTopic, cursor, subscription) + : new PersistentDispatcherMultipleConsumers(isolatedTopic, cursor, subscription); + doAnswer(invocation -> { + dispatcher.consumerFlow(invocation.getArgument(0), invocation.getArgument(1)); + return null; + }).when(subscription).consumerFlow(any(), anyInt()); + + TransportCnx cnx = mock(TransportCnx.class); + when(cnx.isActive()).thenReturn(true); + when(cnx.isWritable()).thenReturn(true); + Consumer remainingConsumer = new Consumer(subscription, Shared, topicName, 1, 0, "remaining", true, + cnx, "role", emptyMap(), false, null, MessageId.latest, DEFAULT_CONSUMER_EPOCH); + Consumer removedConsumer = new Consumer(subscription, Shared, topicName, 2, 0, "removed", true, + cnx, "role", emptyMap(), false, null, MessageId.latest, DEFAULT_CONSUMER_EPOCH); + dispatcher.addConsumer(remainingConsumer).join(); + dispatcher.addConsumer(removedConsumer).join(); + + removedConsumer.flowPermits(100); + + assertThat(removedConsumer.getAvailablePermits()).isEqualTo(100); + assertThat(removedConsumer.getAvailablePermitsForDispatcherRemoval()).isZero(); + assertThat(totalAvailablePermits(dispatcher)).isZero(); + dispatcher.removeConsumer(removedConsumer); + assertThat(dispatcher.getConsumers()).containsExactly(remainingConsumer); + assertThat(removedConsumer.getAvailablePermitsForDispatcherRemoval()).isZero(); + assertThat(totalAvailablePermits(dispatcher)).isZero(); + verify(rejectingLane).execute(any()); + } + + private TestContext createTestContext(boolean classic, SubType subType) throws Exception { + String topicName = newTopicName(); + String subscriptionName = "shared-sub"; + admin.topics().createNonPartitionedTopic(topicName); + PersistentTopic topic = (PersistentTopic) getTopic(topicName, false).join().orElseThrow(); + ManagedCursor cursor = mock(ManagedCursorImpl.class); + when(cursor.getName()).thenReturn(subscriptionName); + when(cursor.isClosed()).thenReturn(true); + Subscription subscription = mock(PersistentSubscription.class); + when(subscription.getName()).thenReturn(subscriptionName); + when(subscription.getTopic()).thenReturn(topic); + + Dispatcher dispatcher; + if (subType == Key_Shared) { + dispatcher = classic + ? spy(new PersistentStickyKeyDispatcherMultipleConsumersClassic( + topic, cursor, subscription, getConfig(), + new KeySharedMeta().setKeySharedMode(AUTO_SPLIT))) + : spy(new PersistentStickyKeyDispatcherMultipleConsumers( + topic, cursor, subscription, getConfig(), + new KeySharedMeta().setKeySharedMode(AUTO_SPLIT))); + } else { + dispatcher = classic + ? spy(new PersistentDispatcherMultipleConsumersClassic(topic, cursor, subscription)) + : spy(new PersistentDispatcherMultipleConsumers(topic, cursor, subscription)); + } + doAnswer(invocation -> { + dispatcher.consumerFlow(invocation.getArgument(0), invocation.getArgument(1)); + return null; + }).when(subscription).consumerFlow(any(), anyInt()); + + TransportCnx cnx = mock(TransportCnx.class); + when(cnx.isActive()).thenReturn(true); + when(cnx.isWritable()).thenReturn(true); + Consumer consumer = new Consumer(subscription, subType, topicName, 1, 0, "consumer", true, cnx, "role", + emptyMap(), false, new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), MessageId.latest, + DEFAULT_CONSUMER_EPOCH); + dispatcher.addConsumer(consumer).join(); + return new TestContext(topic, dispatcher, consumer); + } + + private static void stubReadMoreEntries(Dispatcher dispatcher, Runnable action) { + if (dispatcher instanceof PersistentDispatcherMultipleConsumers pip379Dispatcher) { + doAnswer(invocation -> { + action.run(); + return null; + }).when(pip379Dispatcher).readMoreEntries(); + } else { + PersistentDispatcherMultipleConsumersClassic classicDispatcher = + (PersistentDispatcherMultipleConsumersClassic) dispatcher; + doAnswer(invocation -> { + action.run(); + return null; + }).when(classicDispatcher).readMoreEntries(); + } + } + + private static void drainDispatchMessagesThread(Dispatcher dispatcher) throws Exception { + dispatchMessagesThread(dispatcher).submit(() -> { }).get(5, TimeUnit.SECONDS); + } + + private static ExecutorService dispatchMessagesThread(Dispatcher dispatcher) { + if (dispatcher instanceof PersistentDispatcherMultipleConsumers pip379Dispatcher) { + return pip379Dispatcher.dispatchMessagesThread; + } + return ((PersistentDispatcherMultipleConsumersClassic) dispatcher).dispatchMessagesThread; + } + + private static int totalAvailablePermits(Dispatcher dispatcher) { + if (dispatcher instanceof PersistentDispatcherMultipleConsumers pip379Dispatcher) { + return pip379Dispatcher.totalAvailablePermits; + } + return ((PersistentDispatcherMultipleConsumersClassic) dispatcher).totalAvailablePermits; + } + + private record TestContext(PersistentTopic topic, Dispatcher dispatcher, Consumer consumer) { + } +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java index 5d18b59c76b33..ced7dafacc950 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java @@ -33,7 +33,6 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; import io.netty.util.concurrent.EventExecutor; import io.netty.util.concurrent.ImmediateEventExecutor; @@ -45,8 +44,6 @@ import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; import org.apache.bookkeeper.mledger.Entry; import org.apache.bookkeeper.mledger.ManagedCursor; import org.apache.bookkeeper.mledger.impl.ManagedCursorImpl; @@ -135,15 +132,14 @@ public void testFlowCommandRaceWithConsumerRemovalDoesNotLosePermits(boolean cla remainingConsumer.flowPermits(10); removedConsumer.flowPermits(20); - drainBrokerWorkerGroup(context.topic()); + drainDispatchMessagesThread(context.dispatcher()); assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); - runWithFlowWaitingForDispatcher(context.dispatcher(), removedConsumer, 1_020, + runWithFlowQueuedForDispatcher(context.dispatcher(), removedConsumer, 1_020, () -> removedConsumer.flowPermits(1_000), () -> { assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); context.dispatcher().removeConsumer(removedConsumer); }); - drainBrokerWorkerGroup(context.topic()); assertThat(totalAvailablePermits(context.dispatcher())) .isEqualTo(remainingConsumer.getAvailablePermits()); @@ -187,10 +183,10 @@ public void testRemovalAppliesNegativeAccountedPermitBalance(boolean classic, Su remainingConsumer.flowPermits(10); secondRemainingConsumer.flowPermits(15); removedConsumer.flowPermits(20); - drainBrokerWorkerGroup(context.topic()); + drainDispatchMessagesThread(context.dispatcher()); assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(45); - runWithFlowWaitingForDispatcher(context.dispatcher(), removedConsumer, 120, + runWithFlowQueuedForDispatcher(context.dispatcher(), removedConsumer, 120, () -> removedConsumer.flowPermits(100), () -> { // Dispatch more than the removed consumer's 20 accounted permits. Its removal balance becomes // 70 available - 100 pending = -30 and must be applied as-is. @@ -203,7 +199,6 @@ public void testRemovalAppliesNegativeAccountedPermitBalance(boolean classic, Su context.dispatcher().removeConsumer(removedConsumer); assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(25); }); - drainBrokerWorkerGroup(context.topic()); // Keep two consumers after removal so a single-consumer read floor cannot mask dispatcher permit drift. assertThat(context.dispatcher().getConsumers()) @@ -221,7 +216,7 @@ public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits( Consumer removedConsumer = context.removedConsumer(); remainingConsumer.flowPermits(10); - drainBrokerWorkerGroup(context.topic()); + drainDispatchMessagesThread(context.dispatcher()); assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(10); ConsumerStatsImpl blockedStats = new ConsumerStatsImpl(); @@ -233,87 +228,16 @@ public void testBlockedFlowCommandRaceWithConsumerRemovalDoesNotLosePermits( removedConsumer.flowPermits(1_000); assertThat(removedConsumer.getAvailablePermits()).isZero(); - runWithFlowWaitingForDispatcher(context.dispatcher(), removedConsumer, 1_000, + runWithFlowQueuedForDispatcher(context.dispatcher(), removedConsumer, 1_000, () -> removedConsumer.updateBlockedConsumerOnUnackedMsgs(removedConsumer), () -> { assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(10); context.dispatcher().removeConsumer(removedConsumer); }); - drainBrokerWorkerGroup(context.topic()); assertThat(totalAvailablePermits(context.dispatcher())) .isEqualTo(remainingConsumer.getAvailablePermits()); } - @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) - public void testReadTriggerRunsOutsideDispatcherMonitorAfterConsumerRemoval(boolean classic) throws Exception { - TestContext context = createTestContext(classic, Shared, true); - Consumer remainingConsumer = context.remainingConsumer(); - Consumer removedConsumer = context.removedConsumer(); - AbstractPersistentDispatcherMultipleConsumers dispatcher = - (AbstractPersistentDispatcherMultipleConsumers) context.dispatcher(); - - remainingConsumer.flowPermits(10); - removedConsumer.flowPermits(20); - drainBrokerWorkerGroup(context.topic()); - assertThat(totalAvailablePermits(dispatcher)).isEqualTo(30); - - CountDownLatch readTriggerEntered = new CountDownLatch(1); - CountDownLatch allowReadTrigger = new CountDownLatch(1); - AtomicBoolean readTriggerHeldDispatcherMonitor = new AtomicBoolean(); - AtomicInteger readMoreEntriesCalls = new AtomicInteger(); - if (classic) { - doAnswer(invocation -> { - readMoreEntriesCalls.incrementAndGet(); - return invocation.callRealMethod(); - }).when((PersistentDispatcherMultipleConsumersClassic) dispatcher).readMoreEntries(); - } else { - doAnswer(invocation -> { - readMoreEntriesCalls.incrementAndGet(); - return invocation.callRealMethod(); - }).when((PersistentDispatcherMultipleConsumers) dispatcher).readMoreEntries(); - } - doAnswer(invocation -> { - readTriggerHeldDispatcherMonitor.set(Thread.holdsLock(dispatcher)); - readTriggerEntered.countDown(); - assertThat(allowReadTrigger.await(5, TimeUnit.SECONDS)).isTrue(); - return invocation.callRealMethod(); - }).when(dispatcher).readMoreEntriesAsync(); - - ExecutorService flowExecutor = Executors.newSingleThreadExecutor(); - Future flowFuture = null; - try { - flowFuture = flowExecutor.submit(() -> { - removedConsumer.flowPermits(1_000); - return null; - }); - assertThat(readTriggerEntered.await(5, TimeUnit.SECONDS)).isTrue(); - assertThat(readTriggerHeldDispatcherMonitor).isFalse(); - assertThat(flowFuture.isDone()).isFalse(); - assertThat(totalAvailablePermits(dispatcher)).isEqualTo(1_030); - - dispatcher.removeConsumer(removedConsumer); - assertThat(totalAvailablePermits(dispatcher)).isEqualTo(remainingConsumer.getAvailablePermits()); - // Classic removal invokes readMoreEntries synchronously. Count only the task submitted by the Flow - // trigger after the removal has completed. - readMoreEntriesCalls.set(0); - } finally { - allowReadTrigger.countDown(); - try { - if (flowFuture != null) { - flowFuture.get(5, TimeUnit.SECONDS); - } - } finally { - flowExecutor.shutdownNow(); - assertThat(flowExecutor.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); - } - } - drainBrokerWorkerGroup(context.topic()); - - assertThat(dispatcher.getConsumers()).containsExactly(remainingConsumer); - assertThat(totalAvailablePermits(dispatcher)).isEqualTo(remainingConsumer.getAvailablePermits()); - assertThat(readMoreEntriesCalls).hasValue(1); - } - @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) public void testRemainingConsumerCanContinueAfterFlowAndCloseRace(boolean classic) throws Exception { int receiverQueueSize = 10; @@ -360,13 +284,13 @@ public void testRemainingConsumerCanContinueAfterFlowAndCloseRace(boolean classi assertThat(removedBrokerConsumer.getAvailablePermits()).isEqualTo(receiverQueueSize); }); drainBrokerWorkerGroup(topic); + drainDispatchMessagesThread(dispatcher); assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); - // Follow the production subscription -> dispatcher lock order. The Flow thread first updates the - // consumer permits and then waits for the dispatcher monitor. Closing from this thread removes the - // consumer before the Flow update reaches the dispatcher. + // Follow the production subscription -> dispatcher lock order. Queue the Flow behind a lane barrier + // so closing from this thread removes the consumer before the Flow update reaches the dispatcher. synchronized (subscription) { - runWithFlowWaitingForDispatcher(dispatcher, removedBrokerConsumer, + runWithFlowQueuedForDispatcher(dispatcher, removedBrokerConsumer, receiverQueueSize + additionalFlowPermits, () -> removedBrokerConsumer.flowPermits(additionalFlowPermits), () -> { assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); @@ -383,6 +307,7 @@ public void testRemainingConsumerCanContinueAfterFlowAndCloseRace(boolean classi .containsKey(removedBrokerConsumer.consumerId())).isFalse()); removedClient.close(); drainBrokerWorkerGroup(topic); + drainDispatchMessagesThread(dispatcher); assertThat(dispatcher.getConsumers()).containsExactly(remainingBrokerConsumer); assertThat(totalAvailablePermits(dispatcher)) @@ -407,10 +332,6 @@ private TestContext createTestContext(boolean classic) throws Exception { } private TestContext createTestContext(boolean classic, SubType subType) throws Exception { - return createTestContext(classic, subType, false); - } - - private TestContext createTestContext(boolean classic, SubType subType, boolean spyDispatcher) throws Exception { String topicName = newTopicName(); String subscriptionName = "shared-sub"; admin.topics().createNonPartitionedTopic(topicName); @@ -435,21 +356,17 @@ topic, cursor, subscription, getConfig(), new KeySharedMeta().setKeySharedMode(A ? new PersistentDispatcherMultipleConsumersClassic(topic, cursor, subscription) : new PersistentDispatcherMultipleConsumers(topic, cursor, subscription); } - if (spyDispatcher) { - dispatcher = spy(dispatcher); - } - Dispatcher configuredDispatcher = dispatcher; doAnswer(invocation -> { - configuredDispatcher.consumerFlow(invocation.getArgument(0), invocation.getArgument(1)); + dispatcher.consumerFlow(invocation.getArgument(0), invocation.getArgument(1)); return null; }).when(subscription).consumerFlow(any(), anyInt()); Consumer remainingConsumer = createConsumer(subscription, subType, topicName, 1); Consumer removedConsumer = createConsumer(subscription, subType, topicName, 2); - configuredDispatcher.addConsumer(remainingConsumer).join(); - configuredDispatcher.addConsumer(removedConsumer).join(); + dispatcher.addConsumer(remainingConsumer).join(); + dispatcher.addConsumer(removedConsumer).join(); - return new TestContext(topic, configuredDispatcher, remainingConsumer, removedConsumer); + return new TestContext(topic, dispatcher, remainingConsumer, removedConsumer); } private Consumer createConsumer(Subscription subscription, SubType subType, String topicName, long consumerId) { @@ -497,27 +414,41 @@ private static void decrementTotalAvailablePermits(Dispatcher dispatcher, int pe } } - private static void runWithFlowWaitingForDispatcher(Dispatcher dispatcher, Consumer consumer, - int expectedConsumerPermits, CheckedRunnable flow, - CheckedRunnable whileFlowWaits) throws Exception { + private static void runWithFlowQueuedForDispatcher(Dispatcher dispatcher, Consumer consumer, + int expectedConsumerPermits, CheckedRunnable flow, + CheckedRunnable whileFlowIsQueued) throws Exception { + CountDownLatch dispatchThreadBlocked = new CountDownLatch(1); + CountDownLatch releaseDispatchThread = new CountDownLatch(1); + Future blockingTask = dispatchMessagesThread(dispatcher).submit(() -> { + dispatchThreadBlocked.countDown(); + assertThat(releaseDispatchThread.await(5, TimeUnit.SECONDS)).isTrue(); + return null; + }); ExecutorService flowExecutor = Executors.newSingleThreadExecutor(); try { - Future flowFuture; - synchronized (dispatcher) { - flowFuture = flowExecutor.submit(() -> { - flow.run(); - return null; - }); - Awaitility.await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> - assertThat(consumer.getAvailablePermits()).isEqualTo(expectedConsumerPermits)); - assertThat(flowFuture.isDone()).isFalse(); - whileFlowWaits.run(); - } + assertThat(dispatchThreadBlocked.await(5, TimeUnit.SECONDS)).isTrue(); + Future flowFuture = flowExecutor.submit(() -> { + flow.run(); + return null; + }); + // The Flow caller only queues dispatcher work and must return while the dispatch lane remains blocked. flowFuture.get(5, TimeUnit.SECONDS); + assertThat(consumer.getAvailablePermits()).isEqualTo(expectedConsumerPermits); + whileFlowIsQueued.run(); } finally { - flowExecutor.shutdownNow(); - assertThat(flowExecutor.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); + releaseDispatchThread.countDown(); + try { + blockingTask.get(5, TimeUnit.SECONDS); + } finally { + flowExecutor.shutdownNow(); + assertThat(flowExecutor.awaitTermination(5, TimeUnit.SECONDS)).isTrue(); + } } + drainDispatchMessagesThread(dispatcher); + } + + private static void drainDispatchMessagesThread(Dispatcher dispatcher) throws Exception { + dispatchMessagesThread(dispatcher).submit(() -> { }).get(5, TimeUnit.SECONDS); } private static void drainBrokerWorkerGroup(PersistentTopic topic) throws Exception { @@ -526,6 +457,13 @@ private static void drainBrokerWorkerGroup(PersistentTopic topic) throws Excepti } } + private static ExecutorService dispatchMessagesThread(Dispatcher dispatcher) { + if (dispatcher instanceof PersistentDispatcherMultipleConsumers pip379Dispatcher) { + return pip379Dispatcher.dispatchMessagesThread; + } + return ((PersistentDispatcherMultipleConsumersClassic) dispatcher).dispatchMessagesThread; + } + private static int totalAvailablePermits(Dispatcher dispatcher) { if (dispatcher instanceof PersistentDispatcherMultipleConsumers pip379Dispatcher) { return pip379Dispatcher.totalAvailablePermits; From f81fcae08655e98f69421a4dc202c1b5e4fbbad5 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sat, 29 Aug 2026 23:12:34 +0800 Subject: [PATCH 7/7] [fix][broker] Match queued Flow updates by Consumer identity A queued Flow task can outlive its Consumer and run after a replacement with the same protocol identity has joined. Consumer equality can then make the old task appear connected and add stale permits. Keep the existing ObjectSet field while retaining its ObjectHashSet backing instance for an O(1) identity check. Use that check for modern and classic queued Flow processing and cover same-identity replacement across Shared and Key_Shared. --- .../AbstractDispatcherMultipleConsumers.java | 15 +++++- ...PersistentDispatcherMultipleConsumers.java | 2 +- ...entDispatcherMultipleConsumersClassic.java | 2 +- .../SharedDispatcherFlowThreadingTest.java | 46 +++++++++++++++++++ 4 files changed, 62 insertions(+), 3 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractDispatcherMultipleConsumers.java index bec02e94c79ab..5d1dd116e85bd 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractDispatcherMultipleConsumers.java @@ -32,7 +32,8 @@ public abstract class AbstractDispatcherMultipleConsumers extends AbstractBaseDispatcher { protected final CopyOnWriteArrayList consumerList = new CopyOnWriteArrayList<>(); - protected final ObjectSet consumerSet = new ObjectHashSet<>(); + private final ObjectHashSet consumerSetImpl = new ObjectHashSet<>(); + protected final ObjectSet consumerSet = consumerSetImpl; protected volatile int currentConsumerRoundRobinIndex = 0; protected static final int FALSE = 0; @@ -58,6 +59,18 @@ public synchronized boolean canUnsubscribe(Consumer consumer) { return consumerList.size() == 1 && consumerSet.contains(consumer); } + /** + * Checks whether the exact Consumer instance is still connected. + * + *

This differs from {@link ObjectSet#contains(Object)}, which uses {@link Consumer#equals(Object)} and can + * match a replacement Consumer that reuses the same protocol identity. + * The caller must hold the dispatcher monitor while checking membership and acting on the result. + */ + protected final boolean containsConsumerInstance(Consumer consumer) { + int index = consumerSetImpl.indexOf(consumer); + return consumerSetImpl.indexExists(index) && consumerSetImpl.indexGet(index) == consumer; + } + public boolean isClosed() { return isClosed == TRUE; } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java index 5147e2b3b3096..552a7f952e891 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java @@ -323,7 +323,7 @@ private void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessa int updatedTotalAvailablePermits = 0; synchronized (this) { consumer.completePendingDispatcherFlow(additionalNumberOfMessages); - connected = consumerSet.contains(consumer); + connected = containsConsumerInstance(consumer); if (connected) { totalAvailablePermits += additionalNumberOfMessages; updatedTotalAvailablePermits = totalAvailablePermits; diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java index 31dae2cd8c175..7db4ce7a8e366 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java @@ -303,7 +303,7 @@ private void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessa int updatedTotalAvailablePermits = 0; synchronized (this) { consumer.completePendingDispatcherFlow(additionalNumberOfMessages); - connected = consumerSet.contains(consumer); + connected = containsConsumerInstance(consumer); if (connected) { totalAvailablePermits += additionalNumberOfMessages; updatedTotalAvailablePermits = totalAvailablePermits; diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java index 34f13b117e6b2..53809c34d9151 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java @@ -185,6 +185,52 @@ public void testQueuedFlowCompletesPendingAfterConsumerRemoval(boolean classic, assertThat(readTriggers).hasValue(0); } + @Test(dataProvider = "dispatcherVariants", timeOut = 30_000) + public void testQueuedFlowDoesNotApplyToEqualReplacementConsumer(boolean classic, SubType subType) + throws Exception { + TestContext context = createTestContext(classic, subType); + CountDownLatch dispatchThreadBlocked = new CountDownLatch(1); + CountDownLatch releaseDispatchThread = new CountDownLatch(1); + AtomicInteger readTriggers = new AtomicInteger(); + stubReadMoreEntries(context.dispatcher(), readTriggers::incrementAndGet); + Future dispatchTask = dispatchMessagesThread(context.dispatcher()).submit(() -> { + dispatchThreadBlocked.countDown(); + assertThat(releaseDispatchThread.await(5, TimeUnit.SECONDS)).isTrue(); + return null; + }); + Consumer original = context.consumer(); + Consumer replacement = new Consumer(original.getSubscription(), original.subType(), context.topic().getName(), + original.consumerId(), 0, original.consumerName(), true, original.cnx(), "role", emptyMap(), false, + new KeySharedMeta().setKeySharedMode(AUTO_SPLIT), MessageId.latest, DEFAULT_CONSUMER_EPOCH); + + try { + assertThat(dispatchThreadBlocked.await(5, TimeUnit.SECONDS)).isTrue(); + EventExecutor ioEventLoop = context.topic().getBrokerService().executor().next(); + ioEventLoop.submit(() -> original.flowPermits(100)).get(5, TimeUnit.SECONDS); + + assertThat(original.getAvailablePermits()).isEqualTo(100); + assertThat(original.getAvailablePermitsForDispatcherRemoval()).isZero(); + context.dispatcher().removeConsumer(original); + context.dispatcher().addConsumer(replacement).join(); + assertThat(replacement).isNotSameAs(original).isEqualTo(original); + assertThat(replacement.hashCode()).isEqualTo(original.hashCode()); + } finally { + releaseDispatchThread.countDown(); + try { + dispatchTask.get(5, TimeUnit.SECONDS); + } finally { + drainDispatchMessagesThread(context.dispatcher()); + } + } + + assertThat(context.dispatcher().getConsumers()).containsExactly(replacement); + assertThat(context.dispatcher().getConsumers().get(0)).isSameAs(replacement); + assertThat(original.getAvailablePermitsForDispatcherRemoval()).isEqualTo(100); + assertThat(replacement.getAvailablePermits()).isZero(); + assertThat(totalAvailablePermits(context.dispatcher())).isZero(); + assertThat(readTriggers).hasValue(0); + } + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) public void testRejectedFlowStaysPendingAndIsExcludedFromRemoval(boolean classic) throws Exception { String topicName = newTopicName();