org.apache.samza.system.SystemAdmin Java Examples

Map<String, Integer> getStreamPartitionCounts(String systemName, Set<String> streamNames) {
  Map<String, Integer> streamToPartitionCount = new HashMap<>();

  SystemAdmin systemAdmin = systemAdmins.getSystemAdmin(systemName);
  if (systemAdmin == null) {
    throw new SamzaException(String.format("System %s does not exist.", systemName));
  }

  // retrieve the metadata for the streams in this system
  Map<String, SystemStreamMetadata> streamToMetadata = systemAdmin.getSystemStreamMetadata(streamNames);
  // set the partitions of a stream to its StreamEdge
  streamToMetadata.forEach((stream, data) ->
    streamToPartitionCount.put(stream, data.getSystemStreamPartitionMetadata().size()));

  return streamToPartitionCount;
}
 

private void initializeMocks() {
  SystemAdmin admin = mock(SystemAdmin.class);
  doAnswer(invocation -> {
    String offset1 = invocation.getArgumentAt(0, String.class);
    String offset2 = invocation.getArgumentAt(1, String.class);

    return Long.compare(Long.parseLong(offset1), Long.parseLong(offset2));
  }).when(admin).offsetComparator(any(), any());
  doAnswer(invocation -> {
    Map<SystemStreamPartition, String> sspToOffsets = invocation.getArgumentAt(0, Map.class);

    return sspToOffsets.entrySet()
        .stream()
        .collect(Collectors.toMap(Map.Entry::getKey, entry -> getOffsetAfter(entry.getValue())));
  }).when(admin).getOffsetsAfter(any());
  doReturn(admin).when(systemAdmins).getSystemAdmin(TEST_SYSTEM);
  doReturn(ScalaJavaUtil.toScalaMap(new HashMap<>())).when(streamMetadataCache).getStreamMetadata(any(), anyBoolean());
}
 

@Test
public void testCreateStreams() {
  StreamSpec spec1 = new StreamSpec(STREAM1, STREAM1, SYSTEM1);
  StreamSpec spec2 = new StreamSpec(STREAM2, STREAM2, SYSTEM2);
  List<StreamSpec> specList = new ArrayList<>();
  specList.add(spec1);
  specList.add(spec2);

  SystemAdmin admin1 = mock(SystemAdmin.class);
  SystemAdmin admin2 = mock(SystemAdmin.class);
  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  when(systemAdmins.getSystemAdmin(SYSTEM1)).thenReturn(admin1);
  when(systemAdmins.getSystemAdmin(SYSTEM2)).thenReturn(admin2);
  StreamManager manager = new StreamManager(systemAdmins);
  manager.createStreams(specList);

  ArgumentCaptor<StreamSpec> captor = ArgumentCaptor.forClass(StreamSpec.class);
  verify(admin1).createStream(captor.capture());
  assertEquals(STREAM1, captor.getValue().getPhysicalName());

  captor = ArgumentCaptor.forClass(StreamSpec.class);
  verify(admin2).createStream(captor.capture());
  assertEquals(STREAM2, captor.getValue().getPhysicalName());
}
 

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

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

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

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

/**
 * Create a new {@link SystemAdmin} that returns the provided oldest and newest offsets for its topics
 */
private SystemAdmin newAdmin(String oldestOffset, String newestOffset) {
  SystemStreamMetadata checkpointTopicMetadata = new SystemStreamMetadata(CHECKPOINT_TOPIC,
      ImmutableMap.of(new Partition(0), new SystemStreamPartitionMetadata(oldestOffset,
          newestOffset, Integer.toString(Integer.parseInt(newestOffset) + 1))));
  SystemAdmin mockAdmin = mock(SystemAdmin.class);
  when(mockAdmin.getSystemStreamMetadata(Collections.singleton(CHECKPOINT_TOPIC))).thenReturn(
      ImmutableMap.of(CHECKPOINT_TOPIC, checkpointTopicMetadata));
  return mockAdmin;
}
 

private static CoordinatorStreamStore createCoordinatorStreamStore(Config applicationConfig) {
  SystemStream coordinatorSystemStream = CoordinatorStreamUtil.getCoordinatorSystemStream(applicationConfig);
  SystemAdmins systemAdmins = new SystemAdmins(applicationConfig);
  SystemAdmin coordinatorSystemAdmin = systemAdmins.getSystemAdmin(coordinatorSystemStream.getSystem());
  coordinatorSystemAdmin.start();
  CoordinatorStreamUtil.createCoordinatorStream(coordinatorSystemStream, coordinatorSystemAdmin);
  coordinatorSystemAdmin.stop();
  return new CoordinatorStreamStore(applicationConfig, new NoOpMetricsRegistry());
}
 

