Java MessageCollector类代码示例

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@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);

  try {
    Map<String, Object> parsedJsonObject = parse(event.getRawEvent());

    parsedJsonObject.put("channel", event.getChannel());
    parsedJsonObject.put("source", event.getSource());
    parsedJsonObject.put("time", event.getTime());

    collector.send(new OutgoingMessageEnvelope(new SystemStream("kafka", "wikipedia-edits"), parsedJsonObject));
  } catch (Exception e) {
    System.err.println("Unable to parse line: " + event);
  }
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Override
public Collection<JM> handleMessage(M message, MessageCollector collector, TaskCoordinator coordinator) {
  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());
      return Collections.singletonList(joinResult);
    }
  } catch (Exception e) {
    throw new SamzaException("Error handling message in PartialJoinOperatorImpl " + getOpImplId(), e);
  }
  return Collections.emptyList();
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Override
public Collection<WindowPane<K, Object>> handleTimer(MessageCollector collector, TaskCoordinator coordinator) {
  LOG.trace("Processing timer.");
  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;
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
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();
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
/**
 * Aggregate {@link EndOfStreamMessage} from each ssp of the stream.
 * Invoke onEndOfStream() if the stream reaches the end.
 * @param eos {@link EndOfStreamMessage} object
 * @param ssp system stream partition
 * @param collector message collector
 * @param coordinator task coordinator
 */
public final void aggregateEndOfStream(EndOfStreamMessage eos, SystemStreamPartition ssp, MessageCollector collector,
    TaskCoordinator coordinator) {
  LOG.info("Received end-of-stream message from task {} in {}", eos.getTaskName(), ssp);
  eosStates.update(eos, ssp);

  SystemStream stream = ssp.getSystemStream();
  if (eosStates.isEndOfStream(stream)) {
    LOG.info("Input {} reaches the end for task {}", stream.toString(), taskName.getTaskName());
    onEndOfStream(collector, coordinator);

    if (eosStates.allEndOfStream()) {
      // all inputs have been end-of-stream, shut down the task
      LOG.info("All input streams have reached the end for task {}", taskName.getTaskName());
      coordinator.commit(TaskCoordinator.RequestScope.CURRENT_TASK);
      coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
    }
  }
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
public void joinFnInitAndClose() throws Exception {
  TestJoinFunction joinFn = new TestJoinFunction();
  TestJoinStreamApplication app = new TestJoinStreamApplication(joinFn);
  StreamOperatorTask sot = createStreamOperatorTask(new SystemClock(), app);
  assertEquals(1, joinFn.getNumInitCalls());
  MessageCollector messageCollector = mock(MessageCollector.class);

  // push messages to first stream
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));

  // close should not be called till now
  assertEquals(0, joinFn.getNumCloseCalls());
  sot.close();

  // close should be called from sot.close()
  assertEquals(1, joinFn.getNumCloseCalls());
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
public void joinRetainsLatestMessageForKey() throws Exception {
  TestJoinStreamApplication app = new TestJoinStreamApplication(new TestJoinFunction());
  StreamOperatorTask sot = createStreamOperatorTask(new SystemClock(), app);
  List<Integer> output = new ArrayList<>();
  MessageCollector messageCollector = envelope -> output.add((Integer) envelope.getMessage());

  // push messages to first stream
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));
  // push messages to first stream again with same keys but different values
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, 2 * n), messageCollector, taskCoordinator));
  // push messages to second stream with same key
  numbers.forEach(n -> sot.process(new SecondStreamIME(n, n), messageCollector, taskCoordinator));

  int outputSum = output.stream().reduce(0, (s, m) -> s + m);
  assertEquals(165, outputSum); // should use latest messages in the first stream
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
public void joinRetainsLatestMessageForKeyReverse() throws Exception {
  TestJoinStreamApplication app = new TestJoinStreamApplication(new TestJoinFunction());
  StreamOperatorTask sot = createStreamOperatorTask(new SystemClock(), app);
  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 second stream again with same keys but different values
  numbers.forEach(n -> sot.process(new SecondStreamIME(n, 2 * n), messageCollector, taskCoordinator));
  // push messages to first stream with same key
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));

  int outputSum = output.stream().reduce(0, (s, m) -> s + m);
  assertEquals(165, outputSum); // should use latest messages in the second stream
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
public void joinRetainsMatchedMessages() throws Exception {
  TestJoinStreamApplication app = new TestJoinStreamApplication(new TestJoinFunction());
  StreamOperatorTask sot = createStreamOperatorTask(new SystemClock(), app);
  List<Integer> output = new ArrayList<>();
  MessageCollector messageCollector = envelope -> output.add((Integer) envelope.getMessage());

  // push messages to first stream
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));
  // push messages to second stream with same key
  numbers.forEach(n -> sot.process(new SecondStreamIME(n, n), messageCollector, taskCoordinator));

  int outputSum = output.stream().reduce(0, (s, m) -> s + m);
  assertEquals(110, outputSum);

  output.clear();

  // push messages to first stream with same keys once again.
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));
  int newOutputSum = output.stream().reduce(0, (s, m) -> s + m);
  assertEquals(110, newOutputSum); // should produce the same output as before
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
public void joinRetainsMatchedMessagesReverse() throws Exception {
  TestJoinStreamApplication app = new TestJoinStreamApplication(new TestJoinFunction());
  StreamOperatorTask sot = createStreamOperatorTask(new SystemClock(), app);
  List<Integer> output = new ArrayList<>();
  MessageCollector messageCollector = envelope -> output.add((Integer) envelope.getMessage());

  // push messages to first stream
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));
  // push messages to second stream with same key
  numbers.forEach(n -> sot.process(new SecondStreamIME(n, n), messageCollector, taskCoordinator));

  int outputSum = output.stream().reduce(0, (s, m) -> s + m);
  assertEquals(110, outputSum);

  output.clear();

  // push messages to second stream with same keys once again.
  numbers.forEach(n -> sot.process(new SecondStreamIME(n, n), messageCollector, taskCoordinator));
  int newOutputSum = output.stream().reduce(0, (s, m) -> s + m);
  assertEquals(110, newOutputSum); // should produce the same output as before
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
public void joinRemovesExpiredMessages() throws Exception {
  TestClock testClock = new TestClock();
  TestJoinStreamApplication app = new TestJoinStreamApplication(new TestJoinFunction());
  StreamOperatorTask sot = createStreamOperatorTask(testClock, app);
  List<Integer> output = new ArrayList<>();
  MessageCollector messageCollector = envelope -> output.add((Integer) envelope.getMessage());

  // push messages to first stream
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));

  testClock.advanceTime(JOIN_TTL.plus(Duration.ofMinutes(1))); // 1 minute after ttl
  sot.window(messageCollector, taskCoordinator); // should expire first stream messages

  // push messages to second stream with same key
  numbers.forEach(n -> sot.process(new SecondStreamIME(n, n), messageCollector, taskCoordinator));

  assertTrue(output.isEmpty());
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
public void joinRemovesExpiredMessagesReverse() throws Exception {
  TestClock testClock = new TestClock();
  TestJoinStreamApplication app = new TestJoinStreamApplication(new TestJoinFunction());
  StreamOperatorTask sot = createStreamOperatorTask(testClock, app);
  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));

  testClock.advanceTime(JOIN_TTL.plus(Duration.ofMinutes(1))); // 1 minute after ttl
  sot.window(messageCollector, taskCoordinator); // should expire second stream messages

  // push messages to first stream with same key
  numbers.forEach(n -> sot.process(new FirstStreamIME(n, n), messageCollector, taskCoordinator));

  assertTrue(output.isEmpty());
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
@SuppressWarnings("unchecked")
public void testStreamOperator() {
  StreamOperatorSpec<TestMessageEnvelope, TestOutputMessageEnvelope> mockOp = mock(StreamOperatorSpec.class);
  FlatMapFunction<TestMessageEnvelope, TestOutputMessageEnvelope> txfmFn = mock(FlatMapFunction.class);
  when(mockOp.getTransformFn()).thenReturn(txfmFn);
  Config mockConfig = mock(Config.class);
  TaskContext mockContext = mock(TaskContext.class);
  StreamOperatorImpl<TestMessageEnvelope, TestOutputMessageEnvelope> opImpl =
      new StreamOperatorImpl<>(mockOp, mockConfig, mockContext);
  TestMessageEnvelope inMsg = mock(TestMessageEnvelope.class);
  Collection<TestOutputMessageEnvelope> mockOutputs = mock(Collection.class);
  when(txfmFn.apply(inMsg)).thenReturn(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);
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@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();
}
 

import org.apache.samza.task.MessageCollector; //导入依赖的package包/类
@Test
public void testNonKeyedTumblingWindowsDiscardingMode() throws Exception {

  StreamApplication sgb = new TumblingWindowStreamApplication(AccumulationMode.DISCARDING,
      Duration.ofSeconds(1), Triggers.repeat(Triggers.count(1000)));
  List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();

  TestClock testClock = new TestClock();
  StreamOperatorTask task = new StreamOperatorTask(sgb, runner, testClock);
  task.init(config, taskContext);

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

  integers.forEach(n -> task.process(new IntegerEnvelope(n), messageCollector, taskCoordinator));
  Assert.assertEquals(windowPanes.size(), 0);

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

  task.window(messageCollector, taskCoordinator);
  Assert.assertEquals(windowPanes.size(), 1);
  Assert.assertEquals((windowPanes.get(0).getMessage()).size(), 9);
}