HADOOP-19966. ZKDelegationTokenSecretManager may fail to load tokens and keys from ZooKeeper on startup

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/delegation/ZKDelegationTokenSecretManager.java

@@ -24,6 +24,8 @@
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 
@@ -46,6 +48,7 @@
 import org.apache.hadoop.security.authentication.util.ZookeeperClient;
 import org.apache.hadoop.security.token.Token;
 import org.apache.hadoop.security.token.delegation.web.DelegationTokenManager;
+import org.apache.hadoop.util.Preconditions;
 
 import static org.apache.hadoop.security.SecurityUtil.getServerPrincipal;
 import static org.apache.hadoop.util.Time.now;
@@ -79,6 +82,10 @@ public abstract class ZKDelegationTokenSecretManager<TokenIdent extends Abstract
       + "zkConnectionTimeout";
   public static final String ZK_DTSM_ZK_SHUTDOWN_TIMEOUT = ZK_CONF_PREFIX
       + "zkShutdownTimeout";
+  // Max time in ms to wait for the key/token cache initial load on startup,
+  // 0 or negative waits indefinitely.
+  public static final String ZK_DTSM_ZK_CACHE_INIT_TIMEOUT = ZK_CONF_PREFIX
+      + "zkCacheInitTimeout";
   public static final String ZK_DTSM_ZNODE_WORKING_PATH = ZK_CONF_PREFIX
       + "znodeWorkingPath";
   public static final String ZK_DTSM_ZK_AUTH_TYPE = ZK_CONF_PREFIX
@@ -111,6 +118,7 @@ public abstract class ZKDelegationTokenSecretManager<TokenIdent extends Abstract
   public static final int ZK_DTSM_ZK_SESSION_TIMEOUT_DEFAULT = 10000;
   public static final int ZK_DTSM_ZK_CONNECTION_TIMEOUT_DEFAULT = 10000;
   public static final int ZK_DTSM_ZK_SHUTDOWN_TIMEOUT_DEFAULT = 10000;
+  public static final int ZK_DTSM_ZK_CACHE_INIT_TIMEOUT_DEFAULT = 0;
   public static final String ZK_DTSM_ZNODE_WORKING_PATH_DEAFULT = "zkdtsm";
   // By default, increase seq number by 100 each time to reduce overflow
   // speed of znode dataVersion which is 32-integer now.
