Adding multi-node ZKCluster test util class

bookkeeper-server/src/test/java/org/apache/bookkeeper/bookie/BookieInitializationTest.java

@@ -568,7 +568,7 @@ public void testBookieServerStartupOnEphemeralPorts() throws Exception {
      */
     @Test
     public void testStartBookieWithoutZKServer() throws Exception {
-        zkUtil.killServer();
+        zkUtil.killCluster();
 
         File tmpDir = createTempDir("bookie", "test");
 

bookkeeper-server/src/test/java/org/apache/bookkeeper/bookie/LedgerStorageCheckpointTest.java

@@ -143,7 +143,7 @@ public void tearDown() throws Exception {
      * @throws Exception
      */
     protected void startZKCluster() throws Exception {
-        zkUtil.startServer();
+        zkUtil.startCluster();
     }
 
     /**
@@ -152,7 +152,7 @@ protected void startZKCluster() throws Exception {
      * @throws Exception
      */
     protected void stopZKCluster() throws Exception {
-        zkUtil.killServer();
+        zkUtil.killCluster();
     }
 
     protected void cleanupTempDirs() throws Exception {

bookkeeper-server/src/test/java/org/apache/bookkeeper/client/BookKeeperTest.java

@@ -72,7 +72,7 @@ public void testConstructionZkDelay() throws Exception {
             .setZkTimeout(20000);
 
         CountDownLatch l = new CountDownLatch(1);
-        zkUtil.sleepServer(200, TimeUnit.MILLISECONDS, l);
+        zkUtil.sleepCluster(200, TimeUnit.MILLISECONDS, l);
         l.await();
 
         BookKeeper bkc = new BookKeeper(conf);
@@ -87,7 +87,7 @@ public void testConstructionNotConnectedExplicitZk() throws Exception {
             .setZkTimeout(20000);
 
         CountDownLatch l = new CountDownLatch(1);
-        zkUtil.sleepServer(200, TimeUnit.MILLISECONDS, l);
+        zkUtil.sleepCluster(200, TimeUnit.MILLISECONDS, l);
         l.await();
 
         ZooKeeper zk = new ZooKeeper(

bookkeeper-server/src/test/java/org/apache/bookkeeper/meta/TestLongZkLedgerIdGenerator.java

@@ -58,7 +58,7 @@ public void setUp() throws Exception {
         super.setUp();
 
         zkutil = new ZooKeeperUtil();
-        zkutil.startServer();
+        zkutil.startCluster();
         zk = zkutil.getZooKeeperClient();
 
         ZkLedgerIdGenerator shortLedgerIdGenerator = new ZkLedgerIdGenerator(zk,
@@ -73,7 +73,7 @@ public void tearDown() throws Exception {
         LOG.info("Tearing down test");
         ledgerIdGenerator.close();
         zk.close();
-        zkutil.killServer();
+        zkutil.killCluster();
 
         super.tearDown();
     }

bookkeeper-server/src/test/java/org/apache/bookkeeper/meta/TestZkLedgerIdGenerator.java

@@ -55,7 +55,7 @@ public void setUp() throws Exception {
         super.setUp();
 
         zkutil = new ZooKeeperUtil();
-        zkutil.startServer();
+        zkutil.startCluster();
         zk = zkutil.getZooKeeperClient();
 
         ledgerIdGenerator = new ZkLedgerIdGenerator(zk,
@@ -68,7 +68,7 @@ public void tearDown() throws Exception {
         LOG.info("Tearing down test");
         ledgerIdGenerator.close();
         zk.close();
-        zkutil.killServer();
+        zkutil.killCluster();
 
         super.tearDown();
     }

bookkeeper-server/src/test/java/org/apache/bookkeeper/replication/TestLedgerUnderreplicationManager.java

@@ -93,7 +93,7 @@ public class TestLedgerUnderreplicationManager {
     @Before
     public void setupZooKeeper() throws Exception {
         zkUtil = new ZooKeeperUtil();
-        zkUtil.startServer();
+        zkUtil.startCluster();
 
         conf = TestBKConfiguration.newServerConfiguration();
         conf.setMetadataServiceUri(zkUtil.getMetadataServiceUri());
@@ -134,7 +134,7 @@ public void setupZooKeeper() throws Exception {
     @After
     public void teardownZooKeeper() throws Exception {
         if (zkUtil != null) {
-            zkUtil.killServer();
+            zkUtil.killCluster();
             zkUtil = null;
         }
         if (executor != null) {

bookkeeper-server/src/test/java/org/apache/bookkeeper/test/BookKeeperClusterTestCase.java

@@ -79,7 +79,7 @@ public abstract class BookKeeperClusterTestCase {
     public final Timeout globalTimeout;
 
     // Metadata service related variables
-    protected final ZooKeeperUtil zkUtil = new ZooKeeperUtil();
+    protected final ZooKeeperCluster zkUtil;
     protected ZooKeeper zkc;
     protected String metadataServiceUri;
 
@@ -118,8 +118,21 @@ public BookKeeperClusterTestCase(int numBookies) {
     }
 
     public BookKeeperClusterTestCase(int numBookies, int testTimeoutSecs) {
+        this(numBookies, 1, 120);
+    }
+
+    public BookKeeperClusterTestCase(int numBookies, int numOfZKNodes, int testTimeoutSecs) {
         this.numBookies = numBookies;
         this.globalTimeout = Timeout.seconds(testTimeoutSecs);
+        if (numOfZKNodes == 1) {
+            zkUtil = new ZooKeeperUtil();
+        } else {
+            try {
+                zkUtil = new ZooKeeperClusterUtil(numOfZKNodes);
+            } catch (IOException | KeeperException | InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }
     }
 
     @Before
@@ -202,7 +215,7 @@ protected File createTempDir(String prefix, String suffix) throws IOException {
      * @throws Exception
      */
     protected void startZKCluster() throws Exception {
-        zkUtil.startServer();
+        zkUtil.startCluster();
         zkc = zkUtil.getZooKeeperClient();
     }
 
@@ -212,7 +225,7 @@ protected void startZKCluster() throws Exception {
      * @throws Exception
      */
     protected void stopZKCluster() throws Exception {
-        zkUtil.killServer();
+        zkUtil.killCluster();
     }
 
     /**

bookkeeper-server/src/test/java/org/apache/bookkeeper/test/ZooKeeperCluster.java

@@ -0,0 +1,79 @@
+/*
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+package org.apache.bookkeeper.test;
+
+import static org.apache.bookkeeper.util.BookKeeperConstants.AVAILABLE_NODE;
+import static org.apache.bookkeeper.util.BookKeeperConstants.READONLY;
+
+import java.io.IOException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.bookkeeper.zookeeper.ZooKeeperWatcherBase;
+import org.apache.zookeeper.CreateMode;
+import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.Transaction;
+import org.apache.zookeeper.ZooDefs.Ids;
+import org.apache.zookeeper.ZooKeeper;
+
+/**
+ * Interface for ZooKeeperCluster.
+ */
+public interface ZooKeeperCluster {
+    ZooKeeper getZooKeeperClient();
+
+    String getZooKeeperConnectString();
+
+    String getMetadataServiceUri();
+
+    String getMetadataServiceUri(String zkLedgersRootPath);
+
+    String getMetadataServiceUri(String zkLedgersRootPath, String type);
+
+    void startCluster() throws Exception;
+
+    void stopCluster() throws Exception;
+
+    void restartCluster() throws Exception;
+
+    void killCluster() throws Exception;
+
+    void sleepCluster(final int time, final TimeUnit timeUnit, final CountDownLatch l)
+            throws InterruptedException, IOException;
+
+    default void expireSession(ZooKeeper zk) throws Exception {
+        long id = zk.getSessionId();
+        byte[] password = zk.getSessionPasswd();
+        ZooKeeperWatcherBase w = new ZooKeeperWatcherBase(10000);
+        ZooKeeper zk2 = new ZooKeeper(getZooKeeperConnectString(), zk.getSessionTimeout(), w, id, password);
+        w.waitForConnection();
+        zk2.close();
+    }
+
+    default void createBKEnsemble(String ledgersPath) throws KeeperException, InterruptedException {
+        Transaction txn = getZooKeeperClient().transaction();
+        txn.create(ledgersPath, new byte[0], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
+        txn.create(ledgersPath + "/" + AVAILABLE_NODE, new byte[0], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
+        txn.create(ledgersPath + "/" + AVAILABLE_NODE + "/" + READONLY, new byte[0], Ids.OPEN_ACL_UNSAFE,
+                CreateMode.PERSISTENT);
+        txn.commit();
+    }
+}

bookkeeper-server/src/test/java/org/apache/bookkeeper/test/ZooKeeperClusterUtil.java

@@ -0,0 +1,140 @@
+/*
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+package org.apache.bookkeeper.test;
+
+import com.google.common.io.Files;
+import java.io.IOException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.bookkeeper.meta.LongHierarchicalLedgerManagerFactory;
+import org.apache.bookkeeper.zookeeper.ZooKeeperClient;
+import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.ZooKeeper;
+import org.apache.zookeeper.test.QuorumUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Provides multi node zookeeper cluster.
+ */
+@Slf4j
+public class ZooKeeperClusterUtil implements ZooKeeperCluster {
+
+    static {
+        enableZookeeperTestEnvVariables();
+    }
+
+    static final Logger LOG = LoggerFactory.getLogger(ZooKeeperClusterUtil.class);
+    private final int numOfZKNodes;
+    public QuorumUtil quorumUtil;
+    String connectString;
+    protected ZooKeeper zkc; // zookeeper client
+
+    public static void enableZookeeperTestEnvVariables() {
+        /*
+         * org.apache.zookeeper.test.ClientBase uses FourLetterWordMain, from
+         * 3.5.3 four letter words are disabled by default due to security
+         * reasons
+         */
+        System.setProperty("zookeeper.4lw.commands.whitelist", "*");
+        System.setProperty("zookeeper.admin.enableServer", "false");
+        try {
+            System.setProperty("build.test.dir", Files.createTempDir().getCanonicalPath());
+        } catch (IOException e) {
+            log.error("Failed to create temp dir, so setting build.test.dir system property to /tmp");
+            System.setProperty("build.test.dir", "/tmp");
+        }
+    }
+
+    public ZooKeeperClusterUtil(int numOfZKNodes) throws IOException, KeeperException, InterruptedException {
+        if ((numOfZKNodes < 3) || (numOfZKNodes % 2 == 0)) {
+            throw new IllegalArgumentException("numOfZKNodes should be atleast 3 and it should not be even number");
+        }
+        this.numOfZKNodes = numOfZKNodes;
+    }
+
+    @Override
+    public String getZooKeeperConnectString() {
+        return connectString;
+    }
+
+    @Override
+    public String getMetadataServiceUri() {
+        return getMetadataServiceUri("/ledgers");
+    }
+
+    @Override
+    public String getMetadataServiceUri(String zkLedgersRootPath) {
+        return getMetadataServiceUri(zkLedgersRootPath, LongHierarchicalLedgerManagerFactory.NAME);
+    }
+
+    @Override
+    public String getMetadataServiceUri(String zkLedgersRootPath, String type) {
+        /*
+         * URI doesn't accept ',', for more info. check
+         * AbstractConfiguration.getMetadataServiceUri()
+         */
+        return "zk+" + type + "://" + connectString.replace(",", ";") + zkLedgersRootPath;
+    }
+
+    @Override
+    public ZooKeeper getZooKeeperClient() {
+        return zkc;
+    }
+
+    @Override
+    public void startCluster() throws Exception {
+        // QuorumUtil will start 2*n+1 nodes.
+        quorumUtil = new QuorumUtil(numOfZKNodes / 2);
+        quorumUtil.startAll();
+        connectString = quorumUtil.getConnString();
+        // create a zookeeper client
+        LOG.debug("Instantiate ZK Client");
+        zkc = ZooKeeperClient.newBuilder().connectString(getZooKeeperConnectString()).sessionTimeoutMs(10000).build();
+
+        // create default bk ensemble
+        createBKEnsemble("/ledgers");
+    }
+
+    @Override
+    public void stopCluster() throws Exception {
+        if (zkc != null) {
+            zkc.close();
+        }
+        quorumUtil.shutdownAll();
+    }
+
+    @Override
+    public void restartCluster() throws Exception {
+        quorumUtil.startAll();
+    }
+
+    @Override
+    public void killCluster() throws Exception {
+        quorumUtil.tearDown();
+    }
+
+    @Override
+    public void sleepCluster(int time, TimeUnit timeUnit, CountDownLatch l) throws InterruptedException, IOException {
+        throw new UnsupportedOperationException("sleepServer operation is not supported for ZooKeeperClusterUtil");
+    }
+}