From 2ce09a61cfbee9306511155f44ff126e79139765 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sun, 6 Sep 2026 08:25:27 +0800 Subject: [PATCH 1/2] [fix][broker] Prevent NPE when the last ACK races with sticky hash reassignment Assisted-by: Codex --- .../broker/service/DrainingHashesTracker.java | 84 +++-- .../DrainingHashesTrackerConcurrencyTest.java | 345 ++++++++++++++++++ 2 files changed, 403 insertions(+), 26 deletions(-) create mode 100644 pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerConcurrencyTest.java diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/DrainingHashesTracker.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/DrainingHashesTracker.java index 87ec1cf147055..61e4de3fe33bf 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/DrainingHashesTracker.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/DrainingHashesTracker.java @@ -106,6 +106,22 @@ boolean decrementRefCount() { return REF_COUNT_UPDATER.decrementAndGet(this) == 0; } + /** + * Decrements the reference count only when doing so cannot remove the entry. + * + * @return true if the reference count was decremented, false if the last reference must be handled separately + */ + boolean decrementRefCountIfGreaterThanOne() { + int current = refCount; + while (current > 1) { + if (REF_COUNT_UPDATER.compareAndSet(this, current, current - 1)) { + return true; + } + current = refCount; + } + return false; + } + /** * Increments the blocked count. */ @@ -271,11 +287,11 @@ public void addEntry(Consumer consumer, int stickyHash) { .attr("consumerName", consumer.consumerName()) .log("Draining hash incrementing consumer id: name"); } + // Publish the entry and increment its reference count atomically with respect to removal. + entry.incrementRefCount(); } finally { lock.writeLock().unlock(); } - // increment the reference count of the entry (applies to both new and existing entries) - entry.incrementRefCount(); // perform side-effects outside of the lock to reduce chances for deadlocks if (addedStatsForNewEntry != null) { @@ -333,35 +349,46 @@ public void reduceRefCount(Consumer consumer, int stickyHash, boolean closing) { if (entry == null) { return; } - if (entry.getConsumer() != consumer) { - throw new IllegalStateException( - "Consumer " + entry.getConsumer() + " is already draining hash " + stickyHash - + " in dispatcher " + dispatcherName + ". Same hash being used for consumer " + consumer - + "."); - } - if (entry.decrementRefCount()) { - log.debug() - .attr("dispatcher", dispatcherName) - .attr("hash", stickyHash) - .attr("consumerId", consumer.consumerId()) - .attr("consumerName", consumer.consumerName()) - .log("Draining hash removing consumer id: name"); - - DrainingHashEntry removed; - boolean notifyUnblocking = false; + boolean removed = false; + boolean notifyUnblocking = false; + // A non-final ACK can update the captured entry without serializing on the tracker write lock. + // If another path removes the entry concurrently, changing the detached old object is harmless. + if (entry.getConsumer() != consumer || !entry.decrementRefCountIfGreaterThanOne()) { lock.writeLock().lock(); try { - removed = drainingHashes.remove(stickyHash); - if (!closing && removed.isBlocking()) { - if (batchLevel > 0) { - unblockedWhileBatching = true; - } else { - notifyUnblocking = true; + // Serialize the final decrement with removal and verify that this is still the mapped generation. + if (drainingHashes.get(stickyHash) != entry) { + return; + } + if (entry.getConsumer() != consumer) { + throw new IllegalStateException( + "Consumer " + entry.getConsumer() + " is already draining hash " + stickyHash + + " in dispatcher " + dispatcherName + ". Same hash being used for consumer " + + consumer + "."); + } + removed = entry.decrementRefCount(); + if (removed) { + drainingHashes.remove(stickyHash); + if (!closing && entry.isBlocking()) { + if (batchLevel > 0) { + unblockedWhileBatching = true; + } else { + notifyUnblocking = true; + } } } } finally { lock.writeLock().unlock(); } + } + + if (removed) { + log.debug() + .attr("dispatcher", dispatcherName) + .attr("hash", stickyHash) + .attr("consumerId", consumer.consumerId()) + .attr("consumerName", consumer.consumerName()) + .log("Draining hash removing consumer id: name"); // perform side-effects outside of the lock to reduce chances for deadlocks @@ -413,12 +440,17 @@ public boolean shouldBlockStickyKeyHash(Consumer consumer, int stickyKeyHash) { .attr("consumer", entry.getConsumer()) .attr("refCount", entry.getRefCount()) .log("Hash has been reassigned to consumer. The draining hash entry will be removed."); + boolean removed; lock.writeLock().lock(); try { - drainingHashes.remove(stickyKeyHash, entry); + removed = drainingHashes.remove(stickyKeyHash, entry); } finally { lock.writeLock().unlock(); } + if (!removed) { + // Only the thread that removed this entry is responsible for clearing its stats. + return false; + } // update the consumer specific stats ConsumerDrainingHashesStats drainingHashesStats = @@ -492,4 +524,4 @@ public void updateConsumerStats(Consumer consumer, ConsumerStatsImpl consumerSta public void consumerRemoved(Consumer consumer) { consumerDrainingHashesStatsMap.remove(new ConsumerIdentityWrapper(consumer)); } -} \ No newline at end of file +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerConcurrencyTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerConcurrencyTest.java new file mode 100644 index 0000000000000..39ec48f90839e --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerConcurrencyTest.java @@ -0,0 +1,345 @@ +/* + * 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; + +import static org.apache.pulsar.broker.BrokerTestUtil.createMockConsumer; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.verifyNoMoreInteractions; +import io.netty.util.concurrent.DefaultThreadFactory; +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.AtomicReference; +import org.apache.pulsar.broker.service.DrainingHashesTracker.DrainingHashEntry; +import org.apache.pulsar.broker.service.DrainingHashesTracker.UnblockingHandler; +import org.apache.pulsar.common.policies.data.stats.ConsumerStatsImpl; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +/** + * Exercises concurrent tracker calls directly. Dispatcher topology changes have additional locking, + * so these tests alone do not establish that consumer churn can produce these interleavings. + */ +public class DrainingHashesTrackerConcurrencyTest { + @Test(timeOut = 30000) + public void lastAckShouldCompleteWhenOwnerReclaimsHash() throws Exception { + Consumer owner = createMockConsumer("owner"); + Consumer other = createMockConsumer("other"); + UnblockingHandler handler = mock(UnblockingHandler.class); + PausingTracker tracker = new PausingTracker(handler); + int hash = 1; + tracker.addEntry(owner, hash); + DrainingHashEntry entry = tracker.getEntry(hash); + assertThat(entry.getRefCount()).isEqualTo(1); + assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); + + LookupPause pause = tracker.pauseNextLookup(); + ExecutorService executor = newExecutor(); + try { + Future ack = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); + pause.awaitCaptured(); + assertThat(tracker.shouldBlockStickyKeyHash(owner, hash)).isFalse(); + assertThat(tracker.getEntry(hash)).isNull(); + pause.resume(); + + // Unpatched code throws ExecutionException caused by the NPE at removed.isBlocking(). + ack.get(10, TimeUnit.SECONDS); + + assertThat(entry.getRefCount()).isNotNegative(); + assertThat(tracker.getEntry(hash)).isNull(); + assertStats(tracker, owner, 0, 1, 0); + verifyNoInteractions(handler); + } finally { + pause.resume(); + shutdown(executor); + } + } + + @Test(timeOut = 30000) + public void ownerReassignmentShouldNotClearStatsAfterLastAckAlreadyRemovedHash() throws Exception { + Consumer owner = createMockConsumer("owner"); + Consumer other = createMockConsumer("other"); + UnblockingHandler handler = mock(UnblockingHandler.class); + PausingTracker tracker = new PausingTracker(handler); + int hash = 1; + tracker.addEntry(owner, hash); + assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); + + LookupPause pause = tracker.pauseNextLookup(); + ExecutorService executor = newExecutor(); + try { + Future reassignment = executor.submit(() -> tracker.shouldBlockStickyKeyHash(owner, hash)); + pause.awaitCaptured(); + tracker.reduceRefCount(owner, hash, false); + assertStats(tracker, owner, 0, 1, 0); + pause.resume(); + + assertThat(reassignment.get(10, TimeUnit.SECONDS)).isFalse(); + assertThat(tracker.getEntry(hash)).isNull(); + assertStats(tracker, owner, 0, 1, 0); + verify(handler).stickyKeyHashUnblocked(hash); + verifyNoMoreInteractions(handler); + } finally { + pause.resume(); + shutdown(executor); + } + } + + @DataProvider(name = "replacementOwners") + public Object[][] replacementOwners() { + return new Object[][] {{true}, {false}}; + } + + @Test(dataProvider = "replacementOwners", timeOut = 30000) + public void staleAckShouldNotRemoveReplacementEntry(boolean sameOwner) throws Exception { + Consumer owner = createMockConsumer("owner"); + Consumer replacementOwner = sameOwner ? owner : createMockConsumer("replacement-owner"); + UnblockingHandler handler = mock(UnblockingHandler.class); + PausingTracker tracker = new PausingTracker(handler); + int hash = 1; + tracker.addEntry(owner, hash); + DrainingHashEntry oldEntry = tracker.getEntry(hash); + + LookupPause pause = tracker.pauseNextLookup(); + ExecutorService executor = newExecutor(); + try { + Future ack = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); + pause.awaitCaptured(); + assertThat(tracker.shouldBlockStickyKeyHash(owner, hash)).isFalse(); + tracker.addEntry(replacementOwner, hash); + DrainingHashEntry replacement = tracker.getEntry(hash); + assertThat(replacement).isNotSameAs(oldEntry); + pause.resume(); + ack.get(10, TimeUnit.SECONDS); + + assertThat(tracker.getEntry(hash)).isSameAs(replacement); + assertThat(replacement.getConsumer()).isSameAs(replacementOwner); + assertThat(replacement.getRefCount()).isEqualTo(1); + assertThat(oldEntry.getRefCount()).isNotNegative(); + assertStats(tracker, owner, sameOwner ? 1 : 0, 1, sameOwner ? 1 : 0); + if (!sameOwner) { + assertStats(tracker, replacementOwner, 1, 0, 1); + } + verifyNoInteractions(handler); + } finally { + pause.resume(); + shutdown(executor); + } + } + + @Test(dataProvider = "replacementOwners", timeOut = 30000) + public void staleNonFinalAckShouldNotAffectReplacementEntry(boolean sameOwner) throws Exception { + Consumer owner = createMockConsumer("owner"); + Consumer replacementOwner = sameOwner ? owner : createMockConsumer("replacement-owner"); + UnblockingHandler handler = mock(UnblockingHandler.class); + PausingTracker tracker = new PausingTracker(handler); + int hash = 1; + tracker.addEntry(owner, hash); + tracker.addEntry(owner, hash); + DrainingHashEntry oldEntry = tracker.getEntry(hash); + assertThat(oldEntry.getRefCount()).isEqualTo(2); + + LookupPause pause = tracker.pauseNextLookup(); + ExecutorService executor = newExecutor(); + try { + Future ack = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); + pause.awaitCaptured(); + assertThat(tracker.shouldBlockStickyKeyHash(owner, hash)).isFalse(); + tracker.addEntry(replacementOwner, hash); + DrainingHashEntry replacement = tracker.getEntry(hash); + assertThat(replacement).isNotSameAs(oldEntry); + pause.resume(); + ack.get(10, TimeUnit.SECONDS); + + assertThat(oldEntry.getRefCount()).isEqualTo(1); + assertThat(tracker.getEntry(hash)).isSameAs(replacement); + assertThat(replacement.getConsumer()).isSameAs(replacementOwner); + assertThat(replacement.getRefCount()).isEqualTo(1); + assertStats(tracker, owner, sameOwner ? 1 : 0, 1, sameOwner ? 1 : 0); + if (!sameOwner) { + assertStats(tracker, replacementOwner, 1, 0, 1); + } + verifyNoInteractions(handler); + } finally { + pause.resume(); + shutdown(executor); + } + } + + @Test(timeOut = 30000) + public void concurrentReductionsShouldRemoveAndNotifyOnce() throws Exception { + Consumer owner = createMockConsumer("owner"); + Consumer other = createMockConsumer("other"); + UnblockingHandler handler = mock(UnblockingHandler.class); + DrainingHashesTracker tracker = new DrainingHashesTracker("dispatcher", handler); + int hash = 1; + int workerCount = 8; + int reductionsPerWorker = 8; + for (int i = 0; i < workerCount * reductionsPerWorker; i++) { + tracker.addEntry(owner, hash); + } + DrainingHashEntry entry = tracker.getEntry(hash); + assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); + + ExecutorService executor = newExecutor(workerCount); + CountDownLatch start = new CountDownLatch(1); + Future[] workers = new Future[workerCount]; + try { + for (int i = 0; i < workerCount; i++) { + workers[i] = executor.submit(() -> { + await(start); + for (int j = 0; j < reductionsPerWorker; j++) { + tracker.reduceRefCount(owner, hash, false); + } + }); + } + start.countDown(); + for (Future worker : workers) { + worker.get(10, TimeUnit.SECONDS); + } + + assertThat(entry.getRefCount()).isZero(); + assertThat(tracker.getEntry(hash)).isNull(); + assertStats(tracker, owner, 0, 1, 0); + verify(handler).stickyKeyHashUnblocked(hash); + verifyNoMoreInteractions(handler); + } finally { + start.countDown(); + shutdown(executor); + } + } + + @Test(timeOut = 30000) + public void twoCapturedLastAcksShouldOnlyRemoveAndNotifyOnce() throws Exception { + Consumer owner = createMockConsumer("owner"); + Consumer other = createMockConsumer("other"); + UnblockingHandler handler = mock(UnblockingHandler.class); + PausingTracker tracker = new PausingTracker(handler); + int hash = 1; + tracker.addEntry(owner, hash); + DrainingHashEntry entry = tracker.getEntry(hash); + assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); + + LookupPause firstPause = tracker.pauseNextLookup(); + LookupPause secondPause = null; + ExecutorService executor = newExecutor(); + try { + Future firstAck = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); + firstPause.awaitCaptured(); + secondPause = tracker.pauseNextLookup(); + Future secondAck = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); + secondPause.awaitCaptured(); + firstPause.resume(); + firstAck.get(10, TimeUnit.SECONDS); + secondPause.resume(); + secondAck.get(10, TimeUnit.SECONDS); + + assertThat(entry.getRefCount()).isZero(); + assertThat(tracker.getEntry(hash)).isNull(); + assertStats(tracker, owner, 0, 1, 0); + verify(handler).stickyKeyHashUnblocked(hash); + verifyNoMoreInteractions(handler); + } finally { + firstPause.resume(); + if (secondPause != null) { + secondPause.resume(); + } + shutdown(executor); + } + } + + private static ExecutorService newExecutor() { + return newExecutor(2); + } + + private static ExecutorService newExecutor(int threadCount) { + return Executors.newFixedThreadPool(threadCount, new DefaultThreadFactory("draining-hash-test")); + } + + private static void await(CountDownLatch latch) { + try { + assertThat(latch.await(10, TimeUnit.SECONDS)).as("Workers must start together").isTrue(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while waiting to start", e); + } + } + + private static void shutdown(ExecutorService executor) throws InterruptedException { + executor.shutdownNow(); + assertThat(executor.awaitTermination(10, TimeUnit.SECONDS)).as("Test threads must terminate").isTrue(); + } + + private static void assertStats(DrainingHashesTracker tracker, Consumer consumer, + int count, long clearedTotal, int unackedMessages) { + ConsumerStatsImpl stats = new ConsumerStatsImpl(); + tracker.updateConsumerStats(consumer, stats); + assertThat(stats.drainingHashesCount).isEqualTo(count); + assertThat(stats.drainingHashesClearedTotal).isEqualTo(clearedTotal); + assertThat(stats.drainingHashesUnackedMessages).isEqualTo(unackedMessages); + } + + private static class PausingTracker extends DrainingHashesTracker { + private final AtomicReference nextPause = new AtomicReference<>(); + + PausingTracker(UnblockingHandler handler) { + super("dispatcher", handler); + } + + LookupPause pauseNextLookup() { + LookupPause pause = new LookupPause(); + assertThat(nextPause.compareAndSet(null, pause)).as("Previous lookup pause must be consumed").isTrue(); + return pause; + } + + @Override + public DrainingHashEntry getEntry(int stickyKeyHash) { + DrainingHashEntry entry = super.getEntry(stickyKeyHash); + LookupPause pause = nextPause.getAndSet(null); + if (pause != null) { + pause.captured.countDown(); + try { + assertThat(pause.resumed.await(10, TimeUnit.SECONDS)).as("Paused lookup must resume").isTrue(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while waiting to resume lookup", e); + } + } + return entry; + } + } + + private static class LookupPause { + private final CountDownLatch captured = new CountDownLatch(1); + private final CountDownLatch resumed = new CountDownLatch(1); + + void awaitCaptured() throws InterruptedException { + assertThat(captured.await(10, TimeUnit.SECONDS)).as("Worker must capture the real entry").isTrue(); + } + + void resume() { + resumed.countDown(); + } + } +} From 873dbd703c325fe853d1ce76d171c5d27a389639 Mon Sep 17 00:00:00 2001 From: void-ptr974 Date: Sun, 6 Sep 2026 13:19:05 +0800 Subject: [PATCH 2/2] [test][broker] Cover draining hash reference lifecycle races Add deterministic coverage for entry publication, concurrent reference additions, and final ACK slow-path rechecks. Exercise nested batching and closing during concurrent removals using an injectable test lock. Validation: 57 tracker and pending-ack test invocations, 4 broker regression test invocations, and offline quickCheck passed. Targeted implementation mutations fail the three new tests as expected. Assisted-by: Codex --- .../broker/service/DrainingHashesTracker.java | 9 +- .../DrainingHashesTrackerConcurrencyTest.java | 250 +++++++++++++++--- 2 files changed, 226 insertions(+), 33 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/DrainingHashesTracker.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/DrainingHashesTracker.java index 61e4de3fe33bf..37414dbbfa32f 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/DrainingHashesTracker.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/DrainingHashesTracker.java @@ -19,6 +19,7 @@ package org.apache.pulsar.broker.service; import static org.apache.pulsar.broker.service.StickyKeyConsumerSelector.STICKY_KEY_HASH_NOT_SET; +import com.google.common.annotations.VisibleForTesting; import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap; import java.util.ArrayList; import java.util.Collections; @@ -49,7 +50,7 @@ public class DrainingHashesTracker { private final UnblockingHandler unblockingHandler; // optimize the memory consumption of the map by using primitive int keys private final Int2ObjectOpenHashMap drainingHashes = new Int2ObjectOpenHashMap<>(); - private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + private final ReentrantReadWriteLock lock; int batchLevel; boolean unblockedWhileBatching; private final Map consumerDrainingHashesStatsMap = @@ -241,8 +242,14 @@ public interface UnblockingHandler { } public DrainingHashesTracker(String dispatcherName, UnblockingHandler unblockingHandler) { + this(dispatcherName, unblockingHandler, new ReentrantReadWriteLock()); + } + + @VisibleForTesting + DrainingHashesTracker(String dispatcherName, UnblockingHandler unblockingHandler, ReentrantReadWriteLock lock) { this.dispatcherName = dispatcherName; this.unblockingHandler = unblockingHandler; + this.lock = lock; } /** diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerConcurrencyTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerConcurrencyTest.java index 39ec48f90839e..88b6faf85a8e8 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerConcurrencyTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/DrainingHashesTrackerConcurrencyTest.java @@ -31,6 +31,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantReadWriteLock; import org.apache.pulsar.broker.service.DrainingHashesTracker.DrainingHashEntry; import org.apache.pulsar.broker.service.DrainingHashesTracker.UnblockingHandler; import org.apache.pulsar.common.policies.data.stats.ConsumerStatsImpl; @@ -42,6 +43,119 @@ * so these tests alone do not establish that consumer churn can produce these interleavings. */ public class DrainingHashesTrackerConcurrencyTest { + @Test(timeOut = 30000) + public void newlyPublishedEntryShouldAlreadyHaveItsFirstReference() throws Exception { + Consumer owner = createMockConsumer("owner"); + UnblockingHandler handler = mock(UnblockingHandler.class); + PausingReadWriteLock lock = new PausingReadWriteLock(); + DrainingHashesTracker tracker = new DrainingHashesTracker("dispatcher", handler, lock); + int hash = 1; + OperationPause pause = lock.pauseNextWriteUnlock(); + ExecutorService executor = newExecutor(); + try { + Future addition = executor.submit(() -> tracker.addEntry(owner, hash)); + pause.awaitCaptured(); + + // The map is now readable, but addEntry has not returned from writeLock().unlock(). + // Moving the increment outside the lock would publish an entry with zero references. + DrainingHashEntry entry = tracker.getEntry(hash); + assertThat(entry).isNotNull(); + assertThat(entry.getRefCount()).isEqualTo(1); + pause.resume(); + addition.get(10, TimeUnit.SECONDS); + assertStats(tracker, owner, 1, 0, 1); + + tracker.reduceRefCount(owner, hash, false); + assertThat(entry.getRefCount()).isZero(); + assertThat(tracker.getEntry(hash)).isNull(); + assertStats(tracker, owner, 0, 1, 0); + verifyNoInteractions(handler); + } finally { + pause.resume(); + shutdown(executor); + } + } + + @Test(timeOut = 30000) + public void addedReferenceShouldSurviveAckAfterWriteUnlock() throws Exception { + Consumer owner = createMockConsumer("owner"); + Consumer other = createMockConsumer("other"); + UnblockingHandler handler = mock(UnblockingHandler.class); + PausingReadWriteLock lock = new PausingReadWriteLock(); + DrainingHashesTracker tracker = new DrainingHashesTracker("dispatcher", handler, lock); + int hash = 1; + tracker.addEntry(owner, hash); + DrainingHashEntry entry = tracker.getEntry(hash); + assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); + + OperationPause pause = lock.pauseNextWriteUnlock(); + ExecutorService executor = newExecutor(); + try { + Future addition = executor.submit(() -> tracker.addEntry(owner, hash)); + pause.awaitCaptured(); + // ACK while addEntry is still returning from unlock. The new reference must already be counted. + tracker.reduceRefCount(owner, hash, false); + pause.resume(); + addition.get(10, TimeUnit.SECONDS); + + assertThat(tracker.getEntry(hash)).isSameAs(entry); + assertThat(entry.getRefCount()).isEqualTo(1); + assertStats(tracker, owner, 1, 0, 1); + verifyNoInteractions(handler); + + tracker.reduceRefCount(owner, hash, false); + assertThat(entry.getRefCount()).isZero(); + assertThat(tracker.getEntry(hash)).isNull(); + assertStats(tracker, owner, 0, 1, 0); + verify(handler).stickyKeyHashUnblocked(hash); + verifyNoMoreInteractions(handler); + } finally { + pause.resume(); + shutdown(executor); + } + } + + @Test(timeOut = 30000) + public void slowPathAckShouldPreserveReferenceAddedBeforeWriteLock() throws Exception { + Consumer owner = createMockConsumer("owner"); + Consumer other = createMockConsumer("other"); + UnblockingHandler handler = mock(UnblockingHandler.class); + PausingReadWriteLock lock = new PausingReadWriteLock(); + DrainingHashesTracker tracker = new DrainingHashesTracker("dispatcher", handler, lock); + int hash = 1; + tracker.addEntry(owner, hash); + DrainingHashEntry entry = tracker.getEntry(hash); + assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); + + OperationPause pause = lock.pauseNextWriteLock(); + ExecutorService executor = newExecutor(); + try { + Future ack = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); + pause.awaitCaptured(); + // The fast path observed one reference and declined to decrement it, but the ACK has not locked yet. + assertThat(entry.getRefCount()).isEqualTo(1); + tracker.addEntry(owner, hash); + assertThat(entry.getRefCount()).isEqualTo(2); + pause.resume(); + ack.get(10, TimeUnit.SECONDS); + + assertThat(tracker.getEntry(hash)).isSameAs(entry); + assertThat(entry.getRefCount()).isEqualTo(1); + assertStats(tracker, owner, 1, 0, 1); + verifyNoInteractions(handler); + + tracker.reduceRefCount(owner, hash, false); + assertThat(entry.getRefCount()).isZero(); + assertThat(tracker.getEntry(hash)).isNull(); + assertStats(tracker, owner, 0, 1, 0); + verify(handler).stickyKeyHashUnblocked(hash); + verifyNoMoreInteractions(handler); + } finally { + pause.resume(); + shutdown(executor); + } + } + @Test(timeOut = 30000) public void lastAckShouldCompleteWhenOwnerReclaimsHash() throws Exception { Consumer owner = createMockConsumer("owner"); @@ -54,7 +168,7 @@ public void lastAckShouldCompleteWhenOwnerReclaimsHash() throws Exception { assertThat(entry.getRefCount()).isEqualTo(1); assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); - LookupPause pause = tracker.pauseNextLookup(); + OperationPause pause = tracker.pauseNextLookup(); ExecutorService executor = newExecutor(); try { Future ack = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); @@ -86,7 +200,7 @@ public void ownerReassignmentShouldNotClearStatsAfterLastAckAlreadyRemovedHash() tracker.addEntry(owner, hash); assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); - LookupPause pause = tracker.pauseNextLookup(); + OperationPause pause = tracker.pauseNextLookup(); ExecutorService executor = newExecutor(); try { Future reassignment = executor.submit(() -> tracker.shouldBlockStickyKeyHash(owner, hash)); @@ -121,7 +235,7 @@ public void staleAckShouldNotRemoveReplacementEntry(boolean sameOwner) throws Ex tracker.addEntry(owner, hash); DrainingHashEntry oldEntry = tracker.getEntry(hash); - LookupPause pause = tracker.pauseNextLookup(); + OperationPause pause = tracker.pauseNextLookup(); ExecutorService executor = newExecutor(); try { Future ack = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); @@ -160,7 +274,7 @@ public void staleNonFinalAckShouldNotAffectReplacementEntry(boolean sameOwner) t DrainingHashEntry oldEntry = tracker.getEntry(hash); assertThat(oldEntry.getRefCount()).isEqualTo(2); - LookupPause pause = tracker.pauseNextLookup(); + OperationPause pause = tracker.pauseNextLookup(); ExecutorService executor = newExecutor(); try { Future ack = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); @@ -187,8 +301,13 @@ public void staleNonFinalAckShouldNotAffectReplacementEntry(boolean sameOwner) t } } - @Test(timeOut = 30000) - public void concurrentReductionsShouldRemoveAndNotifyOnce() throws Exception { + @DataProvider(name = "removalModes") + public Object[][] removalModes() { + return new Object[][] {{false, false}, {true, false}, {false, true}, {true, true}}; + } + + @Test(dataProvider = "removalModes", timeOut = 30000) + public void concurrentReductionsShouldRemoveAndNotifyOnce(boolean batching, boolean closing) throws Exception { Consumer owner = createMockConsumer("owner"); Consumer other = createMockConsumer("other"); UnblockingHandler handler = mock(UnblockingHandler.class); @@ -202,6 +321,11 @@ public void concurrentReductionsShouldRemoveAndNotifyOnce() throws Exception { DrainingHashEntry entry = tracker.getEntry(hash); assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); + if (batching) { + tracker.startBatch(); + tracker.startBatch(); + } + ExecutorService executor = newExecutor(workerCount); CountDownLatch start = new CountDownLatch(1); Future[] workers = new Future[workerCount]; @@ -210,7 +334,7 @@ public void concurrentReductionsShouldRemoveAndNotifyOnce() throws Exception { workers[i] = executor.submit(() -> { await(start); for (int j = 0; j < reductionsPerWorker; j++) { - tracker.reduceRefCount(owner, hash, false); + tracker.reduceRefCount(owner, hash, closing); } }); } @@ -222,16 +346,15 @@ public void concurrentReductionsShouldRemoveAndNotifyOnce() throws Exception { assertThat(entry.getRefCount()).isZero(); assertThat(tracker.getEntry(hash)).isNull(); assertStats(tracker, owner, 0, 1, 0); - verify(handler).stickyKeyHashUnblocked(hash); - verifyNoMoreInteractions(handler); + assertUnblocking(tracker, handler, hash, batching, closing); } finally { start.countDown(); shutdown(executor); } } - @Test(timeOut = 30000) - public void twoCapturedLastAcksShouldOnlyRemoveAndNotifyOnce() throws Exception { + @Test(dataProvider = "removalModes", timeOut = 30000) + public void twoCapturedLastAcksShouldOnlyRemoveAndNotifyOnce(boolean batching, boolean closing) throws Exception { Consumer owner = createMockConsumer("owner"); Consumer other = createMockConsumer("other"); UnblockingHandler handler = mock(UnblockingHandler.class); @@ -241,14 +364,19 @@ public void twoCapturedLastAcksShouldOnlyRemoveAndNotifyOnce() throws Exception DrainingHashEntry entry = tracker.getEntry(hash); assertThat(tracker.shouldBlockStickyKeyHash(other, hash)).isTrue(); - LookupPause firstPause = tracker.pauseNextLookup(); - LookupPause secondPause = null; + if (batching) { + tracker.startBatch(); + tracker.startBatch(); + } + + OperationPause firstPause = tracker.pauseNextLookup(); + OperationPause secondPause = null; ExecutorService executor = newExecutor(); try { - Future firstAck = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); + Future firstAck = executor.submit(() -> tracker.reduceRefCount(owner, hash, closing)); firstPause.awaitCaptured(); secondPause = tracker.pauseNextLookup(); - Future secondAck = executor.submit(() -> tracker.reduceRefCount(owner, hash, false)); + Future secondAck = executor.submit(() -> tracker.reduceRefCount(owner, hash, closing)); secondPause.awaitCaptured(); firstPause.resume(); firstAck.get(10, TimeUnit.SECONDS); @@ -258,8 +386,7 @@ public void twoCapturedLastAcksShouldOnlyRemoveAndNotifyOnce() throws Exception assertThat(entry.getRefCount()).isZero(); assertThat(tracker.getEntry(hash)).isNull(); assertStats(tracker, owner, 0, 1, 0); - verify(handler).stickyKeyHashUnblocked(hash); - verifyNoMoreInteractions(handler); + assertUnblocking(tracker, handler, hash, batching, closing); } finally { firstPause.resume(); if (secondPause != null) { @@ -269,6 +396,22 @@ public void twoCapturedLastAcksShouldOnlyRemoveAndNotifyOnce() throws Exception } } + private static void assertUnblocking(DrainingHashesTracker tracker, UnblockingHandler handler, + int hash, boolean batching, boolean closing) { + if (batching) { + verifyNoInteractions(handler); + tracker.endBatch(); + verifyNoInteractions(handler); + tracker.endBatch(); + } + if (closing) { + verifyNoInteractions(handler); + } else { + verify(handler).stickyKeyHashUnblocked(batching ? -1 : hash); + verifyNoMoreInteractions(handler); + } + } + private static ExecutorService newExecutor() { return newExecutor(2); } @@ -301,14 +444,14 @@ private static void assertStats(DrainingHashesTracker tracker, Consumer consumer } private static class PausingTracker extends DrainingHashesTracker { - private final AtomicReference nextPause = new AtomicReference<>(); + private final AtomicReference nextPause = new AtomicReference<>(); PausingTracker(UnblockingHandler handler) { super("dispatcher", handler); } - LookupPause pauseNextLookup() { - LookupPause pause = new LookupPause(); + OperationPause pauseNextLookup() { + OperationPause pause = new OperationPause(); assertThat(nextPause.compareAndSet(null, pause)).as("Previous lookup pause must be consumed").isTrue(); return pause; } @@ -316,26 +459,69 @@ LookupPause pauseNextLookup() { @Override public DrainingHashEntry getEntry(int stickyKeyHash) { DrainingHashEntry entry = super.getEntry(stickyKeyHash); - LookupPause pause = nextPause.getAndSet(null); - if (pause != null) { - pause.captured.countDown(); - try { - assertThat(pause.resumed.await(10, TimeUnit.SECONDS)).as("Paused lookup must resume").isTrue(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new AssertionError("Interrupted while waiting to resume lookup", e); - } - } + pauseIfRequested(nextPause); return entry; } } - private static class LookupPause { + private static class PausingReadWriteLock extends ReentrantReadWriteLock { + private final AtomicReference nextWriteLockPause = new AtomicReference<>(); + private final AtomicReference nextWriteUnlockPause = new AtomicReference<>(); + private final WriteLock pausingWriteLock = new WriteLock(this) { + @Override + public void lock() { + pauseIfRequested(nextWriteLockPause); + super.lock(); + } + + @Override + public void unlock() { + super.unlock(); + pauseIfRequested(nextWriteUnlockPause); + } + }; + + @Override + public WriteLock writeLock() { + return pausingWriteLock; + } + + OperationPause pauseNextWriteLock() { + OperationPause pause = new OperationPause(); + assertThat(nextWriteLockPause.compareAndSet(null, pause)).isTrue(); + return pause; + } + + OperationPause pauseNextWriteUnlock() { + OperationPause pause = new OperationPause(); + assertThat(nextWriteUnlockPause.compareAndSet(null, pause)).isTrue(); + return pause; + } + } + + private static void pauseIfRequested(AtomicReference nextPause) { + OperationPause pause = nextPause.getAndSet(null); + if (pause != null) { + pause.pause(); + } + } + + private static class OperationPause { private final CountDownLatch captured = new CountDownLatch(1); private final CountDownLatch resumed = new CountDownLatch(1); void awaitCaptured() throws InterruptedException { - assertThat(captured.await(10, TimeUnit.SECONDS)).as("Worker must capture the real entry").isTrue(); + assertThat(captured.await(10, TimeUnit.SECONDS)).as("Worker must reach the pause").isTrue(); + } + + void pause() { + captured.countDown(); + try { + assertThat(resumed.await(10, TimeUnit.SECONDS)).as("Paused operation must resume").isTrue(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while waiting to resume operation", e); + } } void resume() {