[fix][broker] Prevent NPE when the last ACK races with sticky hash reassignment

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<DrainingHashEntry> drainingHashes = new Int2ObjectOpenHashMap<>();
-    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
+    private final ReentrantReadWriteLock lock;
     int batchLevel;
     boolean unblockedWhileBatching;
     private final Map<ConsumerIdentityWrapper, ConsumerDrainingHashesStats> consumerDrainingHashesStatsMap =
@@ -106,6 +107,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.
          */
@@ -225,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;
     }
 
     /**
@@ -271,11 +294,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 +356,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 +447,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 +531,4 @@ public void updateConsumerStats(Consumer consumer, ConsumerStatsImpl consumerSta
     public void consumerRemoved(Consumer consumer) {
         consumerDrainingHashesStatsMap.remove(new ConsumerIdentityWrapper(consumer));
     }
-}
+}