TEZ-4725: Fix flaky tests in TestAMRecoveryAggregationBroadcast

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

@@ -961,10 +961,37 @@ static DAGClientAMProtocolBlockingPB getAMProxy(FrameworkClient frameworkClient,
       throw new TezException(e);
     }
 
+    // Returning null lets callers (waitForProxy, sendAMHeartbeat) back off
+    // and retry rather than crash.
+    if (isAMRpcEndpointUnavailable(appReport)) {
+      LOG.debug("AM RPC endpoint not yet available for {} (host={}, port={})"
+          + " - will retry", applicationId, appReport.getHost(), appReport.getRpcPort());
+      return null;
+    }
+
     return getAMProxy(conf, appReport.getHost(), appReport.getRpcPort(),
         appReport.getClientToAMToken(), ugi);
   }
 
+  /**
+   * Returns true when appReport does NOT yet expose an AM RPC
+   * endpoint that is safe to hand to org.apache.hadoop.net.NetUtils#createSocketAddrForHost.
+   * 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 the RPC layer (which would
+   * throw IllegalArgumentException: port out of range):
+   *   host == null / "N/A" : RM has not received registerApplicationMaster()
+   *   rpcPort == 0         : protobuf wire default
+   *   rpcPort == -1        : container up but RPC listener not yet bound
+   */
+  @Private
+  public static boolean isAMRpcEndpointUnavailable(ApplicationReport appReport) {
+    String amHost = appReport.getHost();
+    int amRpcPort = appReport.getRpcPort();
+    return amHost == null || amHost.equals("N/A") || amRpcPort <= 0;
+  }
+
   @Private
   public static DAGClientAMProtocolBlockingPB getAMProxy(final Configuration conf, String amHost,
       int amRpcPort, org.apache.hadoop.yarn.api.records.Token clientToAMToken,

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

@@ -280,10 +280,8 @@ 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
+    if (TezClientUtils.isAMRpcEndpointUnavailable(appReport)) {
+      // AM RPC endpoint not yet available - try again on the next poll.
       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,46 @@ 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");
+
+    // Case 4: host == null  (RM has not yet received AM registration)
+    assertNull(TezClientUtils.getAMProxy(
+        newRunningFrameworkClient(appId, null, 8080), conf, appId, ugi),
+        "host == null 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,7 +49,9 @@
 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.hadoop.yarn.util.Records;
 import org.apache.tez.client.FrameworkClient;
 import org.apache.tez.common.CachedEntity;
 import org.apache.tez.dag.api.NoCurrentDAGException;
@@ -725,6 +727,59 @@ 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).
+    try (DAGClientRPCImplWithFakeReport client = newFakeReportClient(tezConf, "somehost", 0)) {
+      assertFalse(client.createAMProxyIfNeeded(), "rpcPort == 0 should return false");
+    }
+
+    // Case: rpcPort == -1 (YARN-808 gap — AM allocated but RPC not bound).
+    try (DAGClientRPCImplWithFakeReport client = newFakeReportClient(tezConf, "somehost", -1)) {
+      assertFalse(client.createAMProxyIfNeeded(), "rpcPort == -1 should return false");
+    }
+
+    // Case: host == null.
+    try (DAGClientRPCImplWithFakeReport client = newFakeReportClient(tezConf, null, 8080)) {
+      assertFalse(client.createAMProxyIfNeeded(), "host == null should return false");
+    }
+
+    // Case: host == "N/A".
+    try (DAGClientRPCImplWithFakeReport client = newFakeReportClient(tezConf, "N/A", 8080)) {
+      assertFalse(client.createAMProxyIfNeeded(), "host == N/A should return false");
+    }
+  }
+
+  private DAGClientRPCImplWithFakeReport newFakeReportClient(TezConfiguration conf,
+      String host, int rpcPort) throws IOException {
+    ApplicationReport report = Records.newRecord(ApplicationReport.class);
+    report.setYarnApplicationState(YarnApplicationState.RUNNING);
+    report.setHost(host);
+    report.setRpcPort(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,7 @@ public class TestAMRecoveryAggregationBroadcast {
 
   private TezConfiguration tezConf;
   private TezClient tezSession;
+  private Path outPath;
 
   @BeforeAll
   public static void setupAll() {
@@ -173,6 +175,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 +207,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 +224,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 +245,8 @@ 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 AM
+    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 +267,8 @@ 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);
@@ -309,8 +314,7 @@ private DAG createDAG(String dagName) throws Exception {
             .create(AggregationProcessor.class.getName()).setUserPayload(payload), 1);
 
     DataSinkDescriptor dataSink = MROutput
-        .createConfigBuilder(new Configuration(tezConf), TextOutputFormat.class,
-            OUT_PATH.toString())
+        .createConfigBuilder(new Configuration(tezConf), TextOutputFormat.class, outPath.toString())
         .build();
     // Broadcast Hash Join
     Vertex mapJoinVertex = Vertex
@@ -339,12 +343,15 @@ 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.
+      for (String vertexName : vertexNamesToWaitFor) {
+        waitForVertexSucceeded(dagClient, vertexName, TimeUnit.SECONDS.toMillis(60));
+      }
       YarnClient yarnClient = YarnClient.createYarnClient();
       yarnClient.init(tezConf);
       yarnClient.start();
@@ -356,14 +363,41 @@ 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 timeoutNanos = TimeUnit.MILLISECONDS.toNanos(timeoutMs);
+    long startNanos = System.nanoTime();
+    while ((System.nanoTime() - startNanos) < timeoutNanos) {
+      // 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();
+        switch (state) {
+          case SUCCEEDED -> {
+            return;
+          }
+          case FAILED, KILLED, ERROR ->
+              throw new AssertionError(
+                  "Vertex " + vertexName + " reached terminal non-success state: " + state);
+          default -> {
+            // Still running / initializing; fall through to sleep and poll again.
+          }
+        }
+      }
+      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());