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

[void-ptr974]

Motivation

AvgShedder records a destination broker while planning bundle unloading. The cache previously used BundleData as its key, but BundleData has value-based equals and hashCode over load statistics that are updated in place.

This causes two correctness problems:

• a key can become unreachable after its hash changes during a load-report refresh;
• different bundles with equal load data can collide and overwrite each other's destination.

Both cases can assign a bundle to a broker different from the destination selected by the shedding plan.

Modifications

• Key AvgShedder pending destinations by the canonical bundle name instead of mutable BundleData.
• Add selectBrokerForBundle(...) as a backward-compatible default method and pass the stable bundle name through ModularLoadManagerImpl.
• Scope pending destinations to one load-shedding attempt and remove them after the manager finishes processing that attempt.
• Store pending destinations in a ConcurrentHashMap because planning and placement can access the strategy concurrently.
• Preserve the original planned destination when selection temporarily falls back because that broker is unavailable; the fallback is not written back over the plan.
• Exclude the first selected broker from the widened candidate set when retrying an overloaded placement.
• Keep the historical four-argument AvgShedder selector as an uncached fallback and mark that AvgShedder override deprecated because it cannot receive a stable bundle name.
• Return Optional.empty() for empty candidate sets.

Compatibility

The new strategy methods are default methods, so existing ModularLoadManagerStrategy and LoadSheddingStrategy implementations remain source- and binary-compatible.

The historical four-argument method remains callable and is not scheduled for removal. Its AvgShedder implementation is intentionally uncached; callers that have a stable bundle name should use selectBrokerForBundle(...).

Verifying this change

• GitHub CI checks have not been reported for the latest commit yet.
• Verified locally with:

./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.testOverloadedBrokerIsExcludedFromRetry

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

./gradlew --no-daemon \
  :pulsar-broker:checkstyleMain \
  :pulsar-broker:checkstyleTest

The tests cover mutable and equal-but-distinct BundleData, stable-name wiring, default-method compatibility, empty candidates, temporary fallback without replacing the original plan, cleanup of all pending bundles, overloaded broker retry exclusion, and final bundle affinity behavior.

The overload retry regression test was also checked against the unpatched branch: it fails because the retry candidate set still contains the first overloaded broker.

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

• Dependencies
• Public API
  • Adds backward-compatible default methods.
  • Deprecates only the historical AvgShedder override; it remains 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.

[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.

[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.