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/Consumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java index a7f296536881e..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 @@ -117,6 +117,18 @@ public class Consumer { private static final AtomicIntegerFieldUpdater MESSAGE_PERMITS_UPDATER = AtomicIntegerFieldUpdater.newUpdater(Consumer.class, "messagePermits"); 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 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. + */ + 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 +928,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 +955,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 +967,57 @@ 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) { + 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 + * its total or ignore them because the consumer has already been removed. + */ + public void completePendingDispatcherFlow(int additionalNumberOfPermits) { + synchronized (flowPermitAccountingLock) { + // Preserve the accounting delta across signed int wrap, matching the other permit counters. + pendingDispatcherFlowPermits -= 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 requires every dispatcher + * 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 + * 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) { + 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..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 @@ -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; @@ -258,9 +259,12 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE notifyAddedToReplay.setTrue(); } }); - totalAvailablePermits -= consumer.getAvailablePermits(); + // 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() - .attr("diffAvailablePermits", consumer.getAvailablePermits()) + .attr("availablePermits", availablePermits) .attr("totalAvailablePermits", totalAvailablePermits) .log("Decreased totalAvailablePermits"); if (notifyAddedToReplay.booleanValue()) { @@ -300,27 +304,45 @@ protected synchronized void clearComponentsAfterRemovedAllConsumers() { @Override public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { - topic.getBrokerService().executor().execute(() -> { - internalConsumerFlow(consumer, 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 synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { - if (!consumerSet.contains(consumer)) { + private void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { + boolean connected; + int updatedTotalAvailablePermits = 0; + synchronized (this) { + consumer.completePendingDispatcherFlow(additionalNumberOfMessages); + connected = containsConsumerInstance(consumer); + if (connected) { + totalAvailablePermits += additionalNumberOfMessages; + updatedTotalAvailablePermits = totalAvailablePermits; + } + } + + 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(); + 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 3de50042b592d..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 @@ -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; @@ -244,9 +245,12 @@ public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceE consumer.getPendingAcks().forEach((ledgerId, entryId, batchSize, stickyKeyHash) -> { addMessageToReplay(ledgerId, entryId, stickyKeyHash); }); - totalAvailablePermits -= consumer.getAvailablePermits(); + // 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() - .attr("availablePermits", consumer.getAvailablePermits()) + .attr("availablePermits", availablePermits) .attr("totalAvailablePermits", totalAvailablePermits) .log("Decreased totalAvailablePermits by in PersistentDispatcherMultipleConsumers. " + "New dispatcher permit count is"); @@ -280,24 +284,42 @@ private synchronized void clearComponentsAfterRemovedAllConsumers() { @Override public void consumerFlow(Consumer consumer, int additionalNumberOfMessages) { - topic.getBrokerService().executor().execute(() -> { - internalConsumerFlow(consumer, 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 synchronized void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { - if (!consumerSet.contains(consumer)) { + private void internalConsumerFlow(Consumer consumer, int additionalNumberOfMessages) { + boolean connected; + int updatedTotalAvailablePermits = 0; + synchronized (this) { + consumer.completePendingDispatcherFlow(additionalNumberOfMessages); + connected = containsConsumerInstance(consumer); + if (connected) { + totalAvailablePermits += additionalNumberOfMessages; + updatedTotalAvailablePermits = totalAvailablePermits; + } + } + + 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(); @@ -307,6 +329,7 @@ private synchronized void internalConsumerFlow(Consumer consumer, int additional * 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/SharedDispatcherFlowThreadingTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java new file mode 100644 index 0000000000000..53809c34d9151 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherFlowThreadingTest.java @@ -0,0 +1,364 @@ +/* + * 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 = "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(); + 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 new file mode 100644 index 0000000000000..ced7dafacc950 --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/persistent/SharedDispatcherPermitAccountingTest.java @@ -0,0 +1,489 @@ +/* + * 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.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.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +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.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; +import org.apache.pulsar.common.policies.data.stats.ConsumerStatsImpl; +import org.awaitility.Awaitility; +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}}; + } + + @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(); + + remainingConsumer.flowPermits(10); + removedConsumer.flowPermits(20); + drainDispatchMessagesThread(context.dispatcher()); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); + + runWithFlowQueuedForDispatcher(context.dispatcher(), removedConsumer, 1_020, + () -> removedConsumer.flowPermits(1_000), () -> { + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(30); + context.dispatcher().removeConsumer(removedConsumer); + }); + + assertThat(totalAvailablePermits(context.dispatcher())) + .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 = "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(), subType, context.topic().getName(), 3); + context.dispatcher().addConsumer(secondRemainingConsumer).join(); + + remainingConsumer.flowPermits(10); + secondRemainingConsumer.flowPermits(15); + removedConsumer.flowPermits(20); + drainDispatchMessagesThread(context.dispatcher()); + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(45); + + 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. + 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); + }); + + // 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 = "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(); + + remainingConsumer.flowPermits(10); + drainDispatchMessagesThread(context.dispatcher()); + 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(); + + runWithFlowQueuedForDispatcher(context.dispatcher(), removedConsumer, 1_000, + () -> removedConsumer.updateBlockedConsumerOnUnackedMsgs(removedConsumer), () -> { + assertThat(totalAvailablePermits(context.dispatcher())).isEqualTo(10); + context.dispatcher().removeConsumer(removedConsumer); + }); + + assertThat(totalAvailablePermits(context.dispatcher())) + .isEqualTo(remainingConsumer.getAvailablePermits()); + } + + @Test(dataProvider = "dispatcherImplementations", timeOut = 30_000) + public void testRemainingConsumerCanContinueAfterFlowAndCloseRace(boolean classic) throws Exception { + int receiverQueueSize = 10; + int messagesToConsume = receiverQueueSize * 2; + int additionalFlowPermits = 1_000; + String topicName = newTopicName(); + String subscriptionName = "shared-sub"; + 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); + drainDispatchMessagesThread(dispatcher); + assertThat(totalAvailablePermits(dispatcher)).isEqualTo(receiverQueueSize * 2); + + // 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) { + runWithFlowQueuedForDispatcher(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); + drainDispatchMessagesThread(dispatcher); + + 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); + } + } + } finally { + getConfig().setSubscriptionSharedUseClassicPersistentImplementation(previousClassicSetting); + } + } + + 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); + + 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; + 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, 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, 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); + 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); + 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)), stickyKeyHashes, batchSizes, batchIndexesAcks, + permits, 0, 0, redeliveryTracker, DEFAULT_CONSUMER_EPOCH).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 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 { + 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 { + 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 { + for (EventExecutor eventExecutor : topic.getBrokerService().executor()) { + eventExecutor.submit(() -> { }).sync(); + } + } + + 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 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) { + } + + @FunctionalInterface + private interface CheckedRunnable { + void run() throws Exception; + } +}