[fix][broker] Fix AvgShedder assignment cache keying with stable bundle names

[void-ptr974]

Motivation

When AvgShedder decides to move load away from an overloaded broker, it first chooses bundles to unload and a destination broker for each bundle. When the load manager processes those bundles, it must use the same destination that AvgShedder selected.

AvgShedder previously remembered this relationship in a Map<BundleData, String>:

bundle load data -> destination broker

BundleData is not a stable map key. Its load statistics are updated in place, and its equals and hashCode are based on those statistics. This leads to two failures:

• after a load report updates a BundleData object, its hash can change and the existing map entry can no longer be found;
• two different bundles with the same load statistics compare as equal, so one destination can overwrite the other.

In either case, the bundle may be assigned to a broker different from the destination chosen by AvgShedder.

How the fix works

The temporary shedding plan is now keyed by the canonical bundle name:

AvgShedder planning
    bundle name -> planned destination
          |
          v
ModularLoadManager selects that destination
          |
          v
unload request records the destination in the existing bundle affinity map
          |
          v
the next ownership lookup consumes the affinity

The AvgShedder map is needed only while one load-shedding pass is being processed. The existing bundle affinity map carries the selected destination from the unload request to the later ownership lookup. AvgShedder therefore clears its temporary plan when the shedding pass finishes, including when planning fails.

If the planned broker is no longer available, normal placement selects from the remaining candidates without replacing the original plan. If the selected broker is overloaded, the load manager restores the policy-allowed candidates and lets the configured placement strategy choose again without forcing the first broker out of the candidate set.

Modifications

• Store planned destinations as bundle name -> broker in a ConcurrentHashMap.
• Add ModularLoadManagerStrategy.selectBrokerForBundle(...) as a default method so the load manager can pass the stable bundle name during assignment.
• Add an attempt-completion callback to LoadSheddingStrategy and invoke it in finally around both planning and unload processing.
• Keep the historical four-argument selector as the normal fallback when no planned destination can be used.
• Return Optional.empty() when broker selection receives an empty candidate set.

Compatibility

The new strategy methods have default implementations, so existing ModularLoadManagerStrategy and LoadSheddingStrategy implementations do not need to implement them. The historical four-argument selector remains available and is not scheduled for removal.

The change is limited to the Modular load manager and does not modify configuration defaults, metadata or wire formats, protocols, REST APIs, or the Simple and Extensible load managers.

Verification

• GitHub CI checks have not been reported for the latest commit yet.
• Verified locally with focused broker tests and repository checks:

./gradlew --no-daemon :pulsar-broker:test -PtestRetryCount=0 \
  --tests org.apache.pulsar.broker.loadbalance.impl.AvgShedderTest \
  --tests org.apache.pulsar.broker.loadbalance.ModularLoadManagerStrategyTest \
  --tests org.apache.pulsar.broker.loadbalance.impl.ModularLoadManagerImplTest.testLoadSheddingPassesBundleNameAndCompletesAttempt \
  --tests org.apache.pulsar.broker.loadbalance.impl.ModularLoadManagerImplTest.testOverloadRetryBypassesBundleAwareSelectionWithoutRemovingCandidates

./gradlew --no-daemon :pulsar-broker:test -PtestRetryCount=0 \
  --tests org.apache.pulsar.broker.loadbalance.impl.ModularLoadManagerImplTest

./gradlew --no-daemon quickCheck

The tests cover:

• mutation of a BundleData object after planning;
• different bundles with equal load statistics;
• propagation of the canonical bundle name into broker selection;
• cleanup after successful processing and after a planning exception;
• unavailable and overloaded planned destinations;
• the existing unload-destination affinity flow;
• compatibility of the strategy fallback methods.

The complete ModularLoadManagerImplTest class passed with 17 tests and no failures or errors.

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

• Dependencies
• Public API
  • Adds backward-compatible default methods.
  • Keeps the historical selector available.
• Schema or metadata formats
• Configuration defaults
• Threading model
• Binary protocol
• REST endpoints
• Admin CLI
• Metrics
• Deployment

