[fix][meta] Fix ephemeral handling of ZK nodes and fix MockZooKeeper ephemeral and ZK stat handling

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

@@ -285,15 +285,15 @@ public static CompletableFuture<Set<String>> applyNamespacePoliciesAsync(
     public static String getBundleRangeFromBundleName(String bundleName) {
         // the bundle format is property/cluster/namespace/0x00000000_0xFFFFFFFF
         int pos = bundleName.lastIndexOf("/");
-        checkArgument(pos != -1);
+        checkArgument(pos != -1, "Invalid bundle name format: %s", bundleName);
         return bundleName.substring(pos + 1);
     }
 
     // From a full bundle name, extract the namespace name.
     public static String getNamespaceNameFromBundleName(String bundleName) {
         // the bundle format is property/cluster/namespace/0x00000000_0xFFFFFFFF
         int pos = bundleName.lastIndexOf('/');
-        checkArgument(pos != -1);
+        checkArgument(pos != -1, "Invalid bundle name format: %s", bundleName);
         return bundleName.substring(0, pos);
     }
 

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

@@ -583,7 +583,7 @@ private void updateBundleData() {
         for (String bundle : bundleData.keySet()) {
             if (!activeBundles.contains(bundle)){
                 bundleData.remove(bundle);
-                if (pulsar.getLeaderElectionService().isLeader()){
+                if (pulsar.getLeaderElectionService() != null && pulsar.getLeaderElectionService().isLeader()){
                     deleteBundleDataFromMetadataStore(bundle);
                 }
             }

pulsar-broker/src/main/java/org/apache/pulsar/broker/web/PulsarWebResource.java

@@ -1012,7 +1012,7 @@ protected static boolean isLeaderBroker(PulsarService pulsar) {
         if (ExtensibleLoadManagerImpl.isLoadManagerExtensionEnabled(pulsar)) {
             return true;
         }
-        return  pulsar.getLeaderElectionService().isLeader();
+        return pulsar.getLeaderElectionService() != null && pulsar.getLeaderElectionService().isLeader();
     }
 
     public void validateTenantOperation(String tenant, TenantOperation operation) {

pulsar-broker/src/main/java/org/apache/pulsar/common/naming/NamespaceBundles.java

@@ -108,9 +108,11 @@ public int size() {
 
     public void validateBundle(NamespaceBundle nsBundle) throws Exception {
         int idx = Arrays.binarySearch(partitions, nsBundle.getLowerEndpoint());
-        checkArgument(idx >= 0, "Cannot find bundle in the bundles list");
-        checkArgument(nsBundle.getUpperEndpoint().equals(bundles.get(idx).getUpperEndpoint()),
-                "Invalid upper boundary for bundle");
+        checkArgument(idx >= 0, "Cannot find bundle %s in the bundles list", nsBundle);
+        NamespaceBundle foundBundle = bundles.get(idx);
+        Long upperEndpoint = foundBundle.getUpperEndpoint();
+        checkArgument(nsBundle.getUpperEndpoint().equals(upperEndpoint),
+                "Invalid upper boundary for bundle %s. Expected upper boundary of %s", nsBundle, foundBundle);
     }
 
     public NamespaceBundle getFullBundle() {

pulsar-broker/src/test/java/org/apache/pulsar/broker/auth/MockedPulsarServiceBaseTest.java

@@ -150,6 +150,9 @@ public static String getTlsFileForClient(String name) {
 
     private final List<AutoCloseable> closeables = new ArrayList<>();
 
+    // Set to true in test's constructor to use a real Zookeeper (TestZKServer)
+    protected boolean useTestZookeeper;
+
     public MockedPulsarServiceBaseTest() {
         resetConfig();
     }
@@ -363,7 +366,14 @@ protected void afterPulsarStart(PulsarService pulsar) throws Exception {
      * @throws Exception if an error occurs
      */
     protected void restartBroker() throws Exception {
+        restartBroker(null);
+    }
+
+    protected void restartBroker(Consumer<ServiceConfiguration> configurationChanger) throws Exception {
         stopBroker();
+        if (configurationChanger != null) {
+            configurationChanger.accept(conf);
+        }
         startBroker();
         if (pulsarClient == null) {
             pulsarClient = newPulsarClient(lookupUrl.toString(), 0);
@@ -461,7 +471,6 @@ protected PulsarTestContext.Builder createPulsarTestContextBuilder(ServiceConfig
         PulsarTestContext.Builder builder = PulsarTestContext.builder()
                 .spyByDefault()
                 .config(conf)
-                .withMockZookeeper(true)
                 .pulsarServiceCustomizer(pulsarService -> {
                     try {
                         beforePulsarStart(pulsarService);
@@ -470,9 +479,25 @@ protected PulsarTestContext.Builder createPulsarTestContextBuilder(ServiceConfig
                     }
                 })
                 .brokerServiceCustomizer(this::customizeNewBrokerService);
+        configureMetadataStores(builder);
         return builder;
     }
 
+    /**
+     * Configures the metadata stores for the PulsarTestContext.Builder instance.
+     * Set useTestZookeeper to true in the test's constructor to use TestZKServer which is a real ZooKeeper
+     * implementation.
+     *
+     * @param builder the PulsarTestContext.Builder instance to configure
+     */
+    protected void configureMetadataStores(PulsarTestContext.Builder builder) {
+        if (useTestZookeeper) {
+            builder.withTestZookeeper();
+        } else {
+            builder.withMockZookeeper(true);
+        }
+    }
+
     protected PulsarTestContext createAdditionalPulsarTestContext(ServiceConfiguration conf) throws Exception {
         return createAdditionalPulsarTestContext(conf, null);
     }

pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java

@@ -19,15 +19,13 @@
 
 package org.apache.pulsar.broker.testcontext;
 
-import com.google.common.util.concurrent.MoreExecutors;
 import io.netty.channel.EventLoopGroup;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdkBuilder;
 import io.opentelemetry.sdk.testing.exporter.InMemoryMetricReader;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Optional;
@@ -37,6 +35,7 @@
 import lombok.Builder;
 import lombok.Getter;
 import lombok.Singular;
+import lombok.SneakyThrows;
 import lombok.ToString;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.bookkeeper.client.BookKeeper;
@@ -65,16 +64,19 @@
 import org.apache.pulsar.compaction.CompactionServiceFactory;
 import org.apache.pulsar.compaction.Compactor;
 import org.apache.pulsar.compaction.PulsarCompactionServiceFactory;
+import org.apache.pulsar.metadata.TestZKServer;
 import org.apache.pulsar.metadata.api.MetadataStore;
 import org.apache.pulsar.metadata.api.MetadataStoreConfig;
 import org.apache.pulsar.metadata.api.MetadataStoreException;
 import org.apache.pulsar.metadata.api.extended.MetadataStoreExtended;
 import org.apache.pulsar.metadata.impl.MetadataStoreFactoryImpl;
 import org.apache.pulsar.metadata.impl.ZKMetadataStore;
 import org.apache.zookeeper.CreateMode;
+import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.MockZooKeeper;
 import org.apache.zookeeper.MockZooKeeperSession;
-import org.apache.zookeeper.data.ACL;
+import org.apache.zookeeper.ZooDefs;
+import org.apache.zookeeper.ZooKeeper;
 import org.jetbrains.annotations.NotNull;
 import org.mockito.Mockito;
 import org.mockito.internal.util.MockUtil;
@@ -161,6 +163,10 @@ public class PulsarTestContext implements AutoCloseable {
 
     private final MockZooKeeper mockZooKeeperGlobal;
 
+    private final TestZKServer testZKServer;
+
+    private final TestZKServer testZKServerGlobal;
+
     private final SpyConfig spyConfig;
 
     private final boolean startable;
@@ -409,6 +415,11 @@ public Builder reuseMockBookkeeperAndMetadataStores(PulsarTestContext otherConte
                 if (otherContext.getMockZooKeeperGlobal() != null) {
                     mockZooKeeperGlobal(otherContext.getMockZooKeeperGlobal());
                 }
+            } else if (otherContext.getTestZKServer() != null) {
+                testZKServer(otherContext.getTestZKServer());
+                if (otherContext.getTestZKServerGlobal() != null) {
+                    testZKServerGlobal(otherContext.getTestZKServerGlobal());
+                }
             } else {
                 localMetadataStore(NonClosingProxyHandler.createNonClosingProxy(otherContext.getLocalMetadataStore(),
                         MetadataStoreExtended.class
@@ -476,17 +487,56 @@ public Builder withMockZookeeper(boolean useSeparateGlobalZk) {
         }
 
         private MockZooKeeper createMockZooKeeper() throws Exception {
-            MockZooKeeper zk = MockZooKeeper.newInstance(MoreExecutors.newDirectExecutorService());
-            List<ACL> dummyAclList = new ArrayList<>(0);
+            MockZooKeeper zk = MockZooKeeper.newInstance();
+            initializeZookeeper(zk);
+            registerCloseable(zk::shutdown);
+            return zk;
+        }
 
+        private static void initializeZookeeper(ZooKeeper zk) throws KeeperException, InterruptedException {
             ZkUtils.createFullPathOptimistic(zk, "/ledgers/available/192.168.1.1:" + 5000,
-                    "".getBytes(StandardCharsets.UTF_8), dummyAclList, CreateMode.PERSISTENT);
+                    "".getBytes(StandardCharsets.UTF_8), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
 
-            zk.create("/ledgers/LAYOUT", "1\nflat:1".getBytes(StandardCharsets.UTF_8), dummyAclList,
+            zk.create("/ledgers/LAYOUT", "1\nflat:1".getBytes(StandardCharsets.UTF_8), ZooDefs.Ids.OPEN_ACL_UNSAFE,
                     CreateMode.PERSISTENT);
+        }
 
-            registerCloseable(zk::shutdown);
-            return zk;
+        /**
+         * Configure this PulsarTestContext to use a test ZooKeeper instance which is
+         * shared for both the local and configuration metadata stores.
+         *
+         * @return the builder
+         */
+        public Builder withTestZookeeper() {
+            return withTestZookeeper(false);
+        }
+
+        /**
+         * Configure this PulsarTestContext to use a test ZooKeeper instance.
+         *
+         * @param useSeparateGlobalZk if true, the global (configuration) zookeeper will be a separate instance
+         * @return the builder
+         */
+        public Builder withTestZookeeper(boolean useSeparateGlobalZk) {
+            try {
+                testZKServer(createTestZookeeper());
+                if (useSeparateGlobalZk) {
+                    testZKServerGlobal(createTestZookeeper());
+                }
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+            return this;
+        }
+
+        private TestZKServer createTestZookeeper() throws Exception {
+            TestZKServer testZKServer = new TestZKServer();
+            try (ZooKeeper zkc = new ZooKeeper(testZKServer.getConnectionString(), 5000, event -> {
+            })) {
+                initializeZookeeper(zkc);
+            }
+            registerCloseable(testZKServer);
+            return testZKServer;
         }
 
         /**
@@ -676,6 +726,20 @@ private void initializeCommonPulsarServices(SpyConfig spyConfig) {
                             configurationMetadataStore(mockZookeeperMetadataStore);
                         }
                     }
+                } else if (super.testZKServer != null) {
+                    MetadataStoreExtended testZookeeperMetadataStore =
+                            createTestZookeeperMetadataStore(super.testZKServer, MetadataStoreConfig.METADATA_STORE);
+                    if (super.localMetadataStore == null) {
+                        localMetadataStore(testZookeeperMetadataStore);
+                    }
+                    if (super.configurationMetadataStore == null) {
+                        if (super.testZKServerGlobal != null) {
+                            configurationMetadataStore(createTestZookeeperMetadataStore(super.testZKServerGlobal,
+                                    MetadataStoreConfig.CONFIGURATION_METADATA_STORE));
+                        } else {
+                            configurationMetadataStore(testZookeeperMetadataStore);
+                        }
+                    }
                 } else {
                     try {
                         MetadataStoreExtended store = MetadataStoreFactoryImpl.createExtended("memory:local",
@@ -720,6 +784,17 @@ private MetadataStoreExtended createMockZookeeperMetadataStore(MockZooKeeper moc
             return nonClosingProxy;
         }
 
+        @SneakyThrows
+        private MetadataStoreExtended createTestZookeeperMetadataStore(TestZKServer zkServer,
+                                                                       String metadataStoreName) {
+            MetadataStoreExtended store = MetadataStoreExtended.create("zk:" + zkServer.getConnectionString(),
+                    MetadataStoreConfig.builder().metadataStoreName(metadataStoreName).build());
+            registerCloseable(store);
+            MetadataStoreExtended nonClosingProxy =
+                    NonClosingProxyHandler.createNonClosingProxy(store, MetadataStoreExtended.class);
+            return nonClosingProxy;
+        }
+
         protected abstract void initializePulsarServices(SpyConfig spyConfig, Builder builder);
     }