public final CompletionStage<Void> aggregateDrainMessages(DrainMessage drainMessage, SystemStreamPartition ssp,
      MessageCollector collector, TaskCoordinator coordinator) {
    LOG.info("Received drain message from task {} in {}", drainMessage.getTaskName(), ssp);
    drainStates.update(drainMessage, ssp);

    SystemStream stream = ssp.getSystemStream();
    CompletionStage<Void> drainFuture = CompletableFuture.completedFuture(null);

    if (drainStates.isDrained(stream)) {
      LOG.info("Input {} is drained for task {}", stream.toString(), taskName.getTaskName());
      if (drainMessage.getTaskName() != null) {
        // This is the aggregation task which already received all the drain messages from upstream.
        // Broadcast the drain messages to all the peer partitions.
        controlMessageSender.broadcastToOtherPartitions(new DrainMessage(drainMessage.getRunId()), ssp, collector);
      }

      drainFuture = acceptFutureWithExecutor(onDrainOfStream(collector, coordinator), result -> {
        if (drainStates.areAllStreamsDrained()) {
          // All input streams have been drained, shut down the task
          LOG.info("All input streams have been drained for task {}. Requesting shutdown.", taskName.getTaskName());
          coordinator.commit(TaskCoordinator.RequestScope.CURRENT_TASK);
          coordinator.shutdown(TaskCoordinator.RequestScope.CURRENT_TASK);
        }
      });
    }

    return drainFuture;
  }
Aggregate {@link DrainMessage} from each ssp of the stream.
Invoke {@link #onDrainOfStream(MessageCollector, TaskCoordinator)} if the stream reaches the end.
@param drainMessage {@link DrainMessage} object
@param ssp system stream partition
@param collector message collector
@param coordinator task coordinator
samza-core/src/main/java/org/apache/samza/operators/impl/OperatorImpl.java