From 980363e92fbd5f37cf896c4758ff8d4fa74dcc48 Mon Sep 17 00:00:00 2001 From: Denys Kuzmenko Date: Wed, 30 Sep 2026 18:35:53 +0300 Subject: [PATCH] TEZ-4759: An external AM never re-registers after a ZooKeeper session loss - ZkAMRegistry.add() wrote the AMRecord once, as a bare EPHEMERAL node; a session expiry deletes it and nothing re-creates it, so clients see an empty registry while the AM keeps running - Hold a Curator PersistentNode per AMRecord: it re-creates the node after a session loss or a delete; remove() closes it - add() fails when the node is not created within the ZooKeeper connection timeout - TestZkAMRegistry: after the ZooKeeper server is replaced, the AMRecord node is back --- .../registry/zookeeper/ZkAMRegistry.java | 38 +++++++++++-------- .../registry/zookeeper/TestZkAMRegistry.java | 35 +++++++++++++++++ 2 files changed, 57 insertions(+), 16 deletions(-) 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 75caf57583..515a735e63 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 @@ -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 amRecords = Collections.synchronizedList(new ArrayList<>()); + private final Map amNodes = new ConcurrentHashMap<>(); private final String externalId; private CuratorFramework client = null; @@ -80,7 +82,7 @@ public void start() throws Exception { * is logged.

*/ 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: / @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); } 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 acc6f2beb7..6d8a18a421 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 @@ -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();