org.apache.samza.system.SystemConsumer Java Examples

/**
 *  Creates SystemConsumer objects for store restoration, creating one consumer per system.
 */
private static Map<String, SystemConsumer> createConsumers(Set<String> storeSystems,
    Map<String, SystemFactory> systemFactories, Config config, MetricsRegistry registry) {
  // Create one consumer for each system in use, map with one entry for each such system
  Map<String, SystemConsumer> consumers = new HashMap<>();

  // Iterate over the list of storeSystems and create one sysConsumer per system
  for (String storeSystemName : storeSystems) {
    SystemFactory systemFactory = systemFactories.get(storeSystemName);
    if (systemFactory == null) {
      throw new SamzaException("System " + storeSystemName + " does not exist in config");
    }
    consumers.put(storeSystemName, systemFactory.getConsumer(storeSystemName, config, registry));
  }

  return consumers;
}
 

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> consumeUntilTimeoutOrWatermark(
    SystemConsumer consumer, SystemStreamPartition ssp, long timeoutMillis)
    throws InterruptedException {
  assertTrue("Expected timeoutMillis (" + timeoutMillis + ") >= 0", timeoutMillis >= 0);

  final List<IncomingMessageEnvelope> accumulator = new ArrayList<>();
  final long start = System.currentTimeMillis();
  long now = start;
  while (timeoutMillis + start >= now) {
    accumulator.addAll(pollOnce(consumer, ssp, now - start - timeoutMillis));
    if (!accumulator.isEmpty()
        && MessageType.of(accumulator.get(accumulator.size() - 1).getMessage())
            == MessageType.WATERMARK) {
      break;
    }
    now = System.currentTimeMillis();
  }
  return accumulator;
}
 

@Override
public void restore() throws InterruptedException {
  Map<String, RestoreOffsets> storesToRestore = storeActions.storesToRestore;

  for (Map.Entry<String, RestoreOffsets> entry : storesToRestore.entrySet()) {
    String storeName = entry.getKey();
    String endOffset = entry.getValue().endingOffset;
    SystemStream systemStream = storeChangelogs.get(storeName);
    SystemAdmin systemAdmin = systemAdmins.getSystemAdmin(systemStream.getSystem());
    SystemConsumer systemConsumer = storeConsumers.get(storeName);
    SystemStreamPartition changelogSSP = new SystemStreamPartition(systemStream, taskModel.getChangelogPartition());

    ChangelogSSPIterator changelogSSPIterator =
        new ChangelogSSPIterator(systemConsumer, changelogSSP, endOffset, systemAdmin, true);
    StorageEngine taskStore = storeEngines.get(storeName);

    LOG.info("Restoring store: {} for task: {}", storeName, taskModel.getTaskName());
    taskStore.restore(changelogSSPIterator);
  }
}
 

private static List<IncomingMessageEnvelope> consumeUntilTimeoutOrEos(
    SystemConsumer consumer, SystemStreamPartition ssp, long timeoutMillis)
    throws InterruptedException {
  assertTrue("Expected timeoutMillis (" + timeoutMillis + ") >= 0", timeoutMillis >= 0);

  final List<IncomingMessageEnvelope> accumulator = new ArrayList<>();
  final long start = System.currentTimeMillis();
  long now = start;
  while (timeoutMillis + start >= now) {
    accumulator.addAll(pollOnce(consumer, ssp, now - start - timeoutMillis));
    if (!accumulator.isEmpty() && accumulator.get(accumulator.size() - 1).isEndOfStream()) {
      break;
    }
    now = System.currentTimeMillis();
  }
  return accumulator;
}
 

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

