org.apache.samza.task.TaskCoordinator Java Examples

public Optional<WindowPane<K, Object>> onMessage(TriggerKey<K> triggerKey, M message,
    MessageCollector collector, TaskCoordinator coordinator) {
  if (!isCancelled) {
    LOG.trace("Forwarding callbacks for {}", message);
    impl.onMessage(message, triggerScheduler);

    if (impl.shouldFire()) {
      // repeating trigger can trigger multiple times, So, clear the state to allow future triggerings.
      if (impl instanceof RepeatingTriggerImpl) {
        ((RepeatingTriggerImpl<M, K>) impl).clear();
      }
      return onTriggerFired(triggerKey, collector, coordinator);
    }
  }
  return Optional.empty();
}
 

@Override
public Collection<WindowPane<K, Object>> handleTimer(MessageCollector collector, TaskCoordinator coordinator) {
  LOG.trace("Processing time triggers");
  List<WindowPane<K, Object>> results = new ArrayList<>();
  List<TriggerKey<K>> keys = triggerScheduler.runPendingCallbacks();

  for (TriggerKey<K> key : keys) {
    TriggerImplHandler triggerImplHandler = triggers.get(key);
    if (triggerImplHandler != null) {
      Optional<WindowPane<K, Object>> maybeTriggeredPane = triggerImplHandler.onTimer(key, collector, coordinator);
      maybeTriggeredPane.ifPresent(results::add);
    }
  }
  LOG.trace("Triggered panes: " + results.size());
  return results;
}
 

@Override
protected CompletionStage<Collection<JM>> handleMessageAsync(M message, MessageCollector collector,
    TaskCoordinator coordinator) {
  Collection<JM> output = Collections.emptyList();

  try {
    KeyValueStore<K, TimestampedValue<M>> thisState = thisPartialJoinFn.getState();
    KeyValueStore<K, TimestampedValue<OM>> otherState = otherPartialJoinFn.getState();

    K key = thisPartialJoinFn.getKey(message);
    thisState.put(key, new TimestampedValue<>(message, clock.currentTimeMillis()));
    TimestampedValue<OM> otherMessage = otherState.get(key);

    long now = clock.currentTimeMillis();
    if (otherMessage != null && otherMessage.getTimestamp() > now - ttlMs) {
      JM joinResult = thisPartialJoinFn.apply(message, otherMessage.getValue());
      output = Collections.singletonList(joinResult);
    }
  } catch (Exception e) {
    throw new SamzaException("Error handling message in PartialJoinOperatorImpl " + getOpImplId(), e);
  }

  return CompletableFuture.completedFuture(output);
}
 

@Override
public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) {
  if (envelope.getSystemStreamPartition().getStream().equals("epoch")) {
    int newEpoch = Integer.parseInt((String) envelope.getMessage());
    logger.info("New epoch in message - " + newEpoch);

    Integer epoch = getInt(EPOCH);
    if (epoch == null || newEpoch == epoch)
      return;
    if (newEpoch < epoch)
      throw new IllegalArgumentException("Got new epoch " + newEpoch + " which is less than current epoch " + epoch);

    // it's a new era, reset current epoch and count
    logger.info("Epoch: " + newEpoch);
    this.state.put(EPOCH, Integer.toString(newEpoch));
    this.state.put(COUNT, "0");
    coordinator.commit(RequestScope.ALL_TASKS_IN_CONTAINER);
  }
}
 

