TEZ-4725: Fix flaky tests in TestAMRecoveryAggregationBroadcast

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

@@ -961,7 +961,26 @@ static DAGClientAMProtocolBlockingPB getAMProxy(FrameworkClient frameworkClient,
       throw new TezException(e);
     }
 
-    return getAMProxy(conf, appReport.getHost(), appReport.getRpcPort(),
+    // YARN-808 gap: when the AM container is first allocated YARN briefly
+    // reports state=RUNNING before the AM has registered with the RM or bound
+    // its RPC listener. During that window the ApplicationReport contains
+    // sentinel values that must not be passed to
+    // NetUtils.createSocketAddrForHost() (which would throw
+    // IllegalArgumentException: port out of range) or to RPC.getProxy():
+    //   host == null / "N/A" : RM has not received registerApplicationMaster()
+    //   rpcPort == 0         : protobuf wire default
+    //   rpcPort == -1        : container up but RPC listener not yet bound
+    // Returning null lets callers (waitForProxy, sendAMHeartbeat) back off
+    // and retry rather than crash.
+    String amHost = appReport.getHost();
+    int amRpcPort = appReport.getRpcPort();
+    if (amHost == null || amHost.equals("N/A") || amRpcPort <= 0) {
+      LOG.debug("AM RPC endpoint not yet available for {} (host={}, port={})"
+          + " - will retry", applicationId, amHost, amRpcPort);
+      return null;
+    }
+
+    return getAMProxy(conf, amHost, amRpcPort,
         appReport.getClientToAMToken(), ugi);
   }
 

tez-api/src/main/java/org/apache/tez/dag/api/client/rpc/DAGClientRPCImpl.java

@@ -280,10 +280,14 @@ boolean createAMProxyIfNeeded() throws IOException, TezException,
     }
 
     // YARN-808. Cannot ascertain if AM is ready until we connect to it.
-    // workaround check the default string set by YARN
-    if(appReport.getHost() == null || appReport.getHost().equals("N/A") ||
-        appReport.getRpcPort() == 0){
-      // attempt not running
+    // Workaround: check the default strings/sentinels set by YARN.
+    //   port == 0  : protobuf wire default, AM has not called
+    //                registerApplicationMaster() yet.
+    //   port == -1 : AM container is allocated (state=RUNNING) but the
+    //                RPC listener has not been bound yet.
+    if (appReport.getHost() == null || appReport.getHost().equals("N/A") ||
+        appReport.getRpcPort() <= 0) {
+      // AM RPC endpoint not yet available
       return false;
     }
 

tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java

@@ -26,6 +26,8 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.io.File;
 import java.io.FileNotFoundException;
@@ -54,16 +56,19 @@
 import org.apache.hadoop.io.DataInputByteBuffer;
 import org.apache.hadoop.io.Text;
 import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.security.token.Token;
 import org.apache.hadoop.security.token.TokenIdentifier;
 import org.apache.hadoop.yarn.api.ApplicationConstants;
 import org.apache.hadoop.yarn.api.ApplicationConstants.Environment;
 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.api.records.ContainerLaunchContext;
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
 import org.apache.hadoop.yarn.api.records.Resource;
+import org.apache.hadoop.yarn.api.records.YarnApplicationState;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.util.Records;
@@ -1003,4 +1008,41 @@ public void testSessionCredentialsMergedBeforeAmConfigCredentials() throws Excep
     // session token should be applied while creating ContainerLaunchContext
     assertEquals(sessionToken, amLaunchCredentials.getToken(tokenType));
   }
