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

[SongOf]

Motivation

OwnershipCache.removeOwnership(bundle) removes whichever ResourceLock happens to be in
locallyAcquiredLocks at that moment, without checking that it belongs to the ownership
"generation" the caller was working on. Three defects follow from this:

1. A stale unload chain can destroy a re-acquired ownership. When a bundle's resource lock
   expires, the lock-expiry callback in the cache loader fires unloadNamespaceBundle(...)
   asynchronously (dropping its future) and invalidates the local owner cache immediately. A
   concurrent lookup can then re-acquire the bundle (new lock, new OwnedBundle) while the stale
   unload chain is still running. When that chain reaches its final removeOwnership(bundle), it
   removes and releases the newly acquired lock, deleting the new owner's znode while the local
   cache still claims active ownership. This opens a transient dual-serving window (fenced-ledger
   errors on persistent topics, silent message loss on non-persistent topics) and causes ownership
   flapping. The same stale chain also deactivates the newer OwnedBundle through the cache lookup
   in updateBundleState(bundle, false).

2. removeOwnership reports success while an acquisition is in flight, leaving zombie
   ownership. The lock is only put into locallyAcquiredLocks after acquireLock completes, so
   a removeOwnership call racing with an in-flight acquisition finds the map empty, returns a
   completed future ("released"), and the lock installed afterwards is never cleaned up: the broker
   silently keeps the lock and the cache entry for a bundle the caller believes released (e.g. a
   deleted bundle/namespace), polluting load reports until restart.

3. The lock-expiry callback never removes the expired lock from locallyAcquiredLocks
   (asymmetric add/remove bookkeeping), so a dead lock handle lingers in the map and drifts from
   the owned-bundle cache. The callback also drops the future returned by
   unloadNamespaceBundle(...), so failures of the triggered unload are silent.

Modifications

• Bind each OwnedBundle to the ResourceLock acquired for it, so an OwnedBundle instance
  represents a single ownership generation. Add removeOwnership(OwnedBundle) which releases only
  that generation via ConcurrentHashMap.remove(key, value); if the bundle has since been
  re-acquired, the newer ownership is left untouched. OwnedBundle.handleUnloadRequest now uses it
  and no longer calls updateBundleState (its own isActive flag is already flipped via CAS
  earlier in the method; the cache lookup could deactivate a newer generation).
• removeOwnership(NamespaceBundle) now waits for an in-flight acquisition to settle and then
  releases whatever it installed, instead of reporting success while the pending lock leaks. Its
  "release the current ownership, whatever generation" semantics are intentionally kept for the
  bundle-deletion and split paths.
