[EFM] Recoverable Random Beacon State Machine

engine/consensus/dkg/reactor_engine.go

@@ -110,7 +110,7 @@ func (e *ReactorEngine) Ready() <-chan struct{} {
 		if phase == flow.EpochPhaseSetup {
 			e.startDKGForEpoch(currentCounter, first)
 		} else if phase == flow.EpochPhaseCommitted {
-			// If we start up in EpochCommitted phase, ensure the DKG end state is set correctly.
+			// If we start up in EpochCommitted phase, ensure the DKG current state is set correctly.
 			e.handleEpochCommittedPhaseStarted(currentCounter, first)
 		}
 	})
@@ -249,10 +249,10 @@ func (e *ReactorEngine) startDKGForEpoch(currentEpochCounter uint64, first *flow
 //
 // This function checks that the local DKG completed and that our locally computed
 // key share is consistent with the canonical key vector. When this function returns,
-// an end state for the just-completed DKG is guaranteed to be stored (if not, the
+// an current state for the just-completed DKG is guaranteed to be stored (if not, the
 // program will crash). Since this function is invoked synchronously before the end
 // of the current epoch, this guarantees that when we reach the end of the current epoch
-// we will either have a usable beacon key (successful DKG) or a DKG failure end state
+// we will either have a usable beacon key (successful DKG) or a DKG failure current state
 // stored, so we can safely fall back to using our staking key.
 //
 // CAUTION: This function is not safe for concurrent use. This is not enforced within
@@ -267,13 +267,18 @@ func (e *ReactorEngine) handleEpochCommittedPhaseStarted(currentEpochCounter uin
 		Uint64("next_epoch", nextEpochCounter).   // the epoch the just-finished DKG was preparing for
 		Logger()
 
-	// Check whether we have already set the end state for this DKG.
+	// Check whether we have already set the current state for this DKG.
 	// This can happen if the DKG failed locally, if we failed to generate
 	// a local private beacon key, or if we crashed while performing this
 	// check previously.
 	currentState, err := e.dkgState.GetDKGState(nextEpochCounter)
 	if err != nil {
-		log.Fatal().Err(err).Msg("failed to get dkg state, by this point it should have been set")
+		if errors.Is(err, storage.ErrNotFound) {
+			log.Warn().Msg("failed to get dkg state, assuming this node has skipped epoch setup phase")
+		} else {
+			log.Fatal().Err(err).Msg("failed to get dkg state")
+		}
+
 		return
 	}
 	if currentState != flow.DKGStateCompleted {
@@ -299,7 +304,7 @@ func (e *ReactorEngine) handleEpochCommittedPhaseStarted(currentEpochCounter uin
 		err := e.dkgState.SetDKGState(nextEpochCounter, flow.DKGStateFailure)
 		if err != nil {
 			// TODO use irrecoverable context
-			log.Fatal().Err(err).Msg("failed to set dkg end state")
+			log.Fatal().Err(err).Msg("failed to set dkg state")
 		}
 		return
 	} else if err != nil {
@@ -316,7 +321,7 @@ func (e *ReactorEngine) handleEpochCommittedPhaseStarted(currentEpochCounter uin
 	}
 	localPubKey := myBeaconPrivKey.PublicKey()
 
-	// we computed a local beacon key but it is inconsistent with our canonical
+	// we computed a local beacon key, but it is inconsistent with our canonical
 	// public key - therefore it is unsafe for use
 	if !nextDKGPubKey.Equals(localPubKey) {
 		log.Warn().
@@ -326,15 +331,15 @@ func (e *ReactorEngine) handleEpochCommittedPhaseStarted(currentEpochCounter uin
 		err := e.dkgState.SetDKGState(nextEpochCounter, flow.DKGStateFailure)
 		if err != nil {
 			// TODO use irrecoverable context
-			log.Fatal().Err(err).Msg("failed to set dkg end state")
+			log.Fatal().Err(err).Msg("failed to set dkg current state")
 		}
 		return
 	}
 
 	err = e.dkgState.SetDKGState(nextEpochCounter, flow.RandomBeaconKeyCommitted)
 	if err != nil {
 		// TODO use irrecoverable context
-		e.log.Fatal().Err(err).Msg("failed to set dkg end state")
+		e.log.Fatal().Err(err).Msg("failed to set dkg current state")
 	}
 	log.Info().Msgf("successfully ended DKG, my beacon pub key for epoch %d is %s", nextEpochCounter, localPubKey)
 }
@@ -427,16 +432,17 @@ func (e *ReactorEngine) end(nextEpochCounter uint64) func() error {
 		if crypto.IsDKGFailureError(err) {
 			// Failing to complete the DKG protocol is a rare but expected scenario, which we must handle.
 			// By convention, if we are leaving the happy path, we want to persist the _first_ failure symptom
-			// in the `dkgState`. If the write yields a `storage.ErrAlreadyExists`, we know the overall protocol
-			// has already abandoned the happy path, because on the happy path the ReactorEngine is the only writer.
-			// Then this function just stops and returns without error.
+			// in the `dkgState`. If the write yields a [storage.InvalidDKGStateTransitionError], it means that the state machine
+			// is in the terminal state([flow.RandomBeaconKeyCommitted]) as all other transitions(even to [flow.DKGStateFailure] -> [flow.DKGStateFailure])
+			// are allowed. If the protocol is in terminal state, and we have a failure symptom, then it means that recovery has happened
+			// before ending the DKG. In this case, we want to ignore the error and return without error.
 			e.log.Warn().Err(err).Msgf("node %s with index %d failed DKG locally", e.me.NodeID(), e.controller.GetIndex())
 			err := e.dkgState.SetDKGState(nextEpochCounter, flow.DKGStateFailure)
 			if err != nil {
-				if errors.Is(err, storage.ErrAlreadyExists) {
-					return nil // DKGState already being set is expected in case of epoch recovery
+				if storage.IsInvalidDKGStateTransitionError(err) {
+					return nil
 				}
-				return fmt.Errorf("failed to set dkg end state following dkg end error: %w", err)
+				return fmt.Errorf("failed to set dkg current state following dkg end error: %w", err)
 			}
 			return nil // local DKG protocol has failed (the expected scenario)
 		} else if err != nil {

storage/badger/dkg_state.go

@@ -23,7 +23,7 @@ var allowedStateTransitions = map[flow.DKGState][]flow.DKGState{
 	flow.DKGStateStarted:          {flow.DKGStateCompleted, flow.DKGStateFailure, flow.RandomBeaconKeyCommitted},
 	flow.DKGStateCompleted:        {flow.RandomBeaconKeyCommitted, flow.DKGStateFailure},
 	flow.RandomBeaconKeyCommitted: {flow.RandomBeaconKeyCommitted},
-	flow.DKGStateFailure:          {flow.RandomBeaconKeyCommitted},
+	flow.DKGStateFailure:          {flow.RandomBeaconKeyCommitted, flow.DKGStateFailure},
 	flow.DKGStateUninitialized:    {flow.DKGStateStarted, flow.DKGStateFailure, flow.RandomBeaconKeyCommitted},
 }
 

storage/badger/dkg_state_test.go

@@ -275,10 +275,9 @@ func TestDKGState_FailureState(t *testing.T) {
 			require.True(t, storage.IsInvalidDKGStateTransitionError(err))
 		})
 
-		t.Run("-> flow.DKGStateFailure, not allowed", func(t *testing.T) {
+		t.Run("-> flow.DKGStateFailure, should be allowed", func(t *testing.T) {
 			err = store.SetDKGState(setupState(), flow.DKGStateFailure)
-			require.Error(t, err)
-			require.True(t, storage.IsInvalidDKGStateTransitionError(err))
+			require.NoError(t, err)
 		})
 
 		t.Run("-> flow.DKGStateCompleted, not allowed", func(t *testing.T) {