+
+  /**
+   * Covers the YARN-808 guard in TezClientUtils#getAMProxy
+   */
+  @Test
+  @Timeout(value = 5000, unit = TimeUnit.MILLISECONDS)
+  public void testGetAMProxyReturnsNullWhenRpcEndpointNotAvailable() throws Exception {
+    TezConfiguration conf = new TezConfiguration();
+    ApplicationId appId = ApplicationId.newInstance(1L, 1);
+    UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
+
+    // Case 1: rpcPort == -1  (YARN-808 gap — AM container up, RPC not bound)
+    assertNull(TezClientUtils.getAMProxy(
+        newRunningFrameworkClient(appId, "somehost", -1), conf, appId, ugi),
+        "rpcPort == -1 should return null");
+
+    // Case 2: rpcPort == 0  (protobuf wire default)
+    assertNull(TezClientUtils.getAMProxy(
+        newRunningFrameworkClient(appId, "somehost", 0), conf, appId, ugi),
+        "rpcPort == 0 should return null");
+
+    // Case 3: host == "N/A"  (RM has not yet received AM registration)
+    assertNull(TezClientUtils.getAMProxy(
+        newRunningFrameworkClient(appId, "N/A", 8080), conf, appId, ugi),
+        "host == N/A should return null");
+  }
+
+  private static FrameworkClient newRunningFrameworkClient(ApplicationId appId,
+      String host, int rpcPort) throws IOException, YarnException {
+    ApplicationReport report = mock(ApplicationReport.class);
+    when(report.getYarnApplicationState()).thenReturn(YarnApplicationState.RUNNING);
+    when(report.getHost()).thenReturn(host);
+    when(report.getRpcPort()).thenReturn(rpcPort);
+    FrameworkClient frameworkClient = mock(FrameworkClient.class);
+    when(frameworkClient.getApplicationReport(appId)).thenReturn(report);
+    return frameworkClient;
+  }
 }

tez-api/src/test/java/org/apache/tez/dag/api/client/rpc/TestDAGClient.java

@@ -49,6 +49,7 @@
 import org.apache.hadoop.security.ssl.KeyStoreTestUtil;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ApplicationReport;
+import org.apache.hadoop.yarn.api.records.YarnApplicationState;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.tez.client.FrameworkClient;
 import org.apache.tez.common.CachedEntity;
@@ -725,6 +726,55 @@ public void testGetDagStatusWithCachedStatusExpiration() throws Exception {
     }
   }
 
