[fix][broker] Fix ownership-generation races in OwnershipCache removeOwnership and lock-expiry cleanup

pulsar-broker/src/main/java/org/apache/pulsar/broker/namespace/OwnedBundle.java

@@ -29,13 +29,23 @@
 import org.apache.pulsar.broker.PulsarService;
 import org.apache.pulsar.common.naming.NamespaceBundle;
 import org.apache.pulsar.common.util.FutureUtil;
+import org.apache.pulsar.metadata.api.coordination.ResourceLock;
 
 @EqualsAndHashCode
 @ToString
 @CustomLog
 public class OwnedBundle {
     private final NamespaceBundle bundle;
 
+    /**
+     * The resource lock acquired for this ownership. Ties this instance to a single ownership generation: a
+     * bundle can be re-acquired (new lock, new {@link OwnedBundle}) after this instance's lock expired, and
+     * cleanup done through this instance must never touch the newer generation.
+     */
+    @ToString.Exclude
+    @EqualsAndHashCode.Exclude
+    private final ResourceLock<NamespaceEphemeralData> resourceLock;
+
     /**
      * {@link #nsLock} is used to protect read/write access to {@link #isActive} flag and the corresponding code section
      * based on {@link #isActive} flag.
@@ -55,9 +65,8 @@ public class OwnedBundle {
      * @param suName
      */
     public OwnedBundle(NamespaceBundle suName) {
-        this.bundle = suName;
-        IS_ACTIVE_UPDATER.set(this, TRUE);
-    };
+        this(suName, true);
+    }
 
     /**
      * Constructor to allow set initial active flag.
@@ -67,9 +76,20 @@ public OwnedBundle(NamespaceBundle suName) {
      */
     public OwnedBundle(NamespaceBundle suName, boolean active) {
         this.bundle = suName;
+        this.resourceLock = null;
         IS_ACTIVE_UPDATER.set(this, active ? TRUE : FALSE);
     }
 
