org.apache.samza.Partition Java Examples

/**
 * Sends a specified message envelope from a specified Samza source.

 * @param source String representing the source of the message.
 * @param envelope Aggregate object representing the serialized message to send from the source.
 */
@Override
public void send(String source, OutgoingMessageEnvelope envelope) {
  Object key = envelope.getKey();
  Object message = envelope.getMessage();

  Object partitionKey;
  // We use the partition key from message if available, if not fallback to message key or use message as partition
  // key as the final resort.
  if (envelope.getPartitionKey() != null) {
    partitionKey = envelope.getPartitionKey();
  } else if (key != null) {
    partitionKey = key;
  } else {
    partitionKey = message;
  }

  Preconditions.checkNotNull(partitionKey, "Failed to compute partition key for the message: " + envelope);

  int partition =
      Math.abs(hashCode(partitionKey)) % memoryManager.getPartitionCountForSystemStream(envelope.getSystemStream());

  SystemStreamPartition ssp = new SystemStreamPartition(envelope.getSystemStream(), new Partition(partition));
  memoryManager.put(ssp, key, message);
}
 

@Test
public void testStop() {
  final String storeName = "test-stop-store";
  final String taskName = "test-stop-task";
  final SystemStreamPartition ssp = new SystemStreamPartition("test-system", "test-stream", new Partition(0));
  final String offset = "123";
  final ImmutableMap<SystemStreamPartition, String> processedOffsets = ImmutableMap.of(ssp, offset);

  TaskSideInputStorageManager testSideInputStorageManager = new MockTaskSideInputStorageManagerBuilder(taskName, NON_LOGGED_STORE_DIR)
      .addInMemoryStore(storeName, ImmutableSet.of())
      .build();

  initializeSideInputStorageManager(testSideInputStorageManager);
  testSideInputStorageManager.stop(processedOffsets);

  verify(testSideInputStorageManager.getStore(storeName)).stop();
  verify(testSideInputStorageManager).writeFileOffsets(eq(processedOffsets));
}
 

/**
 * Given some SSPs, but missing metadata for one of the streams, getSSPMetadata should delegate to
 * getSystemStreamMetadata and only fill in results for the SSPs corresponding to streams with metadata.
 */
@Test
public void testGetSSPMetadataMissingStream() {
  SystemStreamPartition streamPartition0 = new SystemStreamPartition(SYSTEM, STREAM, new Partition(0));
  SystemStreamPartition otherStreamPartition0 = new SystemStreamPartition(SYSTEM, OTHER_STREAM, new Partition(0));
  SystemAdmin systemAdmin = mock(MySystemAdmin.class);
  SystemStreamMetadata.SystemStreamPartitionMetadata streamPartition0Metadata =
      new SystemStreamMetadata.SystemStreamPartitionMetadata("1", "2", "3");
  when(systemAdmin.getSystemStreamMetadata(ImmutableSet.of(STREAM, OTHER_STREAM))).thenReturn(ImmutableMap.of(
      STREAM, new SystemStreamMetadata(STREAM, ImmutableMap.of(new Partition(0), streamPartition0Metadata))));
  Set<SystemStreamPartition> ssps = ImmutableSet.of(streamPartition0, otherStreamPartition0);
  when(systemAdmin.getSSPMetadata(ssps)).thenCallRealMethod();
  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> expected =
      ImmutableMap.of(streamPartition0, streamPartition0Metadata);
  assertEquals(expected, systemAdmin.getSSPMetadata(ssps));
  verify(systemAdmin).getSystemStreamMetadata(ImmutableSet.of(STREAM, OTHER_STREAM));
}
 

protected void runApplication(Config config) {
  // Use MockSystemFactory for the coordinator system
  MockSystemFactory.MSG_QUEUES.put(new SystemStreamPartition(MOCK_METADATA_SYSTEM,
      CoordinatorStreamUtil.getCoordinatorStreamName(SamzaSqlTestConfig.SQL_JOB, SamzaSqlTestConfig.SQL_JOB_PROCESSOR_ID),
      new Partition(0)), new ArrayList<>());
  HashMap<String, String> mapConfig = new HashMap<>();
  mapConfig.put(JobConfig.JOB_COORDINATOR_SYSTEM, MOCK_METADATA_SYSTEM);
  mapConfig.put(String.format(SystemConfig.SYSTEM_FACTORY_FORMAT, MOCK_METADATA_SYSTEM), MockSystemFactory.class.getName());

  // configs for using in-memory system as the default system
  mapConfig.putAll(baseInMemorySystemConfigs());
  mapConfig.putAll(config);

  SamzaSqlApplicationRunner runner = new SamzaSqlApplicationRunner(true, new MapConfig(mapConfig));
  executeRun(runner, config);
  runner.waitForFinish();
}
 

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