• The lock-expiry callback removes the expired lock from locallyAcquiredLocks with a
  generation-safe two-arg remove (mirroring LockManagerImpl's own lock bookkeeping), invalidates
  the owned-bundle cache entry only if it still holds this generation, and logs failures of the
  triggered unload instead of dropping the future.

Verifying this change

• Make sure that the change passes the CI checks.

This change added tests and can be verified as follows:

- *Added `OwnershipCacheTest.testStaleUnloadDoesNotReleaseReacquiredOwnership`: simulates the
  expiry-callback interleaving (cache invalidated, bundle re-acquired, then the stale unload
  chain completes) and asserts the re-acquired lock, znode, and active state survive. Fails
  before the fix (the new lock is released and the znode deleted).*
- *Added `OwnershipCacheTest.testRemoveOwnershipWithAcquisitionInFlight`: gates `acquireLock`
  behind a controllable future to call `removeOwnership` deterministically inside the in-flight
  window, then asserts no lock/cache entry/znode remains once both settle. Fails before the fix
  (zombie ownership remains).*
- *Added `OwnershipCacheTest.testExpiredLockIsRemovedFromLocallyAcquiredLocks`: releases the lock
  without going through `removeOwnership` (same `expiredFuture` path as a session expiry) and
  asserts `locallyAcquiredLocks` is cleaned up together with the cache. Fails before the fix
  (dead lock handle lingers).*
- *All 8 pre-existing `OwnershipCacheTest` tests, `NamespaceServiceTest` (27),
  `NamespaceUnloadingTest`, `NamespaceOwnershipListenerTest` and
  `OwnerShipCacheForCurrentServerTest` pass locally.*

Does this pull request potentially affect one of the following parts:

If the box was checked, please highlight the changes

• Dependencies (add or upgrade a dependency)
• The public API
• The schema
• The default values of configurations
• The threading model
• The binary protocol
• The REST endpoints
• The admin CLI options
• The metrics
• Anything that affects deployment

[Denovo1998]

The lock release is now generation-safe, but the stale unload is still bundle-wide at the topic layer.

BrokerService.unloadServiceUnit operates on the broker's current topic map, and cleanUnloadedTopicFromCache later scans the current map again and eventually calls topics.remove(topic) because no create future is supplied. If this unload belongs to an expired OwnedBundle and the bundle is re-acquired while the old close futures are still running, the stale cleanup can remove a topic future installed by the newer ownership generation.

The new test mocks BrokerService, so it only verifies that the new lock/znode/OwnedBundle survives and cannot detect topic-cache corruption.

Could we either prevent re-acquisition until the previous generation's unload cleanup finishes, or make the cleanup identity-safe by capturing the topic futures at unload start and removing them with topics.remove(topicName, capturedFuture)?

[SongOf]

@Denovo1998
Fixed — thanks for the detailed repro path, that's exactly what was happening.
We went with the identity-safe cleanup option, but implemented it as a single shared snapshot rather than two independent scans (capturing once for unload, again for cleanup would leave its own gap — see below):

• BrokerService.getTopicFuturesInBundle(NamespaceBundle) captures (topicName → topicFuture) for the bundle in one pass, at unload start.
• unloadServiceUnit now has an overload that consumes that snapshot directly instead of re-scanning topics, so the set of topics it closes and the set cleanUnloadedTopicFromCache later validates against are guaranteed to be identical — no second scan, no window for a topic to be missed by one side and caught by the other.
• cleanUnloadedTopicFromCache now takes that same snapshot and removes via topics.remove(topic, capturedFuture) instead of the old null-future call, so a stale cleanup can only ever touch the exact future it observed at unload start. A newer generation's topic future is a different object and simply won't match the CAS.
• Both call sites — OwnedBundle.handleUnloadRequest and ServiceUnitStateChannelImpl.closeServiceUnit — capture the snapshot once and thread it through both calls.

On re-acquisition: we didn't block it. Blocking re-acquisition until the previous generation's cleanup finishes would add a synchronization point across OwnershipCache and BrokerService and slow down bundle flapping recovery; identity-safe removal gets full correctness without that coupling.

On the test: agreed the mocked-BrokerService test can't see this class of bug. Added BrokerServiceTest#testCleanUnloadedTopicFromCacheIsGenerationSafe against a real embedded broker — it loads a topic, closes+reloads it to install a new generation's future under the same name (standing in for the re-acquire), then invokes the stale cleanup with the old snapshot and asserts the new generation's topic is still present. Confirmed it fails without the identity check (topics.remove(topic, null)-equivalent) and passes with it.

While implementing this we also caught a follow-on issue in review: unloadServiceUnit's topic-closing scan and the pre-captured cleanup snapshot were originally two separate scans of topics, which reopened a smaller version of the same gap for any topic created in between. Fixed by making unloadServiceUnit consume the caller-supplied snapshot directly (see above) instead of doing its own scan — now there's exactly one point-in-time view shared by both steps.

[Denovo1998]

The future.isDone() guard creates an expiry-before-publication race.

An expiry callback is registered inside the cache loader's thenApply, before the load future completes. If the lock has already expired, or expires concurrently during callback registration, the callback may execute while the cache future is still incomplete:

1. The lock is removed from locallyAcquiredLocks.
2. unloadNamespaceBundle cannot find a completed OwnedBundle.
3. This method skips invalidation because future.isDone() is false.
4. The loader then returns and publishes an active OwnedBundle whose lock has already expired.

This results in the local cache claiming active ownership without a corresponding lock.

Incomplete futures should be handled by attaching a completion callback and removing the future via ownedBundlesCache.asMap().remove(bundle, future) if it resolves to expectedOwnedBundle. Unloading should also be deferred until that generation is published, or the acquisition should fail if expiry wins the race. A deterministic test could use a ResourceLock whose expired future has already completed when the listener is attached.

[SongOf]

@Denovo1998
Confirmed, and thanks for spelling out the exact mechanism — traced it and it's real: the listener is registered inside .thenApply, so if rl.getLockExpiredFuture() is already done (or completes concurrently) at that point, .thenRun(...) fires inline, before the loader's own future is published. invalidateLocalOwnerCache's future.isDone() guard is false at that moment, so it no-ops, and the loader goes on to return ownedBundle, publishing an "active" OwnedBundle with no backing lock.

Did both things you suggested, plus one more:

1. invalidateLocalOwnerCache — replaced the isDone() guard with future.whenComplete(...):
   private void invalidateLocalOwnerCache(NamespaceBundle namespaceBundle, OwnedBundle expectedOwnedBundle) { CompletableFuture<OwnedBundle> future = ownedBundlesCache.getIfPresent(namespaceBundle); if (future == null) { return; } future.whenComplete((ownedBundle, ex) -> { if (ex == null && ownedBundle == expectedOwnedBundle) { ownedBundlesCache.asMap().remove(namespaceBundle, future); } }); }
   Runs immediately if already done, deferred otherwise — same identity check as before, just no longer gated on timing.

2. Went with "fail the acquisition" over "defer unload." Deferring unload still leaves unloadNamespaceBundle's caller-visible behavior wrong for that window, and — more importantly — it leaves onNamespaceBundleUnload firing before onNamespaceBundleOwned for this generation (the unload listener call happens synchronously inside the racing callback, ownership-owned fires only once the future publishes afterward), so every NamespaceBundleOwnershipListener sees "unloaded" before it ever saw "owned," and tryAcquiringOwnership's caller gets a fleeting, already-invalid success. Failing the load outright avoids all three:

AtomicBoolean expiredBeforePublication = new AtomicBoolean(false); rl.getLockExpiredFuture() .thenRun(() -> { ... expiredBeforePublication.set(true); ... }); if (expiredBeforePublication.get()) { throw new IllegalStateException( "Lock for bundle " + namespaceBundle + " expired before ownership could be published"); } return ownedBundle;

The AtomicBoolean check right after .thenRun(...) returns is race-free by construction — thenRun guarantees the action runs inline, before it returns to the caller, exactly when the upstream future is already complete at attach time. No polling, no separate timing check. Caffeine auto-evicts a failed load, so there's no zombie cache entry to separately clean up here; (1) above still matters for the normal case where expiry happens strictly after publication.

4. Test — used your suggested approach almost exactly: a ResourceLock whose getLockExpiredFuture() is pre-completed, plus a gated LockManager.acquireLock so the acquisition is genuinely in-flight (registered in the cache, not yet published) when the listener attaches — otherwise everything resolves synchronously and Caffeine hasn't even registered a placeholder yet, which doesn't reproduce your scenario. testExpiryBeforePublicationDoesNotLeaveActiveZombieOwnership asserts tryAcquiringOwnership itself fails and leaves no trace in locallyAcquiredLocks/ownedBundlesCache. Verified it fails without both fixes and passes with them.

[Denovo1998]

The new wait handles an acquisition that is already visible when removeOwnership starts, but it does not establish an acquire/remove barrier. A new cache load can begin after getIfPresent is checked and install another lock before or after this method completes.

Also, NamespaceService#splitAndOwnBundleOnceAndRetry currently ignores the returned future and completes the split immediately, so a delayed or failed release is not observed by the caller.

Would it be safer to serialize acquisition and removal per bundle, for example with an explicit ACQUIRING / OWNED / UNLOADING / RELEASED state or an unload barrier, and make the split path compose the removal future?

[SongOf]

@Denovo1998
Agreed on both points — the wait was reactive (only caught an already-visible in-flight future) rather than an actual barrier, and yes, the split path was firing-and-forgetting the release. Went with your "serialize per
bundle" suggestion, implemented as an unload/acquire barrier rather than an explicit state enum — reasoning below.

1. Per-bundle async barrier. Added bundleOperationBarriers (Map<NamespaceBundle, CompletableFuture>) and a private serialize(bundle, operation) helper that both tryAcquiringOwnership and removeOwnership(NamespaceBundle) now route through, so the two can never interleave for the same bundle. serialize uses ConcurrentHashMap.compute to atomically chain each new operation behind whatever's currently queued for that bundle, and cleans its own entry back out once nothing's queued behind it. removeOwnership(NamespaceBundle)'s body is now just the blind release, wrapped in serialize — that subsumes the old "wait if visible" logic, since any acquire already queued on
   the barrier is guaranteed to settle before a subsequent remove runs, whether or not it was visible at call time.

Went with this over an explicit ACQUIRING/OWNED/UNLOADING/RELEASED enum because the enum would still need the same kind of per-bundle mutex underneath its transitions (it doesn't replace serialize, only adds a label on top), and every existing external caller of checkOwnershipAsync/isNamespaceBundleOwned/getOwnerAsync would need to be individually re-audited for what an in-between state should mean to it — that's a much bigger surface (15 call sites in NamespaceService alone) than this bug warrants. Happy to revisit as a separate proposal if the intermediate-state visibility is independently wanted.

removeOwnership(OwnedBundle) (the generation-identity CAS release used by the normal unload path) deliberately isn't routed through serialize — it's already race-safe by construction, so serializing it would only add latency to the main unload path for no correctness gain.

2. Split path composes the release. In splitAndOwnBundleOnceAndRetry, the removeOwnership(bundle) call is now thenComposed into the chain (with an .exceptionally that logs and swallows) before completionFuture.complete(null) runs, instead of being fired and discarded. A release failure is logged rather than failing the split — the split itself already succeeded by that point, so a release hiccup is cleanup-level, not split-level, failure. While rewriting this block I also caught that its adjacent exception handler referenced the outer, always-null ex instead of its own parameter — would have NPE'd and left completionFuture never completing on a real updateBundleState failure; fixed alongside since it's the same block.

Tests: testRemoveOwnershipWithAcquisitionInFlight covers the barrier (gates the lock acquisition so it's genuinely in-flight when removeOwnership is called concurrently, asserts no zombie lock survives once both settle); existing NamespaceServiceTest#testSplitAndOwnBundles / #testSplitBundleAndRemoveOldBundleFromOwnerShipCache continue to pass against the composed split path.

[lhotari]

I ran an AI-assisted review of this PR (Claude as the local reviewer plus OpenAI Codex gpt-5.6-sol as a second independent pass; both sets of findings were cross-verified against the actual sources before posting). Overall assessment first: this is a solid, well-tested fix for real races — the generation-bound OwnedBundle, the per-bundle acquire/release barrier, and the symmetric expiry-callback cleanup each address a genuine defect, and the new tests are deterministic and fail before the fix. Two findings are worth addressing before merge; the rest are optional.

1. cleanUnloadedTopicFromCache is only partially generation-safe (BrokerService.removeTopicFromCache)

In removeTopicFromCache(String, NamespaceBundle, CompletableFuture), only the final topics.remove(topic, createTopicFuture) is identity-guarded. The TopicEvent.UNLOAD BEFORE/SUCCESS notifications, the multiLayerTopicsMap removal (including replication-metrics teardown when a namespace empties), the compactor-stats removal, and forgetSegmentLoad all run unconditionally. A stale cleanup whose captured future no longer matches therefore still strips the newer generation's stats/load-report bookkeeping and fires unload events for a topic that stays cached and serving. This guarded-remove-with-unconditional-bookkeeping shape is admittedly pre-existing (removeTopicFromCache(AbstractTopic) has the same issue), but the new javadoc claims a stale call "can never evict a newer generation's topic", which currently only holds for the topics map itself, and testCleanUnloadedTopicFromCacheIsGenerationSafe only asserts getTopicReference() presence. Suggestion: perform the conditional topics.remove first, and only run the bookkeeping/events when that remove actually won.

2. The "race-free" thenRun claim in the expiry-before-publication guard isn't a JDK guarantee (OwnershipCache cache loader)

The comment on expiredBeforePublication asserts thenRun "only runs its action inline, before returning, when the future it is attached to is already complete." When the expiry future completes concurrently with registration, the JDK allows the action to be claimed and run by the completing thread after thenRun returns — so the flag can still read false and the loader publishes an OwnedBundle whose lock is already dead; tryAcquiringOwnership reports success and onNamespaceBundleOwned fires. Tracing the fallback shows the damage is tightly bounded: the listener has already done the two-arg locallyAcquiredLocks.remove, and the deferred invalidateLocalOwnerCache(bundle, ownedBundle) fires as part of the load future's completion, so the cache entry is dropped essentially at publication time. The residual effect is indistinguishable from a lock expiring immediately after a legitimately successful acquire, and state converges with no zombie — so low severity in practice, but the comment (and the matching claim in the invalidateLocalOwnerCache javadoc) should be corrected, and could optionally be hardened by re-checking rl.getLockExpiredFuture().isDone() after registering the listener.

3. (optional) serialize() invokes the operation inside ConcurrentHashMap.compute

In the common uncontended case (previous == null), thenCompose(ignore -> operation.get()) runs the supplier synchronously inside the compute lambda, under the map's bin lock. The suppliers only initiate async work today, but if any synchronous completion chain ever reaches serialize() for the same bundle on the same thread (plausible with already-completed metadata futures, e.g. a fast-failing acquire, or mocks in tests), that is a reentrant compute — behavior ConcurrentHashMap leaves undefined, and it could silently break the very mutual exclusion the barrier provides. Safer shape: inside compute, only capture the previous barrier and install the new one; chain precedingOp → operation.get() → completion after compute returns.

4. (optional) Straggler-topic semantics change

A topic installed in topics for the bundle after the snapshot is captured is now neither closed by the unload nor evicted by the cleanup. The old behavior was arguably worse (force-evicting live topics without closing them, enabling a second instance on reload), so this looks like a deliberate trade — but convergence for stragglers now relies entirely on per-connection ownership checks and lookup redirects. Worth a sentence in the PR description confirming it's intended.

5. (optional) Barrier latency notes, for awareness only

removeOwnership(NamespaceBundle) callers (bundle deletion, split, shutdown) now wait behind an in-flight acquire (bounded by metadata operation timeouts); tryAcquiringOwnership queues per bundle even on a cache hit while another operation is pending; and concurrent acquires that previously coalesced onto one shared failed load now retry sequentially. All per-bundle and bounded.

Also: nice catch on the ex → e fix in the split path — the old code could NPE by dereferencing the enclosing handler's null throwable.

[SongOf]

@lhotari
Thanks for the thorough, cross-verified review — all addressed:

1. Fixed. removeTopicFromCache now runs the identity-guarded topics.remove first and only proceeds to the multiLayerTopicsMap/replication-metrics/compactor-stats/segment-load bookkeeping and UNLOAD events if that remove actually wins. Added testCleanUnloadedTopicFromCacheIsGenerationSafeForBookkeeping to cover the case the existing test missed.
2. Fixed. Corrected both the thenRun comment and the matching invalidateLocalOwnerCache javadoc — they no longer claim synchronous execution on a concurrent completion. Also added the suggested rl.getLockExpiredFuture().isDone() recheck alongside the flag to narrow the window further.
3. Fixed. serialize() now only captures/installs the barrier inside compute(); the precedingOp → operation.get() → completion chain is built after compute() returns, so a fully synchronous, self-reentrant operation only ever produces ordinary, non-nested compute() calls.
4. Documented. Added a comment at both snapshot sites (OwnedBundle and ServiceUnitStateChannelImpl, which share the same pattern) confirming this is intentional, why it's safer than the old force-evict behavior, and that BookKeeper ledger fencing — not this unload path — is the actual backstop against a stale writer.
5. Noted, no action needed — though while addressing the above we found and fixed a related gap in the same spirit: removeOwnership(OwnedBundle), the path every normal unload actually goes through, wasn't wrapped in the barrier at all, so a concurrent acquire could still observe and report success for a generation mid-release. It's now serialized too, with a regression test (testTryAcquiringOwnershipWaitsForInFlightOwnedBundleRelease) — so this latency note extends there as well.

[lhotari]

please merge origin/master (don't rebase) to the PR branch and resolve conflicts

[SongOf]

please merge origin/master (don't rebase) to the PR branch and resolve conflicts

merged @lhotari