@Test
@SuppressWarnings("unchecked")
public void testAsyncFlatMapOperator() {
  AsyncFlatMapOperatorSpec<TestMessageEnvelope, TestOutputMessageEnvelope> mockOp = mock(AsyncFlatMapOperatorSpec.class);
  AsyncFlatMapFunction<TestMessageEnvelope, TestOutputMessageEnvelope> txfmFn = mock(AsyncFlatMapFunction.class);
  when(mockOp.getTransformFn()).thenReturn(txfmFn);
  AsyncFlatmapOperatorImpl<TestMessageEnvelope, TestOutputMessageEnvelope> opImpl =
      new AsyncFlatmapOperatorImpl<>(mockOp);
  TestMessageEnvelope inMsg = mock(TestMessageEnvelope.class);
  Collection<TestOutputMessageEnvelope> mockOutputs = mock(Collection.class);
  when(txfmFn.apply(inMsg)).thenReturn(CompletableFuture.supplyAsync(() -> mockOutputs));
  MessageCollector mockCollector = mock(MessageCollector.class);
  TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);
  Collection<TestOutputMessageEnvelope> results = opImpl
      .handleMessage(inMsg, mockCollector, mockCoordinator);
  verify(txfmFn, times(1)).apply(inMsg);
  assertEquals(results, mockOutputs);
}
 

@Test
public void testOnMessageUpdatesMetrics() {
  ReadableMetricsRegistry mockMetricsRegistry = mock(ReadableMetricsRegistry.class);
  when(this.context.getContainerContext().getContainerMetricsRegistry()).thenReturn(mockMetricsRegistry);
  Counter mockCounter = mock(Counter.class);
  Timer mockTimer = mock(Timer.class);
  when(mockMetricsRegistry.newCounter(anyString(), anyString())).thenReturn(mockCounter);
  when(mockMetricsRegistry.newTimer(anyString(), anyString())).thenReturn(mockTimer);

  Object mockTestOpImplOutput = mock(Object.class);
  OperatorImpl<Object, Object> opImpl = new TestOpImpl(mockTestOpImplOutput);
  opImpl.init(this.internalTaskContext);

  // 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 updates message count and timer metrics
  verify(mockCounter, times(1)).inc();
  verify(mockTimer, times(1)).update(anyLong());
}
 

@Override
protected CompletionStage<Collection<Object>> handleMessageAsync(IncomingMessageEnvelope message,
    MessageCollector collector, TaskCoordinator coordinator) {
  Object result;
  InputTransformer transformer = inputOpSpec.getTransformer();
  if (transformer != null) {
    result = transformer.apply(message);
  } else {
    result = this.inputOpSpec.isKeyed() ? KV.of(message.getKey(), message.getMessage()) : message.getMessage();
  }

  Collection<Object> output = Optional.ofNullable(result)
      .map(Collections::singletonList)
      .orElse(Collections.emptyList());

  return CompletableFuture.completedFuture(output);
}
 

@Override
public void window(MessageCollector collector,
  TaskCoordinator coordinator) {

  KeyValueIterator<String, String> entries = store.all();
  while (entries.hasNext()) {                                        // c
    Entry<String, String> entry = entries.next();
    String key = entry.getKey();
    String value = entry.getValue();
    if (isTimestampKey(key) && Cart.isAbandoned(value)) {            // d
      String shopper = extractShopper(key);
      String cart = store.get(asCartKey(shopper));
      
      AbandonedCartEvent event =
        new AbandonedCartEvent(shopper, cart);
      collector.send(new OutgoingMessageEnvelope(
        new SystemStream("kafka", "derived-events-ch04"), event));    // e
      
      resetShopper(shopper);
    }
  }
}
 

@Test
public void testEndOfStreamFlushesWithNoTriggerFirings() throws Exception {

  OperatorSpecGraph sgb =
      this.getKeyedSessionWindowStreamGraph(AccumulationMode.DISCARDING, Duration.ofMillis(500)).getOperatorSpecGraph();
  TestClock testClock = new TestClock();
  List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
  StreamOperatorTask task = new StreamOperatorTask(sgb, testClock);
  task.init(this.context);

  MessageCollector messageCollector =
    envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());

  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);

  final IncomingMessageEnvelope endOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(
      new SystemStreamPartition("kafka", "integers", new Partition(0)));
  task.processAsync(endOfStream, messageCollector, taskCoordinator, taskCallback);
  Assert.assertEquals(windowPanes.size(), 1);
  Assert.assertEquals(windowPanes.get(0).getMessage().size(), 4);
  verify(taskCoordinator, times(1)).commit(TaskCoordinator.RequestScope.CURRENT_TASK);
  verify(taskCoordinator, times(1)).shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
}
 

