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

Example 21 with TaskCoordinator

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

the class TestOperatorImpl method testOnMessagePropagatesResults.

@Test
public void testOnMessagePropagatesResults() {
    TaskContext mockTaskContext = mock(TaskContext.class);
    when(mockTaskContext.getMetricsRegistry()).thenReturn(new MetricsRegistryMap());
    Object mockTestOpImplOutput = mock(Object.class);
    OperatorImpl<Object, Object> opImpl = new TestOpImpl(mockTestOpImplOutput);
    opImpl.init(mock(Config.class), mockTaskContext);
    // register a couple of operators
    OperatorImpl mockNextOpImpl1 = mock(OperatorImpl.class);
    when(mockNextOpImpl1.getOperatorSpec()).thenReturn(new TestOpSpec());
    when(mockNextOpImpl1.handleMessage(anyObject(), anyObject(), anyObject())).thenReturn(Collections.emptyList());
    mockNextOpImpl1.init(mock(Config.class), mockTaskContext);
    opImpl.registerNextOperator(mockNextOpImpl1);
    OperatorImpl mockNextOpImpl2 = mock(OperatorImpl.class);
    when(mockNextOpImpl2.getOperatorSpec()).thenReturn(new TestOpSpec());
    when(mockNextOpImpl2.handleMessage(anyObject(), anyObject(), anyObject())).thenReturn(Collections.emptyList());
    mockNextOpImpl2.init(mock(Config.class), mockTaskContext);
    opImpl.registerNextOperator(mockNextOpImpl2);
    // send a message to this operator
    MessageCollector mockCollector = mock(MessageCollector.class);
    TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);
    opImpl.onMessage(mock(Object.class), mockCollector, mockCoordinator);
    // verify that it propagates its handleMessage results to next operators
    verify(mockNextOpImpl1, times(1)).handleMessage(mockTestOpImplOutput, mockCollector, mockCoordinator);
    verify(mockNextOpImpl2, times(1)).handleMessage(mockTestOpImplOutput, mockCollector, mockCoordinator);
}

Example 22 with TaskCoordinator

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

the class TestWindowOperator method testCancellationOfOnceTrigger.

@Test
public void testCancellationOfOnceTrigger() throws Exception {
    StreamApplication sgb = new KeyedTumblingWindowStreamApplication(AccumulationMode.ACCUMULATING, Duration.ofSeconds(1), Triggers.count(2));
    TestClock testClock = new TestClock();
    StreamOperatorTask task = new StreamOperatorTask(sgb, runner, testClock);
    task.init(config, taskContext);
    List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
    MessageCollector messageCollector = envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 1);
    Assert.assertEquals(windowPanes.get(0).getKey().getPaneId(), "0");
    Assert.assertEquals(windowPanes.get(0).getKey().getKey(), new Integer(1));
    Assert.assertEquals(windowPanes.get(0).getFiringType(), FiringType.EARLY);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 1);
    testClock.advanceTime(Duration.ofSeconds(1));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 2);
    Assert.assertEquals(windowPanes.get(0).getKey().getPaneId(), "0");
    Assert.assertEquals(windowPanes.get(1).getKey().getPaneId(), "0");
    Assert.assertEquals(windowPanes.get(1).getFiringType(), FiringType.DEFAULT);
    task.process(new IntegerEnvelope(3), messageCollector, taskCoordinator);
    testClock.advanceTime(Duration.ofSeconds(1));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 3);
    Assert.assertEquals(windowPanes.get(2).getKey().getKey(), new Integer(3));
    Assert.assertEquals(windowPanes.get(2).getKey().getPaneId(), "1000");
    Assert.assertEquals(windowPanes.get(2).getFiringType(), FiringType.DEFAULT);
    Assert.assertEquals((windowPanes.get(2).getMessage()).size(), 1);
}

Example 23 with TaskCoordinator

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

the class TestWindowOperator method testCancellationOfAnyTrigger.

