How to use

SamzaContainer

in

org.apache.samza.container

/**
 * This method returns the number of containers in the job. It works by querying the server, and getting the job model.
 * Then it extracts the number of containers from the job model
 *
 * @return current number of containers in the job
 */
public int getCurrentNumContainers() {
 return SamzaContainer.readJobModel(coordinatorServerURL, defaultReadJobModelDelayMs).getContainers().values().size();
}

@Override
public void onNewJobModel(String processorId, JobModel jobModel) {
 synchronized (lock) {
  if (state == State.IN_REBALANCE) {
   containerShutdownLatch = new CountDownLatch(1);
   container = createSamzaContainer(processorId, jobModel);
   container.setContainerListener(new ContainerListener());
   LOGGER.info("Starting the container: {} for the stream processor: {}.", container, processorId);
   containerExcecutorService.submit(container);
  } else {
   LOGGER.info("Ignoring onNewJobModel invocation since the current state is {} and not {}.", state, State.IN_REBALANCE);
  }
 }
}

/**
 * Stops the {@link SamzaContainer}.
 * @return true if {@link SamzaContainer} had shutdown within task.shutdown.ms. false otherwise.
 */
private boolean stopSamzaContainer() {
 boolean hasContainerShutdown = true;
 if (container != null) {
  try {
   container.shutdown();
   LOGGER.info("Waiting {} ms for the container: {} to shutdown.", taskShutdownMs, container);
   hasContainerShutdown = containerShutdownLatch.await(taskShutdownMs, TimeUnit.MILLISECONDS);
  } catch (InterruptedException e) {
   LOGGER.error("Exception occurred when shutting down the container: {}.", container, e);
   hasContainerShutdown = false;
   if (containerException != null) {
    containerException = e;
   }
  }
  LOGGER.info(String.format("Shutdown status of container: %s for stream processor: %s is: %b.", container, processorId, hasContainerShutdown));
 }
 // We want to propagate TimeoutException when container shutdown times out. It is possible that the timeout exception
 // we propagate to the application runner maybe overwritten by container failure cause in case of interleaved execution.
 // It is acceptable since container exception is much more useful compared to timeout exception.
 // We can infer from the logs about the fact that container shutdown timed out or not for additional inference.
 if (!hasContainerShutdown && containerException == null) {
  containerException = new TimeoutException("Container shutdown timed out after " + taskShutdownMs + " ms.");
 }
 return hasContainerShutdown;
}

@Test
public void testOnNewJobModelShouldResultInValidStateTransitions() throws Exception {
 JobCoordinator mockJobCoordinator = Mockito.mock(JobCoordinator.class);
 ProcessorLifecycleListener lifecycleListener = Mockito.mock(ProcessorLifecycleListener.class);
 SamzaContainer mockSamzaContainer = Mockito.mock(SamzaContainer.class);
 MapConfig config = new MapConfig(ImmutableMap.of("task.shutdown.ms", "0"));
 StreamProcessor streamProcessor = new TestableStreamProcessor(config, new HashMap<>(), null,
   lifecycleListener, mockJobCoordinator, mockSamzaContainer);
 streamProcessor.state = State.IN_REBALANCE;
 Mockito.doNothing().when(mockSamzaContainer).run();
 streamProcessor.jobCoordinatorListener.onNewJobModel("TestProcessorId", new JobModel(new MapConfig(), new HashMap<>()));
 Mockito.verify(mockSamzaContainer, Mockito.times(1)).setContainerListener(any());
 Mockito.verify(mockSamzaContainer, Mockito.atMost(1)).run();
}