@Test
public void testSyncTaskWithMultiplePartitionMultithreadedWithCustomIME() throws Exception {
  Map<Integer, List<KV>> inputPartitionData = new HashMap<>();
  Map<Integer, List<KV>> inputPartitionIME = new HashMap<>();
  Map<Integer, List<Integer>> expectedOutputPartitionData = new HashMap<>();
  genData(inputPartitionData, expectedOutputPartitionData);

  for (Map.Entry<Integer, List<KV>> entry: inputPartitionData.entrySet()) {
    Integer partitionId = entry.getKey();
    List<KV> messages = entry.getValue();
    SystemStreamPartition ssp = new SystemStreamPartition("test", "input", new Partition(partitionId));
    inputPartitionIME.put(partitionId, new ArrayList<>());
    int offset = 0;
    for (KV message: messages) {
      IncomingMessageEnvelope ime = new IncomingMessageEnvelope(ssp, String.valueOf(offset++), message.key, message.getValue());
      inputPartitionIME.get(partitionId).add(KV.of(message.key, ime));
    }
  }
  syncTaskWithMultiplePartitionMultithreadedHelper(inputPartitionIME, expectedOutputPartitionData);
}
 

public TestZkJobCoordinator() {
  Map<String, String> configMap = ImmutableMap.of(
      "job.coordinator.system", "kafka",
      "job.name", "test-job",
      "systems.kafka.samza.factory", "org.apache.samza.system.MockSystemFactory");
  config = new MapConfig(configMap);

  Set<SystemStreamPartition> ssps = ImmutableSet.of(
      new SystemStreamPartition("system1", "stream1_1", new Partition(0)),
      new SystemStreamPartition("system1", "stream1_2", new Partition(0)));
  Map<TaskName, TaskModel> tasksForContainer = ImmutableMap.of(
      new TaskName("t1"), new TaskModel(new TaskName("t1"), ssps, new Partition(0)));
  ContainerModel containerModel = new ContainerModel("0", tasksForContainer);
  jobModel = new JobModel(config, ImmutableMap.of("0", containerModel));
  zkMetadataStore = Mockito.mock(MetadataStore.class);
  coordinatorStreamStore = Mockito.mock(CoordinatorStreamStore.class);
}
 

@Test
public void testPartitionCountMonitorWithDurableStates() {
  configMap.put("stores.mystore.changelog", "mychangelog");
  configMap.put(JobConfig.JOB_CONTAINER_COUNT, "1");
  when(CoordinatorStreamUtil.readConfigFromCoordinatorStream(anyObject())).thenReturn(new MapConfig(configMap));
  Config config = new MapConfig(configMap);

  // mimic job runner code to write the config to coordinator stream
  CoordinatorStreamSystemProducer producer = new CoordinatorStreamSystemProducer(config, mock(MetricsRegistry.class));
  producer.writeConfig("test-job", config);

  ClusterBasedJobCoordinator clusterCoordinator = ClusterBasedJobCoordinator.createFromMetadataStore(config);

  // change the input system stream metadata
  MockSystemFactory.MSG_QUEUES.put(new SystemStreamPartition("kafka", "topic1", new Partition(1)), new ArrayList<>());

  StreamPartitionCountMonitor monitor = clusterCoordinator.getPartitionMonitor();
  monitor.updatePartitionCountMetric();
  assertEquals(clusterCoordinator.getAppStatus(), SamzaApplicationState.SamzaAppStatus.FAILED);
}
 

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

  produceMessages(event1, event2);

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

  List<PageViewEvent> results = consumeMessages(sspsToPoll);

  assertEquals(2, results.size());
  assertTrue(results.contains(event1));
  assertTrue(results.contains(event2));
}
 

@Test
public void testCoordinatorStreamSystemConsumer() {
  Map<String, String> expectedConfig = new LinkedHashMap<String, String>();
  expectedConfig.put("job.id", "1234");
  SystemStream systemStream = new SystemStream("system", "stream");
  MockSystemConsumer systemConsumer = new MockSystemConsumer(new SystemStreamPartition(systemStream, new Partition(0)));
  CoordinatorStreamSystemConsumer consumer = new CoordinatorStreamSystemConsumer(systemStream, systemConsumer, new SinglePartitionWithoutOffsetsSystemAdmin());
  assertEquals(0, systemConsumer.getRegisterCount());
  consumer.register();
  assertEquals(1, systemConsumer.getRegisterCount());
  assertFalse(systemConsumer.isStarted());
  consumer.start();
  assertTrue(systemConsumer.isStarted());
  try {
    consumer.getConfig();
    fail("Should have failed when retrieving config before bootstrapping.");
  } catch (SamzaException e) {
    // Expected.
  }
  consumer.bootstrap();
  assertEquals(expectedConfig, consumer.getConfig());
  assertFalse(systemConsumer.isStopped());
  consumer.stop();
  assertTrue(systemConsumer.isStopped());
}
 

