org.apache.samza.system.SystemProducer Java Examples

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

@Test
public void testDiagnosticsManagerStart() {
  SystemProducer mockSystemProducer = Mockito.mock(SystemProducer.class);
  DiagnosticsManager diagnosticsManager =
      new DiagnosticsManager(jobName, jobId, containerModels, containerMb, containerNumCores, numPersistentStores,
          maxHeapSize, containerThreadPoolSize, "0", executionEnvContainerId, taskClassVersion, samzaVersion,
          hostname, diagnosticsSystemStream, mockSystemProducer, Duration.ofSeconds(1), mockExecutorService,
          autosizingEnabled);

  diagnosticsManager.start();

  Mockito.verify(mockSystemProducer, Mockito.times(1)).start();
  Mockito.verify(mockExecutorService, Mockito.times(1))
      .scheduleWithFixedDelay(Mockito.any(Runnable.class), Mockito.anyLong(), Mockito.anyLong(),
          Mockito.any(TimeUnit.class));
}
 

@Test
public void testDiagnosticsManagerStop() throws InterruptedException {
  SystemProducer mockSystemProducer = Mockito.mock(SystemProducer.class);
  Mockito.when(mockExecutorService.isTerminated()).thenReturn(true);
  Duration terminationDuration = Duration.ofSeconds(1);
  DiagnosticsManager diagnosticsManager =
      new DiagnosticsManager(jobName, jobId, containerModels, containerMb, containerNumCores, numPersistentStores,
          maxHeapSize, containerThreadPoolSize, "0", executionEnvContainerId, taskClassVersion, samzaVersion,
          hostname, diagnosticsSystemStream, mockSystemProducer, terminationDuration, mockExecutorService,
          autosizingEnabled);

  diagnosticsManager.stop();

  Mockito.verify(mockExecutorService, Mockito.times(1)).shutdown();
  Mockito.verify(mockExecutorService, Mockito.times(1))
      .awaitTermination(terminationDuration.toMillis(), TimeUnit.MILLISECONDS);
  Mockito.verify(mockExecutorService, Mockito.never()).shutdownNow();
  Mockito.verify(mockSystemProducer, Mockito.times(1)).stop();
}
 

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

public DiagnosticsManager(String jobName,
    String jobId,
    Map<String, ContainerModel> containerModels,
    int containerMemoryMb,
    int containerNumCores,
    int numPersistentStores,
    long maxHeapSizeBytes,
    int containerThreadPoolSize,
    String containerId,
    String executionEnvContainerId,
    String taskClassVersion,
    String samzaVersion,
    String hostname,
    SystemStream diagnosticSystemStream,
    SystemProducer systemProducer,
    Duration terminationDuration, boolean autosizingEnabled) {

  this(jobName, jobId, containerModels, containerMemoryMb, containerNumCores, numPersistentStores, maxHeapSizeBytes, containerThreadPoolSize,
      containerId, executionEnvContainerId, taskClassVersion, samzaVersion, hostname, diagnosticSystemStream, systemProducer,
      terminationDuration, Executors.newSingleThreadScheduledExecutor(
          new ThreadFactoryBuilder().setNameFormat(PUBLISH_THREAD_NAME).setDaemon(true).build()), autosizingEnabled);
}
 

@Test
public void testDiagnosticsManagerForceStop() throws InterruptedException {
  SystemProducer mockSystemProducer = Mockito.mock(SystemProducer.class);
  Mockito.when(mockExecutorService.isTerminated()).thenReturn(false);
  Duration terminationDuration = Duration.ofSeconds(1);
  DiagnosticsManager diagnosticsManager =
      new DiagnosticsManager(jobName, jobId, containerModels, containerMb, containerNumCores, numPersistentStores,
          maxHeapSize, containerThreadPoolSize, "0", executionEnvContainerId, taskClassVersion, samzaVersion,
          hostname, diagnosticsSystemStream, mockSystemProducer, terminationDuration, mockExecutorService,
          autosizingEnabled);

  diagnosticsManager.stop();

  Mockito.verify(mockExecutorService, Mockito.times(1)).shutdown();
  Mockito.verify(mockExecutorService, Mockito.times(1))
      .awaitTermination(terminationDuration.toMillis(), TimeUnit.MILLISECONDS);
  Mockito.verify(mockExecutorService, Mockito.times(1)).shutdownNow();
  Mockito.verify(mockSystemProducer, Mockito.times(1)).stop();
}
 

