From 37f881e1b4dc31e4b9a396c24b94a78b173f1fbf Mon Sep 17 00:00:00 2001 From: Mahesh Raju Somalaraju Date: Tue, 14 Jul 2026 12:43:24 +0530 Subject: [PATCH] TEZ-4725: Fix flaky tests in TestAMRecoveryAggregationBroadcast --- .../org/apache/tez/client/TezClientUtils.java | 21 ++++++- .../dag/api/client/rpc/DAGClientRPCImpl.java | 12 ++-- .../apache/tez/client/TestTezClientUtils.java | 42 +++++++++++++ .../tez/dag/api/client/rpc/TestDAGClient.java | 50 +++++++++++++++ .../TestAMRecoveryAggregationBroadcast.java | 62 ++++++++++++++++--- 5 files changed, 173 insertions(+), 14 deletions(-) diff --git a/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java b/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java index 9e6338ddd8..c277511af9 100644 --- a/tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java +++ b/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); } diff --git a/tez-api/src/main/java/org/apache/tez/dag/api/client/rpc/DAGClientRPCImpl.java b/tez-api/src/main/java/org/apache/tez/dag/api/client/rpc/DAGClientRPCImpl.java index c0d8836634..821429c8ca 100644 --- a/tez-api/src/main/java/org/apache/tez/dag/api/client/rpc/DAGClientRPCImpl.java +++ b/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; } diff --git a/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java b/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java index d635259a6a..7c3d9c6cbe 100644 --- a/tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java +++ b/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; + } } diff --git a/tez-api/src/test/java/org/apache/tez/dag/api/client/rpc/TestDAGClient.java b/tez-api/src/test/java/org/apache/tez/dag/api/client/rpc/TestDAGClient.java index 69c679947d..71fca6bb4a 100644 --- a/tez-api/src/test/java/org/apache/tez/dag/api/client/rpc/TestDAGClient.java +++ b/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(); diff --git a/tez-tests/src/test/java/org/apache/tez/test/TestAMRecoveryAggregationBroadcast.java b/tez-tests/src/test/java/org/apache/tez/test/TestAMRecoveryAggregationBroadcast.java index dc21f7c133..5e33a09698 100644 --- a/tez-tests/src/test/java/org/apache/tez/test/TestAMRecoveryAggregationBroadcast.java +++ b/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 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 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 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 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 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,7 +374,7 @@ 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(); @@ -364,6 +382,32 @@ TezCounters runDAGAndVerify(DAG dag, boolean killAM) throws Exception { 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 readRecoveryLog(int attemptNum) throws IOException { ApplicationId appId = tezSession.getAppMasterApplicationId(); Path tezSystemStagingDir = TezCommonUtils.getTezSystemStagingPath(tezConf, appId.toString());