Examples with TaskHeartbeatRequest org.apache.tez.serviceplugins.api.TaskHeartbeatRequest used on opensource projects

Example 1 with TaskHeartbeatRequest

use of org.apache.tez.serviceplugins.api.TaskHeartbeatRequest in project tez by apache.

the class TestTaskCommunicatorManager2 method testTaskAttemptFailureViaHeartbeat.

// Tests fatal and non fatal
@SuppressWarnings("unchecked")
@Test(timeout = 5000)
public void testTaskAttemptFailureViaHeartbeat() throws IOException, TezException {
    TaskCommunicatorManagerWrapperForTest wrapper = new TaskCommunicatorManagerWrapperForTest();
    TaskSpec taskSpec1 = wrapper.createTaskSpec();
    AMContainerTask amContainerTask1 = new AMContainerTask(taskSpec1, null, null, false, 10);
    TaskSpec taskSpec2 = wrapper.createTaskSpec();
    AMContainerTask amContainerTask2 = new AMContainerTask(taskSpec2, null, null, false, 10);
    ContainerId containerId1 = wrapper.createContainerId(1);
    wrapper.registerRunningContainer(containerId1);
    wrapper.registerTaskAttempt(containerId1, amContainerTask1);
    ContainerId containerId2 = wrapper.createContainerId(2);
    wrapper.registerRunningContainer(containerId2);
    wrapper.registerTaskAttempt(containerId2, amContainerTask2);
    List<TezEvent> events = new LinkedList<>();
    EventMetaData sourceInfo1 = new EventMetaData(EventMetaData.EventProducerConsumerType.PROCESSOR, "testVertex", null, taskSpec1.getTaskAttemptID());
    TaskAttemptFailedEvent failedEvent1 = new TaskAttemptFailedEvent("non-fatal test error", TaskFailureType.NON_FATAL);
    TezEvent failedEventT1 = new TezEvent(failedEvent1, sourceInfo1);
    events.add(failedEventT1);
    TaskHeartbeatRequest taskHeartbeatRequest1 = new TaskHeartbeatRequest(containerId1.toString(), taskSpec1.getTaskAttemptID(), events, 0, 0, 0);
    wrapper.getTaskCommunicatorManager().heartbeat(taskHeartbeatRequest1);
    ArgumentCaptor<Event> argumentCaptor = ArgumentCaptor.forClass(Event.class);
    verify(wrapper.getEventHandler(), times(1)).handle(argumentCaptor.capture());
    assertTrue(argumentCaptor.getAllValues().get(0) instanceof TaskAttemptEventAttemptFailed);
    TaskAttemptEventAttemptFailed failedEvent = (TaskAttemptEventAttemptFailed) argumentCaptor.getAllValues().get(0);
    assertEquals(TaskFailureType.NON_FATAL, failedEvent.getTaskFailureType());
    assertTrue(failedEvent.getDiagnosticInfo().contains("non-fatal"));
    events.clear();
    reset(wrapper.getEventHandler());
    EventMetaData sourceInfo2 = new EventMetaData(EventMetaData.EventProducerConsumerType.PROCESSOR, "testVertex", null, taskSpec2.getTaskAttemptID());
    TaskAttemptFailedEvent failedEvent2 = new TaskAttemptFailedEvent("-fatal- test error", TaskFailureType.FATAL);
    TezEvent failedEventT2 = new TezEvent(failedEvent2, sourceInfo2);
    events.add(failedEventT2);
    TaskHeartbeatRequest taskHeartbeatRequest2 = new TaskHeartbeatRequest(containerId2.toString(), taskSpec2.getTaskAttemptID(), events, 0, 0, 0);
    wrapper.getTaskCommunicatorManager().heartbeat(taskHeartbeatRequest2);
    argumentCaptor = ArgumentCaptor.forClass(Event.class);
    verify(wrapper.getEventHandler(), times(1)).handle(argumentCaptor.capture());
    assertTrue(argumentCaptor.getAllValues().get(0) instanceof TaskAttemptEventAttemptFailed);
    failedEvent = (TaskAttemptEventAttemptFailed) argumentCaptor.getAllValues().get(0);
    assertEquals(TaskFailureType.FATAL, failedEvent.getTaskFailureType());
    assertTrue(failedEvent.getDiagnosticInfo().contains("-fatal-"));
}

Example 2 with TaskHeartbeatRequest

use of org.apache.tez.serviceplugins.api.TaskHeartbeatRequest in project tez by apache.

the class TestTaskCommunicatorManager1 method generateHeartbeat.

private TaskHeartbeatResponse generateHeartbeat(List<TezEvent> events, int fromEventId, int maxEvents, int nextFromEventId, List<TezEvent> sendEvents) throws IOException, TezException {
    ContainerId containerId = createContainerId(appId, 1);
    Vertex vertex = mock(Vertex.class);
    doReturn(vertex).when(dag).getVertex(vertexID);
    doReturn("test_vertex").when(vertex).getName();
    TaskAttemptEventInfo eventInfo = new TaskAttemptEventInfo(nextFromEventId, sendEvents, 0);
    doReturn(eventInfo).when(vertex).getTaskAttemptTezEvents(taskAttemptID, fromEventId, 0, maxEvents);
    taskAttemptListener.registerRunningContainer(containerId, 0);
    taskAttemptListener.registerTaskAttempt(amContainerTask, containerId, 0);
    TaskHeartbeatRequest request = mock(TaskHeartbeatRequest.class);
    doReturn(containerId.toString()).when(request).getContainerIdentifier();
    doReturn(containerId.toString()).when(request).getContainerIdentifier();
    doReturn(taskAttemptID).when(request).getTaskAttemptId();
    doReturn(events).when(request).getEvents();
    doReturn(maxEvents).when(request).getMaxEvents();
    doReturn(fromEventId).when(request).getStartIndex();
    return taskAttemptListener.heartbeat(request);
}