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

Example 31 with TaskCallback

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

the class TestStartpoint method testStartpointTimestamp.

@Test
public void testStartpointTimestamp() throws InterruptedException {
    Map<Integer, RecordMetadata> sentEvents2 = publishKafkaEventsWithDelayPerEvent(inputKafkaTopic2, 0, NUM_KAFKA_EVENTS, PROCESSOR_IDS[1], Duration.ofMillis(2));
    ConcurrentHashMap<String, IncomingMessageEnvelope> recvEventsInputStartpointTimestamp = new ConcurrentHashMap<>();
    CoordinatorStreamStore coordinatorStreamStore = createCoordinatorStreamStore(applicationConfig1);
    coordinatorStreamStore.init();
    StartpointManager startpointManager = new StartpointManager(coordinatorStreamStore);
    startpointManager.start();
    StartpointTimestamp startpointTimestamp = new StartpointTimestamp(sentEvents2.get(150).timestamp());
    writeStartpoints(startpointManager, inputKafkaTopic2, ZK_TEST_PARTITION_COUNT, startpointTimestamp);
    startpointManager.stop();
    coordinatorStreamStore.close();
    TestTaskApplication.TaskApplicationProcessCallback processedCallback = (IncomingMessageEnvelope ime, TaskCallback callback) -> {
        try {
            String streamName = ime.getSystemStreamPartition().getStream();
            TestKafkaEvent testKafkaEvent = TestKafkaEvent.fromString((String) ime.getMessage());
            String eventIndex = testKafkaEvent.getEventData();
            if (inputKafkaTopic2.equals(streamName)) {
                recvEventsInputStartpointTimestamp.put(eventIndex, ime);
            } else {
                throw new RuntimeException("Unexpected input stream: " + streamName);
            }
            callback.complete();
        } catch (Exception ex) {
            callback.failure(ex);
        }
    };
    // Just fetch a few messages
    CountDownLatch processedMessagesLatchStartpointTimestamp = new CountDownLatch(100);
    CountDownLatch shutdownLatchStartpointTimestamp = new CountDownLatch(1);
    TestTaskApplication testTaskApplicationStartpointTimestamp = new TestTaskApplication(TEST_SYSTEM, inputKafkaTopic2, outputKafkaTopic, processedMessagesLatchStartpointTimestamp, shutdownLatchStartpointTimestamp, Optional.of(processedCallback));
    ApplicationRunner appRunner = ApplicationRunners.getApplicationRunner(testTaskApplicationStartpointTimestamp, applicationConfig2);
    executeRun(appRunner, applicationConfig2);
    assertTrue(processedMessagesLatchStartpointTimestamp.await(1, TimeUnit.MINUTES));
    appRunner.kill();
    appRunner.waitForFinish();
    assertTrue(shutdownLatchStartpointTimestamp.await(1, TimeUnit.MINUTES));
    for (IncomingMessageEnvelope ime : recvEventsInputStartpointTimestamp.values()) {
        Integer eventOffset = Integer.valueOf(ime.getOffset());
        // sanity check
        assertNotEquals(0, ime.getEventTime());
        String assertMsg = String.format("Expecting message timestamp: %d >= Startpoint timestamp: %d", ime.getEventTime(), startpointTimestamp.getTimestampOffset());
        assertTrue(assertMsg, ime.getEventTime() >= startpointTimestamp.getTimestampOffset());
    }
}

Example 32 with TaskCallback

use of org.apache.samza.task.TaskCallback 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 33 with TaskCallback

use of org.apache.samza.task.TaskCallback 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();
}