[lhotari]

the name of the broker is called brokerId, returned by org.apache.pulsar.broker.PulsarService#getBrokerId

[void-ptr974]

Thanks, I’ve updated the Javadoc to use “broker ID” consistently.

[lhotari]

Thanks for digging into this — the diagnosis is correct and I was able to confirm it independently: BundleData is annotated @EqualsAndHashCode (pulsar-common/.../BundleData.java:26) over fields that BundleData.update(NamespaceBundleStats) mutates in place on every load-report refresh, so a HashMap<BundleData, String> key silently becomes unreachable, and two bundles with identical stats collide on one key. Re-keying by bundle name is the right fix.

I also verified the key spaces line up: ModularLoadManagerImpl:902 uses serviceUnit.toString(), which is the same key space as loadData.getBundleData() and getBundleDataForLoadShedding(), so the write in selectBundleForUnloading and the read in the placement path agree. The new selectBrokerForBundle default method is source- and binary-compatible, and LeastLongTermMessageRate, LeastResourceUsageWithWeight and RoundRobinBrokerSelector are correctly unaffected.

Two things I'd like to discuss before this goes in — mainly (1).

1. Planned destinations are now permanent — nothing ever removes an entry from bundleBrokerMap

AvgShedder.java:290-310

selectBrokerWithBundleName returns the cached broker and never removes it. The only writes are the put in selectBundleForUnloading (:196) and the fallback put (:305); there is no removal anywhere, and onActiveBrokersChange (:206) just delegates to the no-op default.

Before this PR the bug provided accidental expiry: the key's hash drifted on the next load-report cycle, the entry became unreachable, and the next assignment re-randomized. After the fix the mapping is sticky for the lifetime of the process, which I think is a bigger behaviour change than "retain a planned destination across load-data refreshes" suggests.

Two consequences:

Manual unload stops moving bundles. A bundle is shed to broker2; the entry is written and consumed correctly (this is the intended behaviour, and it now works). An operator later runs pulsar-admin namespaces unload to move it off broker2. The next lookup hits selectBrokerForBundle, broker2 is still a candidate, so the bundle is assigned straight back to broker2. From the operator's point of view the unload is a no-op, until AvgShedder itself happens to shed the bundle again.

The overload-escape path becomes dead code for AvgShedder. ModularLoadManagerImpl:981-987 re-runs placement over the widened candidate set when the selected broker is above loadBalancerBrokerOverloadedThresholdPercentage. With a sticky entry, the second call returns the identical broker.

I don't have a strong opinion on which way to resolve it — options I can see are dropping the entry once an assignment has consumed it, bounding validity via loadData.getRecentlyUnloadedBundles(), or clearing consumed entries at the start of findBundlesForUnloading. Whichever is chosen needs to keep the entry valid across the retry at ModularLoadManagerImpl:981 within a single assignment.

2. bundleBrokerMap grows without bound, while loadData.getBundleData() is pruned

AvgShedder.java:52

ModularLoadManagerImpl actively removes bundle entries when a bundle goes inactive (:600-604) and when it is split (:796). bundleBrokerMap has no equivalent, so on a long-lived leader with bundle splits, namespace deletion or topic churn, every bundle name ever assigned is retained forever.

This is pre-existing rather than a regression (and it actually grew faster before, since every hash drift orphaned a fresh entry), so I'd be fine with it as a follow-up. But since the PR is already rewriting this field, pruning against loadData.getBundleData().keySet() is only a few lines.

3. bundleBrokerMap is a plain HashMap mutated from two threads under different monitors

AvgShedder.java:52, :196, :305

doLoadShedding() is synchronized on the ModularLoadManagerImpl instance (:637) and reaches bundleBrokerMap.put through findBundlesForUnloading. The assignment path synchronizes on a different monitor — synchronized (brokerCandidateCache) (:902) — and also does get/put. Nothing orders those two, so concurrent put during a resize can corrupt the table or spin.

