fix ns isolation policy for repl namespaces

iosdev747 · Jul 29, 2024 · debf0ff · debf0ff
1 parent fc2e314
commit debf0ff
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 6 deletions.
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/ClustersBase.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/ClustersBase.java
@@ -36,6 +36,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.ws.rs.DELETE;
+import javax.ws.rs.DefaultValue;
 import javax.ws.rs.GET;
 import javax.ws.rs.POST;
 import javax.ws.rs.PUT;
@@ -55,6 +56,7 @@
 import org.apache.pulsar.broker.admin.AdminResource;
 import org.apache.pulsar.broker.web.RestException;
 import org.apache.pulsar.client.admin.PulsarAdmin;
+import org.apache.pulsar.client.admin.PulsarAdminException;
 import org.apache.pulsar.common.naming.Constants;
 import org.apache.pulsar.common.naming.NamedEntity;
 import org.apache.pulsar.common.policies.data.BrokerNamespaceIsolationData;
@@ -705,7 +707,9 @@ public void setNamespaceIsolationPolicy(
         @ApiParam(value = "The namespace isolation policy name", required = true)
         @PathParam("policyName") String policyName,
         @ApiParam(value = "The namespace isolation policy data", required = true)
-        NamespaceIsolationDataImpl policyData
+        NamespaceIsolationDataImpl policyData,
+        @DefaultValue("true")
+        @QueryParam("unloadBundles") boolean unload
     ) {
         validateSuperUserAccessAsync()
                 .thenCompose(__ -> validatePoliciesReadOnlyAccessAsync())
@@ -723,7 +727,13 @@ public void setNamespaceIsolationPolicy(
                     nsIsolationPolicies.setPolicy(policyName, policyData);
                     return namespaceIsolationPolicies()
                                     .setIsolationDataAsync(cluster, old -> nsIsolationPolicies.getPolicies());
-                }).thenCompose(__ -> filterAndUnloadMatchedNamespaceAsync(policyData))
+                }).thenCompose(__ -> {
+                    if (unload) {
+                        return filterAndUnloadMatchedNamespaceAsync(cluster, policyData);
+                    } else {
+                        return CompletableFuture.completedFuture(null);
+                    }
+                })
                 .thenAccept(__ -> {
                     log.info("[{}] Successful to update clusters/{}/namespaceIsolationPolicies/{}.",
                             clientAppId(), cluster, policyName);
@@ -757,7 +767,8 @@ public void setNamespaceIsolationPolicy(
     /**
      * Get matched namespaces; call unload for each namespaces.
      */
-    private CompletableFuture<Void> filterAndUnloadMatchedNamespaceAsync(NamespaceIsolationDataImpl policyData) {
+    private CompletableFuture<Void> filterAndUnloadMatchedNamespaceAsync(String cluster,
+                                                                         NamespaceIsolationDataImpl policyData) {
         PulsarAdmin adminClient;
         try {
             adminClient = pulsar().getAdminClient();
@@ -770,8 +781,13 @@ private CompletableFuture<Void> filterAndUnloadMatchedNamespaceAsync(NamespaceIs
                             .map(tenant -> adminClient.namespaces().getNamespacesAsync(tenant));
                     return FutureUtil.waitForAll(completableFutureStream)
                             .thenApply(namespaces -> {
-                                // if namespace match any policy regex, add it to ns list to be unload.
+                                // Filter namespaces that have current cluster in their replication_clusters
+                                // if namespace match any policy regex, add it to ns list to be unloaded.
                                 return namespaces.stream()
+                                        .filter(namespaceName -> adminClient.namespaces()
+                                                .getPoliciesAsync(namespaceName)
+                                                .thenApply(policies -> policies.replication_clusters.contains(cluster))
+                                                .join())
                                         .filter(namespaceName ->
                                                 policyData.getNamespaces().stream().anyMatch(namespaceName::matches))
                                         .collect(Collectors.toList());
@@ -781,7 +797,33 @@ private CompletableFuture<Void> filterAndUnloadMatchedNamespaceAsync(NamespaceIs
                         return CompletableFuture.completedFuture(null);
                     }
                     List<CompletableFuture<Void>> futures = shouldUnloadNamespaces.stream()
-                            .map(namespaceName -> adminClient.namespaces().unloadAsync(namespaceName))
+                            .map(namespaceName -> {
+                                try {
+                                    return adminClient.namespaces()
+                                            .getPolicies(namespaceName);
+                                } catch (PulsarAdminException e) {
+                                    log.warn("[{}] Failed to get policy for {} namespace.", clientAppId(),
+                                            namespaceName, e);
+                                    throw new RuntimeException(e);
+                                }
+                            })
+                            .map(policies -> {
+                                final List<CompletableFuture<Void>> unloadFutures = new ArrayList<>();
+                                List<String> boundaries = policies.bundles.getBoundaries();
+                                for (int i = 0; i < boundaries.size() - 1; i++) {
+                                    String bundle = String.format("%s_%s", boundaries.get(i), boundaries.get(i + 1));
+                                    try {
+                                        unloadFutures.add(
+                                                pulsar().getAdminClient().namespaces().unloadNamespaceBundleAsync(
+                                                        namespaceName.toString(), bundle));
+                                    } catch (PulsarServerException e) {
+                                        log.error("[{}] Failed to unload namespace {}", clientAppId(), namespaceName,
+                                                e);
+                                        throw new RestException(e);
+                                    }
+                                }
+                                return FutureUtil.waitForAll(unloadFutures);
+                            })
                             .collect(Collectors.toList());
                     return FutureUtil.waitForAll(futures)
                             .thenAccept(__ -> {

diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/AdminTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/AdminTest.java
@@ -309,7 +309,7 @@ public void clusters() throws Exception {
                         .build())
                 .build();
         asyncRequests(ctx -> clusters.setNamespaceIsolationPolicy(ctx,
-                "use", "policy1", policyData));
+                "use", "policy1", policyData, true));
         asyncRequests(ctx -> clusters.getNamespaceIsolationPolicies(ctx, "use"));
 
         try {