public SamoaSystemConsumer(String systemName, Config config) {
  String yarnConfHome = config.get(SamzaConfigFactory.YARN_CONF_HOME_KEY);
  if (yarnConfHome != null && yarnConfHome.length() > 0) // if the property is set, otherwise, assume we are running in local mode and ignore this
    SystemsUtils.setHadoopConfigHome(yarnConfHome);

  String filename = config.get(SamzaConfigFactory.FILE_KEY);
  String filesystem = config.get(SamzaConfigFactory.FILESYSTEM_KEY);
  String name = config.get(SamzaConfigFactory.JOB_NAME_KEY);
  SerializationProxy wrapper = (SerializationProxy) SystemsUtils.deserializeObjectFromFileAndKey(filesystem,
      filename, name);

  this.entranceProcessor = wrapper.processor;
  this.entranceProcessor.onCreate(0);

  // Internal stream from SystemConsumer to EntranceTask, so we
  // need only one partition
  this.systemStreamPartition = new SystemStreamPartition(systemName, wrapper.name, new Partition(0));
}
 

private HashMap<String, ContainerModel> generateContainers() {
  Set<TaskModel> taskModels = ImmutableSet.of(
      new TaskModel(new TaskName("task1"),
                    ImmutableSet.of(new SystemStreamPartition(new SystemStream("system1", "stream1"), new Partition(0))),
                    new Partition(0)),
      new TaskModel(new TaskName("task2"),
          ImmutableSet.of(new SystemStreamPartition(new SystemStream("system1", "stream1"), new Partition(1))),
          new Partition(1)));
  GroupByContainerCount grouper = new GroupByContainerCount(2);
  Set<ContainerModel> containerModels = grouper.group(taskModels);
  HashMap<String, ContainerModel> containers = new HashMap<>();
  for (ContainerModel containerModel : containerModels) {
    containers.put(containerModel.getId(), containerModel);
  }
  return containers;
}
 

@Test
public void testRemoveOldCheckpointsWhenBaseDirContainsRegularFiles() {
  TaskName taskName = new TaskName("Partition 0");
  ContainerStorageManager containerStorageManager = mock(ContainerStorageManager.class);
  Map<String, SystemStream> changelogSystemStreams = mock(Map.class);
  SystemAdmins systemAdmins = mock(SystemAdmins.class);
  File loggedStoreBaseDir = mock(File.class);
  Partition changelogPartition = new Partition(0);
  TaskMode taskMode = TaskMode.Active;
  StorageManagerUtil storageManagerUtil = mock(StorageManagerUtil.class);

  File mockStoreDir = mock(File.class);
  String mockStoreDirName = "notDirectory";

  when(loggedStoreBaseDir.listFiles()).thenReturn(new File[] {mockStoreDir});
  when(mockStoreDir.getName()).thenReturn(mockStoreDirName);
  when(storageManagerUtil.getTaskStoreDir(eq(loggedStoreBaseDir), eq(mockStoreDirName), eq(taskName), eq(taskMode))).thenReturn(mockStoreDir);
  // null here can happen if listFiles is called on a non-directory
  when(mockStoreDir.listFiles(any(FileFilter.class))).thenReturn(null);

  TransactionalStateTaskStorageManager tsm = new TransactionalStateTaskStorageManager(taskName, containerStorageManager,
      changelogSystemStreams, systemAdmins, loggedStoreBaseDir, changelogPartition, taskMode, storageManagerUtil);

  tsm.removeOldCheckpoints(CheckpointId.create());
}
 

private void persistPartitionDescriptor(String streamName,
  Map<Partition, List<String>> partitionDescriptorMap) {
  if (StringUtils.isBlank(stagingDirectory) || StringUtils.isBlank(streamName)) {
    LOG.warn("Staging directory ({}) or stream name ({}) is empty", stagingDirectory, streamName);
    return;
  }
  Path targetPath = PartitionDescriptorUtil.getPartitionDescriptorPath(stagingDirectory, streamName);
  try (FileSystem fs = targetPath.getFileSystem(new Configuration())) {
    // Partition descriptor is supposed to be immutable. So don't override it if it exists.
    if (fs.exists(targetPath)) {
      LOG.warn(targetPath.toString() + " exists. Skip persisting partition descriptor.");
    } else {
      LOG.info("About to persist partition descriptors to path: " + targetPath.toString());
      try (FSDataOutputStream fos = fs.create(targetPath)) {
        fos.write(
          PartitionDescriptorUtil.getJsonFromDescriptorMap(partitionDescriptorMap).getBytes(StandardCharsets.UTF_8));
      }
    }
  } catch (IOException e) {
    throw new SamzaException("Failed to validate/persist partition description on hdfs.", e);
  }
}
 