Again pre-existing and not introduced here, but LeastResourceUsageWithWeight.selectBroker is synchronized for exactly this reason, and switching the declaration to ConcurrentHashMap is a one-word change while that line is already being touched.

4. The identity-scan compatibility path is fragile, and its only callers are this PR's tests

AvgShedder.java:312-323

findBundleNameByIdentity recovers the key by scanning every entry of loadData.getBundleData() for reference equality (entry.getValue() == bundleToAssign). That is O(number of bundles in the cluster) per call, and it only works when the caller passes the exact instance stored in loadData — a contract that appears nowhere in the ModularLoadManagerStrategy javadoc. A caller that passes a defensive copy silently falls into the uncached random branch with no signal.

Since ModularLoadManagerImpl now always uses the name-aware path, the scan exists only for third-party callers and for the tests added here. Leaving the 4-arg selectBroker as a plain uncached fallback (optionally @Deprecated) would be simpler and equally correct; the assertions in AvgShedderTest and testSheddingMultiplePairs would move to selectBrokerForBundle.

5. Undocumented behaviour change: empty candidates no longer throw

AvgShedder.java:278, :292-294

Previously empty candidates reached getExpectedBroker, hit % 0, and the catch (Throwable) fallback threw ArithmeticException again (:335, :343), which propagated out of selectBrokerForAssignment. Both paths now short-circuit to Optional.empty(), which ModularLoadManagerImpl:970 handles as "No brokers available".

That's an improvement, but it's a real behaviour change that isn't called out beyond "keep a valid uncached fallback" — worth an explicit line under Modifications.

6. Tests

The fixture cleanup is genuinely nice: dropping the setNumSamples(i) hacks and adding assertEquals(loadData.getBundleData().get("bundle1-0"), loadData.getBundleData().get("bundle3-0")) in testSheddingMultiplePairs turns the equal-but-distinct BundleData collision into an explicit precondition instead of something the old test worked around. I checked that those are two distinct instances, so the identity scan resolves them unambiguously.

A few gaps:

• Nothing tests the actual integration point of the fix — that ModularLoadManagerImpl passes the correct bundle name through. ModularLoadManagerImpl.selectBroker(ServiceUnitId) is already @VisibleForTesting, so a spy strategy asserting the received name would lock the wiring in.
• Nothing covers the sticky-forever behaviour from (1), which is the riskiest part of the change.
• assertNotEquals(plannedBundleData.hashCode(), originalHashCode) couples the test to Lombok's generated hash including topics. Comparing the objects (or a pre-mutation copy) expresses the same intent without depending on hash internals.
• The new BundleData() → new BundleData(1, 1) and TimeAverageMessageData() → TimeAverageMessageData(1) fixture change in testHitHighThreshold is necessary — with maxSamples == 0, update() can't record the sample — but it's unexplained. A one-line comment would save the next reader the detour.

Nit

AvgShedder.java:291: the BundleData bundleToAssign continuation line is indented 2 columns past the opening paren. Purely cosmetic, checkstyle won't flag it.

---

Intent and implementation match; the only place the description understates things is (1), where "retain a planned destination across load-data refreshes" is doing some heavy lifting for "retain it indefinitely".

I reviewed by reading only — I did not run the build or tests, so I'm relying on your local run plus CI for that. I did confirm that every import in the rewritten ModularLoadManagerStrategyTest is still used (Field and Map are still needed by the untouched LeastResourceUsageWithWeight tests), so there shouldn't be an unused-import checkstyle failure.

_{Assisted-by: Claude Opus 5 (Claude Code), with a second independent pass from Codex gpt-5.6-sol (codex review) that reported no actionable findings. All findings above come from the Claude pass; the code references were checked against the PR head.}

[void-ptr974]

Thanks for the detailed review — addressed in follow-up commits.

• Changed AvgShedder to store pending destinations as bundle name -> broker.
• Added selectBrokerForBundle(...) and wired ModularLoadManagerImpl to pass the bundle name to the placement strategy.
• Added onUnloadAttemptCompleted(...); ModularLoadManagerImpl calls it in finally after processing a shedding plan.
• Replaced the shared map with ConcurrentHashMap. The existing 4-argument selector remains as an uncached compatibility fallback.
• Added tests for mutable/equal BundleData, unavailable planned destinations, cleanup, and the three-broker unload flow.

