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

/**
 * Determines the starting offset for each store SSP (based on {@link #getStartingOffset(SystemStreamPartition, SystemAdmin)}) and
 * registers it with the respective SystemConsumer for starting consumption.
 */
private void registerStartingOffsets() {

  for (Map.Entry<String, SystemStream> changelogSystemStreamEntry : changelogSystemStreams.entrySet()) {
    SystemStreamPartition systemStreamPartition =
        new SystemStreamPartition(changelogSystemStreamEntry.getValue(), taskModel.getChangelogPartition());
    SystemAdmin systemAdmin = systemAdmins.getSystemAdmin(changelogSystemStreamEntry.getValue().getSystem());
    SystemConsumer systemConsumer = storeConsumers.get(changelogSystemStreamEntry.getKey());

    String offset = getStartingOffset(systemStreamPartition, systemAdmin);

    if (offset != null) {
      LOG.info("Registering change log consumer with offset " + offset + " for %" + systemStreamPartition);
      systemConsumer.register(systemStreamPartition, offset);
    } else {
      LOG.info("Skipping change log restoration for {} because stream appears to be empty (offset was null).",
          systemStreamPartition);
      taskStoresToRestore.remove(changelogSystemStreamEntry.getKey());
    }
  }
}
 

@Test
public void testNullMessageWithValidMessageKey() {
  final String messageKey = "validKey";
  SystemProducer systemProducer = systemFactory.getProducer(SYSTEM_NAME, config, mockRegistry);
  systemProducer.send(SOURCE, new OutgoingMessageEnvelope(SYSTEM_STREAM, messageKey, null));

  SystemConsumer consumer = systemFactory.getConsumer(SYSTEM_NAME, config, mockRegistry);

  Set<SystemStreamPartition> sspsToPoll = IntStream.range(0, PARTITION_COUNT)
      .mapToObj(partition -> new SystemStreamPartition(SYSTEM_STREAM, new Partition(partition)))
      .collect(Collectors.toSet());

  // register the consumer for ssps
  for (SystemStreamPartition ssp : sspsToPoll) {
    consumer.register(ssp, "0");
  }

  List<IncomingMessageEnvelope> results = consumeRawMessages(consumer, sspsToPoll);
  assertEquals(1, results.size());
  assertEquals(results.get(0).getKey(), messageKey);
  assertNull(results.get(0).getMessage());
}
 

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

@Test
public void testSinglePartitionConsumptionHappyPath() throws Exception {
  int partitionId = 0;

  TestMetricsRegistry testMetrics = new TestMetricsRegistry();
  SystemStreamPartition ssp = new SystemStreamPartition(SYSTEM_NAME, STREAM_NAME1, new Partition(partitionId));

  Config eventHubConfig = createEventHubConfig();

  EventHubSystemFactory factory = new EventHubSystemFactory();
  SystemConsumer consumer = factory.getConsumer(SYSTEM_NAME, eventHubConfig, testMetrics);
  consumer.register(ssp, EventHubSystemConsumer.START_OF_STREAM);
  consumer.start();

  int numEvents = 0;
  int numRetries = 20;
  while (numRetries-- > 0) {
    List<IncomingMessageEnvelope> result = consumer.poll(Collections.singleton(ssp), 2000).get(ssp);
    numEvents = result == null ? 0 : result.size();
    if (numEvents > 0) {
      EventHubIncomingMessageEnvelope eventData = (EventHubIncomingMessageEnvelope) result.get(0);
      System.out.println("System properties: " + eventData.getEventData().getSystemProperties());
      System.out.println("Key: " + new String((byte[]) eventData.getKey()));
      System.out.println("Message: " + new String((byte[]) eventData.getMessage()));
    }
    System.out.println("Retries left: " + numRetries);
  }
  Assert.assertTrue(numEvents > 0);
}
 

@Test
public void testConsumerRespectsOffset() {
  PageViewEvent event = new PageViewEvent(TEST_MEMBER_X, PAGE_ID_X, System.currentTimeMillis());
  PageViewEvent event1 = new PageViewEvent(TEST_MEMBER_Y, PAGE_ID_Y, System.currentTimeMillis());

  produceMessages(event);

  SystemConsumer consumer = systemFactory.getConsumer(SYSTEM_NAME, config, mockRegistry);

  Set<SystemStreamPartition> sspsToPoll = IntStream.range(0, PARTITION_COUNT)
      .mapToObj(partition -> new SystemStreamPartition(SYSTEM_STREAM, new Partition(partition)))
      .collect(Collectors.toSet());

  // register the consumer for ssps
  for (SystemStreamPartition ssp : sspsToPoll) {
    consumer.register(ssp, "0");
  }

  List<PageViewEvent> results = consumeMessages(consumer, sspsToPoll);
  assertEquals(1, results.size());
  assertTrue(results.contains(event));

  // nothing to poll
  results = consumeMessages(consumer, sspsToPoll);
  assertEquals(0, results.size());

  produceMessages(event1);

  // got new message. check if the offset has progressed
  results = consumeMessages(consumer, sspsToPoll);
  assertEquals(1, results.size());
  assertTrue(results.contains(event1));
}
 

private <T> List<T> consumeMessages(Set<SystemStreamPartition> sspsToPoll) {
  SystemConsumer systemConsumer = systemFactory.getConsumer(SYSTEM_NAME, config, mockRegistry);

  // register the consumer for ssps
  for (SystemStreamPartition ssp : sspsToPoll) {
    systemConsumer.register(ssp, "0");
  }

  return consumeMessages(systemConsumer, sspsToPoll);
}
 

private List<IncomingMessageEnvelope> consumeRawMessages(Set<SystemStreamPartition> sspsToPoll) {
  SystemConsumer systemConsumer = systemFactory.getConsumer(SYSTEM_NAME, config, mockRegistry);

  // register the consumer for ssps
  for (SystemStreamPartition ssp : sspsToPoll) {
    systemConsumer.register(ssp, "0");
  }

  return consumeRawMessages(systemConsumer, sspsToPoll);
}