Merge pull request #8017 from onflow/leo/refactor-execution-fork-evidence

zhangchiqing · web-flow · commit 78c4a1353a16 · 2025-10-10T15:31:26.000Z
[Storage] Refactor execution fork evidence
diff --git a/cmd/consensus/main.go b/cmd/consensus/main.go
@@ -396,6 +396,7 @@ func main() {
 				multipleReceiptsFilterMempool,
 				consensusMempools.LogForkAndCrash(node.Logger),
 				node.ProtocolDB,
+				node.StorageLockMgr,
 				node.Logger,
 			)
 			if err != nil {
diff --git a/module/mempool/consensus/exec_fork_suppressor.go b/module/mempool/consensus/exec_fork_suppressor.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/jordanschalm/lockctx"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"go.uber.org/atomic"
@@ -46,6 +47,7 @@ type ExecForkSuppressor struct {
 	execForkDetected      atomic.Bool
 	onExecFork            ExecForkActor
 	execForkEvidenceStore storage.ExecutionForkEvidence
+	lockManager           storage.LockManager
 	log                   zerolog.Logger
 }
 
@@ -61,6 +63,7 @@ func NewExecStateForkSuppressor(
 	seals mempool.IncorporatedResultSeals,
 	onExecFork ExecForkActor,
 	db storage.DB,
+	lockManager storage.LockManager,
 	log zerolog.Logger,
 ) (*ExecForkSuppressor, error) {
 	executionForkEvidenceStore := store.NewExecutionForkEvidence(db)
@@ -83,6 +86,7 @@ func NewExecStateForkSuppressor(
 		execForkDetected:      *atomic.NewBool(execForkDetectedFlag),
 		onExecFork:            onExecFork,
 		execForkEvidenceStore: executionForkEvidenceStore,
+		lockManager:           lockManager,
 		log:                   log.With().Str("mempool", "ExecForkSuppressor").Logger(),
 	}
 
@@ -365,9 +369,13 @@ func (s *ExecForkSuppressor) filterConflictingSeals(sealsByBlockID map[flow.Iden
 				s.execForkDetected.Store(true)
 				s.Clear()
 				conflictingSeals = append(sealsList{candidateSeal}, conflictingSeals...)
-				err := s.execForkEvidenceStore.StoreIfNotExists(conflictingSeals)
+
+				// Acquire lock and store execution fork evidence
+				err := storage.WithLock(s.lockManager, storage.LockInsertExecutionForkEvidence, func(lctx lockctx.Context) error {
+					return s.execForkEvidenceStore.StoreIfNotExists(lctx, conflictingSeals)
+				})
 				if err != nil {
-					panic("failed to store execution fork evidence")
+					s.log.Fatal().Msg("failed to store execution fork evidence")
 				}
 				s.onExecFork(conflictingSeals)
 				return nil
diff --git a/module/mempool/consensus/exec_fork_suppressor_test.go b/module/mempool/consensus/exec_fork_suppressor_test.go
@@ -255,11 +255,12 @@ func Test_ForkDetectionPersisted(t *testing.T) {
 
 		// This function stores conflicting seals to the underlying database.
 		func(t *testing.T, db storage.DB) {
+			lockManager := storage.NewTestingLockManager()
 
 			// initialize ExecForkSuppressor
 			wrappedMempool := &poolmock.IncorporatedResultSeals{}
 			execForkActor := &actormock.ExecForkActor{}
-			wrapper, _ := NewExecStateForkSuppressor(wrappedMempool, execForkActor.OnExecFork, db, zerolog.New(os.Stderr))
+			wrapper, _ := NewExecStateForkSuppressor(wrappedMempool, execForkActor.OnExecFork, db, lockManager, zerolog.New(os.Stderr))
 
 			// add seal
 			wrappedMempool.On("Add", sealA).Return(true, nil).Once()
@@ -285,14 +286,16 @@ func Test_ForkDetectionPersisted(t *testing.T) {
 
 		// This function retrieves conflicting seals from the same underlying database with a new instance of storage.DB.
 		func(t *testing.T, db storage.DB) {
+			lockManager := storage.NewTestingLockManager()
+
 			wrappedMempool2 := &poolmock.IncorporatedResultSeals{}
 			execForkActor2 := &actormock.ExecForkActor{}
 			execForkActor2.On("OnExecFork", mock.Anything).
 				Run(func(args mock.Arguments) {
 					conflictingSeals := args.Get(0).([]*flow.IncorporatedResultSeal)
 					require.ElementsMatch(t, []*flow.IncorporatedResultSeal{sealA, sealB}, conflictingSeals)
 				}).Return().Once()
-			wrapper2, _ := NewExecStateForkSuppressor(wrappedMempool2, execForkActor2.OnExecFork, db, zerolog.New(os.Stderr))
+			wrapper2, _ := NewExecStateForkSuppressor(wrappedMempool2, execForkActor2.OnExecFork, db, lockManager, zerolog.New(os.Stderr))
 
 			// add another (non-conflicting) seal to ExecForkSuppressor
 			// fail test if seal is added to wrapped mempool
@@ -318,8 +321,10 @@ func Test_AddRemove_SmokeTest(t *testing.T) {
 		require.Fail(t, "no call to onExecFork expected ")
 	}
 	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
+		lockManager := storage.NewTestingLockManager()
+
 		wrappedMempool := stdmap.NewIncorporatedResultSeals(100)
-		wrapper, err := NewExecStateForkSuppressor(wrappedMempool, onExecFork, db, zerolog.New(os.Stderr))
+		wrapper, err := NewExecStateForkSuppressor(wrappedMempool, onExecFork, db, lockManager, zerolog.New(os.Stderr))
 		require.NoError(t, err)
 		require.NotNil(t, wrapper)
 
@@ -355,6 +360,8 @@ func Test_AddRemove_SmokeTest(t *testing.T) {
 // Test adding conflicting seals with different number of matching receipts.
 func Test_ConflictingSeal_SmokeTest(t *testing.T) {
 	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
+		lockManager := storage.NewTestingLockManager()
+
 		executingForkDetected := atomic.NewBool(false)
 		onExecFork := func([]*flow.IncorporatedResultSeal) {
 			executingForkDetected.Store(true)
@@ -363,7 +370,7 @@ func Test_ConflictingSeal_SmokeTest(t *testing.T) {
 		rawMempool := stdmap.NewIncorporatedResultSeals(100)
 		receiptsDB := mockstorage.NewExecutionReceipts(t)
 		wrappedMempool := NewIncorporatedResultSeals(rawMempool, receiptsDB)
-		wrapper, err := NewExecStateForkSuppressor(wrappedMempool, onExecFork, db, zerolog.New(os.Stderr))
+		wrapper, err := NewExecStateForkSuppressor(wrappedMempool, onExecFork, db, lockManager, zerolog.New(os.Stderr))
 		require.NoError(t, err)
 		require.NotNil(t, wrapper)
 
@@ -426,9 +433,11 @@ func Test_ConflictingSeal_SmokeTest(t *testing.T) {
 //  4. executes the `testLogic`
 func WithExecStateForkSuppressor(t *testing.T, testLogic func(wrapper *ExecForkSuppressor, wrappedMempool *poolmock.IncorporatedResultSeals, execForkActor *actormock.ExecForkActor)) {
 	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
+		lockManager := storage.NewTestingLockManager()
+
 		wrappedMempool := &poolmock.IncorporatedResultSeals{}
 		execForkActor := &actormock.ExecForkActor{}
-		wrapper, err := NewExecStateForkSuppressor(wrappedMempool, execForkActor.OnExecFork, db, zerolog.New(os.Stderr))
+		wrapper, err := NewExecStateForkSuppressor(wrappedMempool, execForkActor.OnExecFork, db, lockManager, zerolog.New(os.Stderr))
 		require.NoError(t, err)
 		require.NotNil(t, wrapper)
 		testLogic(wrapper, wrappedMempool, execForkActor)
diff --git a/storage/epoch_protocol_state.go b/storage/epoch_protocol_state.go
@@ -35,15 +35,12 @@ type EpochProtocolStateEntries interface {
 	// 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.
+	//     The lock proof serves as a reminder that the CALLER is responsible to ensure that the DEDUPLICATION CHECK is done elsewhere
+	//     ATOMICALLY within this write operation. Currently it's done by operation.InsertHeader where it performs a check
+	//     to ensure the blockID is new, therefore any data indexed by this blockID is new as well.
 	//
 	// Expected errors during normal operations:
-	// - [storage.ErrDataMismatch] if a _different_ KV store for the given stateID has already been persisted
+	// No expected errors during normal operations.
 	BatchIndex(lctx lockctx.Proof, rw ReaderBatchWriter, blockID flow.Identifier, epochProtocolStateID flow.Identifier) error
 
 	// ByID returns the flow.RichEpochStateEntry by its ID.
diff --git a/storage/execution_fork_evidence.go b/storage/execution_fork_evidence.go
@@ -1,6 +1,10 @@
 package storage
 
-import "github.com/onflow/flow-go/model/flow"
+import (
+	"github.com/jordanschalm/lockctx"
+
+	"github.com/onflow/flow-go/model/flow"
+)
 
 // ExecutionForkEvidence represents persistent storage for execution fork evidence.
 // CAUTION: Not safe for concurrent use by multiple goroutines.
@@ -9,8 +13,9 @@ type ExecutionForkEvidence interface {
 	// if no execution fork evidence is currently stored in the database.
 	// This function is a no-op if evidence is already stored, because
 	// only one execution fork evidence can be stored at a time.
+	// The caller must hold the [storage.LockInsertExecutionForkEvidence] lock.
 	// No errors are expected during normal operations.
-	StoreIfNotExists(conflictingSeals []*flow.IncorporatedResultSeal) error
+	StoreIfNotExists(lctx lockctx.Proof, conflictingSeals []*flow.IncorporatedResultSeal) error
 
 	// Retrieve reads conflicting seals from the database.
 	// No error is returned if database record doesn't exist.
diff --git a/storage/locks.go b/storage/locks.go
@@ -29,8 +29,10 @@ const (
 	LockBootstrapping = "lock_bootstrapping"
 	// LockInsertChunkDataPack protects the insertion of chunk data packs (not yet used anywhere
 	LockInsertChunkDataPack = "lock_insert_chunk_data_pack"
-	LockInsertSafetyData    = "lock_insert_safety_data"
-	LockInsertLivenessData  = "lock_insert_liveness_data"
+	// LockInsertExecutionForkEvidence protects the insertion of execution fork evidence
+	LockInsertExecutionForkEvidence = "lock_insert_execution_fork_evidence"
+	LockInsertSafetyData            = "lock_insert_safety_data"
+	LockInsertLivenessData          = "lock_insert_liveness_data"
 )
 
 // Locks returns a list of all named locks used by the storage layer.
@@ -44,6 +46,7 @@ func Locks() []string {
 		LockInsertCollection,
 		LockBootstrapping,
 		LockInsertChunkDataPack,
+		LockInsertExecutionForkEvidence,
 		LockInsertSafetyData,
 		LockInsertLivenessData,
 	}
diff --git a/storage/mock/execution_fork_evidence.go b/storage/mock/execution_fork_evidence.go
diff --git a/storage/operation/epoch_protocol_state.go b/storage/operation/epoch_protocol_state.go
@@ -28,12 +28,9 @@ func RetrieveEpochProtocolState(r storage.Reader, entryID flow.Identifier, entry
 // 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.
+//     The lock proof serves as a reminder that the CALLER is responsible to ensure that the DEDUPLICATION CHECK is done elsewhere
+//     ATOMICALLY within this write operation. Currently it's done by operation.InsertHeader where it performs a check
+//     to ensure the blockID is new, therefore any data indexed by this blockID is new as well.
 //
 // No error returns are expected during normal operation.
 func IndexEpochProtocolState(lctx lockctx.Proof, w storage.Writer, blockID flow.Identifier, epochProtocolStateEntryID flow.Identifier) error {
diff --git a/storage/operation/execution_fork_evidence.go b/storage/operation/execution_fork_evidence.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"
 )
@@ -28,8 +32,24 @@ func RemoveExecutionForkEvidence(w storage.Writer) error {
 }
 
 // InsertExecutionForkEvidence upserts conflicting seals to the database.
-// If a record already exists, it is overwritten; otherwise a new record is created.
+// If a record already exists, it is NOT overwritten, the new record is ignored.
+// The caller must hold the [storage.LockInsertExecutionForkEvidence] lock.
 // No errors are expected during normal operations.
-func InsertExecutionForkEvidence(w storage.Writer, conflictingSeals []*flow.IncorporatedResultSeal) error {
-	return UpsertByKey(w, MakePrefix(codeExecutionFork), conflictingSeals)
+func InsertExecutionForkEvidence(lctx lockctx.Proof, rw storage.ReaderBatchWriter, conflictingSeals []*flow.IncorporatedResultSeal) error {
+	if !lctx.HoldsLock(storage.LockInsertExecutionForkEvidence) {
+		return fmt.Errorf("InsertExecutionForkEvidence requires LockInsertExecutionForkEvidence to be held")
+	}
+	key := MakePrefix(codeExecutionFork)
+	exist, err := KeyExists(rw.GlobalReader(), key)
+	if err != nil {
+		return fmt.Errorf("failed to check if execution fork evidence exists: %w", err)
+	}
+
+	if exist {
+		// Some evidence about execution fork already stored;
+		// We only keep the first evidence => nothing more to do
+		return nil
+	}
+
+	return UpsertByKey(rw.Writer(), key, conflictingSeals)
 }
diff --git a/storage/operation/execution_fork_evidence_test.go b/storage/operation/execution_fork_evidence_test.go
@@ -3,6 +3,8 @@ package operation_test
 import (
 	"testing"
 
+	"github.com/jordanschalm/lockctx"
+
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/storage"
 	"github.com/onflow/flow-go/storage/operation"
@@ -27,6 +29,7 @@ func Test_ExecutionForkEvidenceOperations(t *testing.T) {
 
 	t.Run("Write evidence and retrieve", func(t *testing.T) {
 		dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
+			lockManager := storage.NewTestingLockManager()
 			exists, err := operation.HasExecutionForkEvidence(db.Reader())
 			require.NoError(t, err)
 			require.False(t, exists)
@@ -41,8 +44,10 @@ func Test_ExecutionForkEvidenceOperations(t *testing.T) {
 							unittest.WithBlock(block))))
 			}
 
-			err = db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.InsertExecutionForkEvidence(rw.Writer(), conflictingSeals)
+			err = unittest.WithLock(t, lockManager, storage.LockInsertExecutionForkEvidence, func(lctx lockctx.Context) error {
+				return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+					return operation.InsertExecutionForkEvidence(lctx, rw, conflictingSeals)
+				})
 			})
 			require.NoError(t, err)
 
diff --git a/storage/operation/protocol_kv_store.go b/storage/operation/protocol_kv_store.go
@@ -11,7 +11,8 @@ import (
 )
 
 // InsertProtocolKVStore inserts a protocol KV store by protocol kv store ID.
-// The caller must ensure the protocolKVStoreID is the hash of the given kvStore,
+// This function can be called, and does not require the caller to hold any lock proof,
+// but the caller must ensure the protocolKVStoreID is the hash of the given kvStore,
 // This is currently true, see makeVersionedModelID in state/protocol/protocol_state/kvstore/models.go
 // No expected error returns during normal operations.
 func InsertProtocolKVStore(rw storage.ReaderBatchWriter, protocolKVStoreID flow.Identifier, kvStore *flow.PSKeyValueStoreData) error {
@@ -27,9 +28,12 @@ func RetrieveProtocolKVStore(r storage.Reader, protocolKVStoreID flow.Identifier
 
 // IndexProtocolKVStore indexes a protocol KV store by block ID.
 //
-// CAUTION: To prevent data corruption, we need to guarantee atomicity of existence-check and the subsequent
-// database write. Hence, we require the caller to acquire [storage.LockInsertBlock] and hold it until the
-// database write has been committed.
+// 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):
+//     The lock proof serves as a reminder that the CALLER is responsible to ensure that the DEDUPLICATION CHECK is done elsewhere
+//     ATOMICALLY within this write operation. Currently it's done by operation.InsertHeader where it performs a check
+//     to ensure the blockID is new, therefore any data indexed by this blockID is new as well.
 //
 // Expected error returns during normal operations:
 //   - [storage.ErrAlreadyExists] if a KV store for the given blockID has already been indexed
diff --git a/storage/store/execution_fork_evidence.go b/storage/store/execution_fork_evidence.go
@@ -4,6 +4,8 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/jordanschalm/lockctx"
+
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/storage"
 	"github.com/onflow/flow-go/storage/operation"
@@ -23,28 +25,11 @@ func NewExecutionForkEvidence(db storage.DB) *ExecutionForkEvidence {
 
 // StoreIfNotExists stores the given conflictingSeals to the database.
 // This is a no-op if there is already a record in the database with the same key.
+// The caller must hold the [storage.LockInsertExecutionForkEvidence] lock.
 // No errors are expected during normal operations.
-// CAUTION: This function is not safe for concurrent use by multiple goroutines.  For safety,
-// we rely on the fact that Execution Fork Evidence has a very small surface area of use
-// in the Execution Fork Suppressor.  The Execution Fork Suppressor is responsible for
-// ensuring mutually exclusive access to this function.
-func (efe *ExecutionForkEvidence) StoreIfNotExists(conflictingSeals []*flow.IncorporatedResultSeal) error {
+func (efe *ExecutionForkEvidence) StoreIfNotExists(lctx lockctx.Proof, conflictingSeals []*flow.IncorporatedResultSeal) error {
 	return efe.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-		found, err := operation.HasExecutionForkEvidence(rw.GlobalReader())
-		if err != nil {
-			return fmt.Errorf("failed to check if evidence about execution fork exists: %w", err)
-		}
-		if found {
-			// Some evidence about execution fork already stored;
-			// We only keep the first evidence => nothing more to do
-			return nil
-		}
-
-		err = operation.InsertExecutionForkEvidence(rw.Writer(), conflictingSeals)
-		if err != nil {
-			return fmt.Errorf("failed to store evidence about execution fork: %w", err)
-		}
-		return nil
+		return operation.InsertExecutionForkEvidence(lctx, rw, conflictingSeals)
 	})
 }
 
diff --git a/storage/store/execution_fork_evidence_test.go b/storage/store/execution_fork_evidence_test.go
diff --git a/storage/store/proposal_signatures.go b/storage/store/proposal_signatures.go
diff --git a/storage/store/protocol_kv_store.go b/storage/store/protocol_kv_store.go

Original file line number	Diff line number	Diff line change
`@@ -396,6 +396,7 @@ func main() {`
`396`	`396`	`multipleReceiptsFilterMempool,`
`397`	`397`	`consensusMempools.LogForkAndCrash(node.Logger),`
`398`	`398`	`node.ProtocolDB,`
	`399`	`+ node.StorageLockMgr,`
`399`	`400`	`node.Logger,`
`400`	`401`	`)`
`401`	`402`	`if err != nil {`