Examples with TaskSpec org.apache.tez.runtime.api.impl.TaskSpec used on opensource projects

Example 6 with TaskSpec

use of org.apache.tez.runtime.api.impl.TaskSpec in project hive by apache.

the class TestLlapTaskCommunicator method testFinishableStateUpdateFailure.

@Test(timeout = 5000)
public void testFinishableStateUpdateFailure() throws Exception {
    LlapTaskCommunicatorWrapperForTest wrapper = null;
    Lock lock = new ReentrantLock();
    Condition condition = lock.newCondition();
    final AtomicBoolean opDone = new AtomicBoolean(false);
    LlapProtocolClientProxy proxy = mock(LlapProtocolClientProxy.class, new FinishableStatusUpdateTestAnswer(lock, condition, opDone));
    try {
        wrapper = new LlapTaskCommunicatorWrapperForTest(proxy);
        // Register tasks on 2 nodes, with a dependency on vertex1 completing.
        ContainerId cId11 = wrapper.registerContainer(1, 0);
        TaskSpec ts11 = wrapper.registerRunningTaskAttemptWithSourceVertex(cId11, 1);
        ContainerId cId12 = wrapper.registerContainer(2, 0);
        TaskSpec ts12 = wrapper.registerRunningTaskAttemptWithSourceVertex(cId12, 2);
        ContainerId cId21 = wrapper.registerContainer(3, 1);
        TaskSpec ts21 = wrapper.registerRunningTaskAttemptWithSourceVertex(cId21, 3);
        // Send a state update for vertex1 completion. This triggers a status update to be sent out.
        VertexStateUpdate vertexStateUpdate = new VertexStateUpdate(LlapTaskCommunicatorWrapperForTest.VERTEX_NAME1, VertexState.SUCCEEDED);
        wrapper.getTaskCommunicator().onVertexStateUpdated(vertexStateUpdate);
        // Wait for all invocations to complete.
        lock.lock();
        try {
            while (!opDone.get()) {
                condition.await();
            }
        } finally {
            lock.unlock();
        }
        // Verify that a task kill went out for all nodes running on the specified host.
        verify(wrapper.getTaskCommunicatorContext(), times(2)).taskKilled(any(TezTaskAttemptID.class), any(TaskAttemptEndReason.class), any(String.class));
        verify(wrapper.getTaskCommunicatorContext()).taskKilled(eq(ts11.getTaskAttemptID()), eq(TaskAttemptEndReason.NODE_FAILED), any(String.class));
        verify(wrapper.getTaskCommunicatorContext()).taskKilled(eq(ts12.getTaskAttemptID()), eq(TaskAttemptEndReason.NODE_FAILED), any(String.class));
        wrapper.getTaskCommunicator().sendStateUpdate(LlapNodeId.getInstance(LlapTaskCommunicatorWrapperForTest.HOSTS[1], LlapTaskCommunicatorWrapperForTest.RPC_PORT), LlapDaemonProtocolProtos.SourceStateUpdatedRequestProto.getDefaultInstance());
        // Verify no more invocations in case of success.
        verify(wrapper.getTaskCommunicatorContext(), times(2)).taskKilled(any(TezTaskAttemptID.class), any(TaskAttemptEndReason.class), any(String.class));
    } finally {
        if (wrapper != null) {
            wrapper.shutdown();
        }
    }
}