public boolean registerDrainCallback(DrainCallback callback) {
    Preconditions.checkNotNull(callback);

    switch (state) {
      case RUNNING:
      case STOPPED:
        LOG.warn("Cannot register callback when it is in {} state. Please register callback before calling start "
            + "on DrainMonitor.", state);
        return false;
      case INIT:
        if (this.callback != null) {
          LOG.warn("Cannot register callback as a callback is already registered.");
          return false;
        }
        this.callback = callback;
        return true;
      default:
        return false;
    }
  }
Register a callback to be executed when DrainNotification is encountered.

@param callback the callback to register.
@return Returns {@code true} if registration was successful and {@code false} if not.
Registration can fail it the DrainMonitor is stopped or a callback is already registered.
samza-core/src/main/java/org/apache/samza/drain/DrainMonitor.java