TEZ-4007: Introduce AmExtensions and Zookeeper-based FrameworkServices

pom.xml

@@ -69,6 +69,7 @@
     <commons-io.version>2.16.0</commons-io.version>
     <commons-lang.version>2.6</commons-lang.version>
     <clover.license>${user.home}/clover.license</clover.license>
+    <curator.version>5.9.0</curator.version>
     <dependency-check-maven.version>3.2.0</dependency-check-maven.version>
     <dependency-maven-plugin.version>3.8.1</dependency-maven-plugin.version>
     <spotbugs.version>4.9.3</spotbugs.version>
@@ -736,13 +737,30 @@
           </exclusion>
         </exclusions>
       </dependency>
+      <dependency>
+        <groupId>org.apache.hadoop</groupId>
+        <artifactId>hadoop-registry</artifactId>
+        <version>${hadoop.version}</version>
+      </dependency>
       <dependency>
         <groupId>org.apache.hadoop</groupId>
         <artifactId>hadoop-hdfs</artifactId>
         <version>${hadoop.version}</version>
         <type>test-jar</type>
         <scope>test</scope>
       </dependency>
+      <dependency>
+        <groupId>org.apache.curator</groupId>
+        <artifactId>curator-test</artifactId>
+        <version>${curator.version}</version>
+        <scope>test</scope>
+        <exclusions>
+          <exclusion>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-api</artifactId>
+          </exclusion>
+        </exclusions>
+      </dependency>
       <dependency>
         <groupId>org.mockito</groupId>
         <artifactId>mockito-core</artifactId>

tez-api/findbugs-exclude.xml

@@ -151,4 +151,37 @@
     <Method name="getApplicationId" />
     <Bug pattern="EI_EXPOSE_REP" />
   </Match>
+
+  <!-- TEZ-4007 -->
+  <Match>
+    <Class name="org.apache.tez.client.registry.zookeeper.ZkFrameworkClient" />
+    <Field name="amPort" />
+    <Bug pattern="AT_STALE_THREAD_WRITE_OF_PRIMITIVE" />
+  </Match>
+
+  <Match>
+    <Class name="org.apache.tez.client.TezYarnClient" />
+    <Method name="&lt;init&gt;" params="org.apache.hadoop.yarn.client.api.YarnClient" 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.yarn.api.records.ApplicationId, java.lang.String, java.lang.String, int, java.lang.String, java.lang.String"
+            returns="void"/>
+    <Bug pattern="EI_EXPOSE_REP2" />
+  </Match>
+
+  <Match>
+    <Class name="org.apache.tez.client.registry.AMRecord" />
+    <Method name="toServiceRecord" />
+    <Bug pattern="EI_EXPOSE_REP" />
+  </Match>
+
+  <Match>
+    <Class name="org.apache.tez.frameworkplugins.FrameworkUtils" />
+    <Method name="get" params="java.lang.Class, org.apache.hadoop.conf.Configuration, java.lang.Class" returns="org.apache.tez.frameworkplugins.FrameworkService" />
+    <Bug pattern="REFLC_REFLECTION_MAY_INCREASE_ACCESSIBILITY_OF_CLASS" />
+  </Match>
 </FindBugsFilter>

tez-api/pom.xml

@@ -76,7 +76,6 @@
     <dependency>
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-registry</artifactId>
-      <version>${hadoop.version}</version>
     </dependency>
     <dependency>
      <groupId>org.apache.commons</groupId>
@@ -124,6 +123,11 @@
         <groupId>org.xerial.snappy</groupId>
         <artifactId>snappy-java</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.curator</groupId>
+      <artifactId>curator-test</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

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

@@ -26,7 +26,6 @@
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ApplicationReport;
 import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext;
-import org.apache.hadoop.yarn.client.api.YarnClient;
 import org.apache.hadoop.yarn.client.api.YarnClientApplication;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.tez.common.RPCUtil;
