PR comments

Author: abstractdog

tez-api/src/main/java/org/apache/tez/client/FrameworkClient.java

@@ -67,9 +67,9 @@ public static FrameworkClient createFrameworkClient(TezConfiguration tezConf) {
         throw new TezUncheckedException("Fail to create LocalClient", e);
       }
     } else {
-      ClientFrameworkService clientFrameworkService = FrameworkUtils.get(ClientFrameworkService.class, tezConf);
-      return clientFrameworkService == null ? new YarnClientFrameworkService().newFrameworkClient()
-          : clientFrameworkService.newFrameworkClient();
+      ClientFrameworkService clientFrameworkService = FrameworkUtils.get(ClientFrameworkService.class, tezConf,
+          YarnClientFrameworkService.class);
+      return clientFrameworkService.newFrameworkClient();
     }
   }
 

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

@@ -19,6 +19,7 @@
 package org.apache.tez.client.registry;
 
 import java.util.Objects;
+import java.util.Optional;
 
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.registry.client.types.ServiceRecord;
@@ -51,6 +52,8 @@ public class AMRecord {
   private final String externalId;
   private final String computeName;
 
+  private ServiceRecord serviceRecord;
+
   /**
    * Creates a new {@code AMRecord} with the given application ID, host, port, and identifier.
    * <p>
@@ -74,8 +77,8 @@ public AMRecord(ApplicationId appId, String hostName, String hostIp, int port, S
     this.hostIp = hostIp;
     this.port = port;
     //externalId is optional, if not provided, convert to empty string
-    this.externalId = (externalId == null) ? "" : externalId;
-    this.computeName = (computeName == null) ? ZkConfig.DEFAULT_COMPUTE_GROUP_NAME : computeName;
+    this.externalId = Optional.ofNullable(externalId).orElse("");
+    this.computeName = Optional.ofNullable(computeName).orElse(ZkConfig.DEFAULT_COMPUTE_GROUP_NAME);
   }
 
   /**
@@ -89,12 +92,15 @@ public AMRecord(ApplicationId appId, String hostName, String hostIp, int port, S
    * @param other the {@code AMRecord} instance to copy
    */
   public AMRecord(AMRecord other) {
-    this.appId = other.getApplicationId();
-    this.hostName = other.getHost();
-    this.hostIp = other.getHostIp();
-    this.port = other.getPort();
-    this.externalId = other.getExternalId();
-    this.computeName = other.getComputeName();
+    this.appId = other.appId;
+    this.hostName = other.hostName;
+    this.hostIp = other.hostIp;
+    this.port = other.port;
+    this.externalId = other.externalId;
+    this.computeName = other.computeName;
+    // all fields are final immutable, we can copy the serviceRecord,
+    // if it's initialized there already, as it won't change
+    this.serviceRecord = other.serviceRecord;
   }
 
   /**
@@ -121,10 +127,6 @@ public ApplicationId getApplicationId() {
     return appId;
   }
 
-  public String getHost() {
-    return hostName;
-  }
-
   public String getHostName() {
     return hostName;
   }
@@ -150,8 +152,7 @@ public boolean equals(Object other) {
     if (this == other) {
       return true;
     }
-    if (other instanceof AMRecord) {
-      AMRecord otherRecord = (AMRecord) other;
+    if (other instanceof AMRecord otherRecord) {
       return appId.equals(otherRecord.appId)
           && hostName.equals(otherRecord.hostName)
           && hostIp.equals(otherRecord.hostIp)
@@ -178,13 +179,17 @@ public boolean equals(Object other) {
    * @return a {@link ServiceRecord} populated with the values of this {@code AMRecord}
    */
   public ServiceRecord toServiceRecord() {
-    ServiceRecord serviceRecord = new ServiceRecord();
+    if (serviceRecord != null) {
+      return serviceRecord;
+    }
+    serviceRecord = new ServiceRecord();
     serviceRecord.set(APP_ID_RECORD_KEY, appId);
     serviceRecord.set(HOST_NAME_RECORD_KEY, hostName);
     serviceRecord.set(HOST_IP_RECORD_KEY, hostIp);
     serviceRecord.set(PORT_RECORD_KEY, port);
     serviceRecord.set(EXTERNAL_ID_KEY, externalId);
     serviceRecord.set(COMPUTE_GROUP_NAME_KEY, computeName);
+
     return serviceRecord;
   }
 

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

@@ -18,16 +18,10 @@
 
 package org.apache.tez.client.registry;
 
-import java.util.ArrayList;
-import java.util.List;
 
-import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.service.AbstractService;
-import org.apache.hadoop.service.ServiceStateException;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /**
  * Base class for {@code AMRegistry} implementations.
@@ -43,62 +37,19 @@
  *   <li>{@code serviceStop}</li>
  * </ul>
  * </p>
- *
- * <p>{@code init} and {@code serviceStart} are invoked during
- * {@code DAGAppMaster.serviceInit()}, while {@code serviceStop} is called
- * when {@code DAGAppMaster} shuts down.</p>
  */
 public abstract class AMRegistry extends AbstractService {
-
-  private static final Logger LOG = LoggerFactory.getLogger(AMRegistry.class);
-  private List<AMRecord> amRecords = new ArrayList<>();
-
-  @Override
-  public void init(Configuration conf) {
-    try {
-      this.serviceInit(conf);
-    } catch (Exception e) {
-      LOG.error("Failed to init AMRegistry: name={}, type={}", getName(), getClass().getName());
-      throw ServiceStateException.convert(e);
-    }
-  }
-
-  @Override
-  public void start() {
-    try {
-      this.serviceStart();
-    } catch(Exception e) {
-      LOG.error("Failed to start AMRegistry: name={}, type={}", getName(), getClass().getName());
-      throw ServiceStateException.convert(e);
-    }
-  }
-
   /* Implementations should provide a public no-arg constructor */
   protected AMRegistry(String name) {
     super(name);
   }
 
-  /* Under typical usage, add will be called once automatically with an AMRecord
-     for the DAGClientServer servicing an AM
-   */
-  public void add(AMRecord server) throws Exception {
-    amRecords.add(server);
-  }
+  public abstract void add(AMRecord server) throws Exception;
 
   public abstract void remove(AMRecord server) throws Exception;
 
-  public ApplicationId generateNewId() throws Exception {
-    return null;
-  }
+  public abstract ApplicationId generateNewId() throws Exception;
 
   public abstract AMRecord createAmRecord(ApplicationId appId, String hostName, String hostIp, int port,
                                           String computeName);
-
-  @Override
-  public void serviceStop() throws Exception {
-    List<AMRecord> records = new ArrayList<>(amRecords);
-    for(AMRecord record : records) {
-      remove(record);
-    }
-  }
 }

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

@@ -23,6 +23,11 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import org.apache.hadoop.yarn.api.records.ApplicationId;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 /**
  * Client-side interface for discovering Application Master (AM) instances
  * registered in the AM registry.
@@ -35,26 +40,18 @@
  * appear or are removed.</p>
  */
 public abstract class AMRegistryClient implements Closeable {
+  private static final Logger LOG = LoggerFactory.getLogger(AMRegistryClient.class);
 
   private final List<AMRegistryClientListener> listeners = new ArrayList<>();
 
-  /**
-   * Returns the current set of registered listeners.
-   *
-   * @return a mutable list of listeners
-   */
-  protected List<AMRegistryClientListener> getListeners() {
-    return listeners;
-  }
-
   /**
    * Lookup AM metadata for the given application ID.
    *
    * @param appId the application ID
    * @return the AM record if found, otherwise {@code null}
    * @throws IOException if the lookup fails
    */
-  public abstract AMRecord getRecord(String appId) throws IOException;
+  public abstract AMRecord getRecord(ApplicationId appId) throws IOException;
 
   /**
    * Retrieve all AM records known in the registry.
@@ -81,7 +78,26 @@ public synchronized void addListener(AMRegistryClientListener listener) {
    */
   protected synchronized void notifyOnAdded(AMRecord record) {
     for (AMRegistryClientListener listener : listeners) {
-      listener.onAdd(record);
+      try {
+        listener.onAdd(record);
+      } catch (Exception e) {
+        LOG.warn("Exception while calling AM add listener, AM record {}", record, e);
+      }
+    }
+  }
+
+  /**
+   * Notify listeners of an updated AM record.
+   *
+   * @param record the updated AM record
+   */
+  protected synchronized void notifyOnUpdated(AMRecord record) {
+    for (AMRegistryClientListener listener : listeners) {
+      try {
+        listener.onUpdate(record);
+      } catch (Exception e) {
+        LOG.warn("Exception while calling AM update listener, AM record {}", record, e);
+      }
     }
   }
 
@@ -92,7 +108,11 @@ protected synchronized void notifyOnAdded(AMRecord record) {
    */
   protected synchronized void notifyOnRemoved(AMRecord record) {
     for (AMRegistryClientListener listener : listeners) {
-      listener.onRemove(record);
+      try {
+        listener.onRemove(record);
+      } catch (Exception e) {
+        LOG.warn("Exception while calling AM remove listener, AM record {}", record, e);
+      }
     }
   }
 }

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

@@ -19,6 +19,23 @@
 package org.apache.tez.client.registry;
 
 public interface AMRegistryClientListener {
+
   void onAdd(AMRecord record);
+
+  /**
+   * Default implementation of {@code onUpdate} delegates to {@code onAdd}.
+   *
+   * <p>This provides a convenient backward-compatible behavior for consumers that
+   * store {@link AMRecord} instances in collections keyed by something stable
+   * (such as {@link ApplicationId}). In such cases, re-adding an {@link AMRecord}
+   * effectively overwrites the previous entry, making an explicit update handler
+   * unnecessary for many implementations.</p>
+   *
+   * @param record the updated {@link AMRecord} instance
+   */
+  default void onUpdate(AMRecord record){
+    onAdd(record);
+  }
+
   void onRemove(AMRecord record);
 }

…meworkplugins/TaskResourceException.java …tez/client/registry/AMRegistryUtils.javatez-dag/src/main/java/org/apache/tez/frameworkplugins/TaskResourceException.java renamed to tez-api/src/main/java/org/apache/tez/client/registry/AMRegistryUtils.java tez-dag/src/main/java/org/apache/tez/frameworkplugins/TaskResourceException.java renamed to tez-api/src/main/java/org/apache/tez/client/registry/AMRegistryUtils.java

@@ -15,10 +15,23 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.tez.frameworkplugins;
+package org.apache.tez.client.registry;
 
-public class TaskResourceException extends Exception {
-  public TaskResourceException(String msg) {
-    super(msg);
+import java.io.IOException;
+
+import org.apache.hadoop.registry.client.binding.RegistryUtils;
+import org.apache.hadoop.registry.client.types.ServiceRecord;
+
+public class AMRegistryUtils {
+
+  public static AMRecord jsonStringToRecord(String json) throws IOException {
+    RegistryUtils.ServiceRecordMarshal marshal = new RegistryUtils.ServiceRecordMarshal();
+    ServiceRecord serviceRecord = marshal.fromJson(json);
+    return new AMRecord(serviceRecord);
+  }
+
+  public static String recordToJsonString(AMRecord amRecord) throws IOException {
+    RegistryUtils.ServiceRecordMarshal marshal = new RegistryUtils.ServiceRecordMarshal();
+    return marshal.toJson(amRecord.toServiceRecord());
   }
 }