TEZ-4759: An external AM never re-registers after a ZooKeeper session loss

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

@@ -18,12 +18,14 @@
  */
 package org.apache.tez.dag.api.client.registry.zookeeper;
 
+import java.io.IOException;
 import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
 
 import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.framework.recipes.nodes.PersistentNode;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
@@ -49,7 +51,7 @@ public class ZkAMRegistry implements AMRegistry {
 
   private static final Logger LOG = LoggerFactory.getLogger(ZkAMRegistry.class);
 
-  private final List<AMRecord> amRecords = Collections.synchronizedList(new ArrayList<>());
+  private final Map<AMRecord, PersistentNode> amNodes = new ConcurrentHashMap<>();
   private final String externalId;
 
   private CuratorFramework client = null;
@@ -80,7 +82,7 @@ public void start() throws Exception {
    * is logged.</p>
    */
   public void close() {
-    for (AMRecord amRecord : new ArrayList<>(amRecords)) {
+    for (AMRecord amRecord : amNodes.keySet()) {
       try {
         remove(amRecord);
       } catch (Exception e) {
@@ -95,23 +97,27 @@ public void close() {
   //zkNode at the path:  <TEZ_AM_REGISTRY_NAMESPACE>/<appId>
   @Override
   public void add(AMRecord server) throws Exception {
-    String json = AMRegistryUtils.recordToJsonString(server);
-    try {
-      final String path = pathFor(server);
-      client.setData().forPath(path, json.getBytes(StandardCharsets.UTF_8));
-      LOG.info("Added AMRecord to zkpath {}", path);
-    } catch (KeeperException.NoNodeException nne) {
-      client.create().creatingParentContainersIfNeeded().withMode(CreateMode.EPHEMERAL)
-          .forPath(pathFor(server), json.getBytes(StandardCharsets.UTF_8));
+    // PersistentNode re-creates the ephemeral node after a session loss or a delete; a plain create would leave
+    // the AM running but unregistered.
+    final String path = pathFor(server);
+    byte[] json = AMRegistryUtils.recordToJsonString(server).getBytes(StandardCharsets.UTF_8);
+    PersistentNode node = new PersistentNode(client, CreateMode.EPHEMERAL, false, path, json);
+    node.start();
+    if (!node.waitForInitialCreate(zkConfig.getConnectionTimeoutMs(), TimeUnit.MILLISECONDS)) {
+      node.close();
+      throw new IOException("Timed out creating AMRecord at zkpath " + path);
     }
-    amRecords.add(server);
+    amNodes.put(server, node);
+    LOG.info("Added AMRecord to zkpath {}", path);
   }
 
   @Override
   public void remove(AMRecord server) throws Exception {
-    amRecords.remove(server);
+    PersistentNode node = amNodes.remove(server);
     final String path = pathFor(server);
-    client.delete().forPath(path);
+    if (node != null) {
+      node.close();   // deletes the node
+    }
     LOG.info("Deleted AMRecord from zkpath {}", path);
   }
 

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

@@ -38,6 +38,7 @@
 import java.util.concurrent.TimeUnit;
 
 import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.test.InstanceSpec;
 import org.apache.curator.test.TestingServer;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.tez.client.registry.AMRecord;
@@ -187,6 +188,40 @@ public void testAddAndRemoveAmRecordUpdatesZooKeeper() throws Exception {
     }
   }
 
+  @Test
+  @Timeout(120)
+  public void testAmRecordSurvivesZooKeeperSessionLoss() throws Exception {
+    // own server: an empty one on the same port expires every session held against the old one
+    int port = InstanceSpec.getRandomPort();
+    TestingServer server = new TestingServer(port);
+    TezConfiguration conf = createTezConf();
+    conf.set(TezConfiguration.TEZ_AM_ZOOKEEPER_QUORUM, "localhost:" + port);
+    ZkConfig zkConfig = new ZkConfig(conf);
+
+    try (ZkAMRegistry registry = new ZkAMRegistry("external-id")) {
+      registry.init(conf);
+      registry.start();
+      AMRecord record = registry.createAmRecord(
+          registry.generateNewId(), "localhost", "127.0.0.1", 10000, "default-compute");
+      registry.add(record);
+
+      server.close();
+      server = new TestingServer(port);
+
+      try (CuratorFramework checkClient = zkConfig.createCuratorFramework()) {
+        checkClient.start();
+        String path = zkConfig.getZkNamespace() + "/" + extractApplicationId(record.getApplicationId());
+        while (checkClient.checkExists().forPath(path) == null) {
+          Thread.sleep(100);
+        }
+        assertEquals(AMRegistryUtils.recordToJsonString(record),
+            new String(checkClient.getData().forPath(path), StandardCharsets.UTF_8));
+      }
+    } finally {
+      server.close();
+    }
+  }
+
   @Test
   public void testGenerateUniqueIds() throws Exception {
     TezConfiguration conf = createTezConf();