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

Example 1 with TezHeartbeatRequest

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

the class TestTaskReporter method testContinuousHeartbeatsOnMaxEvents.

@Test(timeout = 10000)
public void testContinuousHeartbeatsOnMaxEvents() throws Exception {
    final Object lock = new Object();
    final AtomicBoolean hb2Done = new AtomicBoolean(false);
    TezTaskUmbilicalProtocol mockUmbilical = mock(TezTaskUmbilicalProtocol.class);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            TezHeartbeatRequest request = (TezHeartbeatRequest) args[0];
            if (request.getRequestId() == 1 || request.getRequestId() == 2) {
                TezHeartbeatResponse response = new TezHeartbeatResponse(createEvents(5));
                response.setLastRequestId(request.getRequestId());
                return response;
            } else if (request.getRequestId() == 3) {
                TezHeartbeatResponse response = new TezHeartbeatResponse(createEvents(1));
                response.setLastRequestId(request.getRequestId());
                synchronized (lock) {
                    hb2Done.set(true);
                    lock.notify();
                }
                return response;
            } else {
                throw new TezUncheckedException("Invalid request id for test: " + request.getRequestId());
            }
        }
    }).when(mockUmbilical).heartbeat(any(TezHeartbeatRequest.class));
    TezTaskAttemptID mockTaskAttemptId = mock(TezTaskAttemptID.class);
    LogicalIOProcessorRuntimeTask mockTask = mock(LogicalIOProcessorRuntimeTask.class);
    doReturn("vertexName").when(mockTask).getVertexName();
    doReturn(mockTaskAttemptId).when(mockTask).getTaskAttemptID();
    // Setup the sleep time to be way higher than the test timeout
    TaskReporter.HeartbeatCallable heartbeatCallable = new TaskReporter.HeartbeatCallable(mockTask, mockUmbilical, 100000, 100000, 5, new AtomicLong(0), "containerIdStr");
    ExecutorService executor = Executors.newSingleThreadExecutor();
    executor.submit(heartbeatCallable);
    try {
        synchronized (lock) {
            if (!hb2Done.get()) {
                lock.wait();
            }
        }
        verify(mockUmbilical, times(3)).heartbeat(any(TezHeartbeatRequest.class));
        Thread.sleep(2000l);
        // Sleep for 2 seconds, less than the callable sleep time. No more invocations.
        verify(mockUmbilical, times(3)).heartbeat(any(TezHeartbeatRequest.class));
    } finally {
        executor.shutdownNow();
    }
}

Example 2 with TezHeartbeatRequest

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

the class TestTaskReporter method testEventThrottling.

@Test(timeout = 10000)
public void testEventThrottling() throws Exception {
    TezTaskAttemptID mockTaskAttemptId = mock(TezTaskAttemptID.class);
    LogicalIOProcessorRuntimeTask mockTask = mock(LogicalIOProcessorRuntimeTask.class);
    when(mockTask.getMaxEventsToHandle()).thenReturn(10000, 1);
    when(mockTask.getVertexName()).thenReturn("vertexName");
    when(mockTask.getTaskAttemptID()).thenReturn(mockTaskAttemptId);
    TezTaskUmbilicalProtocol mockUmbilical = mock(TezTaskUmbilicalProtocol.class);
    TezHeartbeatResponse resp1 = new TezHeartbeatResponse(createEvents(5));
    resp1.setLastRequestId(1);
    TezHeartbeatResponse resp2 = new TezHeartbeatResponse(createEvents(1));
    resp2.setLastRequestId(2);
    resp2.setShouldDie();
    when(mockUmbilical.heartbeat(isA(TezHeartbeatRequest.class))).thenReturn(resp1, resp2);
    // Setup the sleep time to be way higher than the test timeout
    TaskReporter.HeartbeatCallable heartbeatCallable = new TaskReporter.HeartbeatCallable(mockTask, mockUmbilical, 100000, 100000, 5, new AtomicLong(0), "containerIdStr");
    ExecutorService executor = Executors.newSingleThreadExecutor();
    try {
        Future<Boolean> result = executor.submit(heartbeatCallable);
        Assert.assertFalse(result.get());
    } finally {
        executor.shutdownNow();
    }
    ArgumentCaptor<TezHeartbeatRequest> captor = ArgumentCaptor.forClass(TezHeartbeatRequest.class);
    verify(mockUmbilical, times(2)).heartbeat(captor.capture());
    TezHeartbeatRequest req = captor.getValue();
    Assert.assertEquals(2, req.getRequestId());
    Assert.assertEquals(1, req.getMaxEvents());
}