public static void main(String[] args) throws Exception {
 Thread.setDefaultUncaughtExceptionHandler(
   new SamzaUncaughtExceptionHandler(() -> {
    log.info("Exiting process now.");
    System.exit(1);
   }));
 String containerId = System.getenv(ShellCommandConfig.ENV_CONTAINER_ID());
 log.info(String.format("Got container ID: %s", containerId));
 System.out.println(String.format("Container ID: %s", containerId));
 String coordinatorUrl = System.getenv(ShellCommandConfig.ENV_COORDINATOR_URL());
 log.info(String.format("Got coordinator URL: %s", coordinatorUrl));
 System.out.println(String.format("Coordinator URL: %s", coordinatorUrl));
 int delay = new Random().nextInt(SamzaContainer.DEFAULT_READ_JOBMODEL_DELAY_MS()) + 1;
 JobModel jobModel = SamzaContainer.readJobModel(coordinatorUrl, delay);
 Config config = jobModel.getConfig();
 JobConfig jobConfig = new JobConfig(config);
 if (jobConfig.getName().isEmpty()) {
  throw new SamzaException("can not find the job name");
 }
 String jobName = jobConfig.getName().get();
 String jobId = jobConfig.getJobId();
 MDC.put("containerName", "samza-container-" + containerId);
 MDC.put("jobName", jobName);
 MDC.put("jobId", jobId);
 ApplicationDescriptorImpl<? extends ApplicationDescriptor> appDesc =
   ApplicationDescriptorUtil.getAppDescriptor(ApplicationUtil.fromConfig(config), config);
 run(appDesc, containerId, jobModel, config);
 System.exit(0);
}

@Test
public void testStopShouldBeIdempotent() {
 JobCoordinator mockJobCoordinator = Mockito.mock(JobCoordinator.class);
 ProcessorLifecycleListener lifecycleListener = Mockito.mock(ProcessorLifecycleListener.class);
 SamzaContainer mockSamzaContainer = Mockito.mock(SamzaContainer.class);
 MapConfig config = new MapConfig(ImmutableMap.of("task.shutdown.ms", "0"));
 StreamProcessor streamProcessor = PowerMockito.spy(new StreamProcessor("TestProcessorId", config, new HashMap<>(), null, lifecycleListener, mockJobCoordinator));
 Mockito.doNothing().when(mockJobCoordinator).stop();
 Mockito.doNothing().when(mockSamzaContainer).shutdown();
 Mockito.when(mockSamzaContainer.hasStopped()).thenReturn(false);
 Mockito.when(mockSamzaContainer.getStatus())
     .thenReturn(SamzaContainerStatus.STARTED)
     .thenReturn(SamzaContainerStatus.STOPPED);
 streamProcessor.state = State.RUNNING;
 streamProcessor.stop();
 assertEquals(State.STOPPING, streamProcessor.state);
}

@VisibleForTesting
SamzaContainer createSamzaContainer(String processorId, JobModel jobModel) {
 return SamzaContainer.apply(processorId, jobModel, ScalaJavaUtil.toScalaMap(this.customMetricsReporter),
   this.taskFactory, JobContextImpl.fromConfigWithDefaults(this.config),
   Option.apply(this.applicationDefinedContainerContextFactoryOptional.orElse(null)),
   Option.apply(this.applicationDefinedTaskContextFactoryOptional.orElse(null)));
}

@Test
public void testCoordinatorFailureShouldStopTheStreamProcessor() {
 JobCoordinator mockJobCoordinator = Mockito.mock(JobCoordinator.class);
 ProcessorLifecycleListener lifecycleListener = Mockito.mock(ProcessorLifecycleListener.class);
 SamzaContainer mockSamzaContainer = Mockito.mock(SamzaContainer.class);
 MapConfig config = new MapConfig(ImmutableMap.of("task.shutdown.ms", "0"));
 StreamProcessor streamProcessor = new StreamProcessor("TestProcessorId", config, new HashMap<>(), null, lifecycleListener, mockJobCoordinator);
 Exception failureException = new Exception("dummy exception");
 streamProcessor.container = mockSamzaContainer;
 streamProcessor.state = State.RUNNING;
 streamProcessor.jobCoordinatorListener.onCoordinatorFailure(failureException);
 Mockito.doNothing().when(mockSamzaContainer).shutdown();
 Mockito.when(mockSamzaContainer.hasStopped()).thenReturn(false);
 assertEquals(State.STOPPED, streamProcessor.state);
 Mockito.verify(lifecycleListener).afterFailure(failureException);
 Mockito.verify(mockSamzaContainer).shutdown();
}

  .createInstance(new ProcessorContext() { }, config);