+    OwnedBundle(NamespaceBundle suName, ResourceLock<NamespaceEphemeralData> resourceLock) {
+        this.bundle = suName;
+        this.resourceLock = resourceLock;
+        IS_ACTIVE_UPDATER.set(this, TRUE);
+    }
+
+    ResourceLock<NamespaceEphemeralData> getResourceLock() {
+        return resourceLock;
+    }
+
     /**
      * Access to the namespace name.
      *
@@ -131,10 +151,10 @@ public CompletableFuture<Void> handleUnloadRequest(PulsarService pulsar, long ti
         log.info().attr("ownership", this.bundle).log("Disabling ownership");
 
         // close topics forcefully
-        return pulsar.getNamespaceService().getOwnershipCache()
-                .updateBundleState(this.bundle, false)
-                .thenCompose(v -> pulsar.getBrokerService().unloadServiceUnit(
-                        bundle, true, closeWithoutWaitingClientDisconnect, timeout, timeoutUnit))
+        // isActive was already flipped to false above; looking the bundle up in the ownership cache here could
+        // deactivate a newer OwnedBundle that re-acquired the bundle after this instance's lock expired.
+        return pulsar.getBrokerService().unloadServiceUnit(
+                        bundle, true, closeWithoutWaitingClientDisconnect, timeout, timeoutUnit)
                 .handle((numUnloadedTopics, ex) -> {
                     if (ex != null) {
                         // ignore topic-close failure to unload bundle
@@ -150,8 +170,8 @@ public CompletableFuture<Void> handleUnloadRequest(PulsarService pulsar, long ti
                     return null;
                 })
                 .thenCompose(v -> {
-                    // delete ownership node on zk
-                    return pulsar.getNamespaceService().getOwnershipCache().removeOwnership(bundle);
+                    // delete ownership node on zk, but only for this instance's ownership generation
+                    return pulsar.getNamespaceService().getOwnershipCache().removeOwnership(this);
                 }).whenComplete((ignored, ex) -> {
                     double unloadBundleTime = TimeUnit.NANOSECONDS
                             .toMillis((System.nanoTime() - unloadBundleStartTime));

pulsar-broker/src/main/java/org/apache/pulsar/broker/namespace/OwnershipCache.java

@@ -92,14 +92,24 @@ public CompletableFuture<OwnedBundle> asyncLoad(NamespaceBundle namespaceBundle,
             return lockManager.acquireLock(ServiceUnitUtils.path(namespaceBundle), selfOwnerInfo)
                     .thenApply(rl -> {
                         locallyAcquiredLocks.put(namespaceBundle, rl);
+                        OwnedBundle ownedBundle = new OwnedBundle(namespaceBundle, rl);
                         rl.getLockExpiredFuture()
                                 .thenRun(() -> {
                                     log.info().attr("path", rl.getPath()).log("Resource lock has expired");
-                                    namespaceService.unloadNamespaceBundle(namespaceBundle);
-                                    invalidateLocalOwnerCache(namespaceBundle);
+                                    locallyAcquiredLocks.remove(namespaceBundle, rl);
+                                    namespaceService.unloadNamespaceBundle(namespaceBundle)
+                                            .exceptionally(ex -> {
+                                                log.debug()
+                                                        .attr("bundle", namespaceBundle)
+                                                        .exception(ex)
+                                                        .log("Failed to unload namespace bundle after its"
+                                                                + " resource lock expired");
+                                                return null;
+                                            });
+                                    invalidateLocalOwnerCache(namespaceBundle, ownedBundle);
                                     namespaceService.onNamespaceBundleUnload(namespaceBundle);
                                 });
-                        return new OwnedBundle(namespaceBundle);
+                        return ownedBundle;
                     });
         }
     }
@@ -223,6 +233,15 @@ public CompletableFuture<NamespaceEphemeralData> tryAcquiringOwnership(Namespace
      *
      */
     public CompletableFuture<Void> removeOwnership(NamespaceBundle bundle) {
+        CompletableFuture<OwnedBundle> ownedBundleFuture = ownedBundlesCache.getIfPresent(bundle);
+        if (ownedBundleFuture != null && !ownedBundleFuture.isDone()) {
+            // An acquisition of this bundle is in flight. Wait for it to settle and then release what it
+            // acquired: reporting success while the pending acquisition installs its lock afterwards would
+            // leave the broker silently owning the bundle.
+            return ownedBundleFuture.handle((ownedBundle, ex) -> null)
+                    .thenCompose(ignore -> removeOwnership(bundle));
+        }
+
         ResourceLock<NamespaceEphemeralData> lock = locallyAcquiredLocks.remove(bundle);
         if (lock == null) {
             // We don't own the specified bundle anymore
@@ -232,6 +251,27 @@ public CompletableFuture<Void> removeOwnership(NamespaceBundle bundle) {
         return lock.release();
     }
 
+    /**
+     * Method to remove the ownership that was acquired for the given {@link OwnedBundle} instance only.
+     *
+     * <p>If the bundle has since been re-acquired (the given instance's lock expired and a newer
+     * {@link OwnedBundle} with a newer lock owns the bundle now), the newer ownership is left untouched.
+     */
+    public CompletableFuture<Void> removeOwnership(OwnedBundle ownedBundle) {
+        ResourceLock<NamespaceEphemeralData> lock = ownedBundle.getResourceLock();
+        if (lock == null) {
+            // The instance is not bound to a lock (not created by this cache): fall back to removing whatever
+            // ownership currently exists for the bundle.
+            return removeOwnership(ownedBundle.getNamespaceBundle());
+        }
+        if (!locallyAcquiredLocks.remove(ownedBundle.getNamespaceBundle(), lock)) {
+            // This ownership generation was already released or has expired; a newer acquisition may own the
+            // bundle now and must not be disturbed.
+            return CompletableFuture.completedFuture(null);
+        }
+        return lock.release();
+    }
+
     /**
      * Method to remove ownership of all owned bundles.
      *
@@ -343,6 +383,18 @@ public void invalidateLocalOwnerCache(NamespaceBundle namespaceBundle) {
         this.ownedBundlesCache.synchronous().invalidate(namespaceBundle);
     }
 
+    /**
+     * Invalidate the local owner cache entry only if it still holds the given {@link OwnedBundle} instance,
+     * so that a stale lock-expiry callback cannot drop an entry installed by a newer acquisition.
+     */
+    private void invalidateLocalOwnerCache(NamespaceBundle namespaceBundle, OwnedBundle expectedOwnedBundle) {
+        CompletableFuture<OwnedBundle> future = ownedBundlesCache.getIfPresent(namespaceBundle);
+        if (future != null && future.isDone() && !future.isCompletedExceptionally()
+                && future.getNow(null) == expectedOwnedBundle) {
+            ownedBundlesCache.asMap().remove(namespaceBundle, future);
+        }
+    }
+
     @VisibleForTesting
     public Map<NamespaceBundle, ResourceLock<NamespaceEphemeralData>> getLocallyAcquiredLocks() {
         return locallyAcquiredLocks;

pulsar-broker/src/test/java/org/apache/pulsar/broker/namespace/OwnershipCacheTest.java

@@ -27,12 +27,15 @@
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNotSame;
 import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
 import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
 import com.google.common.collect.Range;
 import com.google.common.hash.Hashing;
 import java.util.EnumSet;
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -50,6 +53,8 @@
 import org.apache.pulsar.metadata.api.MetadataStoreConfig;
 import org.apache.pulsar.metadata.api.MetadataStoreException;
 import org.apache.pulsar.metadata.api.coordination.CoordinationService;
+import org.apache.pulsar.metadata.api.coordination.LockManager;
+import org.apache.pulsar.metadata.api.coordination.ResourceLock;
 import org.apache.pulsar.metadata.api.extended.CreateOption;
 import org.apache.pulsar.metadata.api.extended.MetadataStoreExtended;
 import org.apache.pulsar.metadata.coordination.impl.CoordinationServiceImpl;
@@ -97,6 +102,8 @@ public void setup() throws Exception {
 
         bundleFactory = new NamespaceBundleFactory(pulsar, Hashing.crc32());
         nsService = mock(NamespaceService.class);
+        doReturn(CompletableFuture.completedFuture(null)).when(nsService)
+                .unloadNamespaceBundle(any(NamespaceBundle.class));
         brokerService = mock(BrokerService.class);
         doReturn(CompletableFuture.completedFuture(1)).when(brokerService)
                 .unloadServiceUnit(any(), anyBoolean(), anyBoolean(), anyLong(), any());
@@ -407,4 +414,115 @@ public void testReestablishOwnership() throws Exception {
         assertNotNull(cache.getOwnedBundle(testFullBundle));
     }
 
+    @Test
+    public void testStaleUnloadDoesNotReleaseReacquiredOwnership() throws Exception {
+        OwnershipCache cache = new OwnershipCache(this.pulsar, nsService);
+        doReturn(cache).when(nsService).getOwnershipCache();
+        NamespaceBundle bundle = new NamespaceBundle(NamespaceName.get("pulsar/ns-stale-unload"),
+                Range.closedOpen(0L, (long) Integer.MAX_VALUE),
+                bundleFactory);
+
+        cache.tryAcquiringOwnership(bundle).get();
+        OwnedBundle staleOwnedBundle = cache.getOwnedBundle(bundle);
+        assertNotNull(staleOwnedBundle);
+
+        // Simulate the lock-expiry self-heal path: the expiry callback invalidated the local cache while the
+        // unload it triggered is still in flight, and a concurrent lookup re-acquires the bundle in between.
+        cache.invalidateLocalOwnerCache(bundle);
+        cache.tryAcquiringOwnership(bundle).get();
+        OwnedBundle reacquiredOwnedBundle = cache.getOwnedBundle(bundle);
+        assertNotNull(reacquiredOwnedBundle);
+        assertNotSame(reacquiredOwnedBundle, staleOwnedBundle);
+        ResourceLock<NamespaceEphemeralData> reacquiredLock = cache.getLocallyAcquiredLocks().get(bundle);
+        assertNotNull(reacquiredLock);
+
+        // The stale unload chain now runs its remaining steps against the old OwnedBundle instance.
+        staleOwnedBundle.handleUnloadRequest(pulsar, 5, TimeUnit.SECONDS).join();
+
+        // The re-acquired ownership must survive the stale unload untouched.
+        assertSame(cache.getLocallyAcquiredLocks().get(bundle), reacquiredLock);
+        assertTrue(store.exists(ServiceUnitUtils.path(bundle)).join());
+        assertTrue(reacquiredOwnedBundle.isActive());
+        assertTrue(cache.checkOwnershipAsync(bundle).get());
+    }
+
+    @Test
+    public void testRemoveOwnershipWithAcquisitionInFlight() throws Exception {
+        // Gate the lock acquisition so the test can invoke removeOwnership while the acquisition is in flight
+        LockManager<NamespaceEphemeralData> realLockManager =
+                coordinationService.getLockManager(NamespaceEphemeralData.class);
+        CompletableFuture<Void> gate = new CompletableFuture<>();
+        LockManager<NamespaceEphemeralData> gatedLockManager = new LockManager<NamespaceEphemeralData>() {
+            @Override
+            public CompletableFuture<Optional<NamespaceEphemeralData>> readLock(String path) {
+                return realLockManager.readLock(path);
+            }
+
+            @Override
+            public CompletableFuture<ResourceLock<NamespaceEphemeralData>> acquireLock(String path,
+                                                                                       NamespaceEphemeralData value) {
+                return gate.thenCompose(__ -> realLockManager.acquireLock(path, value));
+            }
+
+            @Override
+            public CompletableFuture<List<String>> listLocks(String path) {
+                return realLockManager.listLocks(path);
+            }
+
+            @Override
+            public CompletableFuture<Void> asyncClose() {
+                return realLockManager.asyncClose();
+            }
+
+            @Override
+            public void close() throws Exception {
+                realLockManager.close();
+            }
+        };
+        CoordinationService gatedCoordinationService = mock(CoordinationService.class);
+        doReturn(gatedLockManager).when(gatedCoordinationService).getLockManager(NamespaceEphemeralData.class);
+        doReturn(gatedCoordinationService).when(pulsar).getCoordinationService();
+
+        OwnershipCache cache = new OwnershipCache(this.pulsar, nsService);
+        NamespaceBundle bundle = new NamespaceBundle(NamespaceName.get("pulsar/ns-inflight-remove"),
+                Range.closedOpen(0L, (long) Integer.MAX_VALUE),
+                bundleFactory);
+
+        CompletableFuture<NamespaceEphemeralData> acquireFuture = cache.tryAcquiringOwnership(bundle);
+        assertFalse(acquireFuture.isDone());
+
+        CompletableFuture<Void> removeFuture = cache.removeOwnership(bundle);
+        gate.complete(null);
+        acquireFuture.join();
+        removeFuture.get(10, TimeUnit.SECONDS);
+
+        // After removeOwnership reported success and the in-flight acquisition settled, the broker must not
+        // silently retain (zombie) ownership.
+        Awaitility.await().untilAsserted(() -> {
+            assertTrue(cache.getLocallyAcquiredLocks().isEmpty());
+            assertTrue(cache.getOwnedBundles().isEmpty());
+            assertFalse(store.exists(ServiceUnitUtils.path(bundle)).join());
+        });
+    }
+
+    @Test
+    public void testExpiredLockIsRemovedFromLocallyAcquiredLocks() throws Exception {
+        OwnershipCache cache = new OwnershipCache(this.pulsar, nsService);
+        NamespaceBundle bundle = new NamespaceBundle(NamespaceName.get("pulsar/ns-expired-lock"),
+                Range.closedOpen(0L, (long) Integer.MAX_VALUE),
+                bundleFactory);
+
+        cache.tryAcquiringOwnership(bundle).get();
+        ResourceLock<NamespaceEphemeralData> lock = cache.getLocallyAcquiredLocks().get(bundle);
+        assertNotNull(lock);
+
+        // The lock dies without going through removeOwnership, like on a metadata session expiry
+        lock.release().join();
+
+        Awaitility.await().untilAsserted(() -> {
+            assertTrue(cache.getOwnedBundles().isEmpty());
+            assertTrue(cache.getLocallyAcquiredLocks().isEmpty());
+        });
+    }
+
 }