From ce391c9673099967b31a519052f6f9f1e3e91e57 Mon Sep 17 00:00:00 2001 From: Mukti Date: Thu, 4 Mar 2021 17:13:14 +0530 Subject: [PATCH 1/2] ZOOKEEPER-1871: Add an option to zkCli to wailt for connection before executing commands --- .../org/apache/zookeeper/ZooKeeperMain.java | 24 +++++++++++++++++++ .../org/apache/zookeeper/ZooKeeperTest.java | 22 +++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java b/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java index 243aca33c39..2f5515bcfb2 100644 --- a/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java +++ b/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java @@ -33,6 +33,9 @@ import java.util.Map; import java.util.Map.Entry; import java.util.NoSuchElementException; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.regex.Matcher; import java.util.regex.Pattern; import org.apache.yetus.audience.InterfaceAudience; @@ -89,6 +92,7 @@ public class ZooKeeperMain { protected ZooKeeper zk; protected String host = ""; + private CountDownLatch connectLatch = null; public boolean getPrintWatches() { return printWatches; @@ -146,6 +150,13 @@ public void process(WatchedEvent event) { ZooKeeperMain.printMessage("WATCHER::"); ZooKeeperMain.printMessage(event.toString()); } + if(connectLatch != null){ + // connection success + if ( event.getType() == Event.EventType.None + && event.getState() == Event.KeeperState.SyncConnected) { + connectLatch.countDown(); + } + } } } @@ -208,6 +219,8 @@ public boolean parseOptions(String[] args) { options.put("readonly", "true"); } else if (opt.equals("-client-configuration")) { options.put("client-configuration", it.next()); + } else if (opt.equals("-waitforconnection")) { + options.put("waitforconnection", "true"); } } catch (NoSuchElementException e) { System.err.println("Error: no argument found for option " + opt); @@ -289,6 +302,17 @@ protected void connectToZK(String newHost) throws InterruptedException, IOExcept System.setProperty(ZKClientConfig.SECURE_CLIENT, "true"); System.out.println("Secure connection is enabled"); } + if (cl.getOption("waitforconnection") != null) { + connectLatch = new CountDownLatch(1); + } + int timeout = Integer.parseInt(cl.getOption("timeout")); + zk = new ZooKeeperAdmin(host, timeout, new MyWatcher(),readOnly); + if (connectLatch != null) { + if (!connectLatch.await(timeout, TimeUnit.MILLISECONDS)) { + zk.close(); + throw new IOException(KeeperException.create(KeeperException.Code.CONNECTIONLOSS)); + } + } ZKClientConfig clientConfig = null; diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java index 4467a1ac471..3a422192915 100644 --- a/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java +++ b/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java @@ -48,6 +48,7 @@ import org.apache.zookeeper.data.Id; import org.apache.zookeeper.data.Stat; import org.apache.zookeeper.test.ClientBase; +import org.apache.zookeeper.PortAssignment; import org.junit.Test; /** @@ -685,4 +686,25 @@ public void testSyncCommand() throws Exception { runCommandExpect(cmd, expected); } + @Test + public void testWaitForConnection() throws Exception { + // get a wrong port number + int invalid_port = PortAssignment.unique(); + long timeout = 3000L; // millisecond + String[] args1 = {"-server", "localhost:" + invalid_port, "-timeout", + Long.toString(timeout), "-waitforconnection", "ls", "/"}; + long start_time = System.currentTimeMillis(); + // try to connect to a non-existing server so as to wait until wait_timeout + try{ + ZooKeeperMain zkMain = new ZooKeeperMain(args1); + fail("IOException was expected"); + }catch (IOException e) { + // do nothing + } + long end_time = System.currentTimeMillis(); + assertTrue("ZooKeeeperMain does not wait until the specified timeout", + end_time - start_time >= timeout); + + } + } From 0a7866e9ce5231e743a173e6ac29507dad7601c0 Mon Sep 17 00:00:00 2001 From: Mukti Date: Tue, 9 Mar 2021 18:37:17 +0530 Subject: [PATCH 2/2] ZOOKEEPER-1871: Used CamelCase naming style and corrected code format --- .../java/org/apache/zookeeper/ZooKeeperMain.java | 6 +++--- .../java/org/apache/zookeeper/ZooKeeperTest.java | 15 +++++++-------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java b/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java index 2f5515bcfb2..bf64f1d1966 100644 --- a/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java +++ b/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java @@ -150,9 +150,9 @@ public void process(WatchedEvent event) { ZooKeeperMain.printMessage("WATCHER::"); ZooKeeperMain.printMessage(event.toString()); } - if(connectLatch != null){ + if (connectLatch != null){ // connection success - if ( event.getType() == Event.EventType.None + if (event.getType() == Event.EventType.None && event.getState() == Event.KeeperState.SyncConnected) { connectLatch.countDown(); } @@ -306,7 +306,7 @@ protected void connectToZK(String newHost) throws InterruptedException, IOExcept connectLatch = new CountDownLatch(1); } int timeout = Integer.parseInt(cl.getOption("timeout")); - zk = new ZooKeeperAdmin(host, timeout, new MyWatcher(),readOnly); + zk = new ZooKeeperAdmin(host, timeout, new MyWatcher(), readOnly); if (connectLatch != null) { if (!connectLatch.await(timeout, TimeUnit.MILLISECONDS)) { zk.close(); diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java index 3a422192915..a31983f615f 100644 --- a/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java +++ b/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java @@ -48,7 +48,6 @@ import org.apache.zookeeper.data.Id; import org.apache.zookeeper.data.Stat; import org.apache.zookeeper.test.ClientBase; -import org.apache.zookeeper.PortAssignment; import org.junit.Test; /** @@ -689,21 +688,21 @@ public void testSyncCommand() throws Exception { @Test public void testWaitForConnection() throws Exception { // get a wrong port number - int invalid_port = PortAssignment.unique(); + int invalidPort = PortAssignment.unique(); long timeout = 3000L; // millisecond - String[] args1 = {"-server", "localhost:" + invalid_port, "-timeout", + String[] args1 = {"-server", "localhost:" + invalidPort, "-timeout", Long.toString(timeout), "-waitforconnection", "ls", "/"}; - long start_time = System.currentTimeMillis(); + long startTime = System.currentTimeMillis(); // try to connect to a non-existing server so as to wait until wait_timeout - try{ + try { ZooKeeperMain zkMain = new ZooKeeperMain(args1); fail("IOException was expected"); - }catch (IOException e) { + } catch (IOException e) { // do nothing } - long end_time = System.currentTimeMillis(); + long endTime = System.currentTimeMillis(); assertTrue("ZooKeeeperMain does not wait until the specified timeout", - end_time - start_time >= timeout); + endTime - startTime >= timeout); }