This is an automated email from the ASF dual-hosted git repository.

arshad pushed a commit to branch branch-3.6
in repository https://gitbox.apache.org/repos/asf/zookeeper.git


The following commit(s) were added to refs/heads/branch-3.6 by this push:
     new 579aa34  ZOOKEEPER-1871: Add an option to zkCli to wait for connection 
before executing commands
579aa34 is described below

commit 579aa3417273ed13a0679344a96538bacb8d2eaf
Author: Mukti Krishnan <[email protected]>
AuthorDate: Mon Mar 15 19:43:43 2021 +0530

    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 <[email protected]>
    
    Reviewers: Mohammad Arshad <[email protected]>
    
    Closes #1639 from MuktiKrishnan/ZOOKEEPER-1871-branch-3.6
---
 .../src/main/resources/markdown/zookeeperCLI.md    |  2 ++
 .../java/org/apache/zookeeper/ZooKeeperMain.java   | 25 +++++++++++++++++++++-
 .../java/org/apache/zookeeper/ZooKeeperTest.java   | 19 ++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/zookeeper-docs/src/main/resources/markdown/zookeeperCLI.md 
b/zookeeper-docs/src/main/resources/markdown/zookeeperCLI.md
index 205c3ba..53da2c4 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 243aca3..435a3f4 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,8 @@ import java.util.List;
 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 +91,7 @@ public class ZooKeeperMain {
 
     protected ZooKeeper zk;
     protected String host = "";
+    private CountDownLatch connectLatch = null;
 
     public boolean getPrintWatches() {
         return printWatches;
@@ -146,6 +149,13 @@ public class ZooKeeperMain {
                 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 +218,8 @@ public class ZooKeeperMain {
                         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);
@@ -301,7 +313,18 @@ public class ZooKeeperMain {
             }
         }
 
-        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));
+            }
+            connectLatch = null;
+        }
     }
 
     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 3e465fc..859d13c 100644
--- a/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java
+++ b/zookeeper-server/src/test/java/org/apache/zookeeper/ZooKeeperTest.java
@@ -700,4 +700,23 @@ public class ZooKeeperTest extends ClientBase {
         assertEquals("Insufficient permission : " + zNodeToBeCreated, 
errorMessage);
     }
 
+    @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 
waitTimeout
+        try {
+            ZooKeeperMain zkMain = new ZooKeeperMain(args1);
+            fail("IOException was expected");
+        } catch (IOException e) {
+            // do nothing
+        }
+        long endTime = System.currentTimeMillis();
+        assertTrue("ZooKeeperMain does not wait until the specified timeout",
+                endTime - startTime >= timeout);
+    }
 }

Reply via email to