From 38c5aaccdf8e4f7db210975021c78b6db556c87f Mon Sep 17 00:00:00 2001
From: Bodor Laszlo <bodorlaszlo0202@gmail.com>
Date: Tue, 28 May 2024 07:45:54 +0200
Subject: [PATCH] TEZ-4564: TezClient to expose Tez AM host:port (#356).
 (Laszlo Bodor, reviewed by Ayush Saxena)

---
 .../apache/tez/client/FrameworkClient.java    |  3 +++
 .../java/org/apache/tez/client/TezClient.java |  8 ++++++++
 .../org/apache/tez/client/TezYarnClient.java  | 15 +++++++++++++++
 .../org/apache/tez/client/TestTezClient.java  | 19 +++++++++++++++++++
 .../org/apache/tez/client/LocalClient.java    | 15 +++++++++++++++
 5 files changed, 60 insertions(+)

diff --git a/tez-api/src/main/java/org/apache/tez/client/FrameworkClient.java b/tez-api/src/main/java/org/apache/tez/client/FrameworkClient.java
index 26c11fd084..a1a0b8a5e8 100644
--- a/tez-api/src/main/java/org/apache/tez/client/FrameworkClient.java
+++ b/tez-api/src/main/java/org/apache/tez/client/FrameworkClient.java
@@ -94,6 +94,9 @@ public abstract ApplicationId submitApplication(ApplicationSubmissionContext app
 
   public abstract ApplicationReport getApplicationReport(ApplicationId appId) throws YarnException, IOException;
 
+  public abstract String getAmHost();
+  public abstract int getAmPort();
+
   public abstract boolean isRunning() throws IOException;
 
   public TezAppMasterStatus getAMStatus(Configuration conf, ApplicationId appId,
diff --git a/tez-api/src/main/java/org/apache/tez/client/TezClient.java b/tez-api/src/main/java/org/apache/tez/client/TezClient.java
index c632700854..937346cacd 100644
--- a/tez-api/src/main/java/org/apache/tez/client/TezClient.java
+++ b/tez-api/src/main/java/org/apache/tez/client/TezClient.java
@@ -1286,4 +1286,12 @@ public static ApplicationId appIdfromString(String appIdStr) {
               + appIdStr, n);
     }
   }
+
+  public String getAmHost() {
+    return frameworkClient == null ? null : frameworkClient.getAmHost();
+  }
+
+  public int getAmPort() {
+    return frameworkClient == null ? -1 : frameworkClient.getAmPort();
+  }
 }
diff --git a/tez-api/src/main/java/org/apache/tez/client/TezYarnClient.java b/tez-api/src/main/java/org/apache/tez/client/TezYarnClient.java
index f89664d863..d109648075 100644
--- a/tez-api/src/main/java/org/apache/tez/client/TezYarnClient.java
+++ b/tez-api/src/main/java/org/apache/tez/client/TezYarnClient.java
@@ -40,6 +40,9 @@ public class TezYarnClient extends FrameworkClient {
 
   private volatile boolean isRunning;
 
+  private String amHost;
+  private int amPort;
+
   protected TezYarnClient(YarnClient yarnClient) {
     this.yarnClient = yarnClient;
   }
@@ -100,6 +103,8 @@ public ApplicationReport getApplicationReport(ApplicationId appId) throws YarnEx
       throw new ApplicationNotFoundException("YARN reports no state for application "
           + appId);
     }
+    this.amHost = report.getHost();
+    this.amPort = report.getRpcPort();
     return report;
   }
 
@@ -107,4 +112,14 @@ public ApplicationReport getApplicationReport(ApplicationId appId) throws YarnEx
   public boolean isRunning() throws IOException {
     return isRunning;
   }
+
+  @Override
+  public String getAmHost() {
+    return amHost;
+  }
+
+  @Override
+  public int getAmPort() {
+    return amPort;
+  }
 }
diff --git a/tez-api/src/test/java/org/apache/tez/client/TestTezClient.java b/tez-api/src/test/java/org/apache/tez/client/TestTezClient.java
index e4d03cc60a..1b93f924a2 100644
--- a/tez-api/src/test/java/org/apache/tez/client/TestTezClient.java
+++ b/tez-api/src/test/java/org/apache/tez/client/TestTezClient.java
@@ -161,6 +161,16 @@ protected DAGClientAMProtocolBlockingPB getProxy(Configuration conf, Application
       }
       return super.getProxy(conf, sessionAppId, ugi);
     }
+
+    @Override
+    public String getAmHost() {
+      return "testhost";
+    }
+
+    @Override
+    public int getAmPort() {
+      return 1234;
+    }
   }
 
   TezClientForTest configureAndCreateTezClient() throws YarnException, IOException, ServiceException {
@@ -1005,4 +1015,13 @@ public void testYarnZkDeprecatedConf() {
     //Test that Exception is not thrown by createFinalConfProtoForApp
     TezClientUtils.createFinalConfProtoForApp(conf, null);
   }
+
+  @Test
+  public void testGetAmHostAndPort() throws Exception {
+    final TezClientForTest client = configureAndCreateTezClient(new TezConfiguration());
+
+    // TezClient exposes AM host and port from the FrameworkClient (now it's a TezYarnClientForTest)
+    assertEquals("testhost", client.getAmHost());
+    assertEquals(1234, client.getAmPort());
+  }
 }
diff --git a/tez-dag/src/main/java/org/apache/tez/client/LocalClient.java b/tez-dag/src/main/java/org/apache/tez/client/LocalClient.java
index bf1cef4c45..73f60bdcc2 100644
--- a/tez-dag/src/main/java/org/apache/tez/client/LocalClient.java
+++ b/tez-dag/src/main/java/org/apache/tez/client/LocalClient.java
@@ -88,6 +88,8 @@ public class LocalClient extends FrameworkClient {
   private TezApiVersionInfo versionInfo = new TezApiVersionInfo();
   private volatile Throwable amFailException = null;
   private boolean isLocalWithoutNetwork;
+  private String amHost;
+  private int amPort;
 
   private static final String localModeDAGSchedulerClassName =
       "org.apache.tez.dag.app.dag.impl.DAGSchedulerNaturalOrderControlled";
@@ -204,6 +206,9 @@ public ApplicationReport getApplicationReport(ApplicationId appId) {
     report.setProgress(dagAppMaster.getProgress());
     report.setAMRMToken(null);
 
+    this.amHost = dagAppMaster.getAppNMHost();
+    this.amPort = dagAppMaster.getRpcPort();
+
     return report;
   }
 
@@ -475,4 +480,14 @@ public boolean shutdownSession(Configuration configuration, ApplicationId sessio
     }
     return super.shutdownSession(configuration, sessionAppId, ugi);
   }
+
+  @Override
+  public String getAmHost() {
+    return amHost;
+  }
+
+  @Override
+  public int getAmPort() {
+    return amPort;
+  }
 }