@Test
public void testCancellationOfAnyTrigger() throws Exception {
    StreamApplication sgb = new KeyedTumblingWindowStreamApplication(AccumulationMode.ACCUMULATING, Duration.ofSeconds(1), Triggers.any(Triggers.count(2), Triggers.timeSinceFirstMessage(Duration.ofMillis(500))));
    TestClock testClock = new TestClock();
    StreamOperatorTask task = new StreamOperatorTask(sgb, runner, testClock);
    task.init(config, taskContext);
    List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
    MessageCollector messageCollector = envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    //assert that the count trigger fired
    Assert.assertEquals(windowPanes.size(), 1);
    //advance the timer to enable the triggering of the inner timeSinceFirstMessage trigger
    testClock.advanceTime(Duration.ofMillis(500));
    //assert that the triggering of the count trigger cancelled the inner timeSinceFirstMessage trigger
    Assert.assertEquals(windowPanes.size(), 1);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    //advance timer by 500 more millis to enable the default trigger
    testClock.advanceTime(Duration.ofMillis(500));
    task.window(messageCollector, taskCoordinator);
    //assert that the default trigger fired
    Assert.assertEquals(windowPanes.size(), 2);
    Assert.assertEquals(windowPanes.get(1).getFiringType(), FiringType.DEFAULT);
    Assert.assertEquals(windowPanes.get(1).getKey().getKey(), new Integer(1));
    Assert.assertEquals(windowPanes.get(1).getKey().getPaneId(), "0");
    Assert.assertEquals((windowPanes.get(1).getMessage()).size(), 5);
    task.process(new IntegerEnvelope(1), messageCollector, taskCoordinator);
    //advance timer by 500 millis to enable the inner timeSinceFirstMessage trigger
    testClock.advanceTime(Duration.ofMillis(500));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 3);
    Assert.assertEquals(windowPanes.get(2).getFiringType(), FiringType.EARLY);
    Assert.assertEquals(windowPanes.get(2).getKey().getKey(), new Integer(1));
    Assert.assertEquals(windowPanes.get(2).getKey().getPaneId(), "1000");
    //advance timer by > 500 millis to enable the default trigger
    testClock.advanceTime(Duration.ofMillis(900));
    task.window(messageCollector, taskCoordinator);
    Assert.assertEquals(windowPanes.size(), 4);
    Assert.assertEquals(windowPanes.get(3).getFiringType(), FiringType.DEFAULT);
    Assert.assertEquals(windowPanes.get(3).getKey().getKey(), new Integer(1));
    Assert.assertEquals(windowPanes.get(3).getKey().getPaneId(), "1000");
}

Example 24 with TaskCoordinator

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

the class TestJoinOperator method joinNoMatchReverse.

@Test
public void joinNoMatchReverse() throws Exception {
    StreamOperatorTask sot = createStreamOperatorTask(new SystemClock(), new TestJoinStreamApplication(new TestJoinFunction()));
    List<Integer> output = new ArrayList<>();
    MessageCollector messageCollector = envelope -> output.add((Integer) envelope.getMessage());
    // push messages to second stream
    numbers.forEach(n -> sot.process(new SecondStreamIME(n, n), messageCollector, taskCoordinator));
    // push messages to first stream with different keys
    numbers.forEach(n -> sot.process(new FirstStreamIME(n + 100, n), messageCollector, taskCoordinator));
    assertTrue(output.isEmpty());
}

Example 25 with TaskCoordinator

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

the class TestMessageStreamImpl method testSink.

@Test
public void testSink() {
    MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph);
    SystemStream testStream = new SystemStream("test-sys", "test-stream");
    SinkFunction<TestMessageEnvelope> xSink = (TestMessageEnvelope m, MessageCollector mc, TaskCoordinator tc) -> {
        mc.send(new OutgoingMessageEnvelope(testStream, m.getMessage()));
        tc.commit(TaskCoordinator.RequestScope.CURRENT_TASK);
    };
    inputStream.sink(xSink);
    Collection<OperatorSpec> subs = inputStream.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> sinkOp = subs.iterator().next();
    assertTrue(sinkOp instanceof SinkOperatorSpec);
    assertEquals(((SinkOperatorSpec) sinkOp).getSinkFn(), xSink);
    TestMessageEnvelope mockTest1 = mock(TestMessageEnvelope.class);
    MessageType mockMsgBody = mock(MessageType.class);
    when(mockTest1.getMessage()).thenReturn(mockMsgBody);
    final List<OutgoingMessageEnvelope> outMsgs = new ArrayList<>();
    MessageCollector mockCollector = mock(MessageCollector.class);
    doAnswer(invocation -> {
        outMsgs.add((OutgoingMessageEnvelope) invocation.getArguments()[0]);
        return null;
    }).when(mockCollector).send(any());
    TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);
    ((SinkOperatorSpec) sinkOp).getSinkFn().apply(mockTest1, mockCollector, mockCoordinator);
    assertEquals(1, outMsgs.size());
    assertEquals(testStream, outMsgs.get(0).getSystemStream());
    assertEquals(mockMsgBody, outMsgs.get(0).getMessage());
}