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

Example 1 with TaskCallback

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

the class TestRunLoop method testProcessCallbacksCompletedOutOfOrder.

@Test
public void testProcessCallbacksCompletedOutOfOrder() {
    int maxMessagesInFlight = 2;
    ExecutorService taskExecutor = Executors.newFixedThreadPool(1);
    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);
            coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
            callback.complete();
            return null;
        });
        return null;
    }).when(task0).process(eq(envelope00), any(), any());
    doAnswer(invocation -> {
        assertEquals(1, task0.metrics().messagesInFlight().getValue());
        assertEquals(0, task0.metrics().asyncCallbackCompleted().getCount());
        TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
        TaskCallback callback = callbackFactory.createCallback();
        callback.complete();
        firstMessageBarrier.countDown();
        return null;
    }).when(task0).process(eq(envelope01), any(), any());
    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, false);
    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());
    verify(offsetManager).update(eq(taskName0), eq(ssp0), eq(envelope00.getOffset()));
    assertEquals(2L, containerMetrics.processes().getCount());
}

Example 2 with TaskCallback

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

the class TestRunLoop method testCommitAllTasks.

@Test
public void testCommitAllTasks() {
    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.ALL_TASKS_IN_CONTAINER);
        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).commit();
}

Example 3 with TaskCallback

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

the class TestRunLoop method testShutdownOnConsensus.

@Test
public void testShutdownOnConsensus() {
    SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
    int maxMessagesInFlight = 1;
    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.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
        callback.complete();
        return null;
    }).when(task0).process(eq(envelope00), any(), any());
    RunLoopTask task1 = getMockRunLoopTask(taskName1, ssp1);
    doAnswer(invocation -> {
        ReadableCoordinator coordinator = invocation.getArgumentAt(1, ReadableCoordinator.class);
        TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
        TaskCallback callback = callbackFactory.createCallback();
        coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
        callback.complete();
        return null;
    }).when(task1).process(eq(envelope11), any(), any());
    Map<TaskName, RunLoopTask> tasks = new HashMap<>();
    tasks.put(taskName0, task0);
    tasks.put(taskName1, task1);
    RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
    // consensus is reached after envelope1 is processed.
    when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope11).thenReturn(null);
    runLoop.run();
    verify(task0).process(any(), any(), any());
    verify(task1).process(any(), any(), any());
    assertEquals(2L, containerMetrics.envelopes().getCount());
    assertEquals(2L, containerMetrics.processes().getCount());
}

Example 4 with TaskCallback

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

the class TestRunLoop method testEndOfStreamWaitsForInFlightMessages.

@Test
public void testEndOfStreamWaitsForInFlightMessages() {
    int maxMessagesInFlight = 2;
    ExecutorService taskExecutor = Executors.newFixedThreadPool(1);
    SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);
    OffsetManager offsetManager = mock(OffsetManager.class);
    RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
    when(task0.offsetManager()).thenReturn(offsetManager);
    CountDownLatch firstMessageBarrier = new CountDownLatch(2);
    doAnswer(invocation -> {
        TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
        TaskCallback callback = callbackFactory.createCallback();
        taskExecutor.submit(() -> {
            firstMessageBarrier.await();
            callback.complete();
            return null;
        });
        return null;
    }).when(task0).process(eq(envelope00), any(), any());
    doAnswer(invocation -> {
        assertEquals(1, task0.metrics().messagesInFlight().getValue());
        TaskCallbackFactory callbackFactory = invocation.getArgumentAt(2, TaskCallbackFactory.class);
        TaskCallback callback = callbackFactory.createCallback();
        callback.complete();
        firstMessageBarrier.countDown();
        return null;
    }).when(task0).process(eq(envelope01), any(), any());
    doAnswer(invocation -> {
        assertEquals(0, task0.metrics().messagesInFlight().getValue());
        assertEquals(2, task0.metrics().asyncCallbackCompleted().getCount());
        return null;
    }).when(task0).endOfStream(any());
    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, false);
    when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(envelope01).thenReturn(ssp0EndOfStream).thenAnswer(invocation -> {
        // this ensures that the end of stream message has passed through run loop BEFORE the last remaining in flight message completes
        firstMessageBarrier.countDown();
        return null;
    });
    runLoop.run();
    verify(task0).endOfStream(any());
}

Example 5 with TaskCallback

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

the class SideInputTask method process.

@Override
public synchronized void process(IncomingMessageEnvelope envelope, ReadableCoordinator coordinator, TaskCallbackFactory callbackFactory) {
    TaskCallback callback = callbackFactory.createCallback();
    this.metrics.processes().inc();
    try {
        this.taskSideInputHandler.process(envelope);
        this.metrics.messagesActuallyProcessed().inc();
        callback.complete();
    } catch (Exception e) {
        callback.failure(e);
    }
}