ZOOKEEPER-1871: Add an option to zkCli to wait for connection before executing commands

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;
 

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);
+
+    }
+
 }