@Test(expected = SamzaException.class)
public void testGetNewestOffsetsThrowsIfNullMetadata() {
  // empty topic == null newest offset
  ContainerStorageManager csm = mock(ContainerStorageManager.class);
  TransactionalStateTaskStorageManager tsm = buildTSM(csm, mock(Partition.class), new StorageManagerUtil());

  TaskName taskName = mock(TaskName.class);
  String changelogSystemName = "systemName";
  String storeName = "storeName";
  String changelogStreamName = "changelogName";
  String newestChangelogSSPOffset = null;
  SystemStream changelogSystemStream = new SystemStream(changelogSystemName, changelogStreamName);
  Partition changelogPartition = new Partition(0);
  SystemStreamPartition changelogSSP = new SystemStreamPartition(changelogSystemStream, changelogPartition);

  Map<String, SystemStream> storeChangelogs =
      ScalaJavaUtil.toScalaMap(ImmutableMap.of(storeName, changelogSystemStream));

  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  SystemAdmin systemAdmin = mock(SystemAdmin.class);
  SystemStreamPartitionMetadata metadata = mock(SystemStreamPartitionMetadata.class);

  when(metadata.getNewestOffset()).thenReturn(newestChangelogSSPOffset);
  when(systemAdmins.getSystemAdmin(changelogSystemName)).thenReturn(systemAdmin);
  when(systemAdmin.getSSPMetadata(eq(ImmutableSet.of(changelogSSP)))).thenReturn(null);

  // invoke the method
  Map<SystemStreamPartition, Option<String>> offsets =
      tsm.getNewestChangelogSSPOffsets(
          taskName, storeChangelogs, changelogPartition, systemAdmins);

  // verify results
  fail("Should have thrown an exception if admin didn't return any metadata");
}
 

@Test
public void testGetNewestOffsetsReturnsNoneForEmptyTopic() {
  // empty topic == null newest offset
  ContainerStorageManager csm = mock(ContainerStorageManager.class);
  TransactionalStateTaskStorageManager tsm = buildTSM(csm, mock(Partition.class), new StorageManagerUtil());

  TaskName taskName = mock(TaskName.class);
  String changelogSystemName = "systemName";
  String storeName = "storeName";
  String changelogStreamName = "changelogName";
  String newestChangelogSSPOffset = null;
  SystemStream changelogSystemStream = new SystemStream(changelogSystemName, changelogStreamName);
  Partition changelogPartition = new Partition(0);
  SystemStreamPartition changelogSSP = new SystemStreamPartition(changelogSystemStream, changelogPartition);

  Map<String, SystemStream> storeChangelogs =
      ScalaJavaUtil.toScalaMap(ImmutableMap.of(storeName, changelogSystemStream));

  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  SystemAdmin systemAdmin = mock(SystemAdmin.class);
  SystemStreamPartitionMetadata metadata = mock(SystemStreamPartitionMetadata.class);

  when(metadata.getNewestOffset()).thenReturn(newestChangelogSSPOffset);
  when(systemAdmins.getSystemAdmin(changelogSystemName)).thenReturn(systemAdmin);
  when(systemAdmin.getSSPMetadata(eq(ImmutableSet.of(changelogSSP)))).thenReturn(ImmutableMap.of(changelogSSP, metadata));

  // invoke the method
  Map<SystemStreamPartition, Option<String>> offsets =
      tsm.getNewestChangelogSSPOffsets(
          taskName, storeChangelogs, changelogPartition, systemAdmins);

  // verify results
  assertEquals(1, offsets.size());
  assertEquals(Option.empty(), offsets.apply(changelogSSP));
}
 

public CoordinatorStreamSystemProducer(SystemStream systemStream, SystemProducer systemProducer, SystemAdmin systemAdmin) {
  this.systemStream = systemStream;
  this.systemProducer = systemProducer;
  this.systemAdmin = systemAdmin;
  this.keySerde = new JsonSerde<>();
  this.messageSerde = new JsonSerde<>();
}
 

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