@@ -46,6 +45,9 @@
 import org.apache.tez.dag.api.client.rpc.DAGClientAMProtocolRPC.ShutdownSessionRequestProto;
 import org.apache.tez.dag.api.client.rpc.DAGClientAMProtocolRPC.SubmitDAGRequestProto;
 import org.apache.tez.dag.api.client.rpc.DAGClientAMProtocolRPC.SubmitDAGResponseProto;
+import org.apache.tez.frameworkplugins.ClientFrameworkService;
+import org.apache.tez.frameworkplugins.FrameworkUtils;
+import org.apache.tez.frameworkplugins.yarn.YarnClientFrameworkService;
 
 import com.google.protobuf.ServiceException;
 
@@ -57,16 +59,18 @@ public abstract class FrameworkClient {
   protected static final Logger LOG = LoggerFactory.getLogger(FrameworkClient.class);
 
   public static FrameworkClient createFrameworkClient(TezConfiguration tezConf) {
-
     boolean isLocal = tezConf.getBoolean(TezConfiguration.TEZ_LOCAL_MODE, TezConfiguration.TEZ_LOCAL_MODE_DEFAULT);
     if (isLocal) {
       try {
         return ReflectionUtils.createClazzInstance("org.apache.tez.client.LocalClient");
       } catch (TezReflectionException e) {
         throw new TezUncheckedException("Fail to create LocalClient", e);
       }
+    } else {
+      ClientFrameworkService clientFrameworkService = FrameworkUtils.get(ClientFrameworkService.class, tezConf,
+          YarnClientFrameworkService.class);
+      return clientFrameworkService.newFrameworkClient();
     }
-    return new TezYarnClient(YarnClient.createYarnClient());
   }
 
   /**

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

@@ -484,7 +484,9 @@ public synchronized TezClient getClient(ApplicationId appId) throws TezException
   }
 
   private void startFrameworkClient() {
-    frameworkClient = createFrameworkClient();
+    if (frameworkClient == null) {
+      frameworkClient = createFrameworkClient();
+    }
     frameworkClient.init(amConfig.getTezConfiguration());
     frameworkClient.start();
   }

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

@@ -43,7 +43,7 @@ public class TezYarnClient extends FrameworkClient {
   private String amHost;
   private int amPort;
 
-  protected TezYarnClient(YarnClient yarnClient) {
+  public TezYarnClient(YarnClient yarnClient) {
     this.yarnClient = yarnClient;
   }
 

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

@@ -19,10 +19,12 @@
 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;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.tez.client.registry.zookeeper.ZkConfig;
 
 
 /**
@@ -37,14 +39,20 @@
 @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 HOST_NAME_RECORD_KEY = "hostName";
+  private static final String HOST_IP_RECORD_KEY = "hostIp";
   private static final String PORT_RECORD_KEY = "port";
-  private static final String OPAQUE_ID_KEY = "id";
+  private static final String EXTERNAL_ID_KEY = "externalId";
+  private static final String COMPUTE_GROUP_NAME_KEY = "computeName";
 
   private final ApplicationId appId;
-  private final String host;
+  private final String hostName;
+  private final String hostIp;
   private final int port;
-  private final String id;
+  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.
@@ -54,17 +62,23 @@ public class AMRecord {
    * 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
+   * @param appId       the {@link ApplicationId} of the Tez application
+   * @param hostName    the hostname where the Application Master is running
+   * @param hostIp      the IP address of the Application Master host
+   * @param port        the RPC port number on which the Application Master is listening
+   * @param externalId  an optional external identifier for the record; if {@code null}, defaults to an empty string
+   * @param computeName the compute group or cluster name; if {@code null},
+   *                    defaults to {@link ZkConfig#DEFAULT_COMPUTE_GROUP_NAME}
    */
-  public AMRecord(ApplicationId appId, String host, int port, String id) {
+  public AMRecord(ApplicationId appId, String hostName, String hostIp, int port, String externalId,
+                  String computeName) {
     this.appId = appId;
-    this.host = host;
+    this.hostName = hostName;
+    this.hostIp = hostIp;
     this.port = port;
-    //If id is not provided, convert to empty string
-    this.id = (id == null) ? "" : id;
+    //externalId is optional, if not provided, convert to empty string
+    this.externalId = Optional.ofNullable(externalId).orElse("");
+    this.computeName = Optional.ofNullable(computeName).orElse(ZkConfig.DEFAULT_COMPUTE_GROUP_NAME);
   }
 
   /**
@@ -78,10 +92,15 @@ public AMRecord(ApplicationId appId, String host, int port, String id) {
    * @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();
+    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;
   }
 
   /**
@@ -97,25 +116,35 @@ public AMRecord(AMRecord other) {
    */
   public AMRecord(ServiceRecord serviceRecord) {
     this.appId = ApplicationId.fromString(serviceRecord.get(APP_ID_RECORD_KEY));
-    this.host = serviceRecord.get(HOST_RECORD_KEY);
+    this.hostName = serviceRecord.get(HOST_NAME_RECORD_KEY);
+    this.hostIp = serviceRecord.get(HOST_IP_RECORD_KEY);
     this.port = Integer.parseInt(serviceRecord.get(PORT_RECORD_KEY));
-    this.id = serviceRecord.get(OPAQUE_ID_KEY);
+    this.externalId = serviceRecord.get(EXTERNAL_ID_KEY);
+    this.computeName = serviceRecord.get(COMPUTE_GROUP_NAME_KEY);
   }
 
   public ApplicationId getApplicationId() {
     return appId;
   }
 
-  public String getHost() {
-    return host;
+  public String getHostName() {
+    return hostName;
+  }
+
+  public String getHostIp() {
+    return hostIp;
   }
 
   public int getPort() {
     return port;
   }
 
-  public String getId() {
-    return id;
+  public String getExternalId() {
+    return externalId;
+  }
+
+  public String getComputeName() {
+    return computeName;
   }
 
   @Override
@@ -125,9 +154,11 @@ public boolean equals(Object other) {
     }
     if (other instanceof AMRecord otherRecord) {
       return appId.equals(otherRecord.appId)
-          && host.equals(otherRecord.host)
+          && hostName.equals(otherRecord.hostName)
+          && hostIp.equals(otherRecord.hostIp)
           && port == otherRecord.port
-          && id.equals(otherRecord.id);
+          && externalId.equals(otherRecord.externalId)
+          && computeName.equals(otherRecord.computeName);
     } else {
       return false;
     }
@@ -148,16 +179,27 @@ 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_RECORD_KEY, host);
+    serviceRecord.set(HOST_NAME_RECORD_KEY, hostName);
+    serviceRecord.set(HOST_IP_RECORD_KEY, hostIp);
     serviceRecord.set(PORT_RECORD_KEY, port);
-    serviceRecord.set(OPAQUE_ID_KEY, id);
+    serviceRecord.set(EXTERNAL_ID_KEY, externalId);
+    serviceRecord.set(COMPUTE_GROUP_NAME_KEY, computeName);
+
     return serviceRecord;
   }
 
+  @Override
+  public String toString() {
+    return toServiceRecord().attributes().toString();
+  }
+
   @Override
   public int hashCode() {
-    return Objects.hash(appId, host, port, id);
+    return Objects.hash(appId, hostName, hostIp, externalId, computeName, port);
   }
 }

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

@@ -0,0 +1,52 @@
+/**
+ * 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 org.apache.hadoop.yarn.api.records.ApplicationId;
+
+
+/**
+ * Base class for {@code AMRegistry} implementations.
+ *
+ * <p>The specific implementation is configured via the
+ * {@code tez.am.registry.class} property.</p>
+ *
+ * <p>Implementations are expected to provide appropriate service lifecycle
+ * behavior, including:
+ * <ul>
+ *   <li>{@code init}</li>
+ *   <li>{@code serviceStart}</li>
+ *   <li>{@code serviceStop}</li>
+ * </ul>
+ * </p>
+ */
+public interface AMRegistry extends AutoCloseable {
+
+  void add(AMRecord server) throws Exception;
+
+  void remove(AMRecord server) throws Exception;
+
+  ApplicationId generateNewId() throws Exception;
+
+  AMRecord createAmRecord(ApplicationId appId, String hostName, String hostIp, int port,
+                                          String computeName);
+
+  void close();
+}