[BFT] Database migration for flow.DKGEndState

cmd/consensus/main.go

@@ -66,6 +66,7 @@ import (
 	"github.com/onflow/flow-go/state/protocol/events/gadgets"
 	protocol_state "github.com/onflow/flow-go/state/protocol/protocol_state/state"
 	bstorage "github.com/onflow/flow-go/storage/badger"
+	"github.com/onflow/flow-go/storage/badger/operation"
 	"github.com/onflow/flow-go/utils/io"
 )
 
@@ -199,6 +200,14 @@ func main() {
 	nodeBuilder.
 		PreInit(cmd.DynamicStartPreInit).
 		ValidateRootSnapshot(badgerState.ValidRootSnapshotContainsEntityExpiryRange).
+		PostInit(func(nodeConfig *cmd.NodeConfig) error {
+			// TODO(EFM, #6794): This function is introduced to implement a backward-compatible upgrade from v1 to v2.
+			// Remove this once we complete the network upgrade.
+			if err := operation.RetryOnConflict(nodeBuilder.DB.Update, operation.MigrateDKGEndStateFromV1()); err != nil {
+				return fmt.Errorf("could not migrate DKG end state from v1 to v2: %w", err)
+			}
+			return nil
+		}).
 		Module("machine account config", func(node *cmd.NodeConfig) error {
 			machineAccountInfo, err = cmd.LoadNodeMachineAccountInfoFile(node.BootstrapDir, node.NodeID)
 			return err

storage/badger/operation/dkg.go

@@ -1,7 +1,9 @@
 package operation
 
 import (
+	"encoding/binary"
 	"errors"
+	"fmt"
 
 	"github.com/dgraph-io/badger/v2"
 
@@ -74,3 +76,54 @@ func UpsertDKGStateForEpoch(epochCounter uint64, newState flow.DKGState) func(*b
 func RetrieveDKGStateForEpoch(epochCounter uint64, currentState *flow.DKGState) func(*badger.Txn) error {
 	return retrieve(makePrefix(codeDKGState, epochCounter), currentState)
 }
+
+// MigrateDKGEndStateFromV1 migrates the database that was used in protocol version v1 to the v2.
+// It reads already stored data by deprecated prefix and writes it to the new prefix with values converted to the new representation.
+// TODO(EFM, #6794): This function is introduced to implement a backward-compatible upgrade from v1 to v2.
+// Remove this once we complete the network upgrade.
+func MigrateDKGEndStateFromV1() func(txn *badger.Txn) error {
+	return func(txn *badger.Txn) error {
+		var ops []func(*badger.Txn) error
+		err := traverse(makePrefix(codeDKGEndState), func() (checkFunc, createFunc, handleFunc) {
+			var epochCounter uint64
+			check := func(key []byte) bool {
+				epochCounter = binary.BigEndian.Uint64(key[1:]) // omit code
+				return true
+			}
+			var oldState uint32
+			create := func() interface{} {
+				return &oldState
+			}
+			handle := func() error {
+				newState := flow.DKGStateUninitialized
+				switch oldState {
+				case 0: // DKGEndStateUnknown
+					newState = flow.DKGStateUninitialized
+				case 1: // DKGEndStateSuccess
+					newState = flow.RandomBeaconKeyCommitted
+				case 2, 3, 4: // DKGEndStateInconsistentKey, DKGEndStateNoKey, DKGEndStateDKGFailure
+					newState = flow.DKGStateFailure
+				}
+
+				// schedule upsert of the new state and removal of the old state
+				// this will be executed after to split collecting and modifying of data.
+				ops = append(ops,
+					UpsertDKGStateForEpoch(epochCounter, newState),
+					remove(makePrefix(codeDKGEndState, epochCounter)))
+
+				return nil
+			}
+			return check, create, handle
+		})(txn)
+		if err != nil {
+			return fmt.Errorf("could not collect deprecated DKG end states: %w", err)
+		}
+
+		for _, op := range ops {
+			if err := op(txn); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+}

storage/badger/operation/dkg_test.go

@@ -1,6 +1,7 @@
 package operation
 
 import (
+	"encoding/binary"
 	"math/rand"
 	"testing"
 
@@ -98,3 +99,76 @@ func TestDKGSetStateForEpoch(t *testing.T) {
 		assert.NoError(t, err)
 	})
 }
+
+// TestMigrateDKGEndStateFromV1 tests the migration of DKG end states from v1 to v2.
+// All possible states in v1 are generated and then checked against the expected states in v2.
+// Afterward the states are then migrated we check that old key was indeed removed and new key was added.
+// This test also checks that the migration is idempotent after the first run.
+func TestMigrateDKGEndStateFromV1(t *testing.T) {
+	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
+		epochCounter := rand.Uint64() % 100
+
+		preMigrationStates := make(map[uint64]uint32)
+		for i := epochCounter; i < epochCounter+100; i++ {
+			state := rand.Uint32() % 5 // [0,4] were supported in v1
+			err := db.Update(insert(makePrefix(codeDKGEndState, i), state))
+			assert.NoError(t, err)
+			preMigrationStates[i] = state
+		}
+
+		assertExpectedState := func(oldState uint32, newState flow.DKGState) {
+			switch oldState {
+			case 0: // DKGEndStateUnknown
+				assert.Equal(t, flow.DKGStateUninitialized, newState)
+			case 1: // DKGEndStateSuccess
+				assert.Equal(t, flow.RandomBeaconKeyCommitted, newState)
+			case 2, 3, 4: // DKGEndStateInconsistentKey, DKGEndStateNoKey, DKGEndStateDKGFailure
+				assert.Equal(t, flow.DKGStateFailure, newState)
+			default:
+				assert.Fail(t, "unexpected state")
+			}
+		}
+
+		// migrate the state
+		err := db.Update(MigrateDKGEndStateFromV1())
+		assert.NoError(t, err)
+
+		assertMigrationSuccessful := func() {
+			// ensure previous keys were removed
+			err = db.View(traverse(makePrefix(codeDKGEndState), func() (checkFunc, createFunc, handleFunc) {
+				assert.Fail(t, "no keys should have been found")
+				return nil, nil, nil
+			}))
+			assert.NoError(t, err)
+
+			migratedStates := make(map[uint64]flow.DKGState)
+			err = db.View(traverse(makePrefix(codeDKGState), func() (checkFunc, createFunc, handleFunc) {
+				var epochCounter uint64
+				check := func(key []byte) bool {
+					epochCounter = binary.BigEndian.Uint64(key[1:]) // omit code
+					return true
+				}
+				var newState flow.DKGState
+				create := func() interface{} {
+					return &newState
+				}
+				handle := func() error {
+					migratedStates[epochCounter] = newState
+					return nil
+				}
+				return check, create, handle
+			}))
+			assert.NoError(t, err)
+			assert.Equal(t, len(preMigrationStates), len(migratedStates))
+			for epochCounter, newState := range migratedStates {
+				assertExpectedState(preMigrationStates[epochCounter], newState)
+			}
+		}
+		assertMigrationSuccessful()
+
+		// migrating again should be no-op
+		err = db.Update(MigrateDKGEndStateFromV1())
+		assert.NoError(t, err)
+		assertMigrationSuccessful()
+	})
+}

storage/badger/operation/prefix.go

@@ -72,8 +72,9 @@ const (
 	codeEpochSetup         = 61 // EpochSetup service event, keyed by ID
 	codeEpochCommit        = 62 // EpochCommit service event, keyed by ID
 	codeBeaconPrivateKey   = 63 // BeaconPrivateKey, keyed by epoch counter
-	_                      = 64 // [DEPRECATED] flag that the DKG for an epoch has been started
-	codeDKGState           = 65 // current state of Recoverable Random Beacon State Machine for given epoch
+	_                      = 64 // [DEPRECATED] flag that the DKG for an epoch has been started, used in protocol version v1
+	codeDKGEndState        = 65 // [DEPRECATED] flag for DKG end state, used in protocol version v1
+	codeDKGState           = 66 // current state of Recoverable Random Beacon State Machine for given epoch
 	codeVersionBeacon      = 67 // flag for storing version beacons
 	codeEpochProtocolState = 68
 	codeProtocolKVStore    = 69