use of org.apache.tez.runtime.api.impl.TezHeartbeatResponse 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();
}
}
use of org.apache.tez.runtime.api.impl.TezHeartbeatResponse 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());
}
Aggregations