private SystemFactory newFactory(SystemProducer producer, SystemConsumer consumer, SystemAdmin admin) {
  SystemFactory factory = mock(SystemFactory.class);
  when(factory.getProducer(anyString(), any(Config.class), any(MetricsRegistry.class))).thenReturn(producer);
  when(factory.getConsumer(anyString(), any(Config.class), any(MetricsRegistry.class))).thenReturn(consumer);
  when(factory.getAdmin(anyString(), any(Config.class))).thenReturn(admin);
  return factory;
}
 

@Test
public void testAllMessagesInTheLogAreRead() 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);

  SystemStreamPartition ssp = new SystemStreamPartition("system-1", "input-topic", new Partition(0));

  int oldestOffset = 0;
  int newestOffset = 10;

  // mock out a consumer that returns ten checkpoint IMEs for the same ssp
  List<List<IncomingMessageEnvelope>> pollOutputs = new ArrayList<>();
  for (int offset = oldestOffset; offset <= newestOffset; offset++) {
    pollOutputs.add(ImmutableList.of(newCheckpointEnvelope(TASK1, ssp, Integer.toString(offset))));
  }

  // return one message at a time from each poll simulating a KafkaConsumer with max.poll.records = 1
  SystemConsumer mockConsumer = newConsumer(pollOutputs);
  SystemAdmin mockAdmin = newAdmin(Integer.toString(oldestOffset), Integer.toString(newestOffset));
  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);
  checkpointManager.start();

  // check that all ten messages are read, and the checkpoint is the newest message
  Checkpoint checkpoint = checkpointManager.readLastCheckpoint(TASK1);
  Assert.assertEquals(checkpoint.getOffsets(), ImmutableMap.of(ssp, Integer.toString(newestOffset)));
}
 

@Test
public void testGetNewestOffsetsReturnsCorrectOffset() {
  ContainerStorageManager csm = mock(ContainerStorageManager.class);
  TransactionalStateTaskStorageManager tsm = buildTSM(csm, mock(Partition.class), new StorageManagerUtil());

  TaskName taskName = mock(TaskName.class);
  String changelogSystemName = "systemName";
  String storeName = "storeName";
  String changelogStreamName = "changelogName";
  String newestChangelogSSPOffset = "1";
  SystemStream changelogSystemStream = new SystemStream(changelogSystemName, changelogStreamName);
  Partition changelogPartition = new Partition(0);
  SystemStreamPartition changelogSSP = new SystemStreamPartition(changelogSystemStream, changelogPartition);

  Map<String, SystemStream> storeChangelogs =
      ScalaJavaUtil.toScalaMap(ImmutableMap.of(storeName, changelogSystemStream));

  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  SystemAdmin systemAdmin = mock(SystemAdmin.class);
  SystemStreamPartitionMetadata metadata = mock(SystemStreamPartitionMetadata.class);

  when(metadata.getNewestOffset()).thenReturn(newestChangelogSSPOffset);
  when(systemAdmins.getSystemAdmin(changelogSystemName)).thenReturn(systemAdmin);
  when(systemAdmin.getSSPMetadata(eq(ImmutableSet.of(changelogSSP)))).thenReturn(ImmutableMap.of(changelogSSP, metadata));

  // invoke the method
  Map<SystemStreamPartition, Option<String>> offsets =
      tsm.getNewestChangelogSSPOffsets(
          taskName, storeChangelogs, changelogPartition, systemAdmins);

  // verify results
  assertEquals(1, offsets.size());
  assertEquals(Option.apply(newestChangelogSSPOffset), offsets.apply(changelogSSP));
}
 

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

/**
 * Get {@link SystemAdmin} instances for all the systems defined in this config.
 *
 * @return map of system name to {@link SystemAdmin}
 */
public Map<String, SystemAdmin> getSystemAdmins() {
  return getSystemFactories().entrySet()
      .stream()
      .collect(Collectors.toMap(Entry::getKey,
        systemNameToFactoryEntry -> systemNameToFactoryEntry.getValue()
            .getAdmin(systemNameToFactoryEntry.getKey(), this)));
}
 