container.setContainerListener(
  new SamzaContainerListener() {
   @Override
container.run();
if (heartbeatMonitor != null) {
 heartbeatMonitor.stop();

public static void main(String[] args) throws Exception {
 Thread.setDefaultUncaughtExceptionHandler(
   new SamzaUncaughtExceptionHandler(() -> {
    log.info("Exiting process now.");
    System.exit(1);
   }));
 String containerId = System.getenv(ShellCommandConfig.ENV_CONTAINER_ID());
 log.info(String.format("Got container ID: %s", containerId));
 System.out.println(String.format("Container ID: %s", containerId));
 String coordinatorUrl = System.getenv(ShellCommandConfig.ENV_COORDINATOR_URL());
 log.info(String.format("Got coordinator URL: %s", coordinatorUrl));
 System.out.println(String.format("Coordinator URL: %s", coordinatorUrl));
 int delay = new Random().nextInt(SamzaContainer.DEFAULT_READ_JOBMODEL_DELAY_MS()) + 1;
 JobModel jobModel = SamzaContainer.readJobModel(coordinatorUrl, delay);
 Config config = jobModel.getConfig();
 JobConfig jobConfig = new JobConfig(config);
 if (jobConfig.getName().isEmpty()) {
  throw new SamzaException("can not find the job name");
 }
 String jobName = jobConfig.getName().get();
 String jobId = jobConfig.getJobId();
 MDC.put("containerName", "samza-container-" + containerId);
 MDC.put("jobName", jobName);
 MDC.put("jobId", jobId);
 ApplicationDescriptorImpl<? extends ApplicationDescriptor> appDesc =
   ApplicationDescriptorUtil.getAppDescriptor(ApplicationUtil.fromConfig(config), config);
 run(appDesc, containerId, jobModel, config);
 System.exit(0);
}

Mockito.doNothing().when(mockSamzaContainer).shutdown();
Mockito.when(mockSamzaContainer.hasStopped()).thenReturn(false);
Mockito.when(mockSamzaContainer.getStatus())
    .thenReturn(SamzaContainerStatus.STARTED)
    .thenReturn(SamzaContainerStatus.STOPPED);
Mockito.verify(mockSamzaContainer, Mockito.times(1)).shutdown();
Mockito.verify(mockJobCoordinator, Mockito.times(1)).stop();

@VisibleForTesting
SamzaContainer createSamzaContainer(String processorId, JobModel jobModel) {
 return SamzaContainer.apply(processorId, jobModel, ScalaJavaUtil.toScalaMap(this.customMetricsReporter),
   this.taskFactory, JobContextImpl.fromConfigWithDefaults(this.config),
   Option.apply(this.applicationDefinedContainerContextFactoryOptional.orElse(null)),
   Option.apply(this.applicationDefinedTaskContextFactoryOptional.orElse(null)));
}

  .createInstance(new ProcessorContext() { }, config);
container.setContainerListener(
  new SamzaContainerListener() {
   @Override
container.run();
if (heartbeatMonitor != null) {
 heartbeatMonitor.stop();

System.out.println(String.format("Coordinator URL: %s", coordinatorUrl));
int delay = new Random().nextInt(SamzaContainer.DEFAULT_READ_JOBMODEL_DELAY_MS()) + 1;
JobModel jobModel = SamzaContainer.readJobModel(coordinatorUrl, delay);
Config config = jobModel.getConfig();
JobConfig jobConfig = new JobConfig(config);

/**
 * Stops the {@link SamzaContainer}.
 * @return true if {@link SamzaContainer} had shutdown within task.shutdown.ms. false otherwise.
 */
private boolean stopSamzaContainer() {
 boolean hasContainerShutdown = true;
 if (container != null) {
  try {
   container.shutdown();
   LOGGER.info("Waiting {} ms for the container: {} to shutdown.", taskShutdownMs, container);
   hasContainerShutdown = containerShutdownLatch.await(taskShutdownMs, TimeUnit.MILLISECONDS);
  } catch (InterruptedException e) {
   LOGGER.error("Exception occurred when shutting down the container: {}.", container, e);
   hasContainerShutdown = false;
   if (containerException != null) {
    containerException = e;
   }
  }
  LOGGER.info(String.format("Shutdown status of container: %s for stream processor: %s is: %b.", container, processorId, hasContainerShutdown));
 }
 // We want to propagate TimeoutException when container shutdown times out. It is possible that the timeout exception
 // we propagate to the application runner maybe overwritten by container failure cause in case of interleaved execution.
 // It is acceptable since container exception is much more useful compared to timeout exception.
 // We can infer from the logs about the fact that container shutdown timed out or not for additional inference.
 if (!hasContainerShutdown && containerException == null) {
  containerException = new TimeoutException("Container shutdown timed out after " + taskShutdownMs + " ms.");
 }
 return hasContainerShutdown;
}

@Override
public void onNewJobModel(String processorId, JobModel jobModel) {
 synchronized (lock) {
  if (state == State.IN_REBALANCE) {
   containerShutdownLatch = new CountDownLatch(1);
   container = createSamzaContainer(processorId, jobModel);
   container.setContainerListener(new ContainerListener());
   LOGGER.info("Starting the container: {} for the stream processor: {}.", container, processorId);
   containerExcecutorService.submit(container);
  } else {
   LOGGER.info("Ignoring onNewJobModel invocation since the current state is {} and not {}.", state, State.IN_REBALANCE);
  }
 }
}

@VisibleForTesting
SamzaContainer createSamzaContainer(String processorId, JobModel jobModel) {
 return SamzaContainer.apply(processorId, jobModel, ScalaJavaUtil.toScalaMap(this.customMetricsReporter),
   this.taskFactory, JobContextImpl.fromConfigWithDefaults(this.config),
   Option.apply(this.applicationDefinedContainerContextFactoryOptional.orElse(null)),
   Option.apply(this.applicationDefinedTaskContextFactoryOptional.orElse(null)));
}

/**
 * This method returns the number of tasks in the job. It works by querying the server, and getting the job model.
 * Then it extracts the number of tasks from the job model
 *
 * @return current number of tasks in the job
 */
public int getCurrentNumTasks() {
 int currentNumTasks = 0;
 for (ContainerModel containerModel : SamzaContainer.readJobModel(coordinatorServerURL, defaultReadJobModelDelayMs).getContainers().values()) {
  currentNumTasks += containerModel.getTasks().size();
 }
 return currentNumTasks;
}

  .createInstance(new ProcessorContext() { }, config);
container.setContainerListener(
  new SamzaContainerListener() {
   @Override
container.run();
if (heartbeatMonitor != null) {
 heartbeatMonitor.stop();

public static void main(String[] args) throws Exception {
 Thread.setDefaultUncaughtExceptionHandler(
   new SamzaUncaughtExceptionHandler(() -> {
    log.info("Exiting process now.");
    System.exit(1);
   }));
 String containerId = System.getenv(ShellCommandConfig.ENV_CONTAINER_ID());
 log.info(String.format("Got container ID: %s", containerId));
 System.out.println(String.format("Container ID: %s", containerId));
 String coordinatorUrl = System.getenv(ShellCommandConfig.ENV_COORDINATOR_URL());
 log.info(String.format("Got coordinator URL: %s", coordinatorUrl));
 System.out.println(String.format("Coordinator URL: %s", coordinatorUrl));
 int delay = new Random().nextInt(SamzaContainer.DEFAULT_READ_JOBMODEL_DELAY_MS()) + 1;
 JobModel jobModel = SamzaContainer.readJobModel(coordinatorUrl, delay);
 Config config = jobModel.getConfig();
 JobConfig jobConfig = new JobConfig(config);
 if (jobConfig.getName().isEmpty()) {
  throw new SamzaException("can not find the job name");
 }
 String jobName = jobConfig.getName().get();
 String jobId = jobConfig.getJobId();
 MDC.put("containerName", "samza-container-" + containerId);
 MDC.put("jobName", jobName);
 MDC.put("jobId", jobId);
 ApplicationDescriptorImpl<? extends ApplicationDescriptor> appDesc =
   ApplicationDescriptorUtil.getAppDescriptor(ApplicationUtil.fromConfig(config), config);
 run(appDesc, containerId, jobModel, config);
 System.exit(0);
}