/**
 * Invoke handleEndOfStream() if all the input streams to the current operator reach the end.
 * Propagate the end-of-stream to downstream operators.
 * @param collector message collector
 * @param coordinator task coordinator
 */
private CompletionStage<Void> onEndOfStream(MessageCollector collector, TaskCoordinator coordinator) {
  CompletionStage<Void> endOfStreamFuture = CompletableFuture.completedFuture(null);

  if (inputStreams.stream().allMatch(input -> eosStates.isEndOfStream(input))) {
    Collection<RM> results = handleEndOfStream(collector, coordinator);

    CompletionStage<Void> resultFuture = CompletableFuture.allOf(
        results.stream()
            .flatMap(r -> this.registeredOperators.stream()
                .map(op -> op.onMessageAsync(r, collector, coordinator)))
            .toArray(CompletableFuture[]::new));

    endOfStreamFuture = resultFuture.thenCompose(x ->
        CompletableFuture.allOf(this.registeredOperators.stream()
            .map(op -> op.onEndOfStream(collector, coordinator))
            .toArray(CompletableFuture[]::new)));
  }

  return endOfStreamFuture;
}
 

@SuppressWarnings("unchecked")
@Override
public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) {
  Map<String, Object> edit = (Map<String, Object>) envelope.getMessage();
  Map<String, Boolean> flags = (Map<String, Boolean>) edit.get("flags");

  Integer editsAllTime = store.get("count-edits-all-time");
  if (editsAllTime == null) editsAllTime = 0;
  store.put("count-edits-all-time", editsAllTime + 1);

  edits += 1;
  byteDiff += (Integer) edit.get("diff-bytes");
  boolean newTitle = titles.add((String) edit.get("title"));

  for (Map.Entry<String, Boolean> flag : flags.entrySet()) {
    if (Boolean.TRUE.equals(flag.getValue())) {
      counts.compute(flag.getKey(), (k, v) -> v == null ? 0 : v + 1);
    }
  }

  if (!newTitle) {
    repeatEdits.inc();
  }
}
 

/**
 * Aggregate the {@link WatermarkMessage} from each ssp into a watermark. Then call onWatermark() if
 * a new watermark exits.
 * @param watermarkMessage a {@link WatermarkMessage} object
 * @param ssp {@link SystemStreamPartition} that the message is coming from.
 * @param collector message collector
 * @param coordinator task coordinator
 */
public final CompletionStage<Void> aggregateWatermark(WatermarkMessage watermarkMessage, SystemStreamPartition ssp,
    MessageCollector collector, TaskCoordinator coordinator) {
  LOG.debug("Received watermark {} from {}", watermarkMessage.getTimestamp(), ssp);
  watermarkStates.update(watermarkMessage, ssp);
  long watermark = watermarkStates.getWatermark(ssp.getSystemStream());
  CompletionStage<Void> watermarkFuture = CompletableFuture.completedFuture(null);

  if (currentWatermark < watermark) {
    LOG.debug("Got watermark {} from stream {}", watermark, ssp.getSystemStream());

    if (watermarkMessage.getTaskName() != null) {
      // This is the aggregation task, which already received all the watermark messages from upstream
      // broadcast the watermark to all the peer partitions
      controlMessageSender.broadcastToOtherPartitions(new WatermarkMessage(watermark), ssp, collector);
    }
    // populate the watermark through the dag
    watermarkFuture = onWatermark(watermark, collector, coordinator)
        .thenAccept(ignored -> watermarkStates.updateAggregateMetric(ssp, watermark));
  }

  return watermarkFuture;
}
 