@Test
public void testMaxChangeLogStreamPartitions() {
  Config config = new MapConfig(ImmutableMap.of("a", "b"));
  Map<TaskName, TaskModel> tasksForContainer1 = ImmutableMap.of(
      new TaskName("t1"), new TaskModel(new TaskName("t1"), ImmutableSet.of(), new Partition(0)),
      new TaskName("t2"), new TaskModel(new TaskName("t2"), ImmutableSet.of(), new Partition(1)));
  Map<TaskName, TaskModel> tasksForContainer2 = ImmutableMap.of(
      new TaskName("t3"), new TaskModel(new TaskName("t3"), ImmutableSet.of(), new Partition(2)),
      new TaskName("t4"), new TaskModel(new TaskName("t4"), ImmutableSet.of(), new Partition(3)),
      new TaskName("t5"), new TaskModel(new TaskName("t5"), ImmutableSet.of(), new Partition(4)));
  ContainerModel containerModel1 = new ContainerModel("0", tasksForContainer1);
  ContainerModel containerModel2 = new ContainerModel("1", tasksForContainer2);
  Map<String, ContainerModel> containers = ImmutableMap.of("0", containerModel1, "1", containerModel2);
  JobModel jobModel = new JobModel(config, containers);
  assertEquals(jobModel.maxChangeLogStreamPartitions, 5);
}
 

@Override
public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
  return streamNames.stream()
      .collect(Collectors.toMap(Function.identity(), streamName -> {
        int messageCount = isBootstrapStream(streamName) ? getMessageCount(streamName) : -1;
        String oldestOffset = messageCount < 0 ? null : "0";
        String newestOffset = messageCount < 0 ? null : String.valueOf(messageCount - 1);
        String upcomingOffset = messageCount < 0 ? null : String.valueOf(messageCount);
        Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> metadataMap = new HashMap<>();
        int partitionCount = config.getInt("streams." + streamName + ".partitionCount", 1);
        for (int i = 0; i < partitionCount; i++) {
          metadataMap.put(new Partition(i), new SystemStreamMetadata.SystemStreamPartitionMetadata(
              oldestOffset, newestOffset, upcomingOffset
          ));
        }
        return new SystemStreamMetadata(streamName, metadataMap);
      }));
}
 

@Test
public void testEndOfStreamFlushesWithNoTriggerFirings() throws Exception {

  OperatorSpecGraph sgb =
      this.getKeyedSessionWindowStreamGraph(AccumulationMode.DISCARDING, Duration.ofMillis(500)).getOperatorSpecGraph();
  TestClock testClock = new TestClock();
  List<WindowPane<Integer, Collection<IntegerEnvelope>>> windowPanes = new ArrayList<>();
  StreamOperatorTask task = new StreamOperatorTask(sgb, testClock);
  task.init(this.context);

  MessageCollector messageCollector =
    envelope -> windowPanes.add((WindowPane<Integer, Collection<IntegerEnvelope>>) envelope.getMessage());

  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);
  task.processAsync(new IntegerEnvelope(1), messageCollector, taskCoordinator, taskCallback);

  final IncomingMessageEnvelope endOfStream = IncomingMessageEnvelope.buildEndOfStreamEnvelope(
      new SystemStreamPartition("kafka", "integers", new Partition(0)));
  task.processAsync(endOfStream, messageCollector, taskCoordinator, taskCallback);
  Assert.assertEquals(windowPanes.size(), 1);
  Assert.assertEquals(windowPanes.get(0).getMessage().size(), 4);
  verify(taskCoordinator, times(1)).commit(TaskCoordinator.RequestScope.CURRENT_TASK);
  verify(taskCoordinator, times(1)).shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
}
 

@Test
public void testPartitionCountMonitorWithoutDurableStates() {
  configMap.put(JobConfig.JOB_CONTAINER_COUNT, "1");
  when(CoordinatorStreamUtil.readConfigFromCoordinatorStream(anyObject())).thenReturn(new MapConfig(configMap));
  Config config = new MapConfig(configMap);

  // mimic job runner code to write the config to coordinator stream
  CoordinatorStreamSystemProducer producer = new CoordinatorStreamSystemProducer(config, mock(MetricsRegistry.class));
  producer.writeConfig("test-job", config);

  ClusterBasedJobCoordinator clusterCoordinator = ClusterBasedJobCoordinator.createFromMetadataStore(config);

  // change the input system stream metadata
  MockSystemFactory.MSG_QUEUES.put(new SystemStreamPartition("kafka", "topic1", new Partition(1)), new ArrayList<>());

  StreamPartitionCountMonitor monitor = clusterCoordinator.getPartitionMonitor();
  monitor.updatePartitionCountMetric();
  assertEquals(clusterCoordinator.getAppStatus(), SamzaApplicationState.SamzaAppStatus.UNDEFINED);
}
 

