private Consumer<Transition<State>> dieOnError(final Consumer<Transition<State>> closure) {
   return transition -> {
     try {
       closure.accept(transition);
     } catch (RuntimeException e) {
       LOG.error("Caught unchecked exception: " + e, e);
       stateMachine.transition(State.DEAD);
       throw e;
     }
   };
 }
 @Override
 public void onDefeated() {
   LOG.error("Lost leadership, committing suicide.");
   stateMachine.transition(State.DEAD);
 }
 @Override
 public void onLeading(LeaderControl control) {
   leaderControl.set(control);
   stateMachine.transition(State.LEADER_AWAITING_REGISTRATION);
 }
 @Subscribe
 public void registered(DriverRegistered event) {
   stateMachine.transition(State.ACTIVE);
 }
 /**
  * Prepares a scheduler to offer itself as a leader candidate. After this call the scheduler will
  * host a live log replica and start syncing data from the leader via the log until it gets called
  * upon to lead.
  *
  * @return A listener that can be offered for leadership of a distributed election.
  */
 public LeadershipListener prepare() {
   stateMachine.transition(State.PREPARING_STORAGE);
   return leadershipListener;
 }
  @VisibleForTesting
  SchedulerLifecycle(
      final NonVolatileStorage storage,
      final Lifecycle lifecycle,
      final Driver driver,
      final DelayedActions delayedActions,
      final ShutdownRegistry shutdownRegistry,
      StatsProvider statsProvider,
      final ServiceManagerIface schedulerActiveServiceManager) {

    requireNonNull(storage);
    requireNonNull(lifecycle);
    requireNonNull(driver);
    requireNonNull(delayedActions);
    requireNonNull(shutdownRegistry);

    statsProvider.makeGauge(
        REGISTERED_GAUGE,
        new Supplier<Integer>() {
          @Override
          public Integer get() {
            return registrationAcked.get() ? 1 : 0;
          }
        });
    for (final State state : State.values()) {
      statsProvider.makeGauge(
          stateGaugeName(state),
          new Supplier<Integer>() {
            @Override
            public Integer get() {
              return (state == stateMachine.getState()) ? 1 : 0;
            }
          });
    }

    shutdownRegistry.addAction(
        new ExceptionalCommand<TimeoutException>() {
          @Override
          public void execute() throws TimeoutException {
            stateMachine.transition(State.DEAD);
            schedulerActiveServiceManager.stopAsync();
            schedulerActiveServiceManager.awaitStopped(5L, TimeUnit.SECONDS);
          }
        });

    final Consumer<Transition<State>> prepareStorage =
        new Consumer<Transition<State>>() {
          @Override
          public void accept(Transition<State> transition) {
            storage.prepare();
            stateMachine.transition(State.STORAGE_PREPARED);
          }
        };

    final Consumer<Transition<State>> handleLeading =
        new Consumer<Transition<State>>() {
          @Override
          public void accept(Transition<State> transition) {
            LOG.info("Elected as leading scheduler!");

            storage.start(
                stores -> {
                  // If storage backfill operations are necessary, they can be done here.
                });

            driver.startAsync().awaitRunning();

            delayedActions.onRegistrationTimeout(
                () -> {
                  if (!registrationAcked.get()) {
                    LOG.error("Framework has not been registered within the tolerated delay.");
                    stateMachine.transition(State.DEAD);
                  }
                });

            delayedActions.onAutoFailover(
                () -> {
                  LOG.info("Triggering automatic failover.");
                  stateMachine.transition(State.DEAD);
                });
          }
        };

    final Consumer<Transition<State>> handleRegistered =
        new Consumer<Transition<State>>() {
          @Override
          public void accept(Transition<State> transition) {
            registrationAcked.set(true);
            delayedActions.blockingDriverJoin(
                () -> {
                  driver.blockUntilStopped();
                  LOG.info("Driver exited, terminating lifecycle.");
                  stateMachine.transition(State.DEAD);
                });

            // TODO(ksweeney): Extract leader advertisement to its own service.
            schedulerActiveServiceManager.startAsync().awaitHealthy();
            try {
              leaderControl.get().advertise();
            } catch (SingletonService.AdvertiseException | InterruptedException e) {
              LOG.error("Failed to advertise leader, shutting down.");
              throw new RuntimeException(e);
            }
          }
        };

    final Consumer<Transition<State>> shutDown =
        new Consumer<Transition<State>>() {
          private final AtomicBoolean invoked = new AtomicBoolean(false);

          @Override
          public void accept(Transition<State> transition) {
            if (!invoked.compareAndSet(false, true)) {
              LOG.info("Shutdown already invoked, ignoring extra call.");
              return;
            }

            // TODO(wfarner): Consider using something like guava's Closer to abstractly tear down
            // resources here.
            try {
              LeaderControl control = leaderControl.get();
              if (control != null) {
                try {
                  control.leave();
                } catch (SingletonService.LeaveException e) {
                  LOG.warn("Failed to leave leadership: " + e, e);
                }
              }

              // TODO(wfarner): Re-evaluate tear-down ordering here.  Should the top-level shutdown
              // be invoked first, or the underlying critical components?
              driver.stopAsync().awaitTerminated();
              storage.stop();
            } finally {
              lifecycle.shutdown();
            }
          }
        };

    stateMachine =
        StateMachine.<State>builder("SchedulerLifecycle")
            .initialState(State.IDLE)
            .logTransitions()
            .addState(
                dieOnError(Consumers.filter(NOT_DEAD, prepareStorage)),
                State.IDLE,
                State.PREPARING_STORAGE,
                State.DEAD)
            .addState(State.PREPARING_STORAGE, State.STORAGE_PREPARED, State.DEAD)
            .addState(
                dieOnError(Consumers.filter(NOT_DEAD, handleLeading)),
                State.STORAGE_PREPARED,
                State.LEADER_AWAITING_REGISTRATION,
                State.DEAD)
            .addState(
                dieOnError(Consumers.filter(NOT_DEAD, handleRegistered)),
                State.LEADER_AWAITING_REGISTRATION,
                State.ACTIVE,
                State.DEAD)
            .addState(State.ACTIVE, State.DEAD)
            .addState(
                State.DEAD,
                // Allow cycles in DEAD to prevent throwing and avoid the need for call-site
                // checking.
                State.DEAD)
            .onAnyTransition(Consumers.filter(IS_DEAD, shutDown))
            .build();

    this.leadershipListener = new SchedulerCandidateImpl(stateMachine, leaderControl);
  }