private CompletionStage<Void> propagateWatermark(Long outputWm, MessageCollector collector, TaskCoordinator coordinator) {
  CompletionStage<Void> watermarkFuture = CompletableFuture.completedFuture(null);

  if (outputWm != null) {
    if (outputWatermark < outputWm) {
      // advance the watermark
      outputWatermark = outputWm;
      LOG.debug("Advance output watermark to {} in operator {}", outputWatermark, getOpImplId());
      watermarkFuture = CompletableFuture.allOf(
          this.registeredOperators
              .stream()
              .map(op -> op.onWatermark(outputWatermark, collector, coordinator))
              .toArray(CompletableFuture[]::new));
    } else if (outputWatermark > outputWm) {
      LOG.warn("Ignore watermark {} that is smaller than the previous watermark {}.", outputWm, outputWatermark);
    }
  }

  return watermarkFuture;
}
 

@Test
public void testSinkOperatorClose() {
  TestOutputMessageEnvelope mockMsg = mock(TestOutputMessageEnvelope.class);
  MessageCollector mockCollector = mock(MessageCollector.class);
  TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);
  SinkFunction<TestOutputMessageEnvelope> sinkFn = mock(SinkFunction.class);

  SinkOperatorImpl<TestOutputMessageEnvelope> sinkImpl = createSinkOperator(sinkFn);
  sinkImpl.handleMessage(mockMsg, mockCollector, mockCoordinator);
  verify(sinkFn, times(1)).apply(mockMsg, mockCollector, mockCoordinator);

  // ensure that close is not called yet
  verify(sinkFn, times(0)).close();

  sinkImpl.handleClose();
  // ensure that close is called once from handleClose()
  verify(sinkFn, times(1)).close();
}
 

@Override
public void window(MessageCollector collector, TaskCoordinator coordinator) {
  for (int i = 0; i < 100 && messagesSent < NewsfeedConfig.NUM_FOLLOW_EVENTS; i++, messagesSent++) {
    String follower = NewsfeedConfig.randomUser();
    String followee = NewsfeedConfig.randomUser();

    HashMap<String, Object> message = new HashMap<String, Object>();
    message.put("event", "follow");
    message.put("follower", follower);
    message.put("followee", followee);
    message.put("time", NewsfeedConfig.currentDateTime());
    collector.send(new OutgoingMessageEnvelope(NewsfeedConfig.FOLLOWS_STREAM, followee, null, message));
  }

  if (messagesSent % 100000 == 0) {
    log.info("Generated " + messagesSent + " follow events");
  }

  if (messagesSent == NewsfeedConfig.NUM_FOLLOW_EVENTS) {
    log.info("Finished generating random follower graph");
    coordinator.shutdown(RequestScope.CURRENT_TASK);
  }
}
 

private void endOfStream() {
  state.complete = true;
  try {
    ReadableCoordinator coordinator = new ReadableCoordinator(task.taskName());

    task.endOfStream(coordinator);
    // issue a request for shutdown of the task
    coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
    coordinatorRequests.update(coordinator);

    // invoke commit on the task - if the endOfStream callback had requested a final commit.
    boolean needFinalCommit = coordinatorRequests.commitRequests().remove(task.taskName());
    if (needFinalCommit) {
      task.commit();
    }
  } finally {
    resume();
  }

}
 

@Override
public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) {
  String key = (String) envelope.getKey();
  String value = (String) envelope.getMessage();
  String[] pieces = value.split("-");
  int epoch = Integer.parseInt(pieces[0]);

  int partition = Integer.parseInt(pieces[1].split(" ")[1]);
  Partitions partitions = loadPartitions(epoch, key);
  logger.info("Joiner got epoch = " + epoch + ", partition = " + partition + ", parts = " + partitions);
  if (partitions.epoch < epoch) {
    // we are in a new era
    if (partitions.partitions.size() != expected)
      throw new IllegalArgumentException("Should have " + expected + " partitions when new epoch starts.");
    logger.info("Reseting epoch to " + epoch);
    this.store.delete(key);
    partitions.epoch = epoch;
    partitions.partitions.clear();
    partitions.partitions.add(partition);
  } else if (partitions.epoch > epoch) {
    logger.info("Ignoring message for epoch " + epoch);
  } else {
    partitions.partitions.add(partition);
    if (partitions.partitions.size() == expected) {
      logger.info("Completed: " + key + " -> " + Integer.toString(epoch));
      collector.send(new OutgoingMessageEnvelope(new SystemStream("kafka", "completed-keys"), key, Integer.toString(epoch)));
    }
  }
  this.store.put(key, partitions.toString());
  logger.info("Join store in Task " + this.taskName + " " + key + " -> " + partitions.toString());
}
 

