TEZ-4564: TezClient to expose Tez AM host:port (#356). (Laszlo Bodor, reviewed by Ayush Saxena)

Author: abstractdog

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,

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();
+  }
 }

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,11 +103,23 @@ 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;
   }
 
   @Override
   public boolean isRunning() throws IOException {
     return isRunning;
   }
+
+  @Override
+  public String getAmHost() {
+    return amHost;
+  }
+
+  @Override
+  public int getAmPort() {
+    return amPort;
+  }
 }

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());
+  }
 }

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;
+  }
 }