public static void createDiagnosticsStream(Config config) {
  if (!new JobConfig(config).getDiagnosticsEnabled()) {
    return;
  }
  // if diagnostics is enabled, create diagnostics stream if it doesnt exist

  SystemAdmins systemAdmins = new SystemAdmins(config);
  String diagnosticsSystemStreamName = new MetricsConfig(config)
      .getMetricsSnapshotReporterStream(MetricsConfig.METRICS_SNAPSHOT_REPORTER_NAME_FOR_DIAGNOSTICS)
      .orElseThrow(() -> new ConfigException("Missing required config: " +
          String.format(MetricsConfig.METRICS_SNAPSHOT_REPORTER_STREAM,
              MetricsConfig.METRICS_SNAPSHOT_REPORTER_NAME_FOR_DIAGNOSTICS)));

  SystemStream diagnosticsSystemStream = StreamUtil.getSystemStreamFromNames(diagnosticsSystemStreamName);
  SystemAdmin diagnosticsSysAdmin = systemAdmins.getSystemAdmin(diagnosticsSystemStream.getSystem());
  StreamSpec diagnosticsStreamSpec = new StreamSpec(DIAGNOSTICS_STREAM_ID, diagnosticsSystemStream.getStream(),
      diagnosticsSystemStream.getSystem(), new StreamConfig(config).getStreamProperties(DIAGNOSTICS_STREAM_ID));

  log.info("Creating diagnostics stream {}", diagnosticsSystemStream.getStream());
  diagnosticsSysAdmin.start();

  if (diagnosticsSysAdmin.createStream(diagnosticsStreamSpec)) {
    log.info("Created diagnostics stream {}", diagnosticsSystemStream.getStream());
  } else {
    log.info("Diagnostics stream {} already exists", diagnosticsSystemStream.getStream());
  }

  diagnosticsSysAdmin.stop();
}
 

/**
 *  Validates each changelog system-stream with its respective SystemAdmin.
 */
private void validateChangelogStreams() {
  LOG.info("Validating change log streams: " + changelogSystemStreams);

  for (SystemStream changelogSystemStream : changelogSystemStreams.values()) {
    SystemAdmin systemAdmin = systemAdmins.getSystemAdmin(changelogSystemStream.getSystem());
    StreamSpec changelogSpec =
        StreamSpec.createChangeLogStreamSpec(changelogSystemStream.getStream(), changelogSystemStream.getSystem(),
            maxChangeLogStreamPartitions);

    systemAdmin.validateStream(changelogSpec);
  }
}
 

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

private static void validateRestoreOffsets(RestoreOffsets restoreOffsets, SystemAdmin systemAdmin) {
  String startingOffset = restoreOffsets.startingOffset;
  String endingOffset = restoreOffsets.endingOffset;
  // cannot enforce that starting offset be non null since SSPMetadata allows oldest offset == null for empty topics.
  if (endingOffset != null) {
    Preconditions.checkState(systemAdmin.offsetComparator(endingOffset, startingOffset) >= 0,
        String.format("Ending offset: %s must be equal to or greater than starting offset: %s",
            endingOffset, startingOffset));
  }
}
 

@Test
public void testGetSystemAdmins() {
  Map<String, String> map = ImmutableMap.of(MOCK_SYSTEM_FACTORY_NAME1, MockSystemFactory.class.getName());
  SystemConfig systemConfig = new SystemConfig(new MapConfig(map));
  Map<String, SystemAdmin> expected = ImmutableMap.of(MOCK_SYSTEM_NAME1, SYSTEM_ADMIN1);
  assertEquals(expected, systemConfig.getSystemAdmins());
}
 

private void checkSideInputCaughtUp(SystemStreamPartition ssp, String offset, SystemStreamMetadata.OffsetType offsetType, boolean isEndOfStream) {

    if (isEndOfStream) {
      this.initialSideInputSSPMetadata.remove(ssp);
      this.sideInputsCaughtUp.countDown();
      LOG.info("Side input ssp {} has caught up to offset {} ({}).", ssp, offset, offsetType);
      return;
    }

    SystemStreamMetadata.SystemStreamPartitionMetadata sspMetadata = this.initialSideInputSSPMetadata.get(ssp);
    String offsetToCheck = sspMetadata == null ? null : sspMetadata.getOffset(offsetType);
    LOG.trace("Checking {} offset {} against {} for {}.", offsetType, offset, offsetToCheck, ssp);

    // Let's compare offset of the chosen message with offsetToCheck.
    Integer comparatorResult;
    if (offset == null || offsetToCheck == null) {
      comparatorResult = -1;
    } else {
      SystemAdmin systemAdmin = systemAdmins.getSystemAdmin(ssp.getSystem());
      comparatorResult = systemAdmin.offsetComparator(offset, offsetToCheck);
    }

    // The SSP is no longer lagging if the envelope's offset is greater than or equal to the
    // latest offset.
    if (comparatorResult != null && comparatorResult.intValue() >= 0) {

      LOG.info("Side input ssp {} has caught up to offset {} ({}).", ssp, offset, offsetType);
      // if its caught up, we remove the ssp from the map, and countDown the latch
      this.initialSideInputSSPMetadata.remove(ssp);
      this.sideInputsCaughtUp.countDown();
      return;
    }
  }
 

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

