Java TaskCoordinator类代码示例

import org.apache.samza.task.TaskCoordinator; //导入依赖的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.TaskCoordinator; //导入依赖的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.TaskCoordinator; //导入依赖的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.TaskCoordinator; //导入依赖的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.TaskCoordinator; //导入依赖的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.TaskCoordinator; //导入依赖的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.TaskCoordinator; //导入依赖的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.TaskCoordinator; //导入依赖的package包/类
@Test
public void testOnMessageUpdatesMetrics() {
  TaskContextImpl mockTaskContext = mock(TaskContextImpl.class);
  ReadableMetricsRegistry mockMetricsRegistry = mock(ReadableMetricsRegistry.class);
  when(mockTaskContext.getMetricsRegistry()).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(mock(Config.class), mockTaskContext);

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

import org.apache.samza.task.TaskCoordinator; //导入依赖的package包/类
@Test
public void testOnTimerUpdatesMetrics() {
  TaskContextImpl mockTaskContext = mock(TaskContextImpl.class);
  ReadableMetricsRegistry mockMetricsRegistry = mock(ReadableMetricsRegistry.class);
  when(mockTaskContext.getMetricsRegistry()).thenReturn(mockMetricsRegistry);
  Counter mockMessageCounter = mock(Counter.class);
  Timer mockTimer = mock(Timer.class);
  when(mockMetricsRegistry.newCounter(anyString(), anyString())).thenReturn(mockMessageCounter);
  when(mockMetricsRegistry.newTimer(anyString(), anyString())).thenReturn(mockTimer);

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

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

  // verify that it updates metrics
  verify(mockMessageCounter, times(0)).inc();
  verify(mockTimer, times(1)).update(anyLong());
}
 

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

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

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

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

import org.apache.samza.task.TaskCoordinator; //导入依赖的package包/类
@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;
  titles.add((String) edit.get("title"));
  byteDiff += (Integer) edit.get("diff-bytes");

  for (Map.Entry<String, Boolean> flag : flags.entrySet()) {
    if (Boolean.TRUE.equals(flag.getValue())) {
      Integer count = counts.get(flag.getKey());

      if (count == null) {
        count = 0;
      }

      count += 1;
      counts.put(flag.getKey(), count);
    }
  }
}
 

import org.apache.samza.task.TaskCoordinator; //导入依赖的package包/类
@Override
public void window(MessageCollector collector, TaskCoordinator coordinator) {
  counts.put("edits", edits);
  counts.put("bytes-added", byteDiff);
  counts.put("unique-titles", titles.size());
  counts.put("edits-all-time", store.get("count-edits-all-time"));

  collector.send(new OutgoingMessageEnvelope(new SystemStream("kafka", "wikipedia-stats"), counts));

  // Reset counts after windowing.
  edits = 0;
  byteDiff = 0;
  titles = new HashSet<String>();
  counts = new HashMap<String, Integer>();
}