Java Code Examples for org.apache.samza.system.SystemConsumer#poll()

@Test
public void testSamzaImpulseSystemConsumer() throws Exception {
  SystemConsumer consumer =
      new SamzaImpulseSystemFactory().getConsumer("default-system", new MapConfig(), null);
  Map<SystemStreamPartition, List<IncomingMessageEnvelope>> result =
      consumer.poll(Collections.singleton(sspForPartition(0)), 100);
  Assert.assertEquals(1, result.size());
  Assert.assertTrue(result.containsKey(sspForPartition(0)));

  List<IncomingMessageEnvelope> messageEnvelopes = result.get(sspForPartition(0));
  Assert.assertEquals(3, messageEnvelopes.size());

  Assert.assertTrue(messageEnvelopes.get(0).getMessage() instanceof OpMessage);
  OpMessage impulseEvent = (OpMessage) messageEnvelopes.get(0).getMessage();
  Assert.assertEquals(OpMessage.Type.ELEMENT, impulseEvent.getType());

  Assert.assertTrue(messageEnvelopes.get(1).getMessage() instanceof WatermarkMessage);

  Assert.assertTrue(messageEnvelopes.get(2).isEndOfStream());
}
 

private List<IncomingMessageEnvelope> consumeRawMessages(SystemConsumer consumer, Set<SystemStreamPartition> sspsToPoll) {
  try {
    Map<SystemStreamPartition, List<IncomingMessageEnvelope>> results = consumer.poll(sspsToPoll, POLL_TIMEOUT_MS);

    return results.entrySet()
        .stream()
        .filter(entry -> entry.getValue().size() != 0)
        .map(Map.Entry::getValue)
        .flatMap(List::stream)
        .collect(Collectors.toList());
  } catch (Exception e) {
    fail("Unable to consume messages");
  }

  return new ArrayList<>();
}
 

private static List<IncomingMessageEnvelope> pollOnce(
    SystemConsumer consumer, SystemStreamPartition ssp, long timeoutMillis)
    throws InterruptedException {
  final Set<SystemStreamPartition> sspSet = Collections.singleton(ssp);
  final Map<SystemStreamPartition, List<IncomingMessageEnvelope>> pollResult =
      consumer.poll(sspSet, timeoutMillis);
  assertEquals(sspSet, pollResult.keySet());
  assertNotNull(pollResult.get(ssp));
  return pollResult.get(ssp);
}
 

private static List<IncomingMessageEnvelope> pollOnce(
    SystemConsumer consumer, SystemStreamPartition ssp, long timeoutMillis)
    throws InterruptedException {
  final Set<SystemStreamPartition> sspSet = Collections.singleton(ssp);
  final Map<SystemStreamPartition, List<IncomingMessageEnvelope>> pollResult =
      consumer.poll(sspSet, timeoutMillis);
  assertEquals(sspSet, pollResult.keySet());
  assertNotNull(pollResult.get(ssp));
  return pollResult.get(ssp);
}
 

public void start() throws IOException, InterruptedException {
  super.start();
  SystemAdmin systemAdmin = factory.getAdmin(systemName, config);
  SystemStreamMetadata ssm =
      systemAdmin.getSystemStreamMetadata(Collections.singleton(physicalStreamName)).get(physicalStreamName);

  NoOpMetricsRegistry metricsRegistry = new NoOpMetricsRegistry();
  Set<SystemStreamPartition> ssps = createSSPs(systemName, physicalStreamName, startPartition, endPartition);
  SystemConsumer consumer = factory.getConsumer(systemName, config, metricsRegistry);
  for (SystemStreamPartition ssp : ssps) {
    consumer.register(ssp, ssm.getSystemStreamPartitionMetadata().get(ssp.getPartition()).getOldestOffset());
  }

  consumer.start();

  System.out.println("starting consumption at " + Instant.now());
  Instant startTime = Instant.now();
  int numEvents = 0;
  while (numEvents < totalEvents) {
    Map<SystemStreamPartition, List<IncomingMessageEnvelope>> pollResult = consumer.poll(ssps, 2000);
    numEvents += pollResult.values().stream().mapToInt(List::size).sum();
  }

  System.out.println("Ending consumption at " + Instant.now());
  System.out.println(String.format("Event Rate is %s Messages/Sec ",
      numEvents * 1000 / Duration.between(startTime, Instant.now()).toMillis()));
  consumer.stop();
  System.exit(0);
}
 