@Override
public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
  return streamNames.stream()
      .collect(
          Collectors.toMap(
              Function.identity(),
              stream -> {
                // Impulse system will always be single partition
                Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata>
                    partitionMetadata =
                        Collections.singletonMap(
                            new Partition(0),
                            new SystemStreamMetadata.SystemStreamPartitionMetadata(
                                DUMMY_OFFSET, DUMMY_OFFSET, DUMMY_OFFSET));
                return new SystemStreamMetadata(stream, partitionMetadata);
              }));
}
 

@Test
public void testFlushOrder() {
  ContainerStorageManager csm = mock(ContainerStorageManager.class);
  StorageEngine mockStore = mock(StorageEngine.class);
  java.util.Map<String, StorageEngine> taskStores = ImmutableMap.of("mockStore", mockStore);
  when(csm.getAllStores(any())).thenReturn(taskStores);

  TransactionalStateTaskStorageManager tsm = spy(buildTSM(csm, mock(Partition.class), new StorageManagerUtil()));
  // stub actual method call
  doReturn(mock(Map.class)).when(tsm).getNewestChangelogSSPOffsets(any(), any(), any(), any());

  // invoke flush
  tsm.flush();

  // ensure that stores are flushed before we get newest changelog offsets
  InOrder inOrder = inOrder(mockStore, tsm);
  inOrder.verify(mockStore).flush();
  inOrder.verify(tsm).getNewestChangelogSSPOffsets(any(), any(), any(), any());
}
 

@Test
public void testStoringAndReadingCheckpointsSamePartition() {
  Partition partition = new Partition(0);
  TaskName taskName = new TaskName("taskName0");
  SystemStreamPartition ssp = new SystemStreamPartition("Azure", "Stream", partition);
  Map<SystemStreamPartition, String> sspMap = new HashMap<>();

  sspMap.put(ssp, "12345");
  Checkpoint cp0 = new Checkpoint(sspMap);

  sspMap.put(ssp, "54321");
  Checkpoint cp1 = new Checkpoint(sspMap);

  checkpointManager.register(taskName);

  checkpointManager.writeCheckpoint(taskName, cp0);
  Checkpoint readCp = checkpointManager.readLastCheckpoint(taskName);
  Assert.assertEquals(cp0, readCp);

  checkpointManager.writeCheckpoint(taskName, cp1);
  Checkpoint readCp1 = checkpointManager.readLastCheckpoint(taskName);
  Assert.assertEquals(cp1, readCp1);
}
 

@Test
public void testGetSystemStreamMetaDataWithValidTopic() {
  System.out.println("STARTING");
  Map<String, SystemStreamMetadata> metadataMap =
      kafkaSystemAdmin.getSystemStreamMetadata(ImmutableSet.of(VALID_TOPIC));

  // verify metadata size
  assertEquals("metadata should return for 1 topic", metadataMap.size(), 1);
  System.out.println("STARTING1");
  // verify the metadata streamName
  assertEquals("the stream name should be " + VALID_TOPIC, metadataMap.get(VALID_TOPIC).getStreamName(), VALID_TOPIC);
  System.out.println("STARTING2");
  // verify the offset for each partition
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> systemStreamPartitionMetadata =
      metadataMap.get(VALID_TOPIC).getSystemStreamPartitionMetadata();
  assertEquals("there are 2 partitions", systemStreamPartitionMetadata.size(), 2);
  System.out.println("STARTING3");
  SystemStreamMetadata.SystemStreamPartitionMetadata partition0Metadata =
      systemStreamPartitionMetadata.get(new Partition(0));
  assertEquals("oldest offset for partition 0", partition0Metadata.getOldestOffset(),
      KAFKA_BEGINNING_OFFSET_FOR_PARTITION0.toString());
  assertEquals("upcoming offset for partition 0", partition0Metadata.getUpcomingOffset(),
      KAFKA_END_OFFSET_FOR_PARTITION0.toString());
  assertEquals("newest offset for partition 0", partition0Metadata.getNewestOffset(),
      Long.toString(KAFKA_END_OFFSET_FOR_PARTITION0 - 1));
  System.out.println("STARTING4");
  SystemStreamMetadata.SystemStreamPartitionMetadata partition1Metadata =
      systemStreamPartitionMetadata.get(new Partition(1));
  assertEquals("oldest offset for partition 1", partition1Metadata.getOldestOffset(),
      KAFKA_BEGINNING_OFFSET_FOR_PARTITION1.toString());
  assertEquals("upcoming offset for partition 1", partition1Metadata.getUpcomingOffset(),
      KAFKA_END_OFFSET_FOR_PARTITION1.toString());
  assertEquals("newest offset for partition 1", partition1Metadata.getNewestOffset(),
      Long.toString(KAFKA_END_OFFSET_FOR_PARTITION1 - 1));
}
 

