synchronized void shutdownTezAM() {
sessionStopped.set(true);
if (currentDAG != null && !currentDAG.isComplete()) {
// send a DAG_KILL message
LOG.info("Sending a kill event to the current DAG" + ", dagId=" + currentDAG.getID());
sendEvent(new DAGEvent(currentDAG.getID(), DAGEventType.DAG_KILL));
} else {
LOG.info("No current running DAG, shutting down the AM");
if (isSession && !state.equals(DAGAppMasterState.ERROR)) {
state = DAGAppMasterState.SUCCEEDED;
}
shutdownHandler.shutdown();
}
}
synchronized String submitDAGToAppMaster(DAGPlan dagPlan) throws TezException {
if (currentDAG != null && !state.equals(DAGAppMasterState.IDLE)) {
throw new TezException("App master already running a DAG");
}
if (state.equals(DAGAppMasterState.ERROR) || sessionStopped.get()) {
throw new TezException(
"AM unable to accept new DAG submissions." + " In the process of shutting down");
}
// RPC server runs in the context of the job user as it was started in
// the job user's UGI context
LOG.info("Starting DAG submitted via RPC");
if (LOG.isDebugEnabled()) {
LOG.debug("Writing DAG plan to: " + TezConfiguration.TEZ_PB_PLAN_TEXT_NAME);
File outFile = new File(TezConfiguration.TEZ_PB_PLAN_TEXT_NAME);
try {
PrintWriter printWriter = new PrintWriter(outFile);
String dagPbString = dagPlan.toString();
printWriter.println(dagPbString);
printWriter.close();
} catch (IOException e) {
throw new TezException("Failed to write TEZ_PLAN to " + outFile.toString(), e);
}
}
submittedDAGs.incrementAndGet();
startDAG(dagPlan);
return currentDAG.getID().toString();
}
private void startDAG(DAG dag) {
currentDAG = dag;
this.state = DAGAppMasterState.RUNNING;
// End of creating the job.
((RunningAppContext) context).setDAG(currentDAG);
// create a job event for job initialization
DAGEvent initDagEvent = new DAGEvent(currentDAG.getID(), DAGEventType.DAG_INIT);
// Send init to the job (this does NOT trigger job execution)
// This is a synchronous call, not an event through dispatcher. We want
// job-init to be done completely here.
dagEventDispatcher.handle(initDagEvent);
// All components have started, start the job.
/** create a job-start event to get this ball rolling */
DAGEvent startDagEvent = new DAGEvent(currentDAG.getID(), DAGEventType.DAG_START);
/** send the job-start event. this triggers the job execution. */
sendEvent(startDagEvent);
}
@Override
public TezDAGID getCurrentDAGID() {
try {
rLock.lock();
if (dag != null) {
return dag.getID();
}
return null;
} finally {
rLock.unlock();
}
}
@Test(timeout = 5000)
public void testDAGClientHandler() throws TezException {
TezDAGID mockTezDAGId = mock(TezDAGID.class);
when(mockTezDAGId.getId()).thenReturn(1);
when(mockTezDAGId.toString()).thenReturn("dag_9999_0001_1");
DAG mockDAG = mock(DAG.class);
when(mockDAG.getID()).thenReturn(mockTezDAGId);
DAGStatusBuilder mockDagStatusBuilder = mock(DAGStatusBuilder.class);
when(mockDAG.getDAGStatus(anySetOf(StatusGetOpts.class))).thenReturn(mockDagStatusBuilder);
VertexStatusBuilder mockVertexStatusBuilder = mock(VertexStatusBuilder.class);
when(mockDAG.getVertexStatus(anyString(), anySetOf(StatusGetOpts.class)))
.thenReturn(mockVertexStatusBuilder);
DAGAppMaster mockDagAM = mock(DAGAppMaster.class);
when(mockDagAM.getState()).thenReturn(DAGAppMasterState.RUNNING);
AppContext mockAppContext = mock(AppContext.class);
when(mockDagAM.getContext()).thenReturn(mockAppContext);
when(mockDagAM.getContext().getCurrentDAG()).thenReturn(mockDAG);
DAGClientHandler dagClientHandler = new DAGClientHandler(mockDagAM);
// getAllDAGs()
assertEquals(1, dagClientHandler.getAllDAGs().size());
assertEquals("dag_9999_0001_1", dagClientHandler.getAllDAGs().get(0));
// getDAGStatus
try {
dagClientHandler.getDAGStatus("dag_9999_0001_2", Sets.newSet(StatusGetOpts.GET_COUNTERS));
fail("should not come here");
} catch (TezException e) {
assertTrue(e.getMessage().contains("Unknown dagId"));
}
DAGStatus dagStatus =
dagClientHandler.getDAGStatus("dag_9999_0001_1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
assertEquals(mockDagStatusBuilder, dagStatus);
// getVertexStatus
try {
dagClientHandler.getVertexStatus(
"dag_9999_0001_2", "v1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
fail("should not come here");
} catch (TezException e) {
assertTrue(e.getMessage().contains("Unknown dagId"));
}
VertexStatus vertexStatus =
dagClientHandler.getVertexStatus(
"dag_9999_0001_1", "v1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
assertEquals(mockVertexStatusBuilder, vertexStatus);
// getTezAppMasterStatus
when(mockDagAM.isSession()).thenReturn(false);
assertEquals(TezAppMasterStatus.RUNNING, dagClientHandler.getTezAppMasterStatus());
when(mockDagAM.isSession()).thenReturn(true);
when(mockDagAM.getState()).thenReturn(DAGAppMasterState.INITED);
assertEquals(TezAppMasterStatus.INITIALIZING, dagClientHandler.getTezAppMasterStatus());
when(mockDagAM.getState()).thenReturn(DAGAppMasterState.ERROR);
assertEquals(TezAppMasterStatus.SHUTDOWN, dagClientHandler.getTezAppMasterStatus());
// tryKillDAG
try {
dagClientHandler.tryKillDAG("dag_9999_0001_2");
fail("should not come here");
} catch (TezException e) {
assertTrue(e.getMessage().contains("Unknown dagId"));
}
dagClientHandler.tryKillDAG("dag_9999_0001_1");
ArgumentCaptor<DAG> eventCaptor = ArgumentCaptor.forClass(DAG.class);
verify(mockDagAM, times(1))
.tryKillDAG(eventCaptor.capture(), contains("Sending client kill from"));
assertEquals(1, eventCaptor.getAllValues().size());
assertTrue(eventCaptor.getAllValues().get(0) instanceof DAG);
assertEquals("dag_9999_0001_1", ((DAG) eventCaptor.getAllValues().get(0)).getID().toString());
// submitDAG
DAGPlan dagPlan = DAGPlan.getDefaultInstance();
Map<String, LocalResource> localResources = new HashMap<String, LocalResource>();
dagClientHandler.submitDAG(dagPlan, localResources);
verify(mockDagAM).submitDAGToAppMaster(dagPlan, localResources);
// shutdown
dagClientHandler.shutdownAM();
verify(mockDagAM).shutdownTezAM(contains("Received message to shutdown AM from"));
}
public void tryKillDAG(String dagIdStr) throws TezException {
DAG dag = getDAG(dagIdStr);
LOG.info("Sending client kill to dag: " + dagIdStr);
// send a DAG_KILL message
sendEvent(new DAGEvent(dag.getID(), DAGEventType.DAG_KILL));
}
private synchronized void handle(DAGAppMasterEvent event) {
switch (event.getType()) {
case INTERNAL_ERROR:
state = DAGAppMasterState.ERROR;
if (currentDAG != null) {
// notify dag to finish which will send the DAG_FINISHED event
LOG.info("Internal Error. Notifying dags to finish.");
sendEvent(new DAGEvent(currentDAG.getID(), DAGEventType.INTERNAL_ERROR));
} else {
LOG.info("Internal Error. Finishing directly as no dag is active.");
shutdownHandler.shutdown();
}
break;
case DAG_FINISHED:
DAGAppMasterEventDAGFinished finishEvt = (DAGAppMasterEventDAGFinished) event;
if (!isSession) {
setStateOnDAGCompletion();
LOG.info("Shutting down on completion of dag:" + finishEvt.getDAGId().toString());
shutdownHandler.shutdown();
} else {
LOG.info(
"DAG completed, dagId="
+ finishEvt.getDAGId().toString()
+ ", dagState="
+ finishEvt.getDAGState());
lastDAGCompletionTime = clock.getTime();
switch (finishEvt.getDAGState()) {
case SUCCEEDED:
successfulDAGs.incrementAndGet();
break;
case ERROR:
case FAILED:
failedDAGs.incrementAndGet();
break;
case KILLED:
killedDAGs.incrementAndGet();
break;
default:
LOG.fatal(
"Received a DAG Finished Event with state="
+ finishEvt.getDAGState()
+ ". Error. Shutting down.");
state = DAGAppMasterState.ERROR;
shutdownHandler.shutdown();
break;
}
if (!state.equals(DAGAppMasterState.ERROR)) {
if (!sessionStopped.get()) {
LOG.info("Waiting for next DAG to be submitted.");
state = DAGAppMasterState.IDLE;
} else {
LOG.info("Session shutting down now.");
state = DAGAppMasterState.SUCCEEDED;
shutdownHandler.shutdown();
}
}
}
break;
default:
throw new TezUncheckedException("AppMaster: No handler for event type: " + event.getType());
}
}