TEZ-4735: Use Ephemeral Sequential Znodes for applicationID generation

tez-dag/src/main/java/org/apache/tez/dag/api/client/registry/zookeeper/ZkAMRegistry.java

@@ -23,8 +23,6 @@
 import java.util.Collections;
 import java.util.List;
 
-import org.apache.curator.RetryLoop;
-import org.apache.curator.RetryPolicy;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.conf.Configuration;
@@ -37,6 +35,8 @@
 import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.data.Stat;
 
+import com.google.common.annotations.VisibleForTesting;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -120,39 +120,17 @@ public ApplicationId generateNewId() throws Exception {
     createNamespaceIfNotExists();
     long namespaceCreationTime = getNamespaceCreationTime();
 
-    boolean success = false;
-    long startTime = System.currentTimeMillis();
-    RetryPolicy retryPolicy = zkConfig.getRetryPolicy();
-    int tryId = 0;
-    for (int i = 0; (i < zkConfig.getCuratorMaxRetries()) && !success; i++) {
-      List<String> children = client.getChildren().forPath(namespace);
-      if (children != null && !children.isEmpty()) {
-        children.sort(Collections.reverseOrder());
-        String last = children.getFirst();
-        ApplicationId lastAppId = ApplicationId.fromString(last);
-        tryId = lastAppId.getId() + 1;
-      }
-      ApplicationId tryAppId = ApplicationId.newInstance(namespaceCreationTime, tryId);
-      try {
-        client
-            .create()
-            .withMode(CreateMode.EPHEMERAL)
-            .forPath(namespace + "/" + tryAppId.toString(), new byte[0]);
-        LOG.debug("Successfully created application id {} for namespace {}", tryAppId, namespace);
-        success = true;
-      } catch (KeeperException.NodeExistsException nodeExists) {
-        LOG.info("Node already exists in ZK for application id {}", tryId);
-        long elapsedTime = System.currentTimeMillis() - startTime;
-        retryPolicy.allowRetry(i + 1, elapsedTime, RetryLoop.getDefaultRetrySleeper());
-        tryId++;
-      }
-    }
-    if (success) {
-      return ApplicationId.newInstance(namespaceCreationTime, tryId);
-    } else {
-      throw new RuntimeException("Could not obtain unique ApplicationId after " +
-          zkConfig.getCuratorMaxRetries() + " tries");
-    }
+    String prefixPath = namespace + "/application_" + namespaceCreationTime + "_";
+    String znodePath = client.create()
+        .withMode(CreateMode.EPHEMERAL_SEQUENTIAL)
+        .forPath(prefixPath, new byte[0]);
+
+    String sequenceStr = znodePath.substring(znodePath.length() - 10);
+    int assignedId = Integer.parseInt(sequenceStr);
+
+    ApplicationId appId = ApplicationId.newInstance(namespaceCreationTime, assignedId);
+    LOG.debug("Successfully created application id {} for namespace {}", appId, namespace);
+    return appId;
   }
 
   @Override
@@ -174,6 +152,11 @@ private void createNamespaceIfNotExists() throws Exception {
   }
 
   private String pathFor(AMRecord record) {
-    return namespace + "/" + record.getApplicationId().toString();
+    return namespace + "/" + extractApplicationId(record.getApplicationId());
+  }
+
+  @VisibleForTesting
+  static String extractApplicationId(ApplicationId appId) {
+    return String.format("application_%d_%010d", appId.getClusterTimestamp(), appId.getId());
   }
 }

tez-dag/src/test/java/org/apache/tez/dag/api/client/registry/zookeeper/TestZkAMRegistry.java

@@ -18,6 +18,7 @@
  */
 package org.apache.tez.dag.api.client.registry.zookeeper;
 
+import static org.apache.tez.dag.api.client.registry.zookeeper.ZkAMRegistry.extractApplicationId;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -172,7 +173,7 @@ public void testAddAndRemoveAmRecordUpdatesZooKeeper() throws Exception {
       // Add record and verify node contents
       registry.add(record);
 
-      String path = zkConfig.getZkNamespace() + "/" + appId.toString();
+      String path = zkConfig.getZkNamespace() + "/" + extractApplicationId(appId);
       byte[] data = checkClient.getData().forPath(path);
 
       assertNotNull(data, "Data should be written to ZooKeeper for AMRecord");
@@ -186,6 +187,31 @@ public void testAddAndRemoveAmRecordUpdatesZooKeeper() throws Exception {
     }
   }
 
+  @Test
+  public void testGenerateUniqueIds() throws Exception {
+    TezConfiguration conf = createTezConf();
+    try {
+      ZkAMRegistry registry1 = new ZkAMRegistry("external-id-1");
+      ZkAMRegistry registry2 = new ZkAMRegistry("external-id-2");
+      registry1.init(conf);
+      registry1.start();
+      registry2.init(conf);
+      registry2.start();
+
+      ApplicationId first = registry1.generateNewId();
+      ApplicationId second = registry2.generateNewId();
+
+      assertNotNull(first);
+      assertNotNull(second);
+      assertEquals(first.getClusterTimestamp(), second.getClusterTimestamp(),
+          "Registries in the same namespace should share cluster timestamp");
+      assertEquals(first.getId() + 1, second.getId(),
+          "Each registry should receive the next sequential id");
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   private TezConfiguration createTezConf() {
     TezConfiguration conf = new TezConfiguration();
     conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "localhost:" + zkServer.getPort());