Examples with TaskName org.apache.samza.container.TaskName used on opensource projects

Example 21 with TaskName

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

Example 22 with TaskName

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

the class ContainerMocks method generateTaskAssignments.

public static Map<String, String> generateTaskAssignments(int numContainers, int taskCount) {
    Map<String, String> mapping = new HashMap<>(taskCount);
    Set<ContainerModel> containers = generateContainerModels(numContainers, taskCount);
    for (ContainerModel container : containers) {
        for (TaskName taskName : container.getTasks().keySet()) {
            mapping.put(taskName.getTaskName(), container.getProcessorId());
        }
    }
    return mapping;
}

Example 23 with TaskName

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

the class TestContainerAllocator method getJobModelReader.

private static JobModelManager getJobModelReader(int containerCount) {
    //Ideally, the JobModelReader should be constructed independent of HttpServer.
    //That way it becomes easier to mock objects. Save it for later.
    HttpServer server = new MockHttpServer("/", 7777, null, new ServletHolder(DefaultServlet.class));
    Map<String, ContainerModel> containers = new java.util.HashMap<>();
    for (int i = 0; i < containerCount; i++) {
        ContainerModel container = new ContainerModel(String.valueOf(i), i, new HashMap<TaskName, TaskModel>());
        containers.put(String.valueOf(i), container);
    }
    JobModel jobModel = new JobModel(getConfig(), containers);
    return new JobModelManager(jobModel, server, null);
}

Example 24 with TaskName

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

the class TestContainerProcessManager method getCoordinator.

private JobModelManager getCoordinator(int containerCount) {
    Map<String, ContainerModel> containers = new java.util.HashMap<>();
    for (int i = 0; i < containerCount; i++) {
        ContainerModel container = new ContainerModel(String.valueOf(i), i, new HashMap<TaskName, TaskModel>());
        containers.put(String.valueOf(i), container);
    }
    Map<String, Map<String, String>> localityMap = new HashMap<>();
    localityMap.put("0", new HashMap<String, String>() {

        {
            put(SetContainerHostMapping.HOST_KEY, "abc");
        }
    });
    LocalityManager mockLocalityManager = mock(LocalityManager.class);
    when(mockLocalityManager.readContainerLocality()).thenReturn(localityMap);
    JobModel jobModel = new JobModel(getConfig(), containers, mockLocalityManager);
    JobModelManager.jobModelRef().getAndSet(jobModel);
    return new JobModelManager(jobModel, this.server, null);
}

Example 25 with TaskName

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

the class TestGroupByPartition method testBroadcastStreamsGroupedCorrectly.

@Test
public void testBroadcastStreamsGroupedCorrectly() {
    HashMap<String, String> configMap = new HashMap<String, String>();
    configMap.put("task.broadcast.inputs", "SystemA.StreamA#0, SystemA.StreamB#1");
    Config config = new MapConfig(configMap);
    GroupByPartition grouper = new GroupByPartition(config);
    HashSet<SystemStreamPartition> allSSPs = new HashSet<SystemStreamPartition>();
    Collections.addAll(allSSPs, aa0, aa1, aa2, ab1, ab2, ac0);
    Map<TaskName, Set<SystemStreamPartition>> result = grouper.group(allSSPs);
    Map<TaskName, Set<SystemStreamPartition>> expectedResult = new HashMap<TaskName, Set<SystemStreamPartition>>();
    HashSet<SystemStreamPartition> partition0 = new HashSet<SystemStreamPartition>();
    // broadcast stream
    partition0.add(aa0);
    partition0.add(ac0);
    // broadcast stream
    partition0.add(ab1);
    expectedResult.put(new TaskName("Partition 0"), partition0);
    HashSet<SystemStreamPartition> partition1 = new HashSet<SystemStreamPartition>();
    partition1.add(aa1);
    // broadcast stream
    partition1.add(ab1);
    // broadcast stream
    partition1.add(aa0);
    expectedResult.put(new TaskName("Partition 1"), partition1);
    HashSet<SystemStreamPartition> partition2 = new HashSet<SystemStreamPartition>();
    partition2.add(aa2);
    partition2.add(ab2);
    // broadcast stream
    partition2.add(aa0);
    // broadcast stream
    partition2.add(ab1);
    expectedResult.put(new TaskName("Partition 2"), partition2);
    assertEquals(expectedResult, result);
}