HIVE-2923 : testAclPositive in TestZooKeeperTokenStore failing in cle…

…an checkout when run on Mac (Thomas Weise via Ashutosh Chauhan)

git-svn-id: https://svn.apache.org/repos/asf/hive/trunk@1310406 13f79535-47bb-0310-9956-ffa450edef68

shims/src/common-secure/java/org/apache/hadoop/hive/thrift/HadoopThriftAuthBridge20S.java

@@ -213,6 +213,8 @@ public static class Server extends HadoopThriftAuthBridge.Server {
        "hive.cluster.delegation.token.store.class";
      public static final String DELEGATION_TOKEN_STORE_ZK_CONNECT_STR =
          "hive.cluster.delegation.token.store.zookeeper.connectString";
+     public static final String DELEGATION_TOKEN_STORE_ZK_CONNECT_TIMEOUTMILLIS =
+         "hive.cluster.delegation.token.store.zookeeper.connectTimeoutMillis";
      public static final String DELEGATION_TOKEN_STORE_ZK_ZNODE =
          "hive.cluster.delegation.token.store.zookeeper.znode";
      public static final String DELEGATION_TOKEN_STORE_ZK_ACL =

shims/src/common-secure/java/org/apache/hadoop/hive/thrift/ZooKeeperTokenStore.java

@@ -23,13 +23,16 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 import org.apache.commons.lang.StringUtils;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.security.token.delegation.AbstractDelegationTokenSecretManager.DelegationTokenInformation;
 import org.apache.hadoop.security.token.delegation.HiveDelegationTokenSupport;
 import org.apache.zookeeper.CreateMode;
 import org.apache.zookeeper.KeeperException;
+import org.apache.zookeeper.WatchedEvent;
 import org.apache.zookeeper.Watcher;
 import org.apache.zookeeper.ZooDefs;
 import org.apache.zookeeper.ZooDefs.Ids;
@@ -56,6 +59,7 @@ public class ZooKeeperTokenStore implements DelegationTokenStore {
   private volatile ZooKeeper zkSession;
   private String zkConnectString;
   private final int zkSessionTimeout = 3000;
+  private long connectTimeoutMillis = -1;
   private List<ACL> newNodeAcl = Ids.OPEN_ACL_UNSAFE;
 
   private class ZooKeeperWatcher implements Watcher {
@@ -70,6 +74,7 @@ public void process(org.apache.zookeeper.WatchedEvent event) {
         }
       }
     }
+
   }
 
   /**
@@ -89,8 +94,8 @@ private ZooKeeper getSession() {
         synchronized (this) {
           if (zkSession == null || zkSession.getState() == States.CLOSED) {
             try {
-            zkSession = new ZooKeeper(this.zkConnectString, this.zkSessionTimeout,
-                new ZooKeeperWatcher());
+              zkSession = createConnectedClient(this.zkConnectString, this.zkSessionTimeout,
+                this.connectTimeoutMillis, new ZooKeeperWatcher());
             } catch (IOException ex) {
               throw new TokenStoreException("Token store error.", ex);
             }
@@ -100,6 +105,49 @@ private ZooKeeper getSession() {
     return zkSession;
   }
 
+  /**
+   * Create a ZooKeeper session that is in connected state.
+   * 
+   * @param connectString ZooKeeper connect String
+   * @param sessionTimeout ZooKeeper session timeout
+   * @param connectTimeout milliseconds to wait for connection, 0 or negative value means no wait
+   * @param watchers
+   * @return
+   * @throws InterruptedException
+   * @throws IOException
+   */
+  public static ZooKeeper createConnectedClient(String connectString,
+      int sessionTimeout, long connectTimeout, final Watcher... watchers)
+      throws IOException {
+    final CountDownLatch connected = new CountDownLatch(1);
+    Watcher connectWatcher = new Watcher() {
+      @Override
+      public void process(WatchedEvent event) {
+        switch (event.getState()) {
+        case SyncConnected:
+          connected.countDown();
+          break;
+        }
+        for (Watcher w : watchers) {
+          w.process(event);
+        }
+      }
+    };
+    ZooKeeper zk = new ZooKeeper(connectString, sessionTimeout, connectWatcher);
+    if (connectTimeout > 0) {
+      try {
+        if (!connected.await(connectTimeout, TimeUnit.MILLISECONDS)) {
+          zk.close();
+          throw new IOException("Timeout waiting for connection after "
+              + connectTimeout + "ms");
+        }
+      } catch (InterruptedException e) {
+        throw new IOException("Error waiting for connection.", e);
+      }
+    }
+    return zk;
+  }
+
   /**
    * Create a path if it does not already exist ("mkdir -p")
    * @param zk ZooKeeper session
@@ -215,6 +263,8 @@ public void setConf(Configuration conf) {
     }
     this.zkConnectString = conf.get(
       HadoopThriftAuthBridge20S.Server.DELEGATION_TOKEN_STORE_ZK_CONNECT_STR, null);
+    this.connectTimeoutMillis = conf.getLong(
+      HadoopThriftAuthBridge20S.Server.DELEGATION_TOKEN_STORE_ZK_CONNECT_TIMEOUTMILLIS, -1);
     this.rootNode = conf.get(
       HadoopThriftAuthBridge20S.Server.DELEGATION_TOKEN_STORE_ZK_ZNODE,
       HadoopThriftAuthBridge20S.Server.DELEGATION_TOKEN_STORE_ZK_ZNODE_DEFAULT);

shims/src/test/org/apache/hadoop/hive/thrift/TestZooKeeperTokenStore.java

@@ -41,7 +41,9 @@ public class TestZooKeeperTokenStore extends TestCase {
   private ZooKeeper zkClient = null;
   private int zkPort = -1;
   private ZooKeeperTokenStore ts;
-
+  // connect timeout large enough for slower test environments
+  private final int connectTimeoutMillis = 30000;
+
   @Override
   protected void setUp() throws Exception {
     File zkDataDir = new File(System.getProperty("java.io.tmpdir"));
@@ -50,8 +52,9 @@ protected void setUp() throws Exception {
     }
     this.zkCluster = new MiniZooKeeperCluster();
     this.zkPort = this.zkCluster.startup(zkDataDir);
-    this.zkClient = new ZooKeeper("localhost:"
-        + zkPort, 300, null);
+
+    this.zkClient = ZooKeeperTokenStore.createConnectedClient("localhost:" + zkPort, 3000,
+        connectTimeoutMillis);
   }
 
   @Override
@@ -72,6 +75,9 @@ private Configuration createConf(String zkPath) {
     conf.set(
         HadoopThriftAuthBridge20S.Server.DELEGATION_TOKEN_STORE_ZK_ZNODE,
         zkPath);
+    conf.setLong(
+        HadoopThriftAuthBridge20S.Server.DELEGATION_TOKEN_STORE_ZK_CONNECT_TIMEOUTMILLIS,
+        connectTimeoutMillis);
     return conf;
   }