@Test
public void testReadSucceedsOnKeySerdeExceptionsWhenValidationIsDisabled() throws Exception {
  KafkaStreamSpec checkpointSpec = new KafkaStreamSpec(CHECKPOINT_TOPIC, CHECKPOINT_TOPIC,
      CHECKPOINT_SYSTEM, 1);
  Config mockConfig = mock(Config.class);
  when(mockConfig.get(JobConfig.SSP_GROUPER_FACTORY)).thenReturn(GROUPER_FACTORY_CLASS);

  // mock out a consumer that returns a single checkpoint IME
  SystemStreamPartition ssp = new SystemStreamPartition("system-1", "input-topic", new Partition(0));
  List<List<IncomingMessageEnvelope>> checkpointEnvelopes = ImmutableList.of(
      ImmutableList.of(newCheckpointEnvelope(TASK1, ssp, "0")));
  SystemConsumer mockConsumer = newConsumer(checkpointEnvelopes);

  SystemAdmin mockAdmin = newAdmin("0", "1");
  SystemFactory factory = newFactory(mock(SystemProducer.class), mockConsumer, mockAdmin);

  // wire up an exception throwing serde with the checkpointmanager
  KafkaCheckpointManager checkpointManager = new KafkaCheckpointManager(checkpointSpec, factory,
      false, mockConfig, mock(MetricsRegistry.class), new ExceptionThrowingCheckpointSerde(),
      new ExceptionThrowingCheckpointKeySerde());
  checkpointManager.register(TASK1);
  checkpointManager.start();

  // expect the read to succeed inspite of the exception from ExceptionThrowingSerde
  checkpointManager.readLastCheckpoint(TASK1);
}
 

@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 TransactionalStateTaskRestoreManager(
    TaskModel taskModel,
    Map<String, StorageEngine> storeEngines,
    Map<String, SystemStream> storeChangelogs,
    SystemAdmins systemAdmins,
    Map<String, SystemConsumer> storeConsumers,
    SSPMetadataCache sspMetadataCache,
    File loggedStoreBaseDirectory,
    File nonLoggedStoreBaseDirectory,
    Config config,
    Clock clock) {
  this.taskModel = taskModel;
  this.storeEngines = storeEngines;
  this.storeChangelogs = storeChangelogs;
  this.systemAdmins = systemAdmins;
  this.storeConsumers = storeConsumers;
  // OK to use SSPMetadataCache here since unlike commit newest changelog ssp offsets will not change
  // between cache init and restore completion
  this.sspMetadataCache = sspMetadataCache;
  this.loggedStoreBaseDirectory = loggedStoreBaseDirectory;
  this.nonLoggedStoreBaseDirectory = nonLoggedStoreBaseDirectory;
  this.config = config;
  this.clock = clock;
  this.storageManagerUtil = new StorageManagerUtil();
  this.fileUtil = new FileUtil();
}
 

@Test(expected = TopicAlreadyMarkedForDeletionException.class)
public void testStartFailsOnTopicCreationErrors() {

  KafkaStreamSpec checkpointSpec = new KafkaStreamSpec(CHECKPOINT_TOPIC, CHECKPOINT_TOPIC,
      CHECKPOINT_SYSTEM, 1);
  // create an admin that throws an exception during createStream
  SystemAdmin mockAdmin = newAdmin("0", "10");
  doThrow(new TopicAlreadyMarkedForDeletionException("invalid stream")).when(mockAdmin).createStream(checkpointSpec);

  SystemFactory factory = newFactory(mock(SystemProducer.class), mock(SystemConsumer.class), mockAdmin);
  KafkaCheckpointManager checkpointManager = new KafkaCheckpointManager(checkpointSpec, factory,
      true, mock(Config.class), mock(MetricsRegistry.class), null, new KafkaCheckpointLogKeySerde());

  // expect an exception during startup
  checkpointManager.createResources();
  checkpointManager.start();
}
 

@Test(expected = StreamValidationException.class)
public void testStartFailsOnTopicValidationErrors() {

  KafkaStreamSpec checkpointSpec = new KafkaStreamSpec(CHECKPOINT_TOPIC, CHECKPOINT_TOPIC,
      CHECKPOINT_SYSTEM, 1);

  // create an admin that throws an exception during validateStream
  SystemAdmin mockAdmin = newAdmin("0", "10");
  doThrow(new StreamValidationException("invalid stream")).when(mockAdmin).validateStream(checkpointSpec);

  SystemFactory factory = newFactory(mock(SystemProducer.class), mock(SystemConsumer.class), mockAdmin);
  KafkaCheckpointManager checkpointManager = new KafkaCheckpointManager(checkpointSpec, factory,
      true, mock(Config.class), mock(MetricsRegistry.class), null, new KafkaCheckpointLogKeySerde());

  // expect an exception during startup
  checkpointManager.createResources();
  checkpointManager.start();
}
 

