org.apache.samza.system.SystemStreamMetadata Java Examples

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

/** The dummy stream created will only be used in Beam tests. */
private static InputDescriptor<OpMessage<String>, ?> createDummyStreamDescriptor(String id) {
  final GenericSystemDescriptor dummySystem =
      new GenericSystemDescriptor(id, InMemorySystemFactory.class.getName());
  final GenericInputDescriptor<OpMessage<String>> dummyInput =
      dummySystem.getInputDescriptor(id, new NoOpSerde<>());
  dummyInput.withOffsetDefault(SystemStreamMetadata.OffsetType.OLDEST);
  final Config config = new MapConfig(dummyInput.toConfig(), dummySystem.toConfig());
  final SystemFactory factory = new InMemorySystemFactory();
  final StreamSpec dummyStreamSpec = new StreamSpec(id, id, id, 1);
  factory.getAdmin(id, config).createStream(dummyStreamSpec);

  final SystemProducer producer = factory.getProducer(id, config, null);
  final SystemStream sysStream = new SystemStream(id, id);
  final Consumer<Object> sendFn =
      (msg) -> {
        producer.send(id, new OutgoingMessageEnvelope(sysStream, 0, null, msg));
      };
  final WindowedValue<String> windowedValue =
      WindowedValue.timestampedValueInGlobalWindow("dummy", new Instant());

  sendFn.accept(OpMessage.ofElement(windowedValue));
  sendFn.accept(new WatermarkMessage(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
  sendFn.accept(new EndOfStreamMessage(null));
  return dummyInput;
}
 

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

public void testAPIUsage() {
  // does not assert anything, but acts as a compile-time check on expected descriptor type parameters
  // and validates that the method calls can be chained.
  ExampleExpandingSystemDescriptor expandingSystem = new ExampleExpandingSystemDescriptor("expandingSystem");
  ExampleExpandingInputDescriptor<Long> input1 = expandingSystem.getInputDescriptor("input1", new IntegerSerde());
  ExampleExpandingOutputDescriptor<Integer> output1 = expandingSystem.getOutputDescriptor("output1", new IntegerSerde());

  input1
      .shouldBootstrap()
      .withOffsetDefault(SystemStreamMetadata.OffsetType.NEWEST)
      .withPriority(1)
      .shouldResetOffset()
      .withStreamConfigs(Collections.emptyMap());

  output1
      .withStreamConfigs(Collections.emptyMap());
}
 

@Test
public void testShouldGetASinglePartition() {
  SinglePartitionWithoutOffsetsSystemAdmin admin = new SinglePartitionWithoutOffsetsSystemAdmin();
  Set<String> streamNames = new HashSet<String>();
  streamNames.add("a");
  streamNames.add("b");

  Map<String, SystemStreamMetadata> metadata = admin.getSystemStreamMetadata(streamNames);
  assertEquals(2, metadata.size());
  SystemStreamMetadata metadata1 = metadata.get("a");
  SystemStreamMetadata metadata2 = metadata.get("b");

  assertEquals(1, metadata1.getSystemStreamPartitionMetadata().size());
  assertEquals(1, metadata2.getSystemStreamPartitionMetadata().size());
  assertNull(metadata.get(new SystemStreamPartition("test-system", "c", new Partition(0))));
}
 

@Override
public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
  return streamNames.stream()
      .collect(
          Collectors.toMap(
              Function.<String>identity(),
              streamName -> {
                try {
                  List<BoundedSource<T>> splits = split(source, pipelineOptions);
                  final Map<Partition, SystemStreamPartitionMetadata> partitionMetaData =
                      new HashMap<>();
                  // we assume that the generated splits are stable,
                  // this is necessary so that the mapping of partition to source is correct
                  // in each container.
                  for (int i = 0; i < splits.size(); i++) {
                    partitionMetaData.put(
                        new Partition(i), new SystemStreamPartitionMetadata(null, null, null));
                  }
                  return new SystemStreamMetadata(streamName, partitionMetaData);
                } catch (Exception e) {
                  throw new SamzaException("Fail to read stream metadata", e);
                }
              }));
}
 

@Test
public void testAPIUsage() {
  // does not assert anything, but acts as a compile-time check on expected descriptor type parameters
  // and validates that the method calls can be chained.
  ExampleTransformingSystemDescriptor imeTransformingSystem =
      new ExampleTransformingSystemDescriptor("imeTransformingSystem")
          .withSystemConfigs(Collections.emptyMap());
  ExampleTransformingInputDescriptor<Long> input1 = imeTransformingSystem.getInputDescriptor("input1", new IntegerSerde());
  ExampleTransformingOutputDescriptor<Integer> output1 = imeTransformingSystem.getOutputDescriptor("output1", new IntegerSerde());

  input1
      .shouldBootstrap()
      .withOffsetDefault(SystemStreamMetadata.OffsetType.NEWEST)
      .withPriority(1)
      .shouldResetOffset()
      .withStreamConfigs(Collections.emptyMap());

  output1
      .withStreamConfigs(Collections.emptyMap());
}
 

@Test
public void testAPIUsage() {
  // does not assert anything, but acts as a compile-time check on expected descriptor type parameters
  // and validates that the method calls can be chained.
  GenericSystemDescriptor mySystem =
      new GenericSystemDescriptor("input-system", "factory.class.name")
          .withSystemConfigs(Collections.emptyMap())
          .withDefaultStreamConfigs(Collections.emptyMap());
  GenericInputDescriptor<Integer> input1 = mySystem.getInputDescriptor("input1", new IntegerSerde());
  GenericOutputDescriptor<Integer> output1 = mySystem.getOutputDescriptor("output1", new IntegerSerde());

  input1
      .withPhysicalName("input-1")
      .shouldBootstrap()
      .withOffsetDefault(SystemStreamMetadata.OffsetType.NEWEST)
      .withPriority(1)
      .shouldResetOffset()
      .isBounded()
      .shouldDeleteCommittedMessages()
      .withStreamConfigs(Collections.emptyMap());

  output1
      .withPhysicalName("output-1")
      .withStreamConfigs(Collections.emptyMap());
}
 

@Test
public void testSDConfigs() {
  GenericSystemDescriptor mySystem =
      new GenericSystemDescriptor("input-system", "factory.class.name")
          .withSystemConfigs(ImmutableMap.of("custom-config-key", "custom-config-value"))
          .withDefaultStreamConfigs(ImmutableMap.of("custom-stream-config-key", "custom-stream-config-value"))
          .withDefaultStreamOffsetDefault(SystemStreamMetadata.OffsetType.UPCOMING);

  Map<String, String> generatedConfigs = mySystem.toConfig();
  Map<String, String> expectedConfigs = ImmutableMap.of(
      "systems.input-system.samza.factory", "factory.class.name",
      "systems.input-system.custom-config-key", "custom-config-value",
      "systems.input-system.default.stream.custom-stream-config-key", "custom-stream-config-value",
      "systems.input-system.default.stream.samza.offset.default", "upcoming"
  );
  assertEquals(expectedConfigs, generatedConfigs);
}
 

@Override
public void validateStream(StreamSpec streamSpec) throws StreamValidationException {
  LOG.info("About to validate stream = " + streamSpec);

  String streamName = streamSpec.getPhysicalName();
  SystemStreamMetadata systemStreamMetadata =
      getSystemStreamMetadata(Collections.singleton(streamName)).get(streamName);
  if (systemStreamMetadata == null) {
    throw new StreamValidationException(
        "Failed to obtain metadata for stream " + streamName + ". Validation failed.");
  }

  int actualPartitionCounter = systemStreamMetadata.getSystemStreamPartitionMetadata().size();
  int expectedPartitionCounter = streamSpec.getPartitionCount();
  LOG.info("actualCount=" + actualPartitionCounter + "; expectedCount=" + expectedPartitionCounter);
  if (actualPartitionCounter != expectedPartitionCounter) {
    throw new StreamValidationException(
        String.format("Mismatch of partitions for stream %s. Expected %d, got %d. Validation failed.", streamName,
            expectedPartitionCounter, actualPartitionCounter));
  }
}
 

@Test
public void testGetSSPMetadataEmptyPartition() {
  SystemStreamPartition ssp = new SystemStreamPartition(TEST_SYSTEM, VALID_TOPIC, new Partition(0));
  SystemStreamPartition otherSSP = new SystemStreamPartition(TEST_SYSTEM, "otherTopic", new Partition(1));
  TopicPartition topicPartition = new TopicPartition(VALID_TOPIC, 0);
  TopicPartition otherTopicPartition = new TopicPartition("otherTopic", 1);
  when(mockKafkaConsumer.beginningOffsets(ImmutableList.of(topicPartition, otherTopicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, 1L));
  when(mockKafkaConsumer.endOffsets(ImmutableList.of(topicPartition, otherTopicPartition))).thenReturn(
      ImmutableMap.of(topicPartition, 11L));

  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> expected =
      ImmutableMap.of(ssp, new SystemStreamMetadata.SystemStreamPartitionMetadata("1", "10", "11"), otherSSP,
          new SystemStreamMetadata.SystemStreamPartitionMetadata(null, null, null));
  assertEquals(expected, kafkaSystemAdmin.getSSPMetadata(ImmutableSet.of(ssp, otherSSP)));
}
 

@Test
public void testAPIUsage() {
  // does not assert anything, but acts as a compile-time check on expected descriptor type parameters
  // and validates that the method calls can be chained.
  ExampleSimpleSystemDescriptor kafkaSystem =
      new ExampleSimpleSystemDescriptor("kafka-system")
          .withSystemConfigs(Collections.emptyMap());
  ExampleSimpleInputDescriptor<Integer> input1 = kafkaSystem.getInputDescriptor("input1", new IntegerSerde());
  ExampleSimpleOutputDescriptor<Integer> output1 = kafkaSystem.getOutputDescriptor("output1", new IntegerSerde());

  input1
      .shouldBootstrap()
      .withOffsetDefault(SystemStreamMetadata.OffsetType.NEWEST)
      .withPriority(1)
      .shouldResetOffset()
      .withStreamConfigs(Collections.emptyMap());

  output1
      .withStreamConfigs(Collections.emptyMap());
}
 

@Test
public void testGetSystemStreamMetaDataWithRetry() {
  final List<PartitionInfo> partitionInfosForTopic = ImmutableList.of(mockPartitionInfo0, mockPartitionInfo1);
  when(mockKafkaConsumer.partitionsFor(VALID_TOPIC)).thenThrow(new RuntimeException())
      .thenReturn(partitionInfosForTopic);

  Map<String, SystemStreamMetadata> metadataMap =
      kafkaSystemAdmin.getSystemStreamMetadata(ImmutableSet.of(VALID_TOPIC));
  assertEquals("metadata should return for 1 topic", metadataMap.size(), 1);

  // retried twice because the first fails and the second succeeds
  Mockito.verify(mockKafkaConsumer, Mockito.times(2)).partitionsFor(VALID_TOPIC);

  final List<TopicPartition> topicPartitions =
      Arrays.asList(new TopicPartition(mockPartitionInfo0.topic(), mockPartitionInfo0.partition()),
          new TopicPartition(mockPartitionInfo1.topic(), mockPartitionInfo1.partition()));
  // the following methods thereafter are only called once
  Mockito.verify(mockKafkaConsumer, Mockito.times(1)).beginningOffsets(topicPartitions);
  Mockito.verify(mockKafkaConsumer, Mockito.times(1)).endOffsets(topicPartitions);
}
 

@Test
public void testGetSSPMetadataWithRetry() {
  SystemStreamPartition oneSSP = new SystemStreamPartition(TEST_SYSTEM, VALID_TOPIC, new Partition(0));
  SystemStreamPartition otherSSP = new SystemStreamPartition(TEST_SYSTEM, "otherTopic", new Partition(1));
  ImmutableSet<SystemStreamPartition> ssps = ImmutableSet.of(oneSSP, otherSSP);
  List<TopicPartition> topicPartitions = ssps.stream()
      .map(ssp -> new TopicPartition(ssp.getStream(), ssp.getPartition().getPartitionId()))
      .collect(Collectors.toList());
  Map<TopicPartition, Long> testBeginningOffsets =
      ImmutableMap.of(testTopicPartition0, KAFKA_BEGINNING_OFFSET_FOR_PARTITION0, testTopicPartition1,
          KAFKA_BEGINNING_OFFSET_FOR_PARTITION1);

  when(mockKafkaConsumer.beginningOffsets(topicPartitions)).thenThrow(new RuntimeException())
      .thenReturn(testBeginningOffsets);
  Map<SystemStreamPartition, SystemStreamMetadata.SystemStreamPartitionMetadata> sspMetadata =
      kafkaSystemAdmin.getSSPMetadata(ssps, new ExponentialSleepStrategy(2,
          1, 1));

  assertEquals("metadata should return for 2 topics", sspMetadata.size(), 2);

  // retried twice because the first fails and the second succeeds
  Mockito.verify(mockKafkaConsumer, Mockito.times(2)).beginningOffsets(topicPartitions);
}
 

/**
 * Default constructor.
 *
 * @param streamsToMonitor  a set of SystemStreams to monitor.
 * @param metadataCache     the metadata cache which will be used to fetch metadata for partition counts.
 * @param metrics           the metrics registry to which the metrics should be added.
 * @param monitorPeriodMs   the period at which the monitor will run in milliseconds.
 * @param monitorCallback   the callback method to be invoked when partition count changes are detected
 */
public StreamPartitionCountMonitor(Set<SystemStream> streamsToMonitor, StreamMetadataCache metadataCache,
    MetricsRegistry metrics, int monitorPeriodMs, Callback monitorCallback) {
  this.streamsToMonitor = streamsToMonitor;
  this.metadataCache = metadataCache;
  this.monitorPeriodMs = monitorPeriodMs;
  this.initialMetadata = getMetadata(streamsToMonitor, metadataCache);
  this.callbackMethod = monitorCallback;

  // Pre-populate the gauges
  Map<SystemStream, Gauge<Integer>> mutableGauges = new HashMap<>();
  for (Map.Entry<SystemStream, SystemStreamMetadata> metadataEntry : initialMetadata.entrySet()) {
    SystemStream systemStream = metadataEntry.getKey();
    Gauge gauge = metrics.newGauge("job-coordinator",
        String.format("%s-%s-partitionCount", systemStream.getSystem(), systemStream.getStream()), 0);
    mutableGauges.put(systemStream, gauge);
  }
  gauges = Collections.unmodifiableMap(mutableGauges);
}
 

/**
 * <p>
 *   Fetches the metadata of the topic partition of coordinator stream. Registers the oldest offset
 *   for the topic partition of coordinator stream with the coordinator system consumer.
 * </p>
 */
private void registerConsumer() {
  LOG.debug("Attempting to register system stream partition: {}", coordinatorSystemStreamPartition);
  String streamName = coordinatorSystemStreamPartition.getStream();
  Map<String, SystemStreamMetadata> systemStreamMetadataMap = systemAdmin.getSystemStreamMetadata(Sets.newHashSet(streamName));

  SystemStreamMetadata systemStreamMetadata = systemStreamMetadataMap.get(streamName);
  Preconditions.checkNotNull(systemStreamMetadata, String.format("System stream metadata does not exist for stream: %s.", streamName));

  SystemStreamPartitionMetadata systemStreamPartitionMetadata = systemStreamMetadata.getSystemStreamPartitionMetadata().get(coordinatorSystemStreamPartition.getPartition());
  Preconditions.checkNotNull(systemStreamPartitionMetadata, String.format("System stream partition metadata does not exist for: %s.", coordinatorSystemStreamPartition));

  String startingOffset = systemStreamPartitionMetadata.getOldestOffset();
  LOG.info("Registering system stream partition: {} with offset: {}.", coordinatorSystemStreamPartition, startingOffset);
  systemConsumer.register(coordinatorSystemStreamPartition, startingOffset);
}
 

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

@Override
public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
  return streamNames.stream()
      .collect(
          Collectors.toMap(
              Function.<String>identity(),
              streamName -> {
                try {
                  final List<UnboundedSource<T, CheckpointMarkT>> splits =
                      split(source, pipelineOptions);
                  final Map<Partition, SystemStreamPartitionMetadata> partitionMetaData =
                      new HashMap<>();
                  // we assume that the generated splits are stable,
                  // this is necessary so that the mapping of partition to source is correct
                  // in each container.
                  for (int i = 0; i < splits.size(); i++) {
                    partitionMetaData.put(
                        new Partition(i), new SystemStreamPartitionMetadata(null, null, null));
                  }
                  return new SystemStreamMetadata(streamName, partitionMetaData);
                } catch (Exception e) {
                  throw new SamzaException("Fail to read stream metadata", e);
                }
              }));
}
 

@Test
public void testSend() {
  SystemStreamMetadata metadata = mock(SystemStreamMetadata.class);
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>();
  partitionMetadata.put(new Partition(0), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(1), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(2), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(3), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  when(metadata.getSystemStreamPartitionMetadata()).thenReturn(partitionMetadata);
  StreamMetadataCache metadataCache = mock(StreamMetadataCache.class);
  when(metadataCache.getSystemStreamMetadata(anyObject(), anyBoolean())).thenReturn(metadata);

  SystemStream systemStream = new SystemStream("test-system", "test-stream");
  Set<Integer> partitions = new HashSet<>();
  MessageCollector collector = mock(MessageCollector.class);
  doAnswer(invocation -> {
    OutgoingMessageEnvelope envelope = (OutgoingMessageEnvelope) invocation.getArguments()[0];
    partitions.add((Integer) envelope.getPartitionKey());
    assertEquals(envelope.getSystemStream(), systemStream);
    return null;
  }).when(collector).send(any());

  ControlMessageSender sender = new ControlMessageSender(metadataCache);
  WatermarkMessage watermark = new WatermarkMessage(System.currentTimeMillis(), "task 0");
  sender.send(watermark, systemStream, collector);
  assertEquals(partitions.size(), 1);
}
 

/**
 * This test is for cases, when calls to systemAdmin (e.g., KafkaSystemAdmin's) get-stream-metadata method return null.
 */
@Test
public void testGetStartingOffsetsWhenStreamMetadataIsNull() {
  final String taskName = "test-get-starting-offset-task";

  Set<SystemStreamPartition> ssps = IntStream.range(1, 2)
      .mapToObj(idx -> new SystemStreamPartition(TEST_SYSTEM, TEST_STREAM, new Partition(idx)))
      .collect(Collectors.toSet());
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = ssps.stream()
      .collect(Collectors.toMap(SystemStreamPartition::getPartition,
        x -> new SystemStreamMetadata.SystemStreamPartitionMetadata(null, "1", "2")));


  TaskSideInputHandler handler = new MockTaskSideInputHandlerBuilder(taskName, TaskMode.Active)
      .addStreamMetadata(Collections.singletonMap(new SystemStream(TEST_SYSTEM, TEST_STREAM),
          new SystemStreamMetadata(TEST_STREAM, partitionMetadata)))
      .addStore(TEST_STORE, ssps)
      .build();

  handler.init();

  ssps.forEach(ssp -> {
    String startingOffset = handler.getStartingOffset(
        new SystemStreamPartition(TEST_SYSTEM, TEST_STREAM, ssp.getPartition()));
    Assert.assertNull("Starting offset should be null", startingOffset);
  });
}
 

@Test
public void testBroadcast() {
  SystemStreamMetadata metadata = mock(SystemStreamMetadata.class);
  Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>();
  partitionMetadata.put(new Partition(0), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(1), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(2), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  partitionMetadata.put(new Partition(3), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
  when(metadata.getSystemStreamPartitionMetadata()).thenReturn(partitionMetadata);
  StreamMetadataCache metadataCache = mock(StreamMetadataCache.class);
  when(metadataCache.getSystemStreamMetadata(anyObject(), anyBoolean())).thenReturn(metadata);

  SystemStream systemStream = new SystemStream("test-system", "test-stream");
  Set<Integer> partitions = new HashSet<>();
  MessageCollector collector = mock(MessageCollector.class);
  doAnswer(invocation -> {
    OutgoingMessageEnvelope envelope = (OutgoingMessageEnvelope) invocation.getArguments()[0];
    partitions.add((Integer) envelope.getPartitionKey());
    assertEquals(envelope.getSystemStream(), systemStream);
    return null;
  }).when(collector).send(any());

  ControlMessageSender sender = new ControlMessageSender(metadataCache);
  WatermarkMessage watermark = new WatermarkMessage(System.currentTimeMillis(), "task 0");
  SystemStreamPartition ssp = new SystemStreamPartition(systemStream, new Partition(0));
  sender.broadcastToOtherPartitions(watermark, ssp, collector);
  assertEquals(partitions.size(), 3);
}
 

@Override
public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
  return streamNames.stream()
      .collect(Collectors.toMap(Function.identity(), streamName -> new SystemStreamMetadata(streamName,
          Collections.singletonMap(new Partition(0),
              new SystemStreamMetadata.SystemStreamPartitionMetadata(null, END_OF_STREAM_OFFSET, null)))));
}
 

/**
 * Builds a map from SystemStreamPartition to oldest offset for changelogs.
 */
private Map<SystemStream, String> getChangeLogOldestOffsetsForPartition(Partition partition,
    Map<SystemStream, SystemStreamMetadata> inputStreamMetadata) {

  Map<SystemStream, String> retVal = new HashMap<>();

  // NOTE: do not use Collectors.Map because of https://bugs.openjdk.java.net/browse/JDK-8148463
  inputStreamMetadata.entrySet()
      .stream()
      .filter(x -> x.getValue().getSystemStreamPartitionMetadata().get(partition) != null)
      .forEach(e -> retVal.put(e.getKey(),
          e.getValue().getSystemStreamPartitionMetadata().get(partition).getOldestOffset()));

  return retVal;
}
 

/**
 * Get the oldest offset for each changelog SSP based on the stream's metadata (obtained from streamMetadataCache).
 */
private void getOldestChangeLogOffsets() {

  Map<SystemStream, SystemStreamMetadata> changeLogMetadata = JavaConverters.mapAsJavaMapConverter(
      streamMetadataCache.getStreamMetadata(
          JavaConverters.asScalaSetConverter(new HashSet<>(changelogSystemStreams.values())).asScala().toSet(),
          false)).asJava();

  LOG.info("Got change log stream metadata: {}", changeLogMetadata);

  changeLogOldestOffsets =
      getChangeLogOldestOffsetsForPartition(taskModel.getChangelogPartition(), changeLogMetadata);
  LOG.info("Assigning oldest change log offsets for taskName {} : {}", taskModel.getTaskName(),
      changeLogOldestOffsets);
}
 

private int getPartitionCount(SystemStream systemStream) {
  return PARTITION_COUNT_CACHE.computeIfAbsent(systemStream, ss -> {
    SystemStreamMetadata metadata = metadataCache.getSystemStreamMetadata(ss, true);
    if (metadata == null) {
      throw new SamzaException("Unable to find metadata for stream " + systemStream);
    }
    return metadata.getSystemStreamPartitionMetadata().size();
  });
}
 

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

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

  consumer.start();

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

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

/**
 * Gets the metadata for all the specified system streams from the provided metadata cache.
 * Handles scala-java conversions.
 *
 * @param streamsToMonitor  the set of system streams for which the metadata is needed.
 * @param metadataCache     the metadata cache which will be used to fetch metadata.
 * @return                  a map from each system stream to its metadata.
 */
private static Map<SystemStream, SystemStreamMetadata> getMetadata(Set<SystemStream> streamsToMonitor,
    StreamMetadataCache metadataCache) {
  return JavaConverters
      .mapAsJavaMapConverter(
          metadataCache.getStreamMetadata(
              JavaConverters.asScalaSetConverter(streamsToMonitor).asScala().toSet(),
              true
          )
      ).asJava();
}
 

@Ignore("Integration Test")
@Test
public void testGetStreamMetadata() {
  EventHubSystemFactory eventHubSystemFactory = new EventHubSystemFactory();
  SystemAdmin eventHubSystemAdmin = eventHubSystemFactory.getAdmin(SYSTEM_NAME,
          MockEventHubConfigFactory.getEventHubConfig(EventHubSystemProducer.PartitioningMethod.EVENT_HUB_HASHING));
  Set<String> streams = new HashSet<>();
  streams.add(STREAM_NAME1);
  streams.add(STREAM_NAME2);
  Map<String, SystemStreamMetadata> metadataMap = eventHubSystemAdmin.getSystemStreamMetadata(streams);

  for (String stream : streams) {
    Assert.assertTrue(metadataMap.containsKey(stream));
    Assert.assertEquals(stream, metadataMap.get(stream).getStreamName());
    Assert.assertNotNull(metadataMap.get(stream).getSystemStreamPartitionMetadata());

    Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadataMap =
            metadataMap.get(stream).getSystemStreamPartitionMetadata();
    Assert.assertTrue(partitionMetadataMap.size() >= MIN_EVENTHUB_ENTITY_PARTITION);
    Assert.assertTrue(partitionMetadataMap.size() <= MAX_EVENTHUB_ENTITY_PARTITION);
    partitionMetadataMap.forEach((partition, metadata) -> {
      Assert.assertEquals(EventHubSystemConsumer.START_OF_STREAM, metadata.getOldestOffset());
      Assert.assertNotSame(EventHubSystemConsumer.END_OF_STREAM, metadata.getNewestOffset());
      String expectedUpcomingOffset = String.valueOf(Long.parseLong(metadata.getNewestOffset()) + 1);
      Assert.assertEquals(expectedUpcomingOffset, metadata.getUpcomingOffset());
    });
  }
}
 

/**
 * Retrieves the oldest offset in the coordinator stream, and registers the
 * coordinator stream with the SystemConsumer using the earliest offset.
 */
public void register() {
  if (isStarted) {
    log.info("Coordinator stream partition {} has already been registered. Skipping.", coordinatorSystemStreamPartition);
    return;
  }
  log.debug("Attempting to register: {}", coordinatorSystemStreamPartition);
  Set<String> streamNames = new HashSet<String>();
  String streamName = coordinatorSystemStreamPartition.getStream();
  streamNames.add(streamName);
  Map<String, SystemStreamMetadata> systemStreamMetadataMap = systemAdmin.getSystemStreamMetadata(streamNames);
  log.info(String.format("Got metadata %s", systemStreamMetadataMap.toString()));

  if (systemStreamMetadataMap == null) {
    throw new SamzaException("Received a null systemStreamMetadataMap from the systemAdmin. This is illegal.");
  }

  SystemStreamMetadata systemStreamMetadata = systemStreamMetadataMap.get(streamName);

  if (systemStreamMetadata == null) {
    throw new SamzaException("Expected " + streamName + " to be in system stream metadata.");
  }

  SystemStreamPartitionMetadata systemStreamPartitionMetadata = systemStreamMetadata.getSystemStreamPartitionMetadata().get(coordinatorSystemStreamPartition.getPartition());

  if (systemStreamPartitionMetadata == null) {
    throw new SamzaException("Expected metadata for " + coordinatorSystemStreamPartition + " to exist.");
  }

  String startingOffset = systemStreamPartitionMetadata.getOldestOffset();
  log.debug("Registering {} with offset {}", coordinatorSystemStreamPartition, startingOffset);
  systemConsumer.register(coordinatorSystemStreamPartition, startingOffset);
}
 

@Override
public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) {
  Map<String, SystemStreamMetadata> requestedMetadata = new HashMap<>();

  try {
    for (String streamName : streamNames) {
      if (!streamPartitions.containsKey(streamName)) {
        LOG.debug(String.format("Partition ids for Stream=%s not found", streamName));

        EventHubClientManager eventHubClientManager = getOrCreateStreamEventHubClient(streamName);
        EventHubClient ehClient = eventHubClientManager.getEventHubClient();

        CompletableFuture<EventHubRuntimeInformation> runtimeInfo = ehClient.getRuntimeInformation();
        long timeoutMs = eventHubConfig.getRuntimeInfoWaitTimeMS(systemName);
        EventHubRuntimeInformation ehInfo = runtimeInfo.get(timeoutMs, TimeUnit.MILLISECONDS);
        LOG.info(String.format("Adding partition ids=%s for stream=%s. EHRuntimetInfo=%s",
            Arrays.toString(ehInfo.getPartitionIds()), streamName, printEventHubRuntimeInfo(ehInfo)));
        streamPartitions.put(streamName, ehInfo.getPartitionIds());
      }
      String[] partitionIds = streamPartitions.get(streamName);
      Map<Partition, SystemStreamPartitionMetadata> sspMetadataMap = getPartitionMetadata(streamName, partitionIds);
      SystemStreamMetadata systemStreamMetadata = new SystemStreamMetadata(streamName, sspMetadataMap);
      requestedMetadata.put(streamName, systemStreamMetadata);
    }
    return requestedMetadata;
  } catch (Exception e) {
    String msg = String.format("Error while fetching EventHubRuntimeInfo for System:%s", systemName);
    LOG.error(msg, e);
    throw new SamzaException(msg, e);
  }
}