TEZ-4008: Pluggable AM FrameworkServices and AmExtensions (2/3)

tez-api/findbugs-exclude.xml

@@ -131,4 +131,24 @@
     <Bug pattern="IS2_INCONSISTENT_SYNC" />
   </Match>
 
+  <!-- TEZ-4008 -->
+  <Match>
+    <Class name="org.apache.tez.client.registry.AMRecord" />
+    <Method name="&lt;init&gt;"
+            params="org.apache.hadoop.yarn.api.records.ApplicationId, java.lang.String, int, java.lang.String"
+            returns="void"/>
+    <Bug pattern="EI_EXPOSE_REP2" />
+  </Match>
+
+  <Match>
+    <Class name="org.apache.tez.client.registry.AMRecord" />
+    <Method name="&lt;init&gt;" params="org.apache.hadoop.registry.client.types.ServiceRecord"/>
+    <Bug pattern="CT_CONSTRUCTOR_THROW" />
+  </Match>
+
+  <Match>
+    <Class name="org.apache.tez.client.registry.AMRecord" />
+    <Method name="getApplicationId" />
+    <Bug pattern="EI_EXPOSE_REP" />
+  </Match>
 </FindBugsFilter>

tez-api/pom.xml

@@ -73,6 +73,11 @@
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-yarn-client</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.hadoop</groupId>
+      <artifactId>hadoop-registry</artifactId>
+      <version>${hadoop.version}</version>
+    </dependency>
     <dependency>
      <groupId>org.apache.commons</groupId>
      <artifactId>commons-collections4</artifactId>

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

@@ -0,0 +1,163 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.tez.client.registry;
+
+import java.util.Objects;
+
+import org.apache.hadoop.classification.InterfaceAudience;
+import org.apache.hadoop.registry.client.types.ServiceRecord;
+import org.apache.hadoop.yarn.api.records.ApplicationId;
+
+
+/**
+ * Record representing an Application Master (AM) instance within Tez.
+ * <p>
+ * This class can be serialized to and from a {@link ServiceRecord}, enabling
+ * storage and retrieval of AM metadata in external systems. Some constructors
+ * and methods are not necessarily used within the Tez codebase itself, but
+ * are part of the Tez API and intended for Tez clients that manage or interact
+ * with Tez unmanaged sessions.
+ */
+@InterfaceAudience.Public
+public class AMRecord {
+  private static final String APP_ID_RECORD_KEY = "appId";
+  private static final String HOST_RECORD_KEY = "host";
+  private static final String PORT_RECORD_KEY = "port";
+  private static final String OPAQUE_ID_KEY = "id";
+
+  private final ApplicationId appId;
+  private final String host;
+  private final int port;
+  private final String id;
+
+  /**
+   * Creates a new {@code AMRecord} with the given application ID, host, port, and identifier.
+   * <p>
+   * If the provided identifier is {@code null}, it will be converted to an empty string.
+   * <p>
+   * Although this constructor may not be used directly within Tez internals,
+   * it is part of the public API for Tez clients that handle unmanaged sessions.
+   *
+   * @param appId the {@link ApplicationId} of the Tez application
+   * @param host the hostname where the Application Master is running
+   * @param port the port number on which the Application Master is listening
+   * @param id an opaque identifier for the record; if {@code null}, defaults to an empty string
+   */
+  public AMRecord(ApplicationId appId, String host, int port, String id) {
+    this.appId = appId;
+    this.host = host;
+    this.port = port;
+    //If id is not provided, convert to empty string
+    this.id = (id == null) ? "" : id;
+  }
+
+  /**
+   * Copy constructor.
+   * <p>
+   * Creates a new {@code AMRecord} by copying the fields of another instance.
+   * <p>
+   * This constructor is mainly useful for client-side logic and session handling,
+   * and may not be invoked directly within the Tez codebase.
+   *
+   * @param other the {@code AMRecord} instance to copy
+   */
+  public AMRecord(AMRecord other) {
+    this.appId = other.getApplicationId();
+    this.host = other.getHost();
+    this.port = other.getPort();
+    this.id = other.getId();
+  }
+
+  /**
+   * Constructs a new {@code AMRecord} from a {@link ServiceRecord}.
+   * <p>
+   * This allows conversion from serialized metadata back into an in-memory {@code AMRecord}.
+   * <p>
+   * While not always used in Tez internals, it exists in the Tez API so
+   * clients can reconstruct AM information when working with unmanaged sessions.
+   *
+   * @param serviceRecord the {@link ServiceRecord} containing AM metadata
+   * @throws IllegalArgumentException if required keys are missing or invalid
+   */
+  public AMRecord(ServiceRecord serviceRecord) {
+    this.appId = ApplicationId.fromString(serviceRecord.get(APP_ID_RECORD_KEY));
+    this.host = serviceRecord.get(HOST_RECORD_KEY);
+    this.port = Integer.parseInt(serviceRecord.get(PORT_RECORD_KEY));
+    this.id = serviceRecord.get(OPAQUE_ID_KEY);
+  }
+
+  public ApplicationId getApplicationId() {
+    return appId;
+  }
+
+  public String getHost() {
+    return host;
+  }
+
+  public int getPort() {
+    return port;
+  }
+
+  public String getId() {
+    return id;
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (this == other) {
+      return true;
+    }
+    if (other instanceof AMRecord otherRecord) {
+      return appId.equals(otherRecord.appId)
+          && host.equals(otherRecord.host)
+          && port == otherRecord.port
+          && id.equals(otherRecord.id);
+    } else {
+      return false;
+    }
+  }
+
+  /**
+   * Converts this {@code AMRecord} into a {@link ServiceRecord}.
+   * <p>
+   * The returned {@link ServiceRecord} contains the Application Master metadata
+   * (application ID, host, port, and opaque identifier) so that it can be stored
+   * in an external registry or retrieved later.
+   * <p>
+   * While this method may not be directly used within Tez internals,
+   * it is part of the Tez public API and is intended for Tez clients
+   * that interact with unmanaged sessions or otherwise need to
+   * persist/reconstruct Application Master information.
+   *
+   * @return a {@link ServiceRecord} populated with the values of this {@code AMRecord}
+   */
+  public ServiceRecord toServiceRecord() {
+    ServiceRecord serviceRecord = new ServiceRecord();
+    serviceRecord.set(APP_ID_RECORD_KEY, appId);
+    serviceRecord.set(HOST_RECORD_KEY, host);
+    serviceRecord.set(PORT_RECORD_KEY, port);
+    serviceRecord.set(OPAQUE_ID_KEY, id);
+    return serviceRecord;
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(appId, host, port, id);
+  }
+}

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