@Test
public void testStartConsumer() {
  final Consumer consumer = Mockito.mock(Consumer.class);
  final KafkaConsumerProxyFactory kafkaConsumerProxyFactory = Mockito.mock(KafkaConsumerProxyFactory.class);

  final KafkaSystemConsumerMetrics kafkaSystemConsumerMetrics = new KafkaSystemConsumerMetrics(TEST_SYSTEM, new NoOpMetricsRegistry());
  final SystemStreamPartition testSystemStreamPartition1 = new SystemStreamPartition(TEST_SYSTEM, TEST_STREAM, new Partition(0));
  final SystemStreamPartition testSystemStreamPartition2 = new SystemStreamPartition(TEST_SYSTEM, TEST_STREAM, new Partition(1));
  final String testOffset = "1";
  final KafkaConsumerProxy kafkaConsumerProxy = Mockito.mock(KafkaConsumerProxy.class);

  Mockito.when(kafkaConsumerProxyFactory.create(Mockito.anyObject())).thenReturn(kafkaConsumerProxy);
  Mockito.doNothing().when(consumer).seek(new TopicPartition(TEST_STREAM, 0), 1);
  Mockito.doNothing().when(consumer).seek(new TopicPartition(TEST_STREAM, 1), 1);

  KafkaSystemConsumer kafkaSystemConsumer = new KafkaSystemConsumer(consumer, TEST_SYSTEM, new MapConfig(), TEST_CLIENT_ID, kafkaConsumerProxyFactory,
                                                                    kafkaSystemConsumerMetrics, new SystemClock());
  kafkaSystemConsumer.register(testSystemStreamPartition1, testOffset);
  kafkaSystemConsumer.register(testSystemStreamPartition2, testOffset);

  kafkaSystemConsumer.startConsumer();

  Mockito.verify(consumer).seek(new TopicPartition(TEST_STREAM, 0), 1);
  Mockito.verify(consumer).seek(new TopicPartition(TEST_STREAM, 1), 1);
  Mockito.verify(kafkaConsumerProxy).start();
  Mockito.verify(kafkaConsumerProxy).addTopicPartition(testSystemStreamPartition1, Long.valueOf(testOffset));
  Mockito.verify(kafkaConsumerProxy).addTopicPartition(testSystemStreamPartition2, Long.valueOf(testOffset));
}
 

@Test
public void testWriteOffsetFilesForPersistedStore() {
  final String storeName = "test-write-offset-persisted-store";
  final String storeName2 = "test-write-offset-persisted-store-2";

  final String taskName = "test-write-offset-for-persisted-task";
  final String offset = "123";
  final SystemStreamPartition ssp = new SystemStreamPartition("test-system", "test-stream", new Partition(0));
  final SystemStreamPartition ssp2 = new SystemStreamPartition("test-system2", "test-stream2", new Partition(0));

  Map<SystemStreamPartition, String> processedOffsets = ImmutableMap.of(ssp, offset, ssp2, offset);

  TaskSideInputStorageManager testSideInputStorageManager = new MockTaskSideInputStorageManagerBuilder(taskName, LOGGED_STORE_DIR)
      .addLoggedStore(storeName, ImmutableSet.of(ssp))
      .addLoggedStore(storeName2, ImmutableSet.of(ssp2))
      .build();

  initializeSideInputStorageManager(testSideInputStorageManager);
  testSideInputStorageManager.writeFileOffsets(processedOffsets);
  File storeDir = testSideInputStorageManager.getStoreLocation(storeName);

  assertTrue("Store directory: " + storeDir.getPath() + " is missing.", storeDir.exists());

  Map<SystemStreamPartition, String> fileOffsets = testSideInputStorageManager.getFileOffsets();

  assertTrue("Failed to get offset for ssp: " + ssp.toString() + " from file.", fileOffsets.containsKey(ssp));
  assertEquals("Mismatch between last processed offset and file offset.", fileOffsets.get(ssp), offset);

  assertTrue("Failed to get offset for ssp: " + ssp2.toString() + " from file.", fileOffsets.containsKey(ssp2));
  assertEquals("Mismatch between last processed offset and file offset.", fileOffsets.get(ssp2), offset);
}
 

