Search in sources :

Example 16 with TaskInstance

use of org.apache.samza.container.TaskInstance in project samza by apache.

the class TestAsyncRunLoop method testEndOfStreamOffsetManagement.

// TODO: Add assertions.
//@Test
public void testEndOfStreamOffsetManagement() throws Exception {
    //explicitly configure to disable commits inside process or window calls and invoke commit from end of stream
    TestTask mockStreamTask1 = new TestTask(true, false, false, null);
    TestTask mockStreamTask2 = new TestTask(true, false, false, null);
    Partition p1 = new Partition(1);
    Partition p2 = new Partition(2);
    SystemStreamPartition ssp1 = new SystemStreamPartition("system1", "stream1", p1);
    SystemStreamPartition ssp2 = new SystemStreamPartition("system1", "stream2", p2);
    IncomingMessageEnvelope envelope1 = new IncomingMessageEnvelope(ssp2, "1", "key1", "message1");
    IncomingMessageEnvelope envelope2 = new IncomingMessageEnvelope(ssp2, "2", "key1", "message1");
    IncomingMessageEnvelope envelope3 = IncomingMessageEnvelope.buildEndOfStreamEnvelope(ssp2);
    Map<SystemStreamPartition, List<IncomingMessageEnvelope>> sspMap = new HashMap<>();
    List<IncomingMessageEnvelope> messageList = new ArrayList<>();
    messageList.add(envelope1);
    messageList.add(envelope2);
    messageList.add(envelope3);
    sspMap.put(ssp2, messageList);
    SystemConsumer mockConsumer = mock(SystemConsumer.class);
    when(mockConsumer.poll(anyObject(), anyLong())).thenReturn(sspMap);
    HashMap<String, SystemConsumer> systemConsumerMap = new HashMap<>();
    systemConsumerMap.put("system1", mockConsumer);
    SystemConsumers consumers = TestSystemConsumers.getSystemConsumers(systemConsumerMap);
    TaskName taskName1 = new TaskName("task1");
    TaskName taskName2 = new TaskName("task2");
    Set<TaskName> taskNames = new HashSet<>();
    taskNames.add(taskName1);
    taskNames.add(taskName2);
    OffsetManager offsetManager = mock(OffsetManager.class);
    when(offsetManager.getLastProcessedOffset(taskName1, ssp1)).thenReturn(Option.apply("3"));
    when(offsetManager.getLastProcessedOffset(taskName2, ssp2)).thenReturn(Option.apply("0"));
    when(offsetManager.getStartingOffset(taskName1, ssp1)).thenReturn(Option.apply(IncomingMessageEnvelope.END_OF_STREAM_OFFSET));
    when(offsetManager.getStartingOffset(taskName2, ssp2)).thenReturn(Option.apply("1"));
    TaskInstance taskInstance1 = createTaskInstance(mockStreamTask1, taskName1, ssp1, offsetManager, consumers);
    TaskInstance taskInstance2 = createTaskInstance(mockStreamTask2, taskName2, ssp2, offsetManager, consumers);
    Map<TaskName, TaskInstance> tasks = new HashMap<>();
    tasks.put(taskName1, taskInstance1);
    tasks.put(taskName2, taskInstance2);
    taskInstance1.registerConsumers();
    taskInstance2.registerConsumers();
    consumers.start();
    int maxMessagesInFlight = 1;
    AsyncRunLoop runLoop = new AsyncRunLoop(tasks, executor, consumers, maxMessagesInFlight, windowMs, commitMs, callbackTimeoutMs, maxThrottlingDelayMs, containerMetrics, () -> 0L, false);
    runLoop.run();
}
Also used : SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Partition(org.apache.samza.Partition) SystemConsumer(org.apache.samza.system.SystemConsumer) TaskInstance(org.apache.samza.container.TaskInstance) SystemConsumers(org.apache.samza.system.SystemConsumers) TestSystemConsumers(org.apache.samza.system.TestSystemConsumers) HashMap(java.util.HashMap) IncomingMessageEnvelope(org.apache.samza.system.IncomingMessageEnvelope) OffsetManager(org.apache.samza.checkpoint.OffsetManager) ArrayList(java.util.ArrayList) TaskName(org.apache.samza.container.TaskName) ArrayList(java.util.ArrayList) List(java.util.List) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) HashSet(java.util.HashSet)

Example 17 with TaskInstance

use of org.apache.samza.container.TaskInstance in project samza by apache.

the class TestContainerStorageManager method addMockedTask.

/**
 * Utility method for creating a mocked taskInstance and taskStorageManager and adding it to the map.
 * @param taskname the desired taskname.
 */
private void addMockedTask(String taskname, int changelogPartition) {
    TaskInstance mockTaskInstance = mock(TaskInstance.class);
    doAnswer(invocation -> {
        return new TaskName(taskname);
    }).when(mockTaskInstance).taskName();
    Gauge testGauge = mock(Gauge.class);
    this.tasks.put(new TaskName(taskname), new TaskModel(new TaskName(taskname), new HashSet<>(), new Partition(changelogPartition)));
    this.taskRestoreMetricGauges.put(new TaskName(taskname), testGauge);
    this.taskInstanceMetrics.put(new TaskName(taskname), mock(TaskInstanceMetrics.class));
}
Also used : TaskInstance(org.apache.samza.container.TaskInstance) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Partition(org.apache.samza.Partition) TaskName(org.apache.samza.container.TaskName) TaskInstanceMetrics(org.apache.samza.container.TaskInstanceMetrics) TaskModel(org.apache.samza.job.model.TaskModel) Gauge(org.apache.samza.metrics.Gauge) HashSet(java.util.HashSet)

Aggregations

TaskInstance (org.apache.samza.container.TaskInstance)17 HashMap (java.util.HashMap)15 TaskName (org.apache.samza.container.TaskName)15 CountDownLatch (java.util.concurrent.CountDownLatch)11 SystemStreamPartition (org.apache.samza.system.SystemStreamPartition)4 HashSet (java.util.HashSet)3 ArrayList (java.util.ArrayList)2 List (java.util.List)2 Partition (org.apache.samza.Partition)2 TaskInstanceMetrics (org.apache.samza.container.TaskInstanceMetrics)2 IncomingMessageEnvelope (org.apache.samza.system.IncomingMessageEnvelope)2 Map (java.util.Map)1 OffsetManager (org.apache.samza.checkpoint.OffsetManager)1 Config (org.apache.samza.config.Config)1 SamzaContainer (org.apache.samza.container.SamzaContainer)1 SamzaContainerContext (org.apache.samza.container.SamzaContainerContext)1 TaskInstanceExceptionHandler (org.apache.samza.container.TaskInstanceExceptionHandler)1 TaskModel (org.apache.samza.job.model.TaskModel)1 Gauge (org.apache.samza.metrics.Gauge)1 MetricsRegistryMap (org.apache.samza.metrics.MetricsRegistryMap)1