From d0669e44008f344a589601765ebcee43eb95b162 Mon Sep 17 00:00:00 2001 From: Mukti Krishnan Date: Fri, 12 Mar 2021 00:05:07 +0530 Subject: [PATCH] ZOOKEEPER-1871: Add an option to zkCli to wait for connection before executing commands -waitforconnection option will make zk client wait for -timeout time to connect to zk server. timeout time is 30ms by default but can be specified explicitly for a session using -timeout option in command line. Author: Mukti Reviewers: maoling , Mohammad Arshad Closes #1626 from MuktiKrishnan/ZOOKEEPER-1871-master and squashes the following commits: 2947514b2 [Mukti] ZOOKEEPER-1871: Removed statement which was re-creating zookeeper admin in ZooKeeperMain.java and added documentation for waitforconnection in zookeeperCLI.md and c475d46ff [Mukti] ZOOKEEPER-1871: Add an option to zkCli to wailt for connection before executing commands --- .../main/resources/markdown/zookeeperCLI.md | 2 ++ .../org/apache/zookeeper/ZooKeeperMain.java | 26 ++++++++++++++++++- .../org/apache/zookeeper/ZooKeeperTest.java | 21 +++++++++++++++ 3 files changed, 48 insertions(+), 1 deletion(-) diff --git a/zookeeper-docs/src/main/resources/markdown/zookeeperCLI.md b/zookeeper-docs/src/main/resources/markdown/zookeeperCLI.md index 668e64e8ba8..7096aa0cc89 100644 --- a/zookeeper-docs/src/main/resources/markdown/zookeeperCLI.md +++ b/zookeeper-docs/src/main/resources/markdown/zookeeperCLI.md @@ -24,6 +24,8 @@ Enter into the ZooKeeper-cli bin/zkCli.sh # connect to the remote host with timeout:3s bin/zkCli.sh -timeout 3000 -server remoteIP:2181 +# connect to the remote host with -waitforconnection option to wait for connection success before executing commands +bin/zkCli.sh -waitforconnection -timeout 3000 -server remoteIP:2181 # connect with a custom client configuration properties file bin/zkCli.sh -client-configuration /path/to/client.properties ``` 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 ff8ce623286..4ee378fac05 100644 --- a/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java +++ b/zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java @@ -32,6 +32,8 @@ import java.util.List; import java.util.Map; 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 java.util.stream.Stream; @@ -68,6 +70,7 @@ public class ZooKeeperMain { protected ZooKeeper zk; protected String host = ""; + private CountDownLatch connectLatch = null; public boolean getPrintWatches() { return printWatches; @@ -106,6 +109,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(); + } + } } } @@ -168,6 +178,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); @@ -261,7 +273,19 @@ protected void connectToZK(String newHost) throws InterruptedException, IOExcept } } - zk = new ZooKeeperAdmin(host, Integer.parseInt(cl.getOption("timeout")), new MyWatcher(), readOnly, clientConfig); + if (cl.getOption("waitforconnection") != null) { + connectLatch = new CountDownLatch(1); + } + + int timeout = Integer.parseInt(cl.getOption("timeout")); + zk = new ZooKeeperAdmin(host, timeout, new MyWatcher(), readOnly, clientConfig); + if (connectLatch != null) { + if (!connectLatch.await(timeout, TimeUnit.MILLISECONDS)) { + zk.close(); + throw new IOException(KeeperException.create(KeeperException.Code.CONNECTIONLOSS)); + } + } + } public static void main(String[] args) throws IOException, InterruptedException { 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 b726c47da06..f69e0c683fb 100644 --- a/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java +++ b/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java @@ -774,4 +774,25 @@ private void assertClientAuthInfo(List expected, String actual) { }); } + @Test + public void testWaitForConnection() throws Exception { + // get a wrong port number + int invalidPort = PortAssignment.unique(); + long timeout = 3000L; // millisecond + String[] args1 = {"-server", "localhost:" + invalidPort, "-timeout", + Long.toString(timeout), "-waitforconnection", "ls", "/"}; + long startTime = 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 endTime = System.currentTimeMillis(); + assertTrue(endTime - startTime >= timeout, + "ZooKeeeperMain does not wait until the specified timeout"); + + } + }