@Test
public void testOnMessagePropagatesResults() {
  Object mockTestOpImplOutput = mock(Object.class);
  OperatorImpl<Object, Object> opImpl = new TestOpImpl(mockTestOpImplOutput);
  opImpl.init(this.internalTaskContext);

  // register a couple of operators
  OperatorImpl mockNextOpImpl1 = mock(OperatorImpl.class);
  when(mockNextOpImpl1.getOperatorSpec()).thenReturn(new TestOpSpec());
  when(mockNextOpImpl1.handleMessageAsync(anyObject(), anyObject(), anyObject()))
      .thenReturn(CompletableFuture.completedFuture(Collections.emptyList()));
  mockNextOpImpl1.init(this.internalTaskContext);
  opImpl.registerNextOperator(mockNextOpImpl1);

  OperatorImpl mockNextOpImpl2 = mock(OperatorImpl.class);
  when(mockNextOpImpl2.getOperatorSpec()).thenReturn(new TestOpSpec());
  when(mockNextOpImpl2.handleMessageAsync(anyObject(), anyObject(), anyObject()))
      .thenReturn(CompletableFuture.completedFuture(Collections.emptyList()));
  mockNextOpImpl2.init(this.internalTaskContext);
  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)).handleMessageAsync(mockTestOpImplOutput, mockCollector, mockCoordinator);
  verify(mockNextOpImpl2, times(1)).handleMessageAsync(mockTestOpImplOutput, mockCollector, mockCoordinator);
}
 

@Test
public void testOnTimerPropagatesResultsAndTimer() {
  Object mockTestOpImplOutput = mock(Object.class);
  OperatorImpl<Object, Object> opImpl = new TestOpImpl(mockTestOpImplOutput);
  opImpl.init(this.internalTaskContext);

  // register a couple of operators
  OperatorImpl mockNextOpImpl1 = mock(OperatorImpl.class);
  when(mockNextOpImpl1.getOperatorSpec()).thenReturn(new TestOpSpec());
  when(mockNextOpImpl1.handleMessageAsync(anyObject(), anyObject(), anyObject()))
      .thenReturn(CompletableFuture.completedFuture(Collections.emptyList()));
  mockNextOpImpl1.init(this.internalTaskContext);
  opImpl.registerNextOperator(mockNextOpImpl1);

  OperatorImpl mockNextOpImpl2 = mock(OperatorImpl.class);
  when(mockNextOpImpl2.getOperatorSpec()).thenReturn(new TestOpSpec());
  when(mockNextOpImpl2.handleMessageAsync(anyObject(), anyObject(), anyObject()))
      .thenReturn(CompletableFuture.completedFuture(Collections.emptyList()));
  mockNextOpImpl2.init(this.internalTaskContext);
  opImpl.registerNextOperator(mockNextOpImpl2);

  // send a timer tick to this operator
  MessageCollector mockCollector = mock(MessageCollector.class);
  TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);
  opImpl.onTimer(mockCollector, mockCoordinator);

  // verify that it propagates its handleTimer results to next operators
  verify(mockNextOpImpl1, times(1)).handleMessageAsync(mockTestOpImplOutput, mockCollector, mockCoordinator);
  verify(mockNextOpImpl2, times(1)).handleMessageAsync(mockTestOpImplOutput, mockCollector, mockCoordinator);

  // verify that it propagates the timer tick to next operators
  verify(mockNextOpImpl1, times(1)).handleTimer(mockCollector, mockCoordinator);
  verify(mockNextOpImpl2, times(1)).handleTimer(mockCollector, mockCoordinator);
}
 