@@ -150,6 +158,7 @@ protected static CuratorFramework getCurator() {
   private CuratorCacheBridge keyCache;
   private CuratorCacheBridge tokenCache;
   private final int seqNumBatchSize;
+  private final int cacheInitTimeoutMs;
   private int currentSeqNum;
   private int currentMaxSeqNum;
 
@@ -168,6 +177,8 @@ public ZKDelegationTokenSecretManager(Configuration conf) {
         ZK_DTSM_TOKEN_SEQNUM_BATCH_SIZE_DEFAULT);
     isTokenWatcherEnabled = conf.getBoolean(ZK_DTSM_TOKEN_WATCHER_ENABLED,
         ZK_DTSM_TOKEN_WATCHER_ENABLED_DEFAULT);
+    cacheInitTimeoutMs = conf.getInt(ZK_DTSM_ZK_CACHE_INIT_TIMEOUT,
+        ZK_DTSM_ZK_CACHE_INIT_TIMEOUT_DEFAULT);
 
     String workPath = conf.get(ZK_DTSM_ZNODE_WORKING_PATH, ZK_DTSM_ZNODE_WORKING_PATH_DEAFULT);
     String nameSpace = workPath + "/" + ZK_DTSM_NAMESPACE;
@@ -240,6 +251,23 @@ static CuratorFramework createCuratorClient(Configuration conf, String namespace
 
   @Override
   public void startThreads() throws IOException {
+    // Fail fast so the cleanup below never tears down a running instance.
+    Preconditions.checkState(!isRunning(), "Secret manager is already running");
+    try {
+      doStartThreads();
+    } catch (IOException | RuntimeException e) {
+      // The caller does not invoke stopThreads() after a failed start, so
+      // release the caches, counters and Curator client here.
+      try {
+        stopThreads();
+      } catch (RuntimeException ce) {
+        e.addSuppressed(ce);
+      }
+      throw e;
+    }
+  }
+
+  private void doStartThreads() throws IOException {
     if (!isExternalClient) {
       try {
         zkClient.start();
@@ -290,8 +318,13 @@ public void startThreads() throws IOException {
     try {
       keyCache = CuratorCache.bridgeBuilder(zkClient, ZK_DTSM_MASTER_KEY_ROOT)
           .build();
+      CountDownLatch keyCacheInitialized = new CountDownLatch(1);
       CuratorCacheListener keyCacheListener = CuratorCacheListener.builder()
           .forCreatesAndChanges((oldNode, node) -> {
+            if (ZK_DTSM_MASTER_KEY_ROOT.equals(node.getPath())) {
+              // The root node itself is a container, not a key.
+              return;
+            }
             try {
               processKeyAddOrUpdate(node.getData());
             } catch (IOException e) {
@@ -301,9 +334,11 @@ public void startThreads() throws IOException {
             }
           })
           .forDeletes(childData -> processKeyRemoved(childData.getPath()))
+          .forInitialized(keyCacheInitialized::countDown)
           .build();
       keyCache.listenable().addListener(keyCacheListener);
       keyCache.start();
+      awaitCacheInitialized(keyCacheInitialized, "key");
       loadFromZKCache(false);
     } catch (Exception e) {
       throw new IOException("Could not start Curator keyCacheListener for keys",
@@ -314,8 +349,13 @@ public void startThreads() throws IOException {
       try {
         tokenCache = CuratorCache.bridgeBuilder(zkClient, ZK_DTSM_TOKENS_ROOT)
             .build();
+        CountDownLatch tokenCacheInitialized = new CountDownLatch(1);
         CuratorCacheListener tokenCacheListener = CuratorCacheListener.builder()
             .forCreatesAndChanges((oldNode, node) -> {
+              if (ZK_DTSM_TOKENS_ROOT.equals(node.getPath())) {
+                // The root node itself is a container, not a token.
+                return;
+              }
               try {
                 processTokenAddOrUpdate(node.getData());
               } catch (IOException e) {
@@ -333,9 +373,11 @@ public void startThreads() throws IOException {
                 throw new UncheckedIOException(e);
               }
             })
+            .forInitialized(tokenCacheInitialized::countDown)
             .build();
         tokenCache.listenable().addListener(tokenCacheListener);
         tokenCache.start();
+        awaitCacheInitialized(tokenCacheInitialized, "token");
         loadFromZKCache(true);
       } catch (Exception e) {
         throw new IOException(
@@ -363,6 +405,11 @@ private void loadFromZKCache(final boolean isTokenCache) {
 
     final AtomicInteger count = new AtomicInteger(0);
     children.forEach(childData -> {
+      if (childData.getPath().equals(
+          isTokenCache ? ZK_DTSM_TOKENS_ROOT : ZK_DTSM_MASTER_KEY_ROOT)) {
+        // The root node itself is a container, not a key or token.
+        return;
+      }
       try {
         if (isTokenCache) {
           processTokenAddOrUpdate(childData.getData());
@@ -386,6 +433,25 @@ private void loadFromZKCache(final boolean isTokenCache) {
     LOG.info("Loaded {} cache.", cacheName);
   }
 
+  private void awaitCacheInitialized(CountDownLatch initialized,
+      String cacheName) throws IOException {
+    try {
+      if (cacheInitTimeoutMs > 0) {
+        if (!initialized.await(cacheInitTimeoutMs, TimeUnit.MILLISECONDS)) {
+          throw new IOException("Timed out after " + cacheInitTimeoutMs
+              + " ms waiting for " + cacheName + " cache initialization, "
+              + "consider increasing " + ZK_DTSM_ZK_CACHE_INIT_TIMEOUT);
+        }
+      } else {
+        initialized.await();
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new IOException("Interrupted while waiting for " + cacheName
+          + " cache initialization", e);
+    }
+  }
+
   private void processKeyAddOrUpdate(byte[] data) throws IOException {
     ByteArrayInputStream bin = new ByteArrayInputStream(data);
     DataInputStream din = new DataInputStream(bin);
@@ -425,6 +491,10 @@ protected TokenIdent processTokenAddOrUpdate(byte[] data) throws IOException {
   }
 
   private void processTokenRemoved(ChildData data) throws IOException {
+    if (ZK_DTSM_TOKENS_ROOT.equals(data.getPath())) {
+      // The root node itself is a container, not a token.
+      return;
+    }
     ByteArrayInputStream bin = new ByteArrayInputStream(data.getData());
     DataInputStream din = new DataInputStream(bin);
     TokenIdent ident = createIdentifier();
@@ -474,7 +544,9 @@ public void stopThreads() {
 
   private void createPersistentNode(String nodePath) throws Exception {
     try {
-      zkClient.create().withMode(CreateMode.PERSISTENT).forPath(nodePath);
+      // Store empty data instead of Curator's default (the local address);
+      // HADOOP-17835 (3.4.0): CuratorCache includes the root node in the cache.
+      zkClient.create().withMode(CreateMode.PERSISTENT).forPath(nodePath, new byte[0]);
     } catch (KeeperException.NodeExistsException ne) {
       LOG.debug(nodePath + " znode already exists !!");
     } catch (Exception e) {

hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/security/token/delegation/TestZKDelegationTokenSecretManager.java

@@ -105,6 +105,7 @@ protected Configuration getSecretConf(String connectString) {
    conf.set(ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH, "testPath");
    conf.set(ZKDelegationTokenSecretManager.ZK_DTSM_ZK_AUTH_TYPE, "none");
    conf.setLong(ZKDelegationTokenSecretManager.ZK_DTSM_ZK_SHUTDOWN_TIMEOUT, 100);
+   conf.setLong(ZKDelegationTokenSecretManager.ZK_DTSM_ZK_CACHE_INIT_TIMEOUT, 10000);
    conf.setLong(DelegationTokenManager.UPDATE_INTERVAL, DAY_IN_SECS);
    conf.setLong(DelegationTokenManager.MAX_LIFETIME, DAY_IN_SECS);
    conf.setLong(DelegationTokenManager.RENEW_INTERVAL, DAY_IN_SECS);
@@ -551,22 +552,29 @@ public void testCreatingParentContainersIfNeeded() throws Exception {
         .build();
     curatorFramework.start();
     ZKDelegationTokenSecretManager.setCurator(curatorFramework);
-    DelegationTokenManager tm1 = new DelegationTokenManager(conf, new Text("foo"));
-
-    // When the init method is called,
-    // the ZKDelegationTokenSecretManager#startThread method will be called,
-    // and the creatingParentContainersIfNeeded will be called to create the nameSpace.
-    tm1.init();
+    DelegationTokenManager tm1 = null;
+    try {
+      tm1 = new DelegationTokenManager(conf, new Text("foo"));
 
-    String workingPath = "/" + conf.get(ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH,
-        ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH_DEAFULT) + "/ZKDTSMRoot";
+      // When the init method is called,
+      // the ZKDelegationTokenSecretManager#startThread method will be called,
+      // and the creatingParentContainersIfNeeded will be called to create the nameSpace.
+      tm1.init();
 
-    // Check if the created NameSpace exists.
-    Stat stat = curatorFramework.checkExists().forPath(workingPath);
-    assertNotNull(stat);
+      String workingPath = "/" + conf.get(ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH,
+          ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH_DEAFULT) + "/ZKDTSMRoot";
 
-    tm1.destroy();
-    curatorFramework.close();
+      // Check if the created NameSpace exists.
+      Stat stat = curatorFramework.checkExists().forPath(workingPath);
+      assertNotNull(stat);
+    } finally {
+      if (tm1 != null) {
+        tm1.destroy();
+      }
+      // Restore the default curator so later tests do not see a closed client.
+      ZKDelegationTokenSecretManager.setCurator(null);
+      curatorFramework.close();
+    }
   }
 
   @Test
@@ -616,36 +624,43 @@ public void testMultipleInit() throws Exception {
 
     DelegationTokenManager tm1 = new DelegationTokenManager(conf, new Text("foo"));
     DelegationTokenManager tm2 = new DelegationTokenManager(conf, new Text("bar"));
-    // When the init method is called,
-    // the ZKDelegationTokenSecretManager#startThread method will be called,
-    // and the creatingParentContainersIfNeeded will be called to create the nameSpace.
-    ExecutorService executorService = Executors.newFixedThreadPool(2);
-
-    Callable<Boolean> tm1Callable = () -> {
-      tm1.init();
-      return true;
-    };
-    Callable<Boolean> tm2Callable = () -> {
-      tm2.init();
-      return true;
-    };
-    List<Future<Boolean>> futures = executorService.invokeAll(
-        Arrays.asList(tm1Callable, tm2Callable));
-    for(Future<Boolean> future : futures) {
-      assertTrue(future.get());
-    }
-    executorService.shutdownNow();
-    assertTrue(executorService.awaitTermination(1, TimeUnit.SECONDS));
-    tm1.destroy();
-    tm2.destroy();
-
-    String workingPath = "/" + conf.get(ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH,
-        ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH_DEAFULT) + "/ZKDTSMRoot";
+    ExecutorService executorService = null;
+    try {
+      // When the init method is called,
+      // the ZKDelegationTokenSecretManager#startThread method will be called,
+      // and the creatingParentContainersIfNeeded will be called to create the nameSpace.
+      executorService = Executors.newFixedThreadPool(2);
+
+      Callable<Boolean> tm1Callable = () -> {
+        tm1.init();
+        return true;
+      };
+      Callable<Boolean> tm2Callable = () -> {
+        tm2.init();
+        return true;
+      };
+      List<Future<Boolean>> futures = executorService.invokeAll(
+          Arrays.asList(tm1Callable, tm2Callable));
+      for (Future<Boolean> future : futures) {
+        assertTrue(future.get());
+      }
 
-    // Check if the created NameSpace exists.
-    Stat stat = curatorFramework.checkExists().forPath(workingPath);
-    assertNotNull(stat);
+      String workingPath = "/" + conf.get(ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH,
+          ZKDelegationTokenSecretManager.ZK_DTSM_ZNODE_WORKING_PATH_DEAFULT) + "/ZKDTSMRoot";
 
-    curatorFramework.close();
+      // Check if the created NameSpace exists.
+      Stat stat = curatorFramework.checkExists().forPath(workingPath);
+      assertNotNull(stat);
+    } finally {
+      if (executorService != null) {
+        executorService.shutdownNow();
+        executorService.awaitTermination(1, TimeUnit.SECONDS);
+      }
+      tm1.destroy();
+      tm2.destroy();
+      // Restore the default curator so later tests do not see a closed client.
+      ZKDelegationTokenSecretManager.setCurator(null);
+      curatorFramework.close();
+    }
   }
 }