Examples with TaskCallbackFactory org.apache.samza.task.TaskCallbackFactory used on opensource projects

Example 6 with TaskCallbackFactory

use of org.apache.samza.task.TaskCallbackFactory in project samza by apache.

the class TestRunLoop method testCommitSingleTask.

@Test
public void testCommitSingleTask() {
    SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
    RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
    doAnswer(invocation -> {
        ReadableCoordinator coordinator = invocation.getArgumentAt(1, ReadableCoordinator.class);
        TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
        TaskCallback callback = callbackFactory.createCallback();
        coordinator.commit(TaskCoordinator.RequestScope.CURRENT_TASK);
        coordinator.shutdown(TaskCoordinator.RequestScope.ALL_TASKS_IN_CONTAINER);
        callback.complete();
        return null;
    }).when(task0).process(eq(envelope00), any(), any());
    RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1);
    Map<TaskName, RunLoopTask> tasks = new HashMap<>();
    tasks.put(this.taskName0, task0);
    tasks.put(taskName1, task1);
    int maxMessagesInFlight = 1;
    RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
    // have a null message in between to make sure task0 finishes processing and invoke the commit
    when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope11).thenReturn(null);
    runLoop.run();
    verify(task0).process(any(), any(), any());
    verify(task1).process(any(), any(), any());
    verify(task0).commit();
    verify(task1, never()).commit();
}

Example 7 with TaskCallbackFactory

use of org.apache.samza.task.TaskCallbackFactory in project samza by apache.

the class TestRunLoop method testCommitWithMessageInFlightWhenAsyncCommitIsEnabled.

@Test
public void testCommitWithMessageInFlightWhenAsyncCommitIsEnabled() {
    int maxMessagesInFlight = 2;
    ExecutorService taskExecutor = Executors.newFixedThreadPool(2);
    SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
    OffsetManager offsetManager = mock(OffsetManager.class);
    RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
    when(task0.offsetManager()).thenReturn(offsetManager);
    CountDownLatch firstMessageBarrier = new CountDownLatch(1);
    doAnswer(invocation -> {
        ReadableCoordinator coordinator = invocation.getArgumentAt(1, ReadableCoordinator.class);
        TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
        TaskCallback callback = callbackFactory.createCallback();
        taskExecutor.submit(() -> {
            firstMessageBarrier.await();
            coordinator.commit(TaskCoordinator.RequestScope.CURRENT_TASK);
            callback.complete();
            return null;
        });
        return null;
    }).when(task0).process(eq(envelope00), any(), any());
    CountDownLatch secondMessageBarrier = new CountDownLatch(1);
    doAnswer(invocation -> {
        ReadableCoordinator coordinator = invocation.getArgumentAt(1, ReadableCoordinator.class);
        TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
        TaskCallback callback = callbackFactory.createCallback();
        taskExecutor.submit(() -> {
            // let the first message proceed to ask for a commit
            firstMessageBarrier.countDown();
            // block this message until commit is executed
            secondMessageBarrier.await();
            coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
            callback.complete();
            return null;
        });
        return null;
    }).when(task0).process(eq(envelope01), any(), any());
    doAnswer(invocation -> {
        assertEquals(1, task0.metrics().asyncCallbackCompleted().getCount());
        assertEquals(1, task0.metrics().messagesInFlight().getValue());
        secondMessageBarrier.countDown();
        return null;
    }).when(task0).commit();
    Map<TaskName, RunLoopTask> tasks = new HashMap<>();
    tasks.put(taskName0, task0);
    RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, true);
    when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(null);
    runLoop.run();
    InOrder inOrder = inOrder(task0);
    inOrder.verify(task0).process(eq(envelope00), any(), any());
    inOrder.verify(task0).process(eq(envelope01), any(), any());
    inOrder.verify(task0).commit();
}