@@ -0,0 +1,24 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@Public
+@Evolving
+package org.apache.tez.client.registry;
+
+import org.apache.hadoop.classification.InterfaceAudience.Public;
+import org.apache.hadoop.classification.InterfaceStability.Evolving;

tez-api/src/main/java/org/apache/tez/dag/api/TezConfiguration.java

@@ -2349,4 +2349,11 @@ static Set<String> getPropertySet() {
   @ConfigurationScope(Scope.AM)
   @ConfigurationProperty
   public static final String TEZ_AM_STANDALONE_CONFS = TEZ_AM_PREFIX + "standalone.confs";
+
+  /**
+  * String value. The class to be used for the AM registry.
+  */
+  @ConfigurationScope(Scope.AM)
+  @ConfigurationProperty
+  public static final String TEZ_AM_REGISTRY_CLASS = TEZ_AM_PREFIX + "registry.class";
 }

tez-api/src/main/java/org/apache/tez/dag/api/TezConstants.java

@@ -102,6 +102,9 @@ public final class TezConstants {
   /// Version-related Environment variables
   public static final String TEZ_CLIENT_VERSION_ENV = "TEZ_CLIENT_VERSION";
 
+  //Arbitrary opaque ID to identify AM instances from AMRegistryClient
+  public static final String TEZ_AM_UUID = "TEZ_AM_UUID";
+
   private static final String TEZ_AM_SERVICE_PLUGIN_NAME_YARN_CONTAINERS = "TezYarn";
   private static final String TEZ_AM_SERVICE_PLUGIN_NAME_IN_AM = "TezUber";
 

tez-dag/src/main/java/org/apache/tez/client/LocalClient.java

@@ -415,7 +415,7 @@ protected DAGAppMaster createDAGAppMaster(ApplicationAttemptId applicationAttemp
           versionInfo.getVersion(), credentials, jobUserName, amPluginDescriptorProto)
       : new DAGAppMaster(applicationAttemptId, cId, currentHost, nmPort, nmHttpPort,
           SystemClock.getInstance(), appSubmitTime, isSession, userDir, localDirs, logDirs,
-          versionInfo.getVersion(), credentials, jobUserName, amPluginDescriptorProto);
+          versionInfo.getVersion(), credentials, jobUserName, amPluginDescriptorProto, null);
   }
 
   @Override

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

@@ -0,0 +1,47 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.tez.dag.api.client.registry;
+
+import org.apache.hadoop.service.AbstractService;
+import org.apache.tez.client.registry.AMRecord;
+
+/**
+ * Base class for AMRegistry implementations.
+ * The specific implementation class is configured by `tez.am.registry.class`.
+ *
+ * Implementations should handle the relevant service lifecycle operations:
+ * `init`, `serviceStart`, `serviceStop`, etc.
+ * - `init` and `serviceStart` are invoked during `DAGAppMaster.serviceInit`.
+ * - `serviceStop` is invoked on `DAGAppMaster` shutdown.
+ */
+public abstract class AMRegistry extends AbstractService {
+
+  /* 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 that services an AM. */
+  public abstract void add(AMRecord server) throws Exception;
+
+  /* Under typical usage, implementations should remove any stale AMRecords upon serviceStop. */
+  public abstract void remove(AMRecord server) throws Exception;
+
+}

tez-dag/src/main/java/org/apache/tez/dag/api/client/registry/package-info.java

@@ -0,0 +1,24 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@Public
+@Evolving
+package org.apache.tez.dag.api.client.registry;
+
+import org.apache.hadoop.classification.InterfaceAudience.Public;
+import org.apache.hadoop.classification.InterfaceStability.Evolving;