[Denovo1998]

The escape path for overload still fails to work for AvgShedder. The first call may return a pending destination; the second call widens the candidate set, but AvgShedder continues to return the same pending broker as long as it remains a candidate.

Even if load data changes after the shedding plan and the pending broker now exceeds the hard overload threshold, we still select it even when a healthier broker is available. Restricting the map to a single unload attempt does not alter this behavior.

Could retrying explicitly bypass or invalidate the pending destination? It would also be valuable to add a test in which the planned broker becomes overloaded between planning and assignment.

[void-ptr974]

Thanks for the careful review. The retry now removes the first selected broker from the widened candidate set before selecting again. A regression test verifies that the retry cannot return the same overloaded broker.

[void-ptr974]

Superseding my previous reply: the first broker is no longer removed. The retry now bypasses only the pending destination.

[Denovo1998]

A cleanup race condition exists here. A selector may read pendingBroker, while onUnloadAttemptCompleted could remove the key, and this put might reinsert a replacement after the attempt has already finished.

While ConcurrentHashMap prevents structural corruption, it does not make this lifecycle transition atomic. Could we use a conditional replacement, such as:

pendingBundleToBroker.replace(bundle, pendingBroker, broker)

and treat a failed replacement as the attempt already being closed? An attempt-id or state object would be even safer. A latch-based concurrency test could cover this race.

[void-ptr974]

I handled this in the same update by removing the fallback write-back. A temporary fallback now leaves the pending plan unchanged, so it cannot reinsert the entry after completion. The test covers fallback, restoration of the original destination, and cleanup.

[Denovo1998]

This cleanup depends on the completion of the unload at the source, not on the completion of the target assignment.

For the legacy load manager, OwnedBundle#handleUnloadRequest finishes after closing topics and removing ownership. The destination affinity is later consumed by ModularLoadManagerWrapper#getLeastLoaded when a lookup actually occurs. As a result, a delayed lookup or destination failure can happen after this callback has removed the pending plan.

Is this intentionally outside the pending-plan contract? If the plan is meant to cover assignment retries, then cleanup should be driven by consumption or expiry, not by the return of the unload RPC.

[void-ptr974]

On this lifecycle point, the AvgShedder entry is only used to select the unload destination. Before the source is unloaded, the destination is stored separately in bundleBrokerAffinityMap and later consumed by lookup. Therefore, this cleanup does not remove the target affinity, and the existing affinity tests cover this flow.

[void-ptr974]

Yes, this is intentional. The pending plan ends after destination selection; directed unload stores the target in the existing one-shot affinity entry consumed by the next lookup. The callback does not indicate that the destination has acquired ownership; the Javadoc now states this explicitly.

[Denovo1998]

The try/finally begins after findBundlesForUnloading, but AvgShedder modifies its shared pending map during that call. If planning throws after inserting one or more entries, onUnloadAttemptCompleted is never called, leaving a partial plan visible to placement.

Could AvgShedder instead construct the destinations in a local plan and only publish it after planning finishes successfully, or at least clear the partially built state when planning fails?

[void-ptr974]

This edge case is theoretically possible if a later broker pair fails after an earlier pair has populated the plan. No unload is issued in that case, the task is rescheduled, and the next AvgShedder pass clears the pending state. To keep this change focused, I suggest handling transactional publication separately if we can reproduce an assignment issue.

[Denovo1998]

Upon re-examining the current threading context, I believe there is a concrete way to reproduce the issue.

updateAll() is not fully synchronized. After cleanupDeadBrokersData() releases the manager monitor, updateAllBrokerData() can mutate loadData.getBrokerData() while doLoadShedding() is still running.

The AvgShedder first snapshots broker names in calculateScoresAndSort(), then later in selectBundleForUnloading() it dereferences:

loadData.getBrokerData().get(underloadedBroker).getLocalData()