@Test(expected = SamzaException.class)
public void testReadFailsOnSerdeExceptions() throws Exception {
  KafkaStreamSpec checkpointSpec = new KafkaStreamSpec(CHECKPOINT_TOPIC, CHECKPOINT_TOPIC,
      CHECKPOINT_SYSTEM, 1);
  Config mockConfig = mock(Config.class);
  when(mockConfig.get(JobConfig.SSP_GROUPER_FACTORY)).thenReturn(GROUPER_FACTORY_CLASS);

  // mock out a consumer that returns a single checkpoint IME
  SystemStreamPartition ssp = new SystemStreamPartition("system-1", "input-topic", new Partition(0));
  List<List<IncomingMessageEnvelope>> checkpointEnvelopes = ImmutableList.of(
      ImmutableList.of(newCheckpointEnvelope(TASK1, ssp, "0")));
  SystemConsumer mockConsumer = newConsumer(checkpointEnvelopes);

  SystemAdmin mockAdmin = newAdmin("0", "1");
  SystemFactory factory = newFactory(mock(SystemProducer.class), mockConsumer, mockAdmin);

  // wire up an exception throwing serde with the checkpointmanager
  KafkaCheckpointManager checkpointManager = new KafkaCheckpointManager(checkpointSpec, factory,
      true, mockConfig, mock(MetricsRegistry.class), new ExceptionThrowingCheckpointSerde(), new KafkaCheckpointLogKeySerde());
  checkpointManager.register(TASK1);
  checkpointManager.start();

  // expect an exception from ExceptionThrowingSerde
  checkpointManager.readLastCheckpoint(TASK1);
}
 

/**
 * Restore each store in taskStoresToRestore sequentially
 */
@Override
public void restore() throws InterruptedException {
  for (String storeName : taskStoresToRestore) {
    LOG.info("Restoring store: {} for task: {}", storeName, taskModel.getTaskName());
    SystemConsumer systemConsumer = storeConsumers.get(storeName);
    SystemStream systemStream = changelogSystemStreams.get(storeName);
    SystemAdmin systemAdmin = systemAdmins.getSystemAdmin(systemStream.getSystem());
    ChangelogSSPIterator changelogSSPIterator = new ChangelogSSPIterator(systemConsumer,
        new SystemStreamPartition(systemStream, taskModel.getChangelogPartition()), null, systemAdmin, false);

    taskStores.get(storeName).restore(changelogSSPIterator);
  }
}
 

@VisibleForTesting
protected CoordinatorStreamStore(Config config, SystemProducer systemProducer, SystemConsumer systemConsumer, SystemAdmin systemAdmin) {
  this.config = config;
  this.systemConsumer = systemConsumer;
  this.systemProducer = systemProducer;
  this.systemAdmin = systemAdmin;
  this.coordinatorSystemStream = CoordinatorStreamUtil.getCoordinatorSystemStream(config);
  this.coordinatorSystemStreamPartition = new SystemStreamPartition(coordinatorSystemStream, new Partition(0));
}
 

NonTransactionalStateTaskRestoreManager(
    TaskModel taskModel,
    Map<String, SystemStream> changelogSystemStreams,
    Map<String, StorageEngine> taskStores,
    SystemAdmins systemAdmins,
    StreamMetadataCache streamMetadataCache,
    Map<String, SystemConsumer> storeConsumers,
    int maxChangeLogStreamPartitions,
    File loggedStoreBaseDirectory,
    File nonLoggedStoreBaseDirectory,
    Config config,
    Clock clock) {
  this.taskStores = taskStores;
  this.taskModel = taskModel;
  this.clock = clock;
  this.changelogSystemStreams = changelogSystemStreams;
  this.systemAdmins = systemAdmins;
  this.fileOffsets = new HashMap<>();
  this.taskStoresToRestore = this.taskStores.entrySet().stream()
      .filter(x -> x.getValue().getStoreProperties().isLoggedStore())
      .map(x -> x.getKey()).collect(Collectors.toSet());
  this.loggedStoreBaseDirectory = loggedStoreBaseDirectory;
  this.nonLoggedStoreBaseDirectory = nonLoggedStoreBaseDirectory;
  this.streamMetadataCache = streamMetadataCache;
  this.storeConsumers = storeConsumers;
  this.maxChangeLogStreamPartitions = maxChangeLogStreamPartitions;
  this.storageConfig = new StorageConfig(config);
  this.storageManagerUtil = new StorageManagerUtil();
}
 

