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

pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/LoadSheddingStrategy.java

@@ -44,4 +44,14 @@ public interface LoadSheddingStrategy {
      * @param activeBrokers active Brokers
      */
     default void onActiveBrokersChange(Set<String> activeBrokers) {}
+
+    /**
+     * Triggered after the load manager has finished processing bundles selected for unloading.
+     *
+     * <p>This is called whether or not individual unloads were issued successfully.
+     *
+     * @param bundles
+     *            The stable names of the bundles that were processed.
+     */
+    default void onUnloadAttemptCompleted(Set<String> bundles) {}
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/ModularLoadManagerStrategy.java

@@ -41,11 +41,31 @@ public interface ModularLoadManagerStrategy {
      *            The load data from the leader broker.
      * @param conf
      *            The service configuration.
-     * @return The name of the selected broker as it appears on ZooKeeper.
+     * @return The selected broker ID.
      */
     Optional<String> selectBroker(Set<String> candidates, BundleData bundleToAssign, LoadData loadData,
             ServiceConfiguration conf);
 
+    /**
+     * Find a suitable broker to assign the given bundle when its stable name is available.
+     *
+     * @param candidates
+     *            The candidates for which the bundle may be assigned.
+     * @param bundle
+     *            The stable bundle name.
+     * @param bundleToAssign
+     *            The data for the bundle to assign.
+     * @param loadData
+     *            The load data from the leader broker.
+     * @param conf
+     *            The service configuration.
+     * @return The selected broker ID.
+     */
+    default Optional<String> selectBrokerForBundle(Set<String> candidates, String bundle, BundleData bundleToAssign,
+                                                   LoadData loadData, ServiceConfiguration conf) {
+        return selectBroker(candidates, bundleToAssign, loadData, conf);
+    }
+
     /**
      * Triggered when active brokers change.
      */

pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/impl/AvgShedder.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pulsar.broker.loadbalance.impl;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.Multimap;
 import com.google.common.hash.Hashing;
@@ -33,6 +34,8 @@
 import java.util.Optional;
 import java.util.Random;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import lombok.CustomLog;
 import org.apache.commons.lang3.mutable.MutableDouble;
 import org.apache.commons.lang3.mutable.MutableInt;
@@ -48,8 +51,8 @@
 
 @CustomLog
 public class AvgShedder implements LoadSheddingStrategy, ModularLoadManagerStrategy {
-    // map bundle to broker.
-    private final Map<BundleData, String> bundleBrokerMap = new HashMap<>();
+    // Pending destinations produced by the current load-shedding attempt, keyed by stable bundle name.
+    private final ConcurrentMap<String, String> pendingBundleToBroker = new ConcurrentHashMap<>();
     // map broker to Scores. scores:0-100
     private final Map<String, Double> brokerScoreMap = new HashMap<>();
     // map broker hit count for high threshold/low threshold
@@ -59,6 +62,8 @@ public class AvgShedder implements LoadSheddingStrategy, ModularLoadManagerStrat
 
     @Override
     public Multimap<String, String> findBundlesForUnloading(LoadData loadData, ServiceConfiguration conf) {
+        // Plans are consumed after each shedding attempt. Clear any state left by a nonstandard caller.
+        pendingBundleToBroker.clear();
         // result returned by shedding, map broker to bundles.
         Multimap<String, String> selectedBundlesCache = ArrayListMultimap.create();
 
@@ -193,7 +198,7 @@ private void selectBundleForUnloading(LoadData loadData, String overloadedBroker
             double traffic = e.getRight();
             if (traffic > 0 && traffic <= trafficMarkedToOffload.doubleValue()) {
                 selectedBundlesCache.put(overloadedBroker, bundle.getKey());
-                bundleBrokerMap.put(bundle.getValue(), underloadedBroker);
+                pendingBundleToBroker.put(bundle.getKey(), underloadedBroker);
                 trafficMarkedToOffload.add(-traffic);
                 log.debug().attr("bundle", bundle).attr("isMsgRateToOffload", isMsgRateToOffload)
                         .attr("traffic", traffic)
@@ -204,7 +209,17 @@ private void selectBundleForUnloading(LoadData loadData, String overloadedBroker
 
     @Override
     public void onActiveBrokersChange(Set<String> activeBrokers) {
-        LoadSheddingStrategy.super.onActiveBrokersChange(activeBrokers);
+        // Keep a stale pending destination so selection can replace it once and reuse the replacement on retry.
+    }
+
+    @Override
+    public void onUnloadAttemptCompleted(Set<String> bundles) {
+        pendingBundleToBroker.keySet().removeAll(bundles);
+    }
+
+    @VisibleForTesting
+    boolean hasPendingDestination(String bundle) {
+        return pendingBundleToBroker.containsKey(bundle);
     }
 
     private List<String> calculateScoresAndSort(LoadData loadData, ServiceConfiguration conf) {
@@ -273,20 +288,39 @@ private List<Pair<String, String>> findBrokerPairs(List<String> brokers,
     @Override
     public Optional<String> selectBroker(Set<String> candidates, BundleData bundleToAssign, LoadData loadData,
                                          ServiceConfiguration conf) {
-        final var brokerToUnload = bundleBrokerMap.getOrDefault(bundleToAssign, null);
-        if (brokerToUnload == null || !candidates.contains(bundleBrokerMap.get(bundleToAssign))) {
-            // cluster initializing or broker is shutdown
-            if (!bundleBrokerMap.containsKey(bundleToAssign)) {
+        return candidates.isEmpty() ? Optional.empty() : Optional.of(getExpectedBroker(candidates, bundleToAssign));
+    }
+
+    @Override
+    public Optional<String> selectBrokerForBundle(Set<String> candidates, String bundle, BundleData bundleToAssign,
+                                                   LoadData loadData, ServiceConfiguration conf) {
+        return bundle == null
+                ? selectBroker(candidates, bundleToAssign, loadData, conf)
+                : selectBrokerWithBundleName(candidates, bundle, bundleToAssign);
+    }
+
+    private Optional<String> selectBrokerWithBundleName(Set<String> candidates, String bundle,
+                                                         BundleData bundleToAssign) {
+        if (candidates.isEmpty()) {
+            return Optional.empty();
+        }
+        final var pendingBroker = pendingBundleToBroker.get(bundle);
+        if (pendingBroker == null || !candidates.contains(pendingBroker)) {
+            // cluster initializing or the pending broker is shutdown or filtered out
+            if (pendingBroker == null) {
                 log.debug("cluster is initializing");
             } else {
-                log.debug().attr("broker", bundleBrokerMap.get(bundleToAssign)).attr("candidates", candidates)
+                log.debug().attr("broker", pendingBroker).attr("candidates", candidates)
                         .log("expected broker is shutdown");
             }
             String broker = getExpectedBroker(candidates, bundleToAssign);
-            bundleBrokerMap.put(bundleToAssign, broker);
+            if (pendingBroker != null) {
+                // Keep a replacement only for the remainder of this unload attempt, including the retry path.
+                pendingBundleToBroker.put(bundle, broker);
+            }
             return Optional.of(broker);
         } else {
-            return Optional.of(brokerToUnload);
+            return Optional.of(pendingBroker);
         }
     }
 

pulsar-broker/src/main/java/org/apache/pulsar/broker/loadbalance/impl/ModularLoadManagerImpl.java

@@ -650,58 +650,63 @@ public synchronized void doLoadShedding() {
 
         Set<String> sheddingExcludedNamespaces = conf.getLoadBalancerSheddingExcludedNamespaces();
         final Multimap<String, String> bundlesToUnload = loadSheddingStrategy.findBundlesForUnloading(loadData, conf);
+        final Set<String> plannedBundles = new HashSet<>(bundlesToUnload.values());
 
-        bundlesToUnload.asMap().forEach((broker, bundles) -> {
-            AtomicBoolean unloadBundleForBroker = new AtomicBoolean(false);
-            bundles.forEach(bundle -> {
-                final String namespaceName = LoadManagerShared.getNamespaceNameFromBundleName(bundle);
-                final String bundleRange = LoadManagerShared.getBundleRangeFromBundleName(bundle);
-                if (sheddingExcludedNamespaces.contains(namespaceName)) {
-                    log.debug().attr("class", loadSheddingStrategy.getClass().getSimpleName())
-                            .attr("namespace", namespaceName)
-                            .log("Skipping load shedding for namespace");
-                    return;
-                }
-                if (!shouldNamespacePoliciesUnload(namespaceName, bundleRange, broker)) {
-                    return;
-                }
+        try {
+            bundlesToUnload.asMap().forEach((broker, bundles) -> {
+                AtomicBoolean unloadBundleForBroker = new AtomicBoolean(false);
+                bundles.forEach(bundle -> {
+                    final String namespaceName = LoadManagerShared.getNamespaceNameFromBundleName(bundle);
+                    final String bundleRange = LoadManagerShared.getBundleRangeFromBundleName(bundle);
+                    if (sheddingExcludedNamespaces.contains(namespaceName)) {
+                        log.debug().attr("class", loadSheddingStrategy.getClass().getSimpleName())
+                                .attr("namespace", namespaceName)
+                                .log("Skipping load shedding for namespace");
+                        return;
+                    }
+                    if (!shouldNamespacePoliciesUnload(namespaceName, bundleRange, broker)) {
+                        return;
+                    }
 
-                if (!shouldAntiAffinityNamespaceUnload(namespaceName, bundleRange, broker)) {
-                    return;
-                }
-                NamespaceBundle bundleToUnload = LoadManagerShared.getNamespaceBundle(pulsar, bundle);
-                Optional<String> destBroker = this.selectBroker(bundleToUnload);
-                if (!destBroker.isPresent()) {
-                    log.info().attr("class", loadSheddingStrategy.getClass().getSimpleName())
-                            .attr("bundle", bundle).attr("broker", broker)
-                            .log("No broker available to unload bundle from broker");
-                    return;
-                }
-                if (destBroker.get().equals(broker)) {
-                    log.warn().attr("class", loadSheddingStrategy.getClass().getSimpleName())
-                            .attr("broker", destBroker.get()).attr("bundle", bundle)
-                            .log("The destination broker is the same as the current owner broker for bundle");
-                    return;
-                }
+                    if (!shouldAntiAffinityNamespaceUnload(namespaceName, bundleRange, broker)) {
+                        return;
+                    }
+                    NamespaceBundle bundleToUnload = LoadManagerShared.getNamespaceBundle(pulsar, bundle);
+                    Optional<String> destBroker = this.selectBroker(bundleToUnload);
+                    if (!destBroker.isPresent()) {
+                        log.info().attr("class", loadSheddingStrategy.getClass().getSimpleName())
+                                .attr("bundle", bundle).attr("broker", broker)
+                                .log("No broker available to unload bundle from broker");
+                        return;
+                    }
+                    if (destBroker.get().equals(broker)) {
+                        log.warn().attr("class", loadSheddingStrategy.getClass().getSimpleName())
+                                .attr("broker", destBroker.get()).attr("bundle", bundle)
+                                .log("The destination broker is the same as the current owner broker for bundle");
+                        return;
+                    }
 
-                log.info().attr("class", loadSheddingStrategy.getClass().getSimpleName())
-                        .attr("bundle", bundle).attr("sourceBroker", broker).attr("destBroker", destBroker.get())
-                        .log("Unloading bundle from source broker to dest broker");
-                try {
-                    pulsar.getAdminClient().namespaces()
-                            .unloadNamespaceBundle(namespaceName, bundleRange, destBroker.get());
-                    loadData.getRecentlyUnloadedBundles().put(bundle, System.currentTimeMillis());
-                    unloadBundleCount++;
-                    unloadBundleForBroker.set(true);
-                } catch (PulsarServerException | PulsarAdminException e) {
-                    log.warn().attr("bundle", bundle).attr("broker", broker).exception(e)
-                            .log("Error when trying to perform load shedding on for broker");
+                    log.info().attr("class", loadSheddingStrategy.getClass().getSimpleName())
+                            .attr("bundle", bundle).attr("sourceBroker", broker).attr("destBroker", destBroker.get())
+                            .log("Unloading bundle from source broker to dest broker");
+                    try {
+                        pulsar.getAdminClient().namespaces()
+                                .unloadNamespaceBundle(namespaceName, bundleRange, destBroker.get());
+                        loadData.getRecentlyUnloadedBundles().put(bundle, System.currentTimeMillis());
+                        unloadBundleCount++;
+                        unloadBundleForBroker.set(true);
+                    } catch (PulsarServerException | PulsarAdminException e) {
+                        log.warn().attr("bundle", bundle).attr("broker", broker).exception(e)
+                                .log("Error when trying to perform load shedding on for broker");
+                    }
+                });
+                if (unloadBundleForBroker.get()) {
+                    unloadBrokerCount++;
                 }
             });
-            if (unloadBundleForBroker.get()) {
-                unloadBrokerCount++;
-            }
-        });
+        } finally {
+            loadSheddingStrategy.onUnloadAttemptCompleted(plannedBundles);
+        }
 
         updateBundleUnloadingMetrics();
     }
@@ -959,10 +964,10 @@ Optional<String> selectBroker(final ServiceUnitId serviceUnit) {
             if (sheddingExcludedNamespaces.contains(namespaceNameFromBundleName)) {
                 log.debug().attr("bundle", bundle).log("Use round robin broker selector for bundle");
                 broker = sheddingExcludedNamespaceSelectionStrategy
-                        .selectBroker(brokerCandidateCache, data, loadData, conf);
+                        .selectBrokerForBundle(brokerCandidateCache, bundle, data, loadData, conf);
             } else {
                 // Choose a broker among the potentially smaller filtered list, when possible
-                broker = placementStrategy.selectBroker(brokerCandidateCache, data, loadData, conf);
+                broker = placementStrategy.selectBrokerForBundle(brokerCandidateCache, bundle, data, loadData, conf);
             }
             log.debug().attr("selectedBroker", broker).attr("candidates", brokerCandidateCache)
                     .log("Selected broker from candidate brokers");
@@ -981,7 +986,7 @@ Optional<String> selectBroker(final ServiceUnitId serviceUnit) {
                         getAvailableBrokers(),
                         brokerTopicLoadingPredicate);
                 Optional<String> brokerTmp =
-                        placementStrategy.selectBroker(brokerCandidateCache, data, loadData, conf);
+                        placementStrategy.selectBrokerForBundle(brokerCandidateCache, bundle, data, loadData, conf);
                 if (brokerTmp.isPresent()) {
                     broker = brokerTmp;
                 }
@@ -990,6 +995,11 @@ Optional<String> selectBroker(final ServiceUnitId serviceUnit) {
         }
     }
 
+    @VisibleForTesting
+    LoadData getLoadData() {
+        return loadData;
+    }
+
     /**
      * As any broker, start the load manager.
      *