[improve][broker] Load topic policies on non-persistent topic load and gate the policy replay (#26134)

(cherry picked from commit 0629dc5)

Author: lhotari

pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java

@@ -1729,6 +1729,15 @@ The max allowed delay for delayed delivery (in milliseconds). If the broker rece
                 + "please enable the system topic first.")
     private boolean topicLevelPoliciesEnabled = true;
 
+    @FieldContext(
+            category = CATEGORY_SERVER,
+            doc = "When enabled, all registered topic-policy listeners in a namespace are re-notified with the current"
+                    + " topic policies after the namespace's topic-policy cache finishes its initial load. Topics load"
+                    + " and apply their own policies when they are loaded, so this broadcast is normally redundant; it"
+                    + " is only needed for custom plugins that register TopicPolicyListeners and depend on it for"
+                    + " backwards compatibility. Disabled by default.")
+    private boolean topicPolicyListenerReplayEnabled = false;
+
     @FieldContext(
             category = CATEGORY_SERVER,
             doc = """

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/AbstractTopic.java

@@ -31,6 +31,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -48,6 +49,7 @@
 import java.util.concurrent.atomic.AtomicLongFieldUpdater;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Function;
 import java.util.function.ToLongFunction;
 import lombok.Getter;
 import lombok.Setter;
@@ -58,6 +60,7 @@
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pulsar.broker.PulsarService;
 import org.apache.pulsar.broker.ServiceConfiguration;
+import org.apache.pulsar.broker.loadbalance.extensions.ExtensibleLoadManagerImpl;
 import org.apache.pulsar.broker.resourcegroup.ResourceGroup;
 import org.apache.pulsar.broker.resourcegroup.ResourceGroupPublishLimiter;
 import org.apache.pulsar.broker.service.BrokerServiceException.ConsumerBusyException;
@@ -114,6 +117,10 @@ public abstract class AbstractTopic implements Topic, TopicPolicyListener {
 
     protected final BrokerService brokerService;
 
+    // Wraps this topic as a TopicPolicyListener so topic-policy updates received while the initial policy is still
+    // loading are buffered and applied in order once initTopicPolicy() completes initialization.
+    protected final TopicPolicyListenerWrapper topicPolicyListener = new TopicPolicyListenerWrapper(this);
+
     // Prefix for replication cursors
     protected final String replicatorPrefix;
 
@@ -173,7 +180,17 @@ public abstract class AbstractTopic implements Topic, TopicPolicyListener {
             AtomicLongFieldUpdater.newUpdater(AbstractTopic.class, "usageCount");
     private volatile long usageCount = 0;
 
-    private Map<String/*subscription*/, SubscriptionPolicies> subscriptionPolicies = Collections.emptyMap();
+    // Effective per-subscription policies, merged from the local and global topic policies with local precedence.
+    // Unlike the PolicyHierarchyValue-backed fields, subscriptionPolicies is a plain map. It is kept as the merge of
+    // the two scopes below (rather than assigned directly) because the local-before-global initialization order
+    // (see TopicPolicyListenerWrapper) would otherwise let the global map -- empty by default in TopicPolicies --
+    // overwrite and clear the local per-subscription policies that were applied just before it.
+    // subscriptionPolicies is volatile because it is read on dispatch threads (getSubscriptionDispatchRate) while it
+    // is updated on the policy-update thread. localSubscriptionPolicies/globalSubscriptionPolicies are only ever read
+    // and written on the (single) policy-update thread, so they don't need to be volatile.
+    private volatile Map<String/*subscription*/, SubscriptionPolicies> subscriptionPolicies = Collections.emptyMap();
+    private Map<String/*subscription*/, SubscriptionPolicies> localSubscriptionPolicies = Collections.emptyMap();
+    private Map<String/*subscription*/, SubscriptionPolicies> globalSubscriptionPolicies = Collections.emptyMap();
 
     protected final LongAdder msgOutFromRemovedSubscriptions = new LongAdder();
     protected final LongAdder bytesOutFromRemovedSubscriptions = new LongAdder();
@@ -310,11 +327,36 @@ protected void updateTopicPolicy(TopicPolicies data) {
         topicPolicies.getEntryFilters().updateTopicValue(data.getEntryFilters(), isGlobalPolicies);
         topicPolicies.getDispatcherPauseOnAckStatePersistentEnabled()
                 .updateTopicValue(data.getDispatcherPauseOnAckStatePersistentEnabled(), isGlobalPolicies);
-        this.subscriptionPolicies = data.getSubscriptionPolicies();
+
+        // Merge instead of assigning directly: keep the local and global per-subscription policies separately and
+        // recompute the effective map with local precedence, so applying the (default-empty) global map does not
+        // clear the local per-subscription policies during the local-before-global initialization.
+        if (isGlobalPolicies) {
+            globalSubscriptionPolicies = data.getSubscriptionPolicies();
+        } else {
+            localSubscriptionPolicies = data.getSubscriptionPolicies();
+        }
+        subscriptionPolicies = mergeSubscriptionPolicies(globalSubscriptionPolicies, localSubscriptionPolicies);
 
         updateEntryFilters();
     }
 
+    // Merges the global and local per-subscription policies with local precedence: a subscription present in the
+    // local policies keeps its local value; otherwise the global value (if any) is used.
+    private static Map<String, SubscriptionPolicies> mergeSubscriptionPolicies(
+            Map<String, SubscriptionPolicies> globalSubscriptionPolicies,
+            Map<String, SubscriptionPolicies> localSubscriptionPolicies) {
+        if (globalSubscriptionPolicies.isEmpty()) {
+            return localSubscriptionPolicies;
+        }
+        if (localSubscriptionPolicies.isEmpty()) {
+            return globalSubscriptionPolicies;
+        }
+        Map<String, SubscriptionPolicies> merged = new HashMap<>(globalSubscriptionPolicies);
+        merged.putAll(localSubscriptionPolicies);
+        return merged;
+    }
+
     protected void updateTopicPolicyByNamespacePolicy(Policies namespacePolicies) {
         if (log.isDebugEnabled()) {
             log.debug("[{}]updateTopicPolicyByNamespacePolicy,data={}", topic, namespacePolicies);
@@ -550,19 +592,68 @@ protected boolean isProducersExceeded(boolean isRemote) {
     }
 
     protected TopicPolicyListener getTopicPolicyListener() {
-        return this;
-    }
-
-    protected void registerTopicPolicyListener() {
-        brokerService.getPulsar().getTopicPoliciesService()
-                .registerListenerAsync(TopicName.getPartitionedTopicName(topic), getTopicPolicyListener());
+        return topicPolicyListener;
     }
 
     protected void unregisterTopicPolicyListener() {
         brokerService.getPulsar().getTopicPoliciesService()
                 .unregisterListener(TopicName.getPartitionedTopicName(topic), getTopicPolicyListener());
     }
 
+    /**
+     * Registers the topic-policy listener and applies the topic's initial policies (global and local) to this topic.
+     * Shared by {@link org.apache.pulsar.broker.service.persistent.PersistentTopic} and
+     * {@link org.apache.pulsar.broker.service.nonpersistent.NonPersistentTopic} so both load their own policies on
+     * topic load, which removes the need to broadcast every topic's policy when a namespace's policy cache finishes
+     * loading (see {@code topicPolicyListenerReplayEnabled}).
+     *
+     * <p>Each call re-initializes the listener wrapper and, whatever the outcome, always completes its initialization
+     * afterwards, so the wrapper never stays in the buffering phase (dropping updates) even if policy loading fails.
+     * This makes the method safe to run again (e.g. a future retry); runs are expected to be serialized.
+     */
+    protected CompletableFuture<Void> initTopicPolicy() {
+        final var topicPoliciesService = brokerService.getPulsar().getTopicPoliciesService();
+        final var partitionedTopicName = TopicName.getPartitionedTopicName(topic);
+
+        // Begin a fresh initialization phase: updates are buffered until initialization completes below. This resets
+        // any previous phase so the method can be run again.
+        topicPolicyListener.startInitialization();
+        CompletableFuture<Void> initTopicPolicyFuture =
+                topicPoliciesService.registerListenerAsync(partitionedTopicName, topicPolicyListener)
+                        .thenCompose(registered -> {
+                            if (!registered) {
+                                return CompletableFuture.completedFuture(null);
+                            }
+                            if (ExtensibleLoadManagerImpl.isInternalTopic(topic)) {
+                                // Internal topics don't load topic-level policies
+                                return CompletableFuture.completedFuture(null);
+                            }
+                            // future for fetching global topic policies
+                            CompletableFuture<Optional<TopicPolicies>> globalPoliciesFuture =
+                                    topicPoliciesService.getTopicPoliciesAsync(partitionedTopicName,
+                                            TopicPoliciesService.GetType.GLOBAL_ONLY);
+                            // future for fetching local topic policies
+                            CompletableFuture<Optional<TopicPolicies>> localPoliciesFuture =
+                                    topicPoliciesService.getTopicPoliciesAsync(partitionedTopicName,
+                                            TopicPoliciesService.GetType.LOCAL_ONLY);
+                            return globalPoliciesFuture.thenCombine(localPoliciesFuture, (global, local) -> {
+                                // finally update the topic policies with the latest value or loaded value
+                                return CompletableFuture.runAsync(() ->
+                                                topicPolicyListener.completeInitialization(global.orElse(null),
+                                                        local.orElse(null)),
+                                        getPoliciesNotifyThread());
+                            }).thenCompose(Function.identity());
+                        });
+        // Whatever the outcome -- success, failure, or the listener not being registered -- make sure the wrapper
+        // leaves the initialization (buffering) phase, so it forwards any buffered value plus all future live updates
+        // instead of dropping them. This is a no-op when the loaded policies were already applied above. Return the
+        // whenComplete stage (not initTopicPolicyFuture) so the returned future completes only after this has run, and
+        // whenComplete's pass-through semantics carry the original success or failure to the caller's initialize().
+        return initTopicPolicyFuture.whenCompleteAsync((v, ex) -> {
+            topicPolicyListener.completeInitializationUnlessAlreadyCompleted();
+        }, getPoliciesNotifyThread());
+    }
+
     protected boolean isSameAddressProducersExceeded(Producer producer) {
         if (isSystemTopic() || producer.isRemote()) {
             return false;

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/SystemTopicBasedTopicPoliciesService.java

@@ -864,18 +864,43 @@ private void initPolicesCache(SystemTopicClient.Reader<PulsarEvent> reader, Comp
                     log.debug("[{}] Reach the end of the system topic.", reader.getSystemTopic().getTopicName());
                 }
 
-                // replay policy message
-                List<CompletableFuture<Void>> notifyFutures = new ArrayList<>();
-                for (Map.Entry<TopicName, TopicPolicies> entry : policiesCache.entrySet()) {
-                    TopicName topicName = entry.getKey();
-                    TopicPolicies policies = entry.getValue();
-                    notifyFutures.add(notifyListenersForTopicAsync(topicName, policies));
+                // Optionally re-notify the topic-policy listeners for this namespace with the cached policies.
+                // Topics load their own policies on load (AbstractTopic#initTopicPolicy), so this replay is only
+                // needed for custom plugins that register TopicPolicyListeners and rely on the broadcast; it is
+                // off by default (topicPolicyListenerReplayEnabled).
+                if (pulsarService.getConfiguration().isTopicPolicyListenerReplayEnabled()) {
+                    NamespaceName namespaceObject = reader.getSystemTopic().getTopicName().getNamespaceObject();
+                    FutureUtil.completeAfter(future, replayTopicPolicyListeners(namespaceObject));
+                } else {
+                    future.complete(null);
                 }
-                FutureUtil.completeAfter(future, FutureUtil.waitForAll(notifyFutures));
             }
         });
     }
 
+    /**
+     * Re-notifies the registered topic-policy listeners for every topic in {@code namespace} with the currently
+     * cached policies (both local and global). Topics apply their own policies on load, so this is only needed for
+     * custom plugins that register {@link TopicPolicyListener}s and rely on the broadcast when a namespace's policy
+     * cache finishes loading (see {@code topicPolicyListenerReplayEnabled}).
+     */
+    @VisibleForTesting
+    CompletableFuture<Void> replayTopicPolicyListeners(NamespaceName namespace) {
+        List<CompletableFuture<Void>> notifyFutures = new ArrayList<>();
+        addNamespacePolicyNotifications(policiesCache, namespace, notifyFutures);
+        addNamespacePolicyNotifications(globalPoliciesCache, namespace, notifyFutures);
+        return FutureUtil.waitForAll(notifyFutures);
+    }
+
+    private void addNamespacePolicyNotifications(Map<TopicName, TopicPolicies> cache, NamespaceName namespace,
+                                                 List<CompletableFuture<Void>> notifyFutures) {
+        for (Map.Entry<TopicName, TopicPolicies> entry : cache.entrySet()) {
+            if (Objects.equals(entry.getKey().getNamespaceObject(), namespace)) {
+                notifyFutures.add(notifyListenersForTopicAsync(entry.getKey(), entry.getValue()));
+            }
+        }
+    }
+
     // Full teardown of a namespace's topic-policies state: removes and closes the reader, the message-handler
     // tracker, the cached policies and the init future. Used when the whole namespace is unloaded.
     @VisibleForTesting