Merge pull request #7964 from onflow/leo/refactor-epoch-protocol-state

[Storage] Refactor epoch protocol state

Author: zhangchiqing

state/protocol/badger/state.go

@@ -256,7 +256,7 @@ func bootstrapProtocolState(
 	for _, proposal := range segment.AllBlocks() {
 		blockID := proposal.Block.ID()
 		protocolStateEntryWrapper := segment.ProtocolStateEntries[proposal.Block.Payload.ProtocolStateID]
-		err := epochProtocolStateSnapshots.BatchIndex(rw, blockID, protocolStateEntryWrapper.EpochEntry.ID())
+		err := epochProtocolStateSnapshots.BatchIndex(lctx, rw, blockID, protocolStateEntryWrapper.EpochEntry.ID())
 		if err != nil {
 			return fmt.Errorf("could not index root protocol state: %w", err)
 		}

state/protocol/protocol_state/epochs/statemachine.go

@@ -242,8 +242,8 @@ func NewEpochStateMachine(
 func (e *EpochStateMachine) Build() (*deferred.DeferredBlockPersist, error) {
 	updatedEpochState, updatedStateID, hasChanges := e.activeStateMachine.Build()
 
-	e.pendingDBUpdates.AddNextOperation(func(_ lockctx.Proof, blockID flow.Identifier, rw storage.ReaderBatchWriter) error {
-		return e.epochProtocolStateDB.BatchIndex(rw, blockID, updatedStateID)
+	e.pendingDBUpdates.AddNextOperation(func(lctx lockctx.Proof, blockID flow.Identifier, rw storage.ReaderBatchWriter) error {
+		return e.epochProtocolStateDB.BatchIndex(lctx, rw, blockID, updatedStateID)
 	})
 
 	if hasChanges {

state/protocol/protocol_state/epochs/statemachine_test.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"testing"
 
+	"github.com/jordanschalm/lockctx"
 	"github.com/stretchr/testify/assert"
 	mocks "github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -16,6 +17,7 @@ import (
 	"github.com/onflow/flow-go/state/protocol/protocol_state/epochs"
 	"github.com/onflow/flow-go/state/protocol/protocol_state/epochs/mock"
 	protocol_statemock "github.com/onflow/flow-go/state/protocol/protocol_state/mock"
+	"github.com/onflow/flow-go/storage"
 	storagemock "github.com/onflow/flow-go/storage/mock"
 	"github.com/onflow/flow-go/utils/unittest"
 )
@@ -40,6 +42,7 @@ type EpochStateMachineSuite struct {
 	happyPathStateMachineFactory    *mock.StateMachineFactoryMethod
 	fallbackPathStateMachineFactory *mock.StateMachineFactoryMethod
 	candidate                       *flow.Header
+	lockManager                     lockctx.Manager
 
 	stateMachine *epochs.EpochStateMachine
 }
@@ -56,6 +59,7 @@ func (s *EpochStateMachineSuite) SetupTest() {
 	s.happyPathStateMachine = mock.NewStateMachine(s.T())
 	s.happyPathStateMachineFactory = mock.NewStateMachineFactoryMethod(s.T())
 	s.fallbackPathStateMachineFactory = mock.NewStateMachineFactoryMethod(s.T())
+	s.lockManager = storage.NewTestingLockManager()
 
 	s.epochStateDB.On("ByBlockID", mocks.Anything).Return(func(_ flow.Identifier) *flow.RichEpochStateEntry {
 		return s.parentEpochState
@@ -97,17 +101,21 @@ func (s *EpochStateMachineSuite) TestBuild_NoChanges() {
 
 	rw := storagemock.NewReaderBatchWriter(s.T())
 
-	s.epochStateDB.On("BatchIndex", rw, s.candidate.ID(), s.parentEpochState.ID()).Return(nil).Once()
-	s.mutator.On("SetEpochStateID", s.parentEpochState.ID()).Return(nil).Once()
+	// Create a proper lock context proof for the BatchIndex operation
+	err = unittest.WithLock(s.T(), s.lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
+		s.epochStateDB.On("BatchIndex", lctx, rw, s.candidate.ID(), s.parentEpochState.ID()).Return(nil).Once()
+		s.mutator.On("SetEpochStateID", s.parentEpochState.ID()).Return(nil).Once()
 
-	dbUpdates, err := s.stateMachine.Build()
-	require.NoError(s.T(), err)
+		dbUpdates, err := s.stateMachine.Build()
+		require.NoError(s.T(), err)
 
-	// Provide the blockID and execute the resulting `dbUpdates`. Thereby, the expected mock methods should be called,
-	// which is asserted by the testify framework. Passing nil lockctx proof because no operations require lock;
-	// operations are deferred only because block ID is not known yet.
-	blockID := s.candidate.ID()
-	err = dbUpdates.Execute(nil, blockID, rw)
+		// Storage operations are deferred, because block ID is not known when the block is newly constructed. Only at the
+		// end after the block is fully constructed, its ID can be computed. We emulate this step here to verify that the
+		// deferred `dbOps` have been correctly constructed. Thereby, the expected mock methods should be called,
+		// which is asserted by the testify framework.
+		blockID := s.candidate.ID()
+		return dbUpdates.Execute(lctx, blockID, rw)
+	})
 	require.NoError(s.T(), err)
 }
 
@@ -139,19 +147,22 @@ func (s *EpochStateMachineSuite) TestBuild_HappyPath() {
 	err := s.stateMachine.EvolveState([]flow.ServiceEvent{epochSetup.ServiceEvent(), epochCommit.ServiceEvent()})
 	require.NoError(s.T(), err)
 
-	// prepare a DB update for epoch state
-	s.epochStateDB.On("BatchIndex", rw, s.candidate.ID(), updatedStateID).Return(nil).Once()
-	s.epochStateDB.On("BatchStore", w, updatedStateID, updatedState.MinEpochStateEntry).Return(nil).Once()
-	s.mutator.On("SetEpochStateID", updatedStateID).Return(nil).Once()
+	// Create a proper lock context proof for the BatchIndex operation
+	err = unittest.WithLock(s.T(), s.lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
+		// prepare a DB update for epoch state
+		s.epochStateDB.On("BatchIndex", lctx, rw, s.candidate.ID(), updatedStateID).Return(nil).Once()
+		s.epochStateDB.On("BatchStore", w, updatedStateID, updatedState.MinEpochStateEntry).Return(nil).Once()
+		s.mutator.On("SetEpochStateID", updatedStateID).Return(nil).Once()
 
-	dbUpdates, err := s.stateMachine.Build()
-	require.NoError(s.T(), err)
+		dbUpdates, err := s.stateMachine.Build()
+		require.NoError(s.T(), err)
 
-	// Provide the blockID and execute the resulting `dbUpdates`. Thereby, the expected mock methods should be called,
-	// which is asserted by the testify framework. Passing nil lockctx proof because no operations require lock;
-	// operations are deferred only because block ID is not known yet.
-	blockID := s.candidate.ID()
-	err = dbUpdates.Execute(nil, blockID, rw)
+		// Provide the blockID and execute the resulting `dbUpdates`. Thereby, the expected mock methods should be called,
+		// which is asserted by the testify framework. The lock context proof is passed to verify that the BatchIndex
+		// operation receives the proper lock context as required by the storage layer.
+		blockID := s.candidate.ID()
+		return dbUpdates.Execute(lctx, blockID, rw)
+	})
 	require.NoError(s.T(), err)
 }
 
@@ -532,29 +543,34 @@ func (s *EpochStateMachineSuite) TestEvolveStateTransitionToNextEpoch_WithInvali
 	err = stateMachine.EvolveState([]flow.ServiceEvent{invalidServiceEvent.ServiceEvent()})
 	require.NoError(s.T(), err)
 
-	s.epochStateDB.On("BatchIndex", mocks.Anything, s.candidate.ID(), mocks.Anything).Return(nil).Once()
+	// Create a proper lock context proof for the BatchIndex operation
+	err = unittest.WithLock(s.T(), s.lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
+		s.epochStateDB.On("BatchIndex", lctx, mocks.Anything, s.candidate.ID(), mocks.Anything).Return(nil).Once()
 
-	expectedEpochState := &flow.MinEpochStateEntry{
-		PreviousEpoch:          s.parentEpochState.CurrentEpoch.Copy(),
-		CurrentEpoch:           *s.parentEpochState.NextEpoch.Copy(),
-		NextEpoch:              nil,
-		EpochFallbackTriggered: true,
-	}
+		expectedEpochState := &flow.MinEpochStateEntry{
+			PreviousEpoch:          s.parentEpochState.CurrentEpoch.Copy(),
+			CurrentEpoch:           *s.parentEpochState.NextEpoch.Copy(),
+			NextEpoch:              nil,
+			EpochFallbackTriggered: true,
+		}
 
-	s.epochStateDB.On("BatchStore", mocks.Anything, expectedEpochState.ID(), expectedEpochState).Return(nil).Once()
-	s.mutator.On("SetEpochStateID", expectedEpochState.ID()).Return().Once()
+		s.epochStateDB.On("BatchStore", mocks.Anything, expectedEpochState.ID(), expectedEpochState).Return(nil).Once()
+		s.mutator.On("SetEpochStateID", expectedEpochState.ID()).Return().Once()
 
-	dbOps, err := stateMachine.Build()
-	require.NoError(s.T(), err)
+		dbOps, err := stateMachine.Build()
+		require.NoError(s.T(), err)
 
-	w := storagemock.NewWriter(s.T())
-	rw := storagemock.NewReaderBatchWriter(s.T())
-	rw.On("Writer").Return(w).Once() // called by epochStateDB.BatchStore
+		w := storagemock.NewWriter(s.T())
+		rw := storagemock.NewReaderBatchWriter(s.T())
+		rw.On("Writer").Return(w).Once() // called by epochStateDB.BatchStore
 
-	// Provide the blockID and execute the resulting `dbUpdates`. Thereby, the expected mock methods should be called,
-	// which is asserted by the testify framework. Passing nil lockctx proof because no operations require lock;
-	// operations are deferred only because block ID is not known yet
-	blockID := s.candidate.ID()
-	err = dbOps.Execute(nil, blockID, rw)
+		// Storage operations are deferred, because block ID is not known when the block is newly constructed. Only at the
+		// end after the block is fully constructed, its ID can be computed. We emulate this step here to verify that the
+		// deferred `dbOps` have been correctly constructed. Thereby, the expected mock methods should be called,
+		// which is asserted by the testify framework.
+		blockID := s.candidate.ID()
+		return dbOps.Execute(lctx, blockID, rw)
+	})
 	require.NoError(s.T(), err)
+
 }

storage/epoch_protocol_state.go

@@ -1,15 +1,23 @@
 package storage
 
 import (
+	"github.com/jordanschalm/lockctx"
+
 	"github.com/onflow/flow-go/model/flow"
 )
 
 // EpochProtocolStateEntries represents persistent, fork-aware storage for the Epoch-related
 // sub-state of the overall of the overall Protocol State (KV Store).
 type EpochProtocolStateEntries interface {
+
 	// BatchStore persists the given epoch protocol state entry as part of a DB batch. Per convention, the identities in
-	// the flow.MinEpochStateEntry must be in canonical order for the current and next epoch (if present),
-	// otherwise an exception is returned.
+	// the flow.MinEpochStateEntry must be in canonical order for the current and next epoch (if present), otherwise an
+	// exception is returned.
+	//
+	// CAUTION: The caller must ensure `epochProtocolStateID` is a collision-resistant hash of the provided
+	// `epochProtocolStateEntry`! This method silently overrides existing data, which is safe only if for the same
+	// key, we always write the same value.
+	//
 	// No errors are expected during normal operation.
 	BatchStore(w Writer, epochProtocolStateID flow.Identifier, epochProtocolStateEntry *flow.MinEpochStateEntry) error
 
@@ -21,11 +29,22 @@ type EpochProtocolStateEntries interface {
 	//     the protocol state changes if we seal some execution results emitting service events.
 	//   - For the key `blockID`, we use the identity of block B which _proposes_ this Protocol State. As value,
 	//     the hash of the resulting protocol state at the end of processing B is to be used.
-	//   - CAUTION: The protocol state requires confirmation by a QC and will only become active at the child block,
+	//   - IMPORTANT: The protocol state requires confirmation by a QC and will only become active at the child block,
 	//     _after_ validating the QC.
 	//
-	// No errors are expected during normal operation.
-	BatchIndex(rw ReaderBatchWriter, blockID flow.Identifier, epochProtocolStateID flow.Identifier) error
+	// CAUTION:
+	//   - The caller must acquire the lock [storage.LockInsertBlock] and hold it until the database write has been committed.
+	//   - OVERWRITES existing data (potential for data corruption):
+	//     This method silently overrides existing data without any sanity checks whether data for the same key already exits.
+	//     Note that the Flow protocol mandates that for a previously persisted key, the data is never changed to a different
+	//     value. Changing data could cause the node to publish inconsistent data and to be slashed, or the protocol to be
+	//     compromised as a whole. This method does not contain any safeguards to prevent such data corruption. The lock proof
+	//     serves as a reminder that the CALLER is responsible to ensure that the DEDUPLICATION CHECK is done elsewhere
+	//     ATOMICALLY with this write operation.
+	//
+	// Expected errors during normal operations:
+	// - [storage.ErrDataMismatch] if a _different_ KV store for the given stateID has already been persisted
+	BatchIndex(lctx lockctx.Proof, rw ReaderBatchWriter, blockID flow.Identifier, epochProtocolStateID flow.Identifier) error
 
 	// ByID returns the flow.RichEpochStateEntry by its ID.
 	// Expected errors during normal operations:

storage/mock/epoch_protocol_state_entries.go

@@ -247,24 +247,31 @@ func TestDirectChildren(t *testing.T) {
 // TestChildrenWrongLockIsRejected verifies that operations fail when called with the wrong lock type.
 // This ensures that IndexNewBlock requires LockInsertBlock and IndexNewClusterBlock requires LockInsertOrFinalizeClusterBlock.
 func TestChildrenWrongLockIsRejected(t *testing.T) {
-	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
-		lockManager := storage.NewTestingLockManager()
-
-		parentID := unittest.IdentifierFixture()
-		childID := unittest.IdentifierFixture()
+	for _, opPair := range getTestOperationPairs() {
+		t.Run(opPair.name, func(t *testing.T) {
+			dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
+				lockManager := storage.NewTestingLockManager()
 
-		err := unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.IndexNewClusterBlock(lctx, rw, childID, parentID)
-			})
-		})
-		require.Error(t, err)
+				parentID := unittest.IdentifierFixture()
+				childID := unittest.IdentifierFixture()
 
-		err = unittest.WithLock(t, lockManager, storage.LockInsertOrFinalizeClusterBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.IndexNewBlock(lctx, rw, childID, parentID)
+				// Use the wrong lock type for each operation
+				var wrongLockType string
+				if opPair.lockType == storage.LockInsertBlock {
+					// For IndexNewBlock, use the cluster block lock (wrong)
+					wrongLockType = storage.LockInsertOrFinalizeClusterBlock
+				} else {
+					// For IndexNewClusterBlock, use the regular block lock (wrong)
+					wrongLockType = storage.LockInsertBlock
+				}
+
+				err := unittest.WithLock(t, lockManager, wrongLockType, func(lctx lockctx.Context) error {
+					return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+						return opPair.indexFunc(lctx, rw, childID, parentID)
+					})
+				})
+				require.Error(t, err)
 			})
 		})
-		require.Error(t, err)
-	})
+	}
 }