@Test
public void testFlush() {
  final String storeName = "test-flush-store";
  final String taskName = "test-flush-task";
  final SystemStreamPartition ssp = new SystemStreamPartition("test-system", "test-stream", new Partition(0));
  final String offset = "123";
  final ImmutableMap<SystemStreamPartition, String> processedOffsets = ImmutableMap.of(ssp, offset);

  TaskSideInputStorageManager testSideInputStorageManager = new MockTaskSideInputStorageManagerBuilder(taskName, LOGGED_STORE_DIR)
      .addLoggedStore(storeName, ImmutableSet.of(ssp))
      .build();
  Map<String, StorageEngine> stores = new HashMap<>();

  initializeSideInputStorageManager(testSideInputStorageManager);
  testSideInputStorageManager.flush(processedOffsets);

  for (StorageEngine storageEngine : stores.values()) {
    verify(storageEngine).flush();
  }

  verify(testSideInputStorageManager).writeFileOffsets(eq(processedOffsets));

  File storeDir = testSideInputStorageManager.getStoreLocation(storeName);
  assertTrue("Store directory: " + storeDir.getPath() + " is missing.", storeDir.exists());

  Map<SystemStreamPartition, String> fileOffsets = testSideInputStorageManager.getFileOffsets();
  assertTrue("Failed to get offset for ssp: " + ssp.toString() + " from file.", fileOffsets.containsKey(ssp));
  assertEquals("Mismatch between last processed offset and file offset.", fileOffsets.get(ssp), offset);
}
 

/**
 * Given some SSPs, getSSPMetadata should delegate to getSystemStreamMetadata and properly extract the results for the
 * requested SSPs.
 */
@Test
public void testGetSSPMetadata() {
  SystemStreamPartition streamPartition0 = new SystemStreamPartition(SYSTEM, STREAM, new Partition(0));
  SystemStreamPartition streamPartition1 = new SystemStreamPartition(SYSTEM, STREAM, new Partition(1));
  SystemStreamPartition otherStreamPartition0 = new SystemStreamPartition(SYSTEM, OTHER_STREAM, new Partition(0));
  SystemAdmin systemAdmin = mock(MySystemAdmin.class);
  SystemStreamMetadata.SystemStreamPartitionMetadata streamPartition0Metadata =
      new SystemStreamMetadata.SystemStreamPartitionMetadata("1", "2", "3");
  SystemStreamMetadata.SystemStreamPartitionMetadata streamPartition1Metadata =
      new SystemStreamMetadata.SystemStreamPartitionMetadata("11", "12", "13");
  SystemStreamMetadata.SystemStreamPartitionMetadata otherStreamPartition0Metadata =
      new SystemStreamMetadata.SystemStreamPartitionMetadata("21", "22", "23");
  when(systemAdmin.getSystemStreamMetadata(ImmutableSet.of(STREAM, OTHER_STREAM))).thenReturn(ImmutableMap.of(
      STREAM, new SystemStreamMetadata(STREAM, ImmutableMap.of(
          new Partition(0), streamPartition0Metadata,
          new Partition(1), streamPartition1Metadata)),
      OTHER_STREAM, new SystemStreamMetadata(OTHER_STREAM, ImmutableMap.of(
          new Partition(0), otherStreamPartition0Metadata))));
  Set<SystemStreamPartition> ssps = ImmutableSet.of(streamPartition0, streamPartition1, otherStreamPartition0);
  when(systemAdmin.getSSPMetadata(ssps)).thenCallRealMethod();
  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> expected = ImmutableMap.of(
      streamPartition0, streamPartition0Metadata,
      streamPartition1, streamPartition1Metadata,
      otherStreamPartition0, otherStreamPartition0Metadata);
  assertEquals(expected, systemAdmin.getSSPMetadata(ssps));
  verify(systemAdmin).getSystemStreamMetadata(ImmutableSet.of(STREAM, OTHER_STREAM));
}
 

/**
 * Get the systemStreamPartitions of the broadcast stream. Specifying
 * one partition for one stream or a range of the partitions for one
 * stream is allowed.
 *
 * @return a Set of SystemStreamPartitions
 */
public Set<SystemStreamPartition> getBroadcastSystemStreamPartitions() {
  HashSet<SystemStreamPartition> systemStreamPartitionSet = new HashSet<>();
  List<String> systemStreamPartitions = getList(BROADCAST_INPUT_STREAMS, Collections.emptyList());

  for (String systemStreamPartition : systemStreamPartitions) {
    int hashPosition = systemStreamPartition.indexOf("#");
    if (hashPosition == -1) {
      throw new IllegalArgumentException("incorrect format in " + systemStreamPartition
          + ". Broadcast stream names should be in the form 'system.stream#partitionId' or 'system.stream#[partitionN-partitionM]'");
    } else {
      String systemStreamName = systemStreamPartition.substring(0, hashPosition);
      String partitionSegment = systemStreamPartition.substring(hashPosition + 1);
      SystemStream systemStream = StreamUtil.getSystemStreamFromNames(systemStreamName);

      if (Pattern.matches(BROADCAST_STREAM_PATTERN, partitionSegment)) {
        systemStreamPartitionSet.add(new SystemStreamPartition(systemStream, new Partition(Integer.valueOf(partitionSegment))));
      } else {
        if (Pattern.matches(BROADCAST_STREAM_RANGE_PATTERN, partitionSegment)) {
          int partitionStart = Integer.valueOf(partitionSegment.substring(1, partitionSegment.lastIndexOf("-")));
          int partitionEnd = Integer.valueOf(partitionSegment.substring(partitionSegment.lastIndexOf("-") + 1, partitionSegment.indexOf("]")));
          if (partitionStart > partitionEnd) {
            LOGGER.warn("The starting partition in stream " + systemStream.toString() + " is bigger than the ending Partition. No partition is added");
          }
          for (int i = partitionStart; i <= partitionEnd; i++) {
            systemStreamPartitionSet.add(new SystemStreamPartition(systemStream, new Partition(i)));
          }
        } else {
          throw new IllegalArgumentException("incorrect format in " + systemStreamPartition
              + ". Broadcast stream names should be in the form 'system.stream#partitionId' or 'system.stream#[partitionN-partitionM]'");
        }
      }
    }
  }
  return systemStreamPartitionSet;
}
 

