diff --git a/tez-dag/src/main/java/org/apache/tez/dag/api/client/registry/zookeeper/ZkAMRegistry.java b/tez-dag/src/main/java/org/apache/tez/dag/api/client/registry/zookeeper/ZkAMRegistry.java index b93f4b90ab..75caf57583 100644 --- a/tez-dag/src/main/java/org/apache/tez/dag/api/client/registry/zookeeper/ZkAMRegistry.java +++ b/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 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()); } } diff --git a/tez-dag/src/test/java/org/apache/tez/dag/api/client/registry/zookeeper/TestZkAMRegistry.java b/tez-dag/src/test/java/org/apache/tez/dag/api/client/registry/zookeeper/TestZkAMRegistry.java index f1d762387c..acc6f2beb7 100644 --- a/tez-dag/src/test/java/org/apache/tez/dag/api/client/registry/zookeeper/TestZkAMRegistry.java +++ b/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());