for each pair.

As a result, an earlier pair can already publish entries into pendingBundleToBroker, and then a broker from a later pair could disappear, causing planning to fail before findBundlesForUnloading() returns. Because the manager’s try/finally block starts after that call, those entries are never cleaned up.

Since LoadSheddingTask only schedules the next run, the aborted plan remains externally visible until the next shedding pass.

I don’t think full transactional publication is required to resolve this. At a minimum, could AvgShedder clear pendingBundleToBroker on an exceptional exit from findBundlesForUnloading()?

[void-ptr974]

Addressed in d1329a0. The try/finally now includes planning, and AvgShedder clears the entire attempt plan on failure. The exceptional path is covered by testLoadSheddingPassesBundleNameAndCompletesAttempt.

[Denovo1998]

Upon re-examining the current threading context, I believe there is a concrete way to reproduce the issue.

updateAll() is not fully synchronized. After cleanupDeadBrokersData() releases the manager monitor, updateAllBrokerData() can mutate loadData.getBrokerData() while doLoadShedding() is still running.

The AvgShedder first snapshots broker names in calculateScoresAndSort(), then later in selectBundleForUnloading() it dereferences:

loadData.getBrokerData().get(underloadedBroker).getLocalData()

for each pair.

As a result, an earlier pair can already publish entries into pendingBundleToBroker, and then a broker from a later pair could disappear, causing planning to fail before findBundlesForUnloading() returns. Because the manager’s try/finally block starts after that call, those entries are never cleaned up.

Since LoadSheddingTask only schedules the next run, the aborted plan remains externally visible until the next shedding pass.

I don’t think full transactional publication is required to resolve this. At a minimum, could AvgShedder clear pendingBundleToBroker on an exceptional exit from findBundlesForUnloading()?

[Denovo1998]

Removing the first broker resolves AvgShedder’s pending-destination retry issue, but it also alters the overload-retry semantics for all ModularLoadManagerStrategy instances.

Previously, the behavior expanded the candidate set and allowed the placement strategy to determine whether a better broker existed. With this change, we now force a different broker even when all alternatives are worse.

For instance, with an overload threshold of 85%, if broker A is at 90% and B at 98%, removing A forces the selection of B. The LeastLongTermMessageRate strategy also treats any broker above the threshold as +INF, so once A is removed, it cannot distinguish between them.

Another issue is that the expanded candidate set is rebuilt only via applyNamespacePolicies(), while earlier filters—such as anti-affinity, topic-count, and broker filters—are not reapplied. If A was the only broker that passed those filters, removing it forces selection from brokers that were intentionally excluded. Overload handling has historically involved anti-affinity considerations (see #9393).

Could we instead adjust the retry to bypass AvgShedder’s pending destination specifically, rather than modifying candidate semantics for all placement strategies? For example, AvgShedder could ignore a pending destination once that broker is confirmed to be overloaded.

[void-ptr974]

Agreed. The retry no longer removes the first broker; it uses the historical selector to bypass the AvgShedder pending destination while preserving candidate semantics. Covered by testOverloadRetryBypassesBundleAwareSelectionWithoutRemovingCandidates.

[Denovo1998]

One compatibility consideration: existing subclasses of AvgShedder that only override the historical selectBroker(...) method will no longer intercept normal placement requests.

ModularLoadManagerImpl now invokes selectBrokerForBundle(...), which AvgShedder overrides. When there is no pending plan, this path directly calls the private getExpectedBroker(...) logic, bypassing the virtual selectBroker(...) method.

For instance, a custom MyAvgShedder extends AvgShedder with an overridden four-argument selector functioned correctly before this change but is silently bypassed now.

Could the case without a pending plan instead delegate to selectBroker(...), reserving the name-aware path only for an actual pending destination? That would preserve existing subclass fallback behavior while still addressing the mutable BundleData key issue.

[void-ptr974]

Addressed in d1329a0. Without a usable pending destination, the bundle-aware path delegates to the historical virtual selector, preserving subclass dispatch.