storage/operation/children_test.go

@@ -1,6 +1,10 @@
 package operation
 
 import (
+	"fmt"
+
+	"github.com/jordanschalm/lockctx"
+
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/storage"
 )
@@ -15,22 +19,34 @@ func InsertEpochProtocolState(w storage.Writer, entryID flow.Identifier, entry *
 // RetrieveEpochProtocolState retrieves an epoch protocol state entry by ID.
 // Error returns:
 //   - storage.ErrNotFound if the key does not exist in the database
-//   - generic error in case of unexpected failure from the database layer
 func RetrieveEpochProtocolState(r storage.Reader, entryID flow.Identifier, entry *flow.MinEpochStateEntry) error {
 	return RetrieveByKey(r, MakePrefix(codeEpochProtocolState, entryID), entry)
 }
 
 // IndexEpochProtocolState indexes an epoch protocol state entry by block ID.
-// Error returns:
-//   - generic error in case of unexpected failure from the database layer or encoding failure.
-func IndexEpochProtocolState(w storage.Writer, blockID flow.Identifier, epochProtocolStateEntryID flow.Identifier) error {
+//
+// CAUTION:
+//   - The caller must acquire the lock [storage.LockInsertBlock] and hold it until the database write has been committed.
+//   - OVERWRITES existing data (potential for data corruption):
+//     This method silently overrides existing data without any sanity checks whether data for the same key already exits.
+//     Note that the Flow protocol mandates that for a previously persisted key, the data is never changed to a different
+//     value. Changing data could cause the node to publish inconsistent data and to be slashed, or the protocol to be
+//     compromised as a whole. This method does not contain any safeguards to prevent such data corruption. The lock proof
+//     serves as a reminder that the CALLER is responsible to ensure that the DEDUPLICATION CHECK is done elsewhere
+//     ATOMICALLY with this write operation.
+//
+// No error returns are expected during normal operation.
+func IndexEpochProtocolState(lctx lockctx.Proof, w storage.Writer, blockID flow.Identifier, epochProtocolStateEntryID flow.Identifier) error {
+	if !lctx.HoldsLock(storage.LockInsertBlock) {
+		return fmt.Errorf("missing required lock: %s", storage.LockInsertBlock)
+	}
+
 	return UpsertByKey(w, MakePrefix(codeEpochProtocolStateByBlockID, blockID), epochProtocolStateEntryID)
 }
 
 // LookupEpochProtocolState finds an epoch protocol state entry ID by block ID.
 // Error returns:
 //   - storage.ErrNotFound if the key does not exist in the database
-//   - generic error in case of unexpected failure from the database layer
 func LookupEpochProtocolState(r storage.Reader, blockID flow.Identifier, epochProtocolStateEntryID *flow.Identifier) error {
 	return RetrieveByKey(r, MakePrefix(codeEpochProtocolStateByBlockID, blockID), epochProtocolStateEntryID)
 }