@SuppressWarnings("unchecked")
@Override
public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) {
  Map<String, Object> jsonObject = (Map<String, Object>) envelope.getMessage();
  WikipediaFeedEvent event = new WikipediaFeedEvent(jsonObject);

  Map<String, Object> parsedJsonObject = WikipediaParser.parseEvent(event);

  if (parsedJsonObject != null) {
    collector.send(new OutgoingMessageEnvelope(OUTPUT_STREAM, parsedJsonObject));
  }
}
 

@Override
public void window(MessageCollector collector, TaskCoordinator coordinator) {
  for (int i = 0; i < NewsfeedConfig.MESSAGES_PER_WINDOW; i++) {
    String sender = NewsfeedConfig.randomUser();

    HashMap<String, Object> message = new HashMap<String, Object>();
    message.put("event", "postMessage");
    message.put("sender", sender);
    message.put("text", "Hello world");
    message.put("time", NewsfeedConfig.currentDateTime());
    collector.send(new OutgoingMessageEnvelope(NewsfeedConfig.MESSAGES_STREAM, sender, null, message));
  }
}
 

@Override
public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) throws Exception {
	for (SamzaStream stream:this.outputStreams) {
		stream.setCollector(collector);
	}
	this.getProcessor().process((ContentEvent) envelope.getMessage());
}
 

@Test
public void testEndOfStreamFlushesWithDefaultTriggerFirings() throws Exception {
  OperatorSpecGraph sgb =
      this.getKeyedSessionWindowStreamGraph(AccumulationMode.DISCARDING, Duration.ofMillis(500)).getOperatorSpecGraph();
  TestClock testClock = new TestClock();
  List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
  StreamOperatorTask task = new StreamOperatorTask(sgb, testClock);
  task.init(this.context);

  MessageCollector messageCollector =
    envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());

  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  testClock.advanceTime(1000);
  task.window(messageCollector, taskCoordinator);
  Assert.assertEquals(windowPanes.size(), 1);

  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);

  final IncomingMessageEnvelope endOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(
      new SystemStreamPartition("kafka", "integers", new Partition(0)));
  task.processAsync(endOfStream, messageCollector, taskCoordinator, taskCallback);
  Assert.assertEquals(windowPanes.size(), 2);
  Assert.assertEquals(windowPanes.get(0).getMessage().size(), 2);
  verify(taskCoordinator, times(1)).commit(TaskCoordinator.RequestScope.CURRENT_TASK);
  verify(taskCoordinator, times(1)).shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
}
 

@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();
}
 

@Test
public void testEndOfStreamFlushesWithEarlyTriggerFirings() throws Exception {

  OperatorSpecGraph sgb = this.getTumblingWindowStreamGraph(AccumulationMode.DISCARDING,
      Duration.ofSeconds(1), Triggers.repeat(Triggers.count(2))).getOperatorSpecGraph();
  List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();

  TestClock testClock = new TestClock();
  StreamOperatorTask task = new StreamOperatorTask(sgb, testClock);
  task.init(this.context);

  MessageCollector messageCollector =
    envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());
  Assert.assertEquals(windowPanes.size(), 0);

  List<Integer> integerList = ImmutableList.of(1, 2, 1, 2, 1);
  integerList.forEach(n -> task.processAsync(new IntegerEnvelope(n), messageCollector, taskCoordinator, taskCallback));

  // early triggers should emit (1,2) and (1,2) in the same window.
  Assert.assertEquals(windowPanes.size(), 2);

  testClock.advanceTime(Duration.ofSeconds(1));
  Assert.assertEquals(windowPanes.size(), 2);

  final IncomingMessageEnvelope endOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(
      new SystemStreamPartition("kafka", "integers", new Partition(0)));
  task.processAsync(endOfStream, messageCollector, taskCoordinator, taskCallback);

  // end of stream flushes the last entry (1)
  Assert.assertEquals(windowPanes.size(), 3);
  Assert.assertEquals((windowPanes.get(0).getMessage()).size(), 2);
  verify(taskCoordinator, times(1)).commit(TaskCoordinator.RequestScope.CURRENT_TASK);
  verify(taskCoordinator, times(1)).shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
}
 