+  /**
+   * Covers the YARN-808 guard in DAGClientRPCImpl#createAMProxyIfNeeded
+   */
+  @Test
+  @Timeout(value = 5000, unit = TimeUnit.MILLISECONDS)
+  public void testCreateAMProxyIfNeededReturnsFalseOnBadEndpoint() throws Exception {
+    TezConfiguration tezConf = new TezConfiguration();
+
+    // Case: rpcPort == 0 (protobuf default sentinel).
+    assertFalse(mockReportClient(tezConf, "somehost", 0).createAMProxyIfNeeded(),
+        "rpcPort == 0 should return false");
+
+    // Case: rpcPort == -1 (YARN-808 gap — AM allocated but RPC not bound).
+    assertFalse(mockReportClient(tezConf, "somehost", -1).createAMProxyIfNeeded(),
+        "rpcPort == -1 should return false");
+
+    // Case: host == null.
+    assertFalse(mockReportClient(tezConf, null, 8080).createAMProxyIfNeeded(),
+        "host == null should return false");
+
+    // Case: host == "N/A".
+    assertFalse(mockReportClient(tezConf, "N/A", 8080).createAMProxyIfNeeded(),
+        "host == N/A should return false");
+  }
+
+  private DAGClientRPCImplWithFakeReport mockReportClient(TezConfiguration conf,
+      String host, int rpcPort) throws IOException {
+    ApplicationReport report = mock(ApplicationReport.class);
+    when(report.getYarnApplicationState()).thenReturn(YarnApplicationState.RUNNING);
+    when(report.getHost()).thenReturn(host);
+    when(report.getRpcPort()).thenReturn(rpcPort);
+    return new DAGClientRPCImplWithFakeReport(mockAppId, dagIdStr, conf, report);
+  }
+
+  private static final class DAGClientRPCImplWithFakeReport extends DAGClientRPCImpl {
+    private final ApplicationReport fakeReport;
+
+    DAGClientRPCImplWithFakeReport(ApplicationId appId, String dagId,
+        TezConfiguration conf, ApplicationReport fakeReport) throws IOException {
+      super(appId, dagId, conf, null, UserGroupInformation.getCurrentUser());
+      this.fakeReport = fakeReport;
+    }
+
+    @Override
+    ApplicationReport getAppReport() {
+      return fakeReport;
+    }
+  }
+
   @Test
   public void testDagClientReturnsFailedDAGOnNoCurrentDAGException() throws Exception {
     TezConfiguration tezConf = new TezConfiguration();

tez-tests/src/test/java/org/apache/tez/test/TestAMRecoveryAggregationBroadcast.java

@@ -26,6 +26,7 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -68,6 +69,7 @@
 import org.apache.tez.dag.api.client.DAGStatus;
 import org.apache.tez.dag.api.client.DAGStatus.State;
 import org.apache.tez.dag.api.client.StatusGetOpts;
+import org.apache.tez.dag.api.client.VertexStatus;
 import org.apache.tez.dag.api.oldrecords.TaskAttemptState;
 import org.apache.tez.dag.app.RecoveryParser;
 import org.apache.tez.dag.history.HistoryEvent;
@@ -105,7 +107,6 @@ public class TestAMRecoveryAggregationBroadcast {
   private static final String TEST_ROOT_DIR = "target" + Path.SEPARATOR
       + TestAMRecoveryAggregationBroadcast.class.getName() + "-tmpDir";
   private static final Path INPUT_FILE = new Path(TEST_ROOT_DIR, "input.csv");
-  private static final Path OUT_PATH = new Path(TEST_ROOT_DIR, "out-groups");
   private static final String EXPECTED_OUTPUT = "1-5\n1-5\n1-5\n1-5\n1-5\n"
       + "2-4\n2-4\n2-4\n2-4\n" + "3-3\n3-3\n3-3\n" + "4-2\n4-2\n" + "5-1\n";
   private static final String TABLE_SCAN_SLEEP = "tez.test.table.scan.sleep";
@@ -119,6 +120,10 @@ public class TestAMRecoveryAggregationBroadcast {
 
   private TezConfiguration tezConf;
   private TezClient tezSession;
+  // Per-test unique output path.
+  // replaces the former static OUT_PATH so that a stale file from a prior run in the same forked JVM
+  // cannot bleed into a later run.
+  private Path outPath;
 
   @BeforeAll
   public static void setupAll() {
@@ -173,6 +178,7 @@ public void setup() throws Exception {
     Path remoteStagingDir = remoteFs.makeQualified(new Path(TEST_ROOT_DIR, String
         .valueOf(new Random().nextInt(100000))));
     TezClientUtils.ensureStagingDirExists(dfsConf, remoteStagingDir);
+    outPath = new Path(TEST_ROOT_DIR, "out-groups-" + new Random().nextInt(100000));
 
     tezConf = new TezConfiguration(tezCluster.getConfig());
     tezConf.setInt(TezConfiguration.DAG_RECOVERY_MAX_UNFLUSHED_EVENTS, 0);
@@ -204,7 +210,7 @@ public void teardown() throws InterruptedException {
   @Timeout(value = 120000, unit = TimeUnit.MILLISECONDS)
   public void testSucceed() throws Exception {
     DAG dag = createDAG("Succeed");
-    TezCounters counters = runDAGAndVerify(dag, false);
+    TezCounters counters = runDAGAndVerify(dag, false, Collections.emptyList());
     assertEquals(3, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue());
 
     List<HistoryEvent> historyEvents1 = readRecoveryLog(1);
@@ -221,7 +227,7 @@ public void testSucceed() throws Exception {
   public void testTableScanTemporalFailure() throws Exception {
     tezConf.setBoolean(TABLE_SCAN_SLEEP, true);
     DAG dag = createDAG("TableScanTemporalFailure");
-    TezCounters counters = runDAGAndVerify(dag, true);
+    TezCounters counters = runDAGAndVerify(dag, true, Collections.emptyList());
     assertEquals(3, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue());
 
     List<HistoryEvent> historyEvents1 = readRecoveryLog(1);
@@ -242,7 +248,9 @@ public void testTableScanTemporalFailure() throws Exception {
   public void testAggregationTemporalFailure() throws Exception {
     tezConf.setBoolean(AGGREGATION_SLEEP, true);
     DAG dag = createDAG("AggregationTemporalFailure");
-    TezCounters counters = runDAGAndVerify(dag, true);
+    // Wait for TableScan to be SUCCEEDED before killing the
+    TezCounters counters = runDAGAndVerify(dag, true,
+        Collections.singletonList(TABLE_SCAN));
     assertEquals(3, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue());
 
     List<HistoryEvent> historyEvents1 = readRecoveryLog(1);
@@ -263,7 +271,9 @@ public void testAggregationTemporalFailure() throws Exception {
   public void testMapJoinTemporalFailure() throws Exception {
     tezConf.setBoolean(MAP_JOIN_SLEEP, true);
     DAG dag = createDAG("MapJoinTemporalFailure");
-    TezCounters counters = runDAGAndVerify(dag, true);
+    // Wait for both upstream vertices to be SUCCEEDED before killing the AM
+    TezCounters counters = runDAGAndVerify(dag, true,
+        Arrays.asList(TABLE_SCAN, AGGREGATION));
     assertEquals(3, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue());
 
     List<HistoryEvent> historyEvents1 = readRecoveryLog(1);
@@ -310,7 +320,7 @@ private DAG createDAG(String dagName) throws Exception {
 
     DataSinkDescriptor dataSink = MROutput
         .createConfigBuilder(new Configuration(tezConf), TextOutputFormat.class,
-            OUT_PATH.toString())
+            outPath.toString())
         .build();
     // Broadcast Hash Join
     Vertex mapJoinVertex = Vertex
@@ -339,12 +349,20 @@ private DAG createDAG(String dagName) throws Exception {
     return dag;
   }
 
-  TezCounters runDAGAndVerify(DAG dag, boolean killAM) throws Exception {
+  TezCounters runDAGAndVerify(DAG dag, boolean killAM,
+      List<String> vertexNamesToWaitFor) throws Exception {
     tezSession.waitTillReady();
     DAGClient dagClient = tezSession.submitDAG(dag);
 
     if (killAM) {
-      TimeUnit.SECONDS.sleep(10);
+      // Deterministic wait: block until every named upstream vertex reaches
+      // SUCCEEDED. Replaces a fixed Thread.sleep(10s) which was too short on
+      // slow CI machines and caused the recovery-log assertions below to
+      // fail intermittently.
+      for (String vertexName : vertexNamesToWaitFor) {
+        waitForVertexSucceeded(dagClient, vertexName,
+            TimeUnit.SECONDS.toMillis(60));
+      }
       YarnClient yarnClient = YarnClient.createYarnClient();
       yarnClient.init(tezConf);
       yarnClient.start();
@@ -356,14 +374,40 @@ TezCounters runDAGAndVerify(DAG dag, boolean killAM) throws Exception {
     LOG.info("Diagnosis: " + dagStatus.getDiagnostics());
     assertEquals(State.SUCCEEDED, dagStatus.getState());
 
-    FSDataInputStream in = remoteFs.open(new Path(OUT_PATH, "part-v002-o000-r-00000"));
+    FSDataInputStream in = remoteFs.open(new Path(outPath, "part-v002-o000-r-00000"));
     ByteBuffer buf = ByteBuffer.allocate(100);
     in.read(buf);
     buf.flip();
     assertEquals(EXPECTED_OUTPUT, StandardCharsets.UTF_8.decode(buf).toString());
     return dagStatus.getDAGCounters();
   }
 
+  private void waitForVertexSucceeded(DAGClient dagClient, String vertexName,
+      long timeoutMs) throws Exception {
+    long deadline = System.currentTimeMillis() + timeoutMs;
+    while (System.currentTimeMillis() < deadline) {
+      // Before the vertex is initialized on the AM, getVertexStatus may
+      // return null - treat that the same as NEW / INITIALIZING and keep
+      // polling.
+      VertexStatus status = dagClient.getVertexStatus(vertexName, null);
+      if (status != null) {
+        VertexStatus.State state = status.getState();
+        if (state == VertexStatus.State.SUCCEEDED) {
+          return;
+        }
+        if (state == VertexStatus.State.FAILED
+            || state == VertexStatus.State.KILLED
+            || state == VertexStatus.State.ERROR) {
+          throw new AssertionError("Vertex " + vertexName
+              + " reached terminal non-success state: " + state);
+        }
+      }
+      TimeUnit.MILLISECONDS.sleep(500);
+    }
+    throw new AssertionError("Timeout waiting for vertex " + vertexName
+        + " to reach SUCCEEDED");
+  }
+
   private List<HistoryEvent> readRecoveryLog(int attemptNum) throws IOException {
     ApplicationId appId = tezSession.getAppMasterApplicationId();
     Path tezSystemStagingDir = TezCommonUtils.getTezSystemStagingPath(tezConf, appId.toString());