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

[void-ptr974]

Motivation

AvgShedder records the destination broker selected while shedding a bundle.
The cache previously used BundleData as its key. BundleData has
value-based equals and hashCode over mutable load statistics, which are
updated in place as load reports are refreshed.

As a result, a planned destination can no longer be found after the key's hash
changes. Different bundles with equal load data can also share the same key and
overwrite each other's planned destination. Both cases can cause AvgShedder to
select a broker different from the planned shedding target.

Modifications

• Key AvgShedder planned destinations by the canonical bundle name.
• Add selectBrokerForBundle to pass the stable bundle name through the
  placement path. Its default delegates to the existing four-argument method,
  preserving existing placement strategy implementations.
• Treat AvgShedder destinations as a pending plan for one load-shedding
  attempt. The plan survives the assignment retry path, including replacement
  of an unavailable destination, and is removed after the manager finishes
  processing that attempt.
• Store pending plans in a ConcurrentHashMap because shedding and placement
  can access the strategy concurrently.
• Keep the legacy four-argument selector as an uncached fallback. Both it and
  the name-aware path return Optional.empty() when there are no candidates.

Verifying this change

• CI checks have not been reported for the latest commit yet.

• Verified with:

  ./gradlew --no-daemon \
    :pulsar-broker:checkstyleMain \
    :pulsar-broker:checkstyleTest \
    :pulsar-broker:compileTestJava
  
  ./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

  • 12 targeted tests passed: 3 AvgShedder, 8 strategy, and 1 manager
    integration test.
  • Covers equal BundleData instances, in-place refresh, unavailable planned
    destinations and retry reuse, pending-plan cleanup, empty candidates,
    default-method delegation, and the manager's stable-name wiring.

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

• Dependencies (add or upgrade a dependency)

• The public API

  Adds backward-compatible default methods to ModularLoadManagerStrategy and
  LoadSheddingStrategy; existing strategy implementations remain valid.

• 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

[lhotari]

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

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