static SystemAdmin createSystemAdmin(Map<String, Integer> streamToPartitions) {

    return new SystemAdmin() {
      @Override
      public Map<SystemStreamPartition, String> getOffsetsAfter(Map<SystemStreamPartition, String> offsets) {
        return null;
      }

      @Override
      public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
        Map<String, SystemStreamMetadata> map = new HashMap<>();
        for (String stream : streamNames) {
          Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> m = new HashMap<>();
          for (int i = 0; i < streamToPartitions.get(stream); i++) {
            m.put(new Partition(i), new SystemStreamMetadata.SystemStreamPartitionMetadata("", "", ""));
          }
          map.put(stream, new SystemStreamMetadata(stream, m));
        }
        return map;
      }

      @Override
      public Integer offsetComparator(String offset1, String offset2) {
        return null;
      }
    };
  }
 

@Test
public void testGetStreamPartitionCounts() {
  SystemAdmin admin1 = mock(SystemAdmin.class);
  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  when(systemAdmins.getSystemAdmin(SYSTEM1)).thenReturn(admin1);

  Map<String, SystemStreamMetadata> map = new HashMap<>();
  SystemStreamMetadata meta1 = mock(SystemStreamMetadata.class);
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitions = new HashMap<>();
  partitions.put(new Partition(0), null);
  when(meta1.getSystemStreamPartitionMetadata()).thenReturn(partitions);
  map.put(STREAM1, meta1);

  SystemStreamMetadata meta2 = mock(SystemStreamMetadata.class);
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitions2 = new HashMap<>();
  partitions2.put(new Partition(0), null);
  partitions2.put(new Partition(1), null);
  when(meta2.getSystemStreamPartitionMetadata()).thenReturn(partitions2);
  map.put(STREAM2, meta2);

  when(admin1.getSystemStreamMetadata(anyObject())).thenReturn(map);

  Set<String> streams = new HashSet<>();
  streams.add(STREAM1);
  streams.add(STREAM2);
  StreamManager manager = new StreamManager(systemAdmins);
  Map<String, Integer> counts = manager.getStreamPartitionCounts(SYSTEM1, streams);

  assertTrue(counts.get(STREAM1).equals(1));
  assertTrue(counts.get(STREAM2).equals(2));
}
 

@Test
public void testClearStreamsFromPreviousRun() {
  SystemAdmin admin1 = mock(SystemAdmin.class);
  SystemAdmin admin2 = mock(SystemAdmin.class);
  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  when(systemAdmins.getSystemAdmin(SYSTEM1)).thenReturn(admin1);
  when(systemAdmins.getSystemAdmin(SYSTEM2)).thenReturn(admin2);

  String runId = "123";
  Map<String, String> config = new HashMap<>();
  config.put(ApplicationConfig.APP_RUN_ID, "123");
  config.put(ApplicationConfig.APP_MODE, ApplicationConfig.ApplicationMode.BATCH.name());

  config.put("streams.stream-1.samza.system", SYSTEM1);
  config.put("streams.stream-1.samza.physical.name", STREAM1 + "-" + runId);
  config.put("streams.stream-1.samza.intermediate", "true");

  config.put("task.checkpoint.factory", MockCheckpointManagerFactory.class.getName());
  config.put("stores.test-store.factory", "dummyfactory");
  config.put("stores.test-store.changelog", SYSTEM2 + "." + STREAM2);

  StreamManager manager = new StreamManager(systemAdmins);
  manager.clearStreamsFromPreviousRun(new MapConfig(config));

  ArgumentCaptor<StreamSpec> captor = ArgumentCaptor.forClass(StreamSpec.class);
  verify(admin1).clearStream(captor.capture());
  assertEquals(captor.getValue().getPhysicalName(), STREAM1 + "-" + runId);

  captor = ArgumentCaptor.forClass(StreamSpec.class);
  verify(admin2).clearStream(captor.capture());
  assertEquals(captor.getValue().getPhysicalName(), STREAM2 + "-" + runId);

  verify(checkpointManager, times(1)).clearCheckpoints();
}