PR comments 2

Author: abstractdog

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

@@ -39,14 +39,12 @@
  */
 public interface AMRegistry extends AutoCloseable {
 
-  void add(AMRecord server) throws Exception;
+  void add(AMRecord record) throws Exception;
 
-  void remove(AMRecord server) throws Exception;
+  void remove(AMRecord record) throws Exception;
 
   ApplicationId generateNewId() throws Exception;
 
   AMRecord createAmRecord(ApplicationId appId, String hostName, String hostIp, int port,
                                           String computeName);
-
-  void close();
 }

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

@@ -81,7 +81,7 @@ protected synchronized void notifyOnAdded(AMRecord record) {
       try {
         listener.onAdd(record);
       } catch (Exception e) {
-        LOG.warn("Exception while calling AM add listener, AM record {}", record, e);
+        LOG.warn("Exception while notifying AM add listener for record {}", record, e);
       }
     }
   }
@@ -96,7 +96,7 @@ protected synchronized void notifyOnUpdated(AMRecord record) {
       try {
         listener.onUpdate(record);
       } catch (Exception e) {
-        LOG.warn("Exception while calling AM update listener, AM record {}", record, e);
+        LOG.warn("Exception while notifying AM update listener for record {}", record, e);
       }
     }
   }
@@ -111,7 +111,7 @@ protected synchronized void notifyOnRemoved(AMRecord record) {
       try {
         listener.onRemove(record);
       } catch (Exception e) {
-        LOG.warn("Exception while calling AM remove listener, AM record {}", record, e);
+        LOG.warn("Exception while notifying AM remove listener for record {}", record, e);
       }
     }
   }

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

@@ -113,6 +113,7 @@ public void start() throws Exception {
     client = zkConf.createCuratorFramework();
     cache = new TreeCache(client, zkConf.getZkNamespace());
     client.start();
+    client.blockUntilConnected();
     cache.start();
     listener = new ZkRegistryListener();
     cache.getListenable().addListener(listener);
@@ -175,7 +176,7 @@ public void childEvent(final CuratorFramework clientParam, final TreeCacheEvent
             if (amRecord != null) {
               LOG.info("AM updated data: {}. Notifying listeners.", amRecord);
               amRecordCache.put(amRecord.getApplicationId(), amRecord);
-              notifyOnAdded(amRecord);
+              notifyOnUpdated(amRecord);
             }
           }
           break;

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

@@ -37,11 +37,9 @@
 
 public class ZkFrameworkClient extends FrameworkClient {
 
-  private AMRecord amRecord;
+  private volatile AMRecord amRecord;
   private ZkAMRegistryClient amRegistryClient = null;
   private volatile boolean isRunning = false;
-  private String amHost;
-  private int amPort;
 
   @Override
   public synchronized void init(TezConfiguration tezConf) {
@@ -56,6 +54,9 @@ public synchronized void init(TezConfiguration tezConf) {
 
   @Override
   public void start() {
+    if (isRunning) {
+      return;
+    }
     try {
       amRegistryClient.start();
       isRunning = true;
@@ -140,8 +141,6 @@ public ApplicationReport getApplicationReport(ApplicationId appId) throws YarnEx
       report.setDiagnostics("AM record not found (likely died) in zookeeper for application id: " + appId);
     } else {
       report.setHost(amRecord.getHostName());
-      amHost = amRecord.getHostName();
-      amPort = amRecord.getPort();
       report.setRpcPort(amRecord.getPort());
       report.setYarnApplicationState(YarnApplicationState.RUNNING);
     }
@@ -155,12 +154,12 @@ public boolean isRunning() {
 
   @Override
   public String getAmHost() {
-    return amHost;
+    return amRecord == null ? null : amRecord.getHostName();
   }
 
   @Override
   public int getAmPort() {
-    return amPort;
+    return amRecord == null ? 0 : amRecord.getPort();
   }
 
   @VisibleForTesting

tez-api/src/main/java/org/apache/tez/frameworkplugins/FrameworkUtils.java

@@ -18,8 +18,11 @@
 package org.apache.tez.frameworkplugins;
 
 
+import java.lang.reflect.InvocationTargetException;
+
 import javax.annotation.Nullable;
 
+import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.tez.common.ReflectionUtils;
 import org.apache.tez.dag.api.TezConfiguration;
@@ -54,17 +57,18 @@ public static <T extends FrameworkService> T get(Class<T> interfaze, @Nullable C
     String modeInConf = conf != null ? conf.get(TezConfiguration.TEZ_FRAMEWORK_MODE) : null;
     String modeInEnv = System.getenv(TezConstants.TEZ_FRAMEWORK_MODE);
     try {
-      if (modeInConf != null) {
+      if (StringUtils.isNotEmpty(modeInConf)) {
         return getByMode(interfaze, modeInConf);
-      } else if (modeInEnv != null) {
+      } else if (StringUtils.isNotEmpty(modeInEnv)) {
         return getByMode(interfaze, modeInEnv);
       } else if (defaultClazz != null) {
-        return (T) defaultClazz.newInstance();
+        return (T) defaultClazz.getDeclaredConstructor().newInstance();
       } else {
         throw new RuntimeException(
             "Framework service not found in any mode: configuration, environment, or default class");
       }
-    } catch (TezReflectionException | InstantiationException | IllegalAccessException e) {
+    } catch (TezReflectionException | InstantiationException | IllegalAccessException | NoSuchMethodException |
+             InvocationTargetException e) {
       throw new RuntimeException("Failed to load framework service for interface: " + interfaze.getName(), e);
     }
   }

tez-api/src/test/java/org/apache/tez/client/TestTezClient.java

@@ -546,6 +546,8 @@ public void testTezClientReconnect(boolean isSession) throws Exception {
     assertEquals(dagClient.getSessionIdentifierString(), appId.toString());
 
     dagClient.close();
+    client.stop();
+    client2.stop();
   }
 
   @Test (timeout=5000)

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

@@ -156,7 +156,7 @@ public ApplicationId generateNewId() throws Exception {
   }
 
   @Override
-  public AMRecord createAmRecord(ApplicationId appId,  String hostName, String hostIp, int port, String computeName) {
+  public AMRecord createAmRecord(ApplicationId appId, String hostName, String hostIp, int port, String computeName) {
     return new AMRecord(appId, hostName, hostIp, port, externalId, computeName);
   }
 

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

@@ -46,7 +46,6 @@
 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;
@@ -1979,7 +1978,17 @@ void stopServices() {
       }
     }
 
-    Optional.ofNullable(frameworkService.getAMRegistry(this.amConf)).ifPresent(AMRegistry::close);
+    AMRegistry registry = frameworkService.getAMRegistry(this.amConf);
+    if (registry != null) {
+      try {
+        registry.close();
+      } catch (Exception e) {
+        LOG.warn("Failed to close registry", e);
+        if (firstException != null) {
+          firstException = e;
+        }
+      }
+    }
 
     //after stopping all services, rethrow the first exception raised
     if (firstException != null) {