Skip to content

TEZ-4725: Fix flaky tests in TestAMRecoveryAggregationBroadcast #520

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maheshrajus wants to merge 1 commit into
base: master
Choose a base branch
from
Open

TEZ-4725: Fix flaky tests in TestAMRecoveryAggregationBroadcast #520

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion tez-api/src/main/java/org/apache/tez/client/TezClientUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Comment on lines +964 to +974

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this long comment can be put to a new method that contains the refactored condition below

except the "return null lets callers" part, which belongs to this method

String amHost = appReport.getHost();
int amRpcPort = appReport.getRpcPort();
if (amHost == null || amHost.equals("N/A") || amRpcPort <= 0) {

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this condition can be refactored to a separate method with appReport parameter and be reused across different codepaths

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

Expand Down
12 changes: 8 additions & 4 deletions tez-api/src/main/java/org/apache/tez/dag/api/client/rpc/DAGClientRPCImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Comment on lines +283 to +287

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this comment can be put to the new method that contains the refactored condition below

if (appReport.getHost() == null || appReport.getHost().equals("N/A") ||
appReport.getRpcPort() <= 0) {
// AM RPC endpoint not yet available
return false;
}

Expand Down
42 changes: 42 additions & 0 deletions tez-api/src/test/java/org/apache/tez/client/TestTezClientUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}
50 changes: 50 additions & 0 deletions tez-api/src/test/java/org/apache/tez/dag/api/client/rpc/TestDAGClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Comment on lines +754 to +761

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's strange that mocking and real classes are mixed: if there is a real DAGClientRPCImplWithFakeReport, it could leverage a real ApplicationReport, couldn't it?


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();
Expand Down
62 changes: 53 additions & 9 deletions tez-tests/src/test/java/org/apache/tez/test/TestAMRecoveryAggregationBroadcast.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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";
Expand All @@ -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() {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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));
Comment on lines +252 to +253

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

line break is not needed here

assertEquals(3, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue());

List<HistoryEvent> historyEvents1 = readRecoveryLog(1);
Expand All @@ -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));
Comment on lines +275 to +276

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

line break is not needed here

assertEquals(3, counters.findCounter(DAGCounter.NUM_SUCCEEDED_TASKS).getValue());

List<HistoryEvent> historyEvents1 = readRecoveryLog(1);
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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));
Comment on lines +363 to +364

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

line break is not needed here

}
YarnClient yarnClient = YarnClient.createYarnClient();
yarnClient.init(tezConf);
yarnClient.start();
Expand All @@ -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) {

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

try to use a monothonic clock instead of System.currentTimeMillis()

// 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);
}
}
Comment on lines +395 to +404

@abstractdog abstractdog Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what about a switch case here, something like

switch (lastState) {
  case SUCCEEDED -> { return; }
  case FAILED, KILLED, ERROR -> throw new AssertionError("Vertex " + vertexName
      + " reached terminal non-success state: " + lastState);
...
}

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());
Expand Down