reworked am registry service lifecycle

Author: abstractdog

tez-api/src/main/java/org/apache/tez/client/registry/AMRegistry.java

@@ -19,7 +19,6 @@
 package org.apache.tez.client.registry;
 
 
-import org.apache.hadoop.service.AbstractService;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 
 
@@ -38,18 +37,16 @@
  * </ul>
  * </p>
  */
-public abstract class AMRegistry extends AbstractService {
-  /* Implementations should provide a public no-arg constructor */
-  protected AMRegistry(String name) {
-    super(name);
-  }
+public interface AMRegistry extends AutoCloseable {
 
-  public abstract void add(AMRecord server) throws Exception;
+  void add(AMRecord server) throws Exception;
 
-  public abstract void remove(AMRecord server) throws Exception;
+  void remove(AMRecord server) throws Exception;
 
-  public abstract ApplicationId generateNewId() throws Exception;
+  ApplicationId generateNewId() throws Exception;
 
-  public abstract AMRecord createAmRecord(ApplicationId appId, String hostName, String hostIp, int port,
+  AMRecord createAmRecord(ApplicationId appId, String hostName, String hostIp, int port,
                                           String computeName);
+
+  void close();
 }

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

@@ -45,7 +45,7 @@
  * Clients should use org.apache.tez.dag.api.client.registry.zookeeper.ZkAMRegistryClient instead.
  */
 @InterfaceAudience.Private
-public class ZkAMRegistry extends AMRegistry {
+public class ZkAMRegistry implements AMRegistry {
 
   private static final Logger LOG = LoggerFactory.getLogger(ZkAMRegistry.class);
 
@@ -57,20 +57,17 @@ public class ZkAMRegistry extends AMRegistry {
   private ZkConfig zkConfig = null;
 
   public ZkAMRegistry(String externalId) {
-    super("ZkAMRegistry");
     this.externalId = externalId;
   }
 
-  @Override
-  public void serviceInit(Configuration conf) {
+  public void init(Configuration conf) {
     zkConfig = new ZkConfig(conf);
     this.client = zkConfig.createCuratorFramework();
     this.namespace = zkConfig.getZkNamespace();
     LOG.info("ZkAMRegistry initialized");
   }
 
-  @Override
-  public void serviceStart() throws Exception {
+  public void start() throws Exception {
     client.start();
     LOG.info("ZkAMRegistry started");
   }
@@ -81,11 +78,8 @@ public void serviceStart() throws Exception {
    *
    * <p>After all removal attempts, the ZooKeeper client is closed and the shutdown
    * is logged.</p>
-   *
-   * @throws Exception if a failure occurs while closing the ZooKeeper client
    */
-  @Override
-  public void serviceStop() throws Exception {
+  public void close() {
     for (AMRecord amRecord : new ArrayList<>(amRecords)) {
       try {
         remove(amRecord);

tez-dag/src/main/java/org/apache/tez/dag/app/DAGAppMaster.java

@@ -25,6 +25,7 @@
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.PrintWriter;
+import java.net.InetSocketAddress;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
@@ -45,6 +46,7 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Random;
 import java.util.Set;
 import java.util.Timer;
@@ -618,11 +620,6 @@ protected void serviceInit(final Configuration conf) throws Exception {
         Executors.newFixedThreadPool(threadCount, new ThreadFactoryBuilder()
             .setDaemon(true).setNameFormat("App Shared Pool - #%d").build());
     execService = MoreExecutors.listeningDecorator(rawExecutor);
-    AMRegistry amRegistry = frameworkService.getAMRegistry(conf);
-    if (amRegistry != null) {
-      initAmRegistry(appAttemptID.getApplicationId(), amRegistry, clientRpcServer);
-      addIfService(amRegistry, false);
-    }
 
     initServices(conf);
     super.serviceInit(conf);
@@ -650,16 +647,35 @@ private static ServerFrameworkService getFrameworkService(Configuration conf) {
     return FrameworkUtils.get(ServerFrameworkService.class, conf, YarnServerFrameworkService.class);
   }
 
-  @VisibleForTesting
-  public void initAmRegistry(ApplicationId appId, AMRegistry amRegistry, DAGClientServer dagClientServer) {
-    dagClientServer.registerServiceListener((service) -> {
+  protected void initClientRpcServer() {
+    clientRpcServer = new DAGClientServer(clientHandler, appAttemptID, recoveryFS);
+    addIfService(clientRpcServer, true);
+
+    initAmRegistryCallbackForRecordAdd();
+  }
+
+  /**
+   * Initializes an AM registry callback when the clientRpcServer is already initialized
+   */
+  private void initAmRegistryCallbackForRecordAdd() {
+    AMRegistry amRegistry = frameworkService.getAMRegistry(this.amConf);
+    if (amRegistry == null) {
+      return;
+    }
+    if (clientRpcServer == null){
+      throw new IllegalStateException(
+          "Client RPC Server has not been initialized before attempting to initialize an AM registry");
+    }
+
+    ApplicationId appId = appAttemptID.getApplicationId();
+
+    clientRpcServer.registerServiceListener((service) -> {
       if (service.isInState(STATE.STARTED)) {
+        InetSocketAddress rpcServerAddress = clientRpcServer.getBindAddress();
+
         final String computeName = System.getenv(ZkConfig.COMPUTE_GROUP_NAME_ENV);
-        AMRecord amRecord = amRegistry.createAmRecord(
-            appId, dagClientServer.getBindAddress().getHostName(),
-            dagClientServer.getBindAddress().getAddress().getHostAddress(),
-            dagClientServer.getBindAddress().getPort(), computeName
-        );
+        AMRecord amRecord = amRegistry.createAmRecord(appId, rpcServerAddress.getHostName(),
+            rpcServerAddress.getAddress().getHostAddress(), rpcServerAddress.getPort(), computeName);
         try {
           amRegistry.add(amRecord);
           LOG.info("Added AMRecord: {} to registry..", amRecord);
@@ -670,11 +686,6 @@ public void initAmRegistry(ApplicationId appId, AMRegistry amRegistry, DAGClient
     });
   }
 
-  protected void initClientRpcServer() {
-    clientRpcServer = new DAGClientServer(clientHandler, appAttemptID, recoveryFS);
-    addIfService(clientRpcServer, true);
-  }
-
   @VisibleForTesting
   protected DAGAppMasterShutdownHandler createShutdownHandler() {
     return new DAGAppMasterShutdownHandler();
@@ -1967,6 +1978,9 @@ void stopServices() {
         firstException = ex;
       }
     }
+
+    Optional.ofNullable(frameworkService.getAMRegistry(this.amConf)).ifPresent(AMRegistry::close);
+
     //after stopping all services, rethrow the first exception raised
     if (firstException != null) {
       throw ServiceStateException.convert(firstException);

tez-dag/src/main/java/org/apache/tez/frameworkplugins/zookeeper/ZkStandaloneServerFrameworkService.java

@@ -59,6 +59,7 @@ public synchronized AMRegistry getAMRegistry(Configuration conf) {
       try {
         amRegistry = new ZkAMRegistry(externalId);
         amRegistry.init(conf);
+        amRegistry.start();
         LOG.info("Created Zookeeper based AM Registry with externalId: {}", externalId);
       } catch (Exception e) {
         throw new RuntimeException(e);

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

@@ -78,8 +78,8 @@ public void teardown() throws Exception {
   public void testGenerateNewIdProducesUniqueIds() throws Exception {
     TezConfiguration conf = createTezConf();
     try (ZkAMRegistry registry = new ZkAMRegistry("external-id")) {
-      registry.serviceInit(conf);
-      registry.serviceStart();
+      registry.init(conf);
+      registry.start();
 
       ApplicationId first = registry.generateNewId();
       ApplicationId second = registry.generateNewId();
@@ -100,8 +100,8 @@ public void testGenerateNewIdFromParallelThreads() throws Exception {
     conf.setInt(TezConfiguration.TEZ_AM_CURATOR_MAX_RETRIES, 29);
 
     try (ZkAMRegistry registry = new ZkAMRegistry("external-id")) {
-      registry.serviceInit(conf);
-      registry.serviceStart();
+      registry.init(conf);
+      registry.start();
 
       ExecutorService executor = Executors.newFixedThreadPool(threadCount);
       CountDownLatch startLatch = new CountDownLatch(1);
@@ -157,8 +157,8 @@ public void testAddAndRemoveAmRecordUpdatesZooKeeper() throws Exception {
 
     try (ZkAMRegistry registry = new ZkAMRegistry("external-id");
          CuratorFramework checkClient = zkConfig.createCuratorFramework()) {
-      registry.serviceInit(conf);
-      registry.serviceStart();
+      registry.init(conf);
+      registry.start();
 
       checkClient.start();