@Override
public SystemConsumer getConsumer(String systemName, Config config, MetricsRegistry registry) {
  final String streamPrefix = "systems." + systemName;
  final Config scopedConfig = config.subset(streamPrefix + ".", true);
  return new Consumer<T, CheckpointMarkT>(
      getUnboundedSource(scopedConfig),
      getPipelineOptions(config),
      new SamzaMetricsContainer((MetricsRegistryMap) registry),
      scopedConfig.get("stepName"));
}
 

private static Map<String, SystemConsumer> createStoreIndexedMap(Map<String, SystemStream> changelogSystemStreams,
    Map<String, SystemConsumer> storeSystemConsumers) {
  // Map of each storeName to its respective systemConsumer
  Map<String, SystemConsumer> storeConsumers = new HashMap<>();

  // Populate the map of storeName to its relevant systemConsumer
  for (String storeName : changelogSystemStreams.keySet()) {
    storeConsumers.put(storeName, storeSystemConsumers.get(changelogSystemStreams.get(storeName).getSystem()));
  }
  return storeConsumers;
}
 

/**
 * Verify that if a particular key-value is written, then another, then the original again,
 * that the original occurs last in the set.
 */
@Test
public void testOrderKeyRewrite() throws InterruptedException {
  final SystemStream systemStream = new SystemStream("system", "stream");
  final SystemStreamPartition ssp = new SystemStreamPartition(systemStream, new Partition(0));
  final SystemConsumer systemConsumer = mock(SystemConsumer.class);

  final List<IncomingMessageEnvelope> list = new ArrayList<>();
  SetConfig setConfig1 = new SetConfig("source", "key1", "value1");
  SetConfig setConfig2 = new SetConfig("source", "key1", "value2");
  SetConfig setConfig3 = new SetConfig("source", "key1", "value1");
  list.add(createIncomingMessageEnvelope(setConfig1, ssp));
  list.add(createIncomingMessageEnvelope(setConfig2, ssp));
  list.add(createIncomingMessageEnvelope(setConfig3, ssp));
  Map<SystemStreamPartition, List<IncomingMessageEnvelope>> messages = new HashMap<SystemStreamPartition, List<IncomingMessageEnvelope>>() {
    {
      put(ssp, list);
    }
  };
  when(systemConsumer.poll(anySet(), anyLong())).thenReturn(messages, Collections.<SystemStreamPartition, List<IncomingMessageEnvelope>>emptyMap());

  CoordinatorStreamSystemConsumer consumer = new CoordinatorStreamSystemConsumer(systemStream, systemConsumer, new SinglePartitionWithoutOffsetsSystemAdmin());

  consumer.bootstrap();

  Set<CoordinatorStreamMessage> bootstrappedMessages = consumer.getBootstrappedStream(SetConfig.TYPE);

  assertEquals(2, bootstrappedMessages.size()); // First message should have been removed as a duplicate
  CoordinatorStreamMessage[] coordinatorStreamMessages = bootstrappedMessages.toArray(new CoordinatorStreamMessage[2]);
  assertEquals(setConfig2, coordinatorStreamMessages[0]);
  assertEquals(setConfig3, coordinatorStreamMessages[1]); //Config 3 MUST be the last message, not config 2
}
 