@Test
public void testBuildDiagnosticsManagerReturnsConfiguredReporter() {
  Config config = new MapConfig(buildTestConfigs());
  JobModel mockJobModel = mock(JobModel.class);
  SystemFactory systemFactory = mock(SystemFactory.class);
  SystemProducer mockProducer = mock(SystemProducer.class);
  MetricsReporterFactory metricsReporterFactory = mock(MetricsReporterFactory.class);
  MetricsSnapshotReporter mockReporter = mock(MetricsSnapshotReporter.class);

  when(systemFactory.getProducer(anyString(), any(Config.class), any(MetricsRegistry.class))).thenReturn(mockProducer);
  when(metricsReporterFactory.getMetricsReporter(anyString(), anyString(), any(Config.class))).thenReturn(
      mockReporter);
  PowerMockito.mockStatic(ReflectionUtil.class);
  when(ReflectionUtil.getObj(REPORTER_FACTORY, MetricsReporterFactory.class)).thenReturn(metricsReporterFactory);
  when(ReflectionUtil.getObj(SYSTEM_FACTORY, SystemFactory.class)).thenReturn(systemFactory);

  Optional<Pair<DiagnosticsManager, MetricsSnapshotReporter>> managerReporterPair =
      DiagnosticsUtil.buildDiagnosticsManager(JOB_NAME, JOB_ID, mockJobModel, CONTAINER_ID, Optional.of(ENV_ID),
          config);

  Assert.assertTrue(managerReporterPair.isPresent());
  Assert.assertEquals(mockReporter, managerReporterPair.get().getValue());
}
 

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

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

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

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

public CoordinatorStreamSystemProducer(Config config, MetricsRegistry registry) {
  SystemStream coordinatorSystemStream = CoordinatorStreamUtil.getCoordinatorSystemStream(config);
  SystemFactory systemFactory = CoordinatorStreamUtil.getCoordinatorSystemFactory(config);
  SystemAdmin systemAdmin = systemFactory.getAdmin(coordinatorSystemStream.getSystem(), config);
  SystemProducer systemProducer = systemFactory.getProducer(coordinatorSystemStream.getSystem(), config, registry);
  this.systemStream = coordinatorSystemStream;
  this.systemProducer = systemProducer;
  this.systemAdmin = systemAdmin;
  this.keySerde = new JsonSerde<>();
  this.messageSerde = new JsonSerde<>();
}
 

@Override
public SystemProducer getProducer(String name, Config config, MetricsRegistry metricsRegistry) {
  ElasticsearchConfig elasticsearchConfig = new ElasticsearchConfig(name, config);
  return new ElasticsearchSystemProducer(name,
                                         getBulkProcessorFactory(elasticsearchConfig),
                                         getClient(elasticsearchConfig),
                                         getIndexRequestFactory(elasticsearchConfig),
                                         new ElasticsearchSystemProducerMetrics(name, metricsRegistry));
}
 

@Override
public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
  EventHubConfig eventHubConfig = new EventHubConfig(config);
  return new EventHubSystemProducer(eventHubConfig, systemName, new EventHubClientManagerFactory(),
          getInterceptorsMap(eventHubConfig, systemName),
          registry);
}
 

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

public void start() throws IOException, InterruptedException {

    super.start();
    String source = "SystemProducerBench";

    int size = Integer.parseInt(cmd.getOptionValue(OPT_SHORT_MESSAGE_SIZE));
    RandomValueGenerator randGenerator = new RandomValueGenerator(System.currentTimeMillis());
    value = randGenerator.getNextString(size, size).getBytes();

    NoOpMetricsRegistry metricsRegistry = new NoOpMetricsRegistry();
    List<SystemStreamPartition> ssps = createSSPs(systemName, physicalStreamName, startPartition, endPartition);
    SystemProducer producer = factory.getProducer(systemName, config, metricsRegistry);
    producer.register(source);
    producer.start();

    System.out.println("starting production at " + Instant.now());
    Instant startTime = Instant.now();
    for (int index = 0; index < totalEvents; index++) {
      SystemStreamPartition ssp = ssps.get(index % ssps.size());
      OutgoingMessageEnvelope messageEnvelope = createMessageEnvelope(ssp, index);
      producer.send(source, messageEnvelope);
    }

    System.out.println("Ending production at " + Instant.now());
    System.out.println(String.format("Event Rate is %s Messages/Sec",
        totalEvents * 1000 / Duration.between(startTime, Instant.now()).toMillis()));

    producer.flush(source);

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

@Override
public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
  return new TestAvroSystemProducer();
}
 

@Override
public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
  return new MockSystemProducer();
}
 

@Override
public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
  return new LoggingSystemProducer();
}
 

@Before
public void setup() {
  producer = mock(SystemProducer.class);
  serializer = new MetricsSnapshotSerdeV2();
}
 

@Override
public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
  return new MockSystemProducer();
}
 

@Override
public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
  throw new RuntimeException("MockSystemProducer not implemented.");
}
 

@Override
public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) {
  return null;
}
 

@Override
public SystemProducer getProducer(String systemName, Config config, MetricsRegistry metricsRegistry) {
  // no producer
  return null;
}
 

@Override
public SystemProducer getProducer(String system, Config config, MetricsRegistry registry) {
  return null;
}