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..4c42bdf34f 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,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, 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..f208121183 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,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; } 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..788c1fbd6e 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,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(); 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..a3e031c0e6 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,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 + TezCounters counters = runDAGAndVerify(dag, true, Collections.singletonList(TABLE_SCAN)); assertEquals(3, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue()); List historyEvents1 = readRecoveryLog(1); @@ -263,7 +270,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 historyEvents1 = readRecoveryLog(1); @@ -309,8 +317,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 +346,18 @@ 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 +369,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 +377,33 @@ TezCounters runDAGAndVerify(DAG dag, boolean killAM) throws Exception { 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 readRecoveryLog(int attemptNum) throws IOException { ApplicationId appId = tezSession.getAppMasterApplicationId(); Path tezSystemStagingDir = TezCommonUtils.getTezSystemStagingPath(tezConf, appId.toString());