public CoordinatorStreamStoreTestUtil(Config config, String systemName) {
  this.config = config;
  this.systemFactory = new MockCoordinatorStreamSystemFactory();
  MockCoordinatorStreamSystemFactory.enableMockConsumerCache();
  SystemConsumer systemConsumer = systemFactory.getConsumer(systemName, config, new NoOpMetricsRegistry());
  SystemProducer systemProducer = systemFactory.getProducer(systemName, config, new NoOpMetricsRegistry());
  SystemAdmin systemAdmin = systemFactory.getAdmin(systemName, config);
  this.coordinatorStreamStore = new CoordinatorStreamStore(config, systemProducer, systemConsumer, systemAdmin);
  this.coordinatorStreamStore.init();
}
 

@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 <T> List<T> consumeMessages(SystemConsumer consumer, Set<SystemStreamPartition> sspsToPoll) {
  return consumeRawMessages(consumer, sspsToPoll)
      .stream()
      .map(IncomingMessageEnvelope::getMessage)
      .map(message -> (T) message)
      .collect(Collectors.toList());
}
 

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

@Override
public SystemConsumer getConsumer(String systemName, Config config, MetricsRegistry registry) {
  String host = config.get("systems." + systemName + ".host");
  int port = config.getInt("systems." + systemName + ".port");
  WikipediaFeed feed = new WikipediaFeed(host, port);

  return new WikipediaConsumer(systemName, feed, registry);
}
 

public CoordinatorStreamSystemConsumer(Config config, MetricsRegistry registry) {
  SystemStream coordinatorSystemStream = CoordinatorStreamUtil.getCoordinatorSystemStream(config);
  SystemFactory systemFactory = CoordinatorStreamUtil.getCoordinatorSystemFactory(config);
  SystemAdmin systemAdmin = systemFactory.getAdmin(coordinatorSystemStream.getSystem(), config);
  SystemConsumer systemConsumer = systemFactory.getConsumer(coordinatorSystemStream.getSystem(), config, registry);

  this.coordinatorSystemStreamPartition = new SystemStreamPartition(coordinatorSystemStream, new Partition(0));
  this.systemConsumer = systemConsumer;
  this.systemAdmin = systemAdmin;
  this.configMap = new HashMap<>();
  this.isBootstrapped = false;
  this.keySerde = new JsonSerde<>();
  this.messageSerde = new JsonSerde<>();
}
 

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

@Test
public void testCheckpointsAreReadFromOldestOffset() throws Exception {
  KafkaStreamSpec checkpointSpec = new KafkaStreamSpec(CHECKPOINT_TOPIC, CHECKPOINT_TOPIC,
      CHECKPOINT_SYSTEM, 1);
  Config mockConfig = mock(Config.class);
  when(mockConfig.get(JobConfig.SSP_GROUPER_FACTORY)).thenReturn(GROUPER_FACTORY_CLASS);

  // mock out a consumer that returns a single checkpoint IME
  SystemStreamPartition ssp = new SystemStreamPartition("system-1", "input-topic", new Partition(0));
  SystemConsumer mockConsumer = newConsumer(ImmutableList.of(
      ImmutableList.of(newCheckpointEnvelope(TASK1, ssp, "0"))));

  String oldestOffset = "0";
  SystemAdmin mockAdmin = newAdmin(oldestOffset, "1");
  SystemFactory factory = newFactory(mock(SystemProducer.class), mockConsumer, mockAdmin);
  KafkaCheckpointManager checkpointManager = new KafkaCheckpointManager(checkpointSpec, factory,
      true, mockConfig, mock(MetricsRegistry.class), new CheckpointSerde(), new KafkaCheckpointLogKeySerde());
  checkpointManager.register(TASK1);

  // 1. verify that consumer.register is called only during checkpointManager.start.
  // 2. verify that consumer.register is called with the oldest offset.
  // 3. verify that no other operation on the CheckpointManager re-invokes register since start offsets are set during
  // register
  verify(mockConsumer, times(0)).register(CHECKPOINT_SSP, oldestOffset);
  checkpointManager.start();
  verify(mockConsumer, times(1)).register(CHECKPOINT_SSP, oldestOffset);

  checkpointManager.readLastCheckpoint(TASK1);
  verify(mockConsumer, times(1)).register(CHECKPOINT_SSP, oldestOffset);
}