@Test
public void testStartpointSpecific() {
  StartpointSpecific startpoint = new StartpointSpecific("123");
  Assert.assertEquals("123", startpoint.getSpecificOffset());
  Assert.assertTrue(startpoint.getCreationTimestamp() <= Instant.now().toEpochMilli());

  MockStartpointVisitor mockStartpointVisitorConsumer = new MockStartpointVisitor();
  startpoint.apply(new SystemStreamPartition("sys", "stream", new Partition(1)), mockStartpointVisitorConsumer);
  Assert.assertEquals(StartpointSpecific.class, mockStartpointVisitorConsumer.visitedClass);
}
 

@Test
public void testRemoveAllFanOuts() throws IOException {
  SystemStreamPartition sspBroadcast = new SystemStreamPartition("mockSystem1", "mockStream1", new Partition(2));
  SystemStreamPartition sspSingle = new SystemStreamPartition("mockSystem2", "mockStream2", new Partition(3));

  TaskName taskWithNonBroadcast = new TaskName("t1");

  List<TaskName> tasks =
      ImmutableList.of(new TaskName("t0"), taskWithNonBroadcast, new TaskName("t2"), new TaskName("t3"), new TaskName("t4"), new TaskName("t5"));

  Map<TaskName, Set<SystemStreamPartition>> taskToSSPs = tasks.stream()
      .collect(Collectors
          .toMap(task -> task, task -> task.equals(taskWithNonBroadcast) ? ImmutableSet.of(sspBroadcast, sspSingle) : ImmutableSet.of(sspBroadcast)));

  StartpointSpecific startpoint42 = new StartpointSpecific("42");

  startpointManager.writeStartpoint(sspBroadcast, startpoint42);
  startpointManager.writeStartpoint(sspSingle, startpoint42);

  // startpoint42 should remap with key sspBroadcast to all tasks + sspBroadcast
  Map<TaskName, Map<SystemStreamPartition, Startpoint>> tasksFannedOutTo = startpointManager.fanOut(taskToSSPs);
  Assert.assertEquals(tasks.size(), tasksFannedOutTo.size());
  Assert.assertTrue(tasksFannedOutTo.keySet().containsAll(tasks));
  Assert.assertFalse("Should be deleted after fan out", startpointManager.readStartpoint(sspBroadcast).isPresent());
  Assert.assertFalse("Should be deleted after fan out", startpointManager.readStartpoint(sspSingle).isPresent());

  startpointManager.removeAllFanOuts();

  // Write back to ensure removing all fan outs doesn't remove all startpoints
  startpointManager.writeStartpoint(sspBroadcast, startpoint42);
  startpointManager.writeStartpoint(sspSingle, startpoint42);

  Assert.assertEquals(0, startpointManager.getFanOutStore().all().size());
  Assert.assertTrue("Should not be deleted after remove all fan outs", startpointManager.readStartpoint(sspBroadcast).isPresent());
  Assert.assertTrue("Should not be deleted after remove all fan outs", startpointManager.readStartpoint(sspSingle).isPresent());
}
 

@Test
public void testStaleStartpoints() {
  SystemStreamPartition ssp = new SystemStreamPartition("mockSystem", "mockStream", new Partition(2));
  TaskName taskName = new TaskName("MockTask");

  long staleTimestamp = Instant.now().toEpochMilli() - StartpointManager.DEFAULT_EXPIRATION_DURATION.toMillis() - 2;
  StartpointTimestamp startpoint = new StartpointTimestamp(staleTimestamp, staleTimestamp);

  startpointManager.writeStartpoint(ssp, startpoint);
  Assert.assertFalse(startpointManager.readStartpoint(ssp).isPresent());

  startpointManager.writeStartpoint(ssp, taskName, startpoint);
  Assert.assertFalse(startpointManager.readStartpoint(ssp, taskName).isPresent());
}