@Test
public void testEndOfStreamCommitBehavior() {
  SystemConsumers consumerMultiplexer = mock(SystemConsumers.class);

  RunLoopTask task0 = getMockRunLoopTask(taskName0, ssp0);
  doAnswer(invocation -> {
    ReadableCoordinator coordinator = invocation.getArgumentAt(0, ReadableCoordinator.class);

    coordinator.commit(TaskCoordinator.RequestScope.CURRENT_TASK);
    return null;
  }).when(task0).endOfStream(any());

  Map<TaskName, RunLoopTask> tasks = new HashMap<>();

  tasks.put(taskName0, task0);

  int maxMessagesInFlight = 1;
  RunLoop runLoop = new RunLoop(tasks, executor, consumerMultiplexer, maxMessagesInFlight, windowMs, commitMs,
                                          callbackTimeoutMs, maxThrottlingDelayMs, maxIdleMs, containerMetrics, () -> 0L, false);
  when(consumerMultiplexer.choose(false)).thenReturn(envelope00).thenReturn(ssp0EndOfStream).thenReturn(null);

  runLoop.run();

  InOrder inOrder = inOrder(task0);

  inOrder.verify(task0).endOfStream(any());
  inOrder.verify(task0).commit();
}
 

@Test
public void testSinkOperatorSinkFunction() {
  SinkFunction<TestOutputMessageEnvelope> sinkFn = mock(SinkFunction.class);
  SinkOperatorImpl<TestOutputMessageEnvelope> sinkImpl = createSinkOperator(sinkFn);
  TestOutputMessageEnvelope mockMsg = mock(TestOutputMessageEnvelope.class);
  MessageCollector mockCollector = mock(MessageCollector.class);
  TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);

  sinkImpl.handleMessage(mockMsg, mockCollector, mockCoordinator);
  verify(sinkFn, times(1)).apply(mockMsg, mockCollector, mockCoordinator);
}
 

@SuppressWarnings("unchecked")
@Override
public void process(IncomingMessageEnvelope envelope,
  MessageCollector collector, TaskCoordinator coordinator) {

  Map<String, Object> event =
    (Map<String, Object>) envelope.getMessage();
  String verb = (String) event.get("verb");
  String shopper = (String) ((Map<String, Object>)
    event.get("subject")).get("shopper");
  
  if (verb.equals("add")) {                                          // a
    String timestamp = (String) ((Map<String, Object>)
      event.get("context")).get("timestamp");

    Map<String, Object> item = (Map<String, Object>)
      ((Map<String, Object>) event.get("directObject")).get("item");
    Cart cart = new Cart(store.get(asCartKey(shopper)));
    cart.addItem(item);

    store.put(asTimestampKey(shopper), timestamp);
    store.put(asCartKey(shopper), cart.asJson());
  
  } else if (verb.equals("place")) {                                 // b
    resetShopper(shopper);
  }
}
 

public void window(MessageCollector collector, TaskCoordinator coordinator) {
  Integer epoch = getInt(EPOCH);
  if (epoch == null) {
    resetEpoch();
    return;
  }
  int counter = getInt(COUNT);
  if (counter < max) {
    logger.info("Emitting: " + counter + ", epoch = " + epoch + ", task = " + taskName);
    OutgoingMessageEnvelope envelope = new OutgoingMessageEnvelope(new SystemStream("kafka", "emitted"), Integer.toString(counter), epoch + "-" + taskName.toString());
    collector.send(envelope);
    this.state.put(COUNT, Integer.toString(getInt(COUNT) + 1));
  }
}
 

public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) {
  store.put((String) envelope.getMessage(), (String) envelope.getMessage());
  count++;
  if (count % LOG_INTERVAL == 0) {
    double ellapsedSecs = (System.currentTimeMillis() - start) / 1000.0;
    System.out.println(String.format("Throughput = %.2f messages/sec.", count / ellapsedSecs));
    start = System.currentTimeMillis();
    count = 0;
    coordinator.commit(RequestScope.ALL_TASKS_IN_CONTAINER);
  }
}