/**
 * Gets the contents of the output stream represented by {@code outputDescriptor} after {@link TestRunner#run(Duration)}
 * has completed
 *
 * @param outputDescriptor describes the stream to be consumed
 * @param timeout timeout for consumption of stream in Ms
 * @param <StreamMessageType> type of message
 *
 * @return a map whose key is {@code partitionId} and value is messages in partition
 * @throws SamzaException Thrown when a poll is incomplete
 */
public static <StreamMessageType> Map<Integer, List<StreamMessageType>> consumeStream(
    InMemoryOutputDescriptor outputDescriptor, Duration timeout) throws SamzaException {
  Preconditions.checkNotNull(outputDescriptor);
  String streamId = outputDescriptor.getStreamId();
  String systemName = outputDescriptor.getSystemName();
  Set<SystemStreamPartition> ssps = new HashSet<>();
  Set<String> streamIds = new HashSet<>();
  streamIds.add(streamId);
  SystemFactory factory = new InMemorySystemFactory();
  Config config = new MapConfig(outputDescriptor.toConfig(), outputDescriptor.getSystemDescriptor().toConfig());
  Map<String, SystemStreamMetadata> metadata = factory.getAdmin(systemName, config).getSystemStreamMetadata(streamIds);
  SystemConsumer consumer = factory.getConsumer(systemName, config, null);
  String name = (String) outputDescriptor.getPhysicalName().orElse(streamId);
  metadata.get(name).getSystemStreamPartitionMetadata().keySet().forEach(partition -> {
    SystemStreamPartition temp = new SystemStreamPartition(systemName, streamId, partition);
    ssps.add(temp);
    consumer.register(temp, "0");
  });

  long t = System.currentTimeMillis();
  Map<SystemStreamPartition, List<IncomingMessageEnvelope>> output = new HashMap<>();
  HashSet<SystemStreamPartition> didNotReachEndOfStream = new HashSet<>(ssps);
  while (System.currentTimeMillis() < t + timeout.toMillis()) {
    Map<SystemStreamPartition, List<IncomingMessageEnvelope>> currentState = null;
    try {
      currentState = consumer.poll(ssps, 10);
    } catch (InterruptedException e) {
      throw new SamzaException("Timed out while consuming stream \n" + e.getMessage());
    }
    for (Map.Entry<SystemStreamPartition, List<IncomingMessageEnvelope>> entry : currentState.entrySet()) {
      SystemStreamPartition ssp = entry.getKey();
      output.computeIfAbsent(ssp, k -> new LinkedList<IncomingMessageEnvelope>());
      List<IncomingMessageEnvelope> currentBuffer = entry.getValue();
      int totalMessagesToFetch = Integer.valueOf(metadata.get(outputDescriptor.getStreamId())
          .getSystemStreamPartitionMetadata()
          .get(ssp.getPartition())
          .getUpcomingOffset());
      if (output.get(ssp).size() + currentBuffer.size() == totalMessagesToFetch) {
        didNotReachEndOfStream.remove(entry.getKey());
        ssps.remove(entry.getKey());
      }
      output.get(ssp).addAll(currentBuffer);
    }
    if (didNotReachEndOfStream.isEmpty()) {
      break;
    }
  }

  if (!didNotReachEndOfStream.isEmpty()) {
    throw new IllegalStateException("Could not poll for all system stream partitions");
  }

  return output.entrySet()
      .stream()
      .collect(Collectors.toMap(entry -> entry.getKey().getPartition().getPartitionId(),
        entry -> entry.getValue().stream().map(e -> (StreamMessageType) e.getMessage()).collect(Collectors.toList())));
}