@Before
public void setup() throws IOException {
YarnConfiguration conf = new YarnConfiguration();
YarnAPIStorageFactory.setConfiguration(conf);
RMStorageFactory.setConfiguration(conf);
RMUtilities.InitializeDB();
}
private RMNodeImpl getRunningNode() throws IOException {
NodeId nodeId = BuilderUtils.newNodeId("localhost", 0);
Resource capability = Resource.newInstance(4096, 4);
int rpcID = HopYarnAPIUtilities.getRPCID();
byte[] allNMRequestData = new byte[10];
try {
RMUtilities.persistAppMasterRPC(rpcID, RPC.Type.RegisterNM, allNMRequestData);
} catch (IOException ex) {
Logger.getLogger(TestRMNodeTransitions.class.getName()).log(Level.SEVERE, null, ex);
}
TransactionState ts =
new TransactionStateImpl(TransactionType.RM); // TransactionStateRM.newInstance(rpcID);
RMNodeImpl node =
new RMNodeImpl(
nodeId,
rmContext,
nodeId.getHost(),
0,
0,
null,
ResourceOption.newInstance(capability, RMNode.OVER_COMMIT_TIMEOUT_MILLIS_DEFAULT),
null);
((TransactionStateImpl) ts).getRMContextInfo().toAddActiveRMNode(nodeId, node, 1);
node.handle(new RMNodeEvent(node.getNodeID(), RMNodeEventType.STARTED, ts));
Assert.assertEquals(NodeState.RUNNING, node.getState());
return node;
}
@Before
public void setUp() throws Exception {
InlineDispatcher rmDispatcher = new InlineDispatcher();
YarnConfiguration conf = new YarnConfiguration();
YarnAPIStorageFactory.setConfiguration(conf);
RMStorageFactory.setConfiguration(conf);
conf.setClass(YarnConfiguration.RM_SCHEDULER, FifoScheduler.class, ResourceScheduler.class);
RMUtilities.InitializeDB();
TransactionStateManager tsm = new TransactionStateManager();
tsm.init(conf);
tsm.start();
rmContext =
new RMContextImpl(
rmDispatcher,
null,
null,
null,
mock(DelegationTokenRenewer.class),
null,
null,
null,
conf,
tsm);
NodesListManager nodesListManager = mock(NodesListManager.class);
HostsFileReader reader = mock(HostsFileReader.class);
when(nodesListManager.getHostsReader()).thenReturn(reader);
((RMContextImpl) rmContext).setNodesListManager(nodesListManager);
scheduler = mock(YarnScheduler.class);
doAnswer(
new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
final SchedulerEvent event = (SchedulerEvent) (invocation.getArguments()[0]);
eventType = event.getType();
if (eventType == SchedulerEventType.NODE_UPDATE) {
List<UpdatedContainerInfo> lastestContainersInfoList =
((NodeUpdateSchedulerEvent) event)
.getRMNode()
.pullContainerUpdates(event.getTransactionState());
for (UpdatedContainerInfo lastestContainersInfo : lastestContainersInfoList) {
completedContainers.addAll(lastestContainersInfo.getCompletedContainers());
}
}
return null;
}
})
.when(scheduler)
.handle(any(SchedulerEvent.class));
rmDispatcher.register(SchedulerEventType.class, new TestSchedulerEventDispatcher());
rmDispatcher.register(
NodesListManagerEventType.class, new TestNodeListManagerEventDispatcher());
NodeId nodeId = BuilderUtils.newNodeId("localhost", 0);
node = new RMNodeImpl(nodeId, rmContext, nodeId.getHost(), 0, 0, null, null, null);
nodesListManagerEvent = null;
}
@Test(timeout = 2000000)
public void testUpdateHeartbeatResponseForCleanup() throws IOException {
RMNodeImpl node = getRunningNode();
NodeId nodeId = node.getNodeID();
int rpcID = HopYarnAPIUtilities.getRPCID();
byte[] allNMRequestData = new byte[1];
allNMRequestData[0] = 0xA;
try {
RMUtilities.persistAppMasterRPC(rpcID, RPC.Type.RegisterNM, allNMRequestData);
} catch (IOException ex) {
Logger.getLogger(TestRMNodeTransitions.class.getName()).log(Level.SEVERE, null, ex);
}
TransactionState ts = new TransactionStateImpl(TransactionType.RM);
// Expire a container
ContainerId completedContainerId =
BuilderUtils.newContainerId(
BuilderUtils.newApplicationAttemptId(BuilderUtils.newApplicationId(0, 0), 0), 0);
node.handle(new RMNodeCleanContainerEvent(nodeId, completedContainerId, ts));
ts.decCounter(TransactionState.TransactionType.INIT);
Assert.assertEquals(1, node.getContainersToCleanUp().size());
rpcID = HopYarnAPIUtilities.getRPCID();
allNMRequestData = new byte[1];
allNMRequestData[0] = 0xA;
try {
RMUtilities.persistAppMasterRPC(rpcID, RPC.Type.FinishApplicationMaster, allNMRequestData);
} catch (IOException ex) {
Logger.getLogger(TestRMNodeTransitions.class.getName()).log(Level.SEVERE, null, ex);
}
TransactionState ts2 =
new TransactionStateImpl(TransactionType.RM); // TransactionStateRM.newInstance(rpcID);
// Finish an application
ApplicationId finishedAppId = BuilderUtils.newApplicationId(0, 1);
node.handle(new RMNodeCleanAppEvent(nodeId, finishedAppId, ts2));
ts2.decCounter(TransactionState.TransactionType.INIT);
Assert.assertEquals(1, node.getAppsToCleanup().size());
rpcID = HopYarnAPIUtilities.getRPCID();
allNMRequestData = new byte[1];
allNMRequestData[0] = 0xA;
try {
RMUtilities.persistAppMasterRPC(rpcID, RPC.Type.FinishApplicationMaster, allNMRequestData);
} catch (IOException ex) {
Logger.getLogger(TestRMNodeTransitions.class.getName()).log(Level.SEVERE, null, ex);
}
TransactionState ts3 =
new TransactionStateImpl(TransactionType.RM); // TransactionStateRM.newInstance(rpcID);
// Verify status update does not clear containers/apps to cleanup
// but updating heartbeat response for cleanup does
RMNodeStatusEvent statusEvent = getMockRMNodeStatusEvent();
RMNodeStatusEvent se =
new RMNodeStatusEvent(
statusEvent.getNodeId(),
statusEvent.getNodeHealthStatus(),
statusEvent.getContainers(),
statusEvent.getKeepAliveAppIds(),
statusEvent.getLatestResponse(),
ts3);
node.handle(se);
ts3.decCounter(TransactionState.TransactionType.INIT);
Assert.assertEquals(1, node.getContainersToCleanUp().size());
Assert.assertEquals(1, node.getAppsToCleanup().size());
NodeHeartbeatResponse hbrsp = Records.newRecord(NodeHeartbeatResponse.class);
node.updateNodeHeartbeatResponseForCleanup(hbrsp, ts);
Assert.assertEquals(0, node.getContainersToCleanUp().size());
Assert.assertEquals(0, node.getAppsToCleanup().size());
Assert.assertEquals(1, hbrsp.getContainersToCleanup().size());
Assert.assertEquals(completedContainerId, hbrsp.getContainersToCleanup().get(0));
Assert.assertEquals(1, hbrsp.getApplicationsToCleanup().size());
Assert.assertEquals(finishedAppId, hbrsp.getApplicationsToCleanup().get(0));
}