Merge pull request #8018 from onflow/leo/refactor-safety-data

zhangchiqing · web-flow · commit da3c9e1f5b47 · 2025-10-10T00:06:27.000Z
[Storage] Refactor safety data operations
diff --git a/cmd/collection/main.go b/cmd/collection/main.go
@@ -636,6 +636,7 @@ func main() {
 				node.Logger,
 				node.Me,
 				node.ProtocolDB,
+				node.StorageLockMgr,
 				node.State,
 				node.Metrics.Engine,
 				node.Metrics.Mempool,
diff --git a/cmd/consensus/main.go b/cmd/consensus/main.go
@@ -621,7 +621,7 @@ func main() {
 			notifier.AddFollowerConsumer(followerDistributor)
 
 			// initialize the persister
-			persist, err := persister.New(node.ProtocolDB, node.RootChainID)
+			persist, err := persister.New(node.ProtocolDB, node.RootChainID, node.StorageLockMgr)
 			if err != nil {
 				return nil, err
 			}
diff --git a/cmd/scaffold.go b/cmd/scaffold.go
@@ -303,7 +303,7 @@ func (fnb *FlowNodeBuilder) EnqueuePingService() {
 		var hotstuffViewFunc func() (uint64, error)
 		// Setup consensus nodes to report their HotStuff view
 		if fnb.BaseConfig.NodeRole == flow.RoleConsensus.String() {
-			hotstuffReader, err := persister.NewReader(node.ProtocolDB, node.RootChainID)
+			hotstuffReader, err := persister.NewReader(node.ProtocolDB, node.RootChainID, node.StorageLockMgr)
 			if err != nil {
 				return nil, err
 			}
diff --git a/cmd/util/cmd/read-hotstuff/cmd/get_liveness.go b/cmd/util/cmd/read-hotstuff/cmd/get_liveness.go
@@ -22,8 +22,10 @@ func init() {
 
 func runGetLivenessData(*cobra.Command, []string) {
 	err := common.WithStorage(flagDatadir, func(db storage.DB) error {
+		lockManager := storage.NewTestingLockManager()
+
 		chainID := flow.ChainID(flagChain)
-		reader, err := persister.NewReader(db, chainID)
+		reader, err := persister.NewReader(db, chainID, lockManager)
 		if err != nil {
 			log.Fatal().Err(err).Msg("could not create reader from db")
 		}
diff --git a/cmd/util/cmd/read-hotstuff/cmd/get_safety.go b/cmd/util/cmd/read-hotstuff/cmd/get_safety.go
@@ -22,9 +22,10 @@ func init() {
 
 func runGetSafetyData(*cobra.Command, []string) {
 	err := common.WithStorage(flagDatadir, func(db storage.DB) error {
+		lockManager := storage.NewTestingLockManager()
 
 		chainID := flow.ChainID(flagChain)
-		reader, err := persister.NewReader(db, chainID)
+		reader, err := persister.NewReader(db, chainID, lockManager)
 		if err != nil {
 			log.Fatal().Err(err).Msg("could not create reader from db")
 		}
diff --git a/consensus/hotstuff/persister/persister.go b/consensus/hotstuff/persister/persister.go
@@ -3,6 +3,8 @@ package persister
 import (
 	"fmt"
 
+	"github.com/jordanschalm/lockctx"
+
 	"github.com/onflow/flow-go/consensus/hotstuff"
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/storage"
@@ -15,8 +17,9 @@ import (
 // SafetyData and LivenessData, for each distinct chain ID. This bootstrapping must be complete
 // before constructing a Persister instance with New (otherwise it will return an error).
 type Persister struct {
-	db      storage.DB
-	chainID flow.ChainID
+	db          storage.DB
+	chainID     flow.ChainID
+	lockManager lockctx.Manager
 }
 
 var _ hotstuff.Persister = (*Persister)(nil)
@@ -26,15 +29,16 @@ var _ hotstuff.PersisterReader = (*Persister)(nil)
 // Persister depends on protocol.State and cluster.State bootstrapping to set initial values for
 // SafetyData and LivenessData, for each distinct chain ID. This bootstrapping must be completed
 // before first using a Persister instance.
-func New(db storage.DB, chainID flow.ChainID) (*Persister, error) {
+func New(db storage.DB, chainID flow.ChainID, lockManager lockctx.Manager) (*Persister, error) {
 	err := ensureSafetyDataAndLivenessDataAreBootstrapped(db, chainID)
 	if err != nil {
 		return nil, fmt.Errorf("fail to check persister was properly bootstrapped: %w", err)
 	}
 
 	p := &Persister{
-		db:      db,
-		chainID: chainID,
+		db:          db,
+		chainID:     chainID,
+		lockManager: lockManager,
 	}
 	return p, nil
 }
@@ -62,8 +66,8 @@ func ensureSafetyDataAndLivenessDataAreBootstrapped(db storage.DB, chainID flow.
 }
 
 // NewReader returns a new Persister as a PersisterReader type (only read methods accessible).
-func NewReader(db storage.DB, chainID flow.ChainID) (hotstuff.PersisterReader, error) {
-	return New(db, chainID)
+func NewReader(db storage.DB, chainID flow.ChainID, lockManager lockctx.Manager) (hotstuff.PersisterReader, error) {
+	return New(db, chainID, lockManager)
 }
 
 // GetSafetyData will retrieve last persisted safety data.
@@ -85,15 +89,19 @@ func (p *Persister) GetLivenessData() (*hotstuff.LivenessData, error) {
 // PutSafetyData persists the last safety data.
 // During normal operations, no errors are expected.
 func (p *Persister) PutSafetyData(safetyData *hotstuff.SafetyData) error {
-	return p.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-		return operation.UpsertSafetyData(rw.Writer(), p.chainID, safetyData)
+	return storage.WithLock(p.lockManager, storage.LockInsertSafetyData, func(lctx lockctx.Context) error {
+		return p.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+			return operation.UpsertSafetyData(lctx, rw, p.chainID, safetyData)
+		})
 	})
 }
 
 // PutLivenessData persists the last liveness data.
 // During normal operations, no errors are expected.
 func (p *Persister) PutLivenessData(livenessData *hotstuff.LivenessData) error {
-	return p.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-		return operation.UpsertLivenessData(rw.Writer(), p.chainID, livenessData)
+	return storage.WithLock(p.lockManager, storage.LockInsertLivenessData, func(lctx lockctx.Context) error {
+		return p.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+			return operation.UpsertLivenessData(lctx, rw, p.chainID, livenessData)
+		})
 	})
 }
diff --git a/consensus/integration/nodes_test.go b/consensus/integration/nodes_test.go
@@ -563,7 +563,7 @@ func createNode(
 
 	signer := verification.NewCombinedSigner(me, beaconKeyStore)
 
-	persist, err := persister.New(db, rootHeader.ChainID)
+	persist, err := persister.New(db, rootHeader.ChainID, lockManager)
 	require.NoError(t, err)
 
 	livenessData, err := persist.GetLivenessData()
diff --git a/engine/collection/epochmgr/factories/hotstuff.go b/engine/collection/epochmgr/factories/hotstuff.go
@@ -32,6 +32,7 @@ type HotStuffFactory struct {
 	baseLogger     zerolog.Logger
 	me             module.Local
 	db             storage.DB
+	lockManager    storage.LockManager
 	protoState     protocol.State
 	engineMetrics  module.EngineMetrics
 	mempoolMetrics module.MempoolMetrics
@@ -43,6 +44,7 @@ func NewHotStuffFactory(
 	log zerolog.Logger,
 	me module.Local,
 	db storage.DB,
+	lockManager storage.LockManager,
 	protoState protocol.State,
 	engineMetrics module.EngineMetrics,
 	mempoolMetrics module.MempoolMetrics,
@@ -54,6 +56,7 @@ func NewHotStuffFactory(
 		baseLogger:     log,
 		me:             me,
 		db:             db,
+		lockManager:    lockManager,
 		protoState:     protoState,
 		engineMetrics:  engineMetrics,
 		mempoolMetrics: mempoolMetrics,
@@ -155,7 +158,7 @@ func (f *HotStuffFactory) CreateModules(
 		return nil, nil, err
 	}
 
-	persist, err := persister.New(f.db, cluster.ChainID())
+	persist, err := persister.New(f.db, cluster.ChainID(), f.lockManager)
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/engine/testutil/nodes.go b/engine/testutil/nodes.go
@@ -377,6 +377,7 @@ func CollectionNode(t *testing.T, hub *stub.Hub, identity bootstrap.NodeInfo, ro
 		node.Log,
 		node.Me,
 		db,
+		node.LockManager,
 		node.State,
 		node.Metrics,
 		node.Metrics,
diff --git a/state/cluster/badger/state.go b/state/cluster/badger/state.go
@@ -35,6 +35,14 @@ func Bootstrap(db storage.DB, lockManager lockctx.Manager, stateRoot *StateRoot)
 	if err != nil {
 		return nil, fmt.Errorf("failed to acquire lock `storage.LockInsertOrFinalizeClusterBlock` for inserting cluster block: %w", err)
 	}
+	err = lctx.AcquireLock(storage.LockInsertSafetyData)
+	if err != nil {
+		return nil, fmt.Errorf("failed to acquire lock `storage.LockInsertSafetyData` for inserting safety data: %w", err)
+	}
+	err = lctx.AcquireLock(storage.LockInsertLivenessData)
+	if err != nil {
+		return nil, fmt.Errorf("failed to acquire lock `storage.LockInsertLivenessData` for inserting liveness data: %w", err)
+	}
 	isBootstrapped, err := IsBootstrapped(db, stateRoot.ClusterID())
 	if err != nil {
 		return nil, fmt.Errorf("failed to determine whether database contains bootstrapped state: %w", err)
@@ -85,12 +93,12 @@ func Bootstrap(db storage.DB, lockManager lockctx.Manager, stateRoot *StateRoot)
 			NewestQC:    rootQC,
 		}
 		// insert safety data
-		err = operation.UpsertSafetyData(rw.Writer(), chainID, safetyData)
+		err = operation.UpsertSafetyData(lctx, rw, chainID, safetyData)
 		if err != nil {
 			return fmt.Errorf("could not insert safety data: %w", err)
 		}
 		// insert liveness data
-		err = operation.UpsertLivenessData(rw.Writer(), chainID, livenessData)
+		err = operation.UpsertLivenessData(lctx, rw, chainID, livenessData)
 		if err != nil {
 			return fmt.Errorf("could not insert liveness data: %w", err)
 		}
diff --git a/state/protocol/badger/state.go b/state/protocol/badger/state.go
@@ -124,6 +124,14 @@ func Bootstrap(
 	if err != nil {
 		return nil, err
 	}
+	err = lctx.AcquireLock(storage.LockInsertSafetyData)
+	if err != nil {
+		return nil, err
+	}
+	err = lctx.AcquireLock(storage.LockInsertLivenessData)
+	if err != nil {
+		return nil, err
+	}
 
 	config := defaultBootstrapConfig()
 	for _, opt := range options {
@@ -585,13 +593,12 @@ func bootstrapStatePointers(lctx lockctx.Proof, rw storage.ReaderBatchWriter, ro
 		NewestQC:    qcForLatestFinalizedBlock,
 	}
 
-	w := rw.Writer()
 	// persist safety and liveness data plus the QuorumCertificate for the latest finalized block for HotStuff/Jolteon consensus
-	err = operation.UpsertSafetyData(w, lastFinalized.ChainID, safetyData)
+	err = operation.UpsertSafetyData(lctx, rw, lastFinalized.ChainID, safetyData)
 	if err != nil {
 		return fmt.Errorf("could not insert safety data: %w", err)
 	}
-	err = operation.UpsertLivenessData(w, lastFinalized.ChainID, livenessData)
+	err = operation.UpsertLivenessData(lctx, rw, lastFinalized.ChainID, livenessData)
 	if err != nil {
 		return fmt.Errorf("could not insert liveness data: %w", err)
 	}
@@ -600,6 +607,7 @@ func bootstrapStatePointers(lctx lockctx.Proof, rw storage.ReaderBatchWriter, ro
 		return fmt.Errorf("could not insert quorum certificate for the latest finalized block: %w", err)
 	}
 
+	w := rw.Writer()
 	// insert height pointers
 	err = operation.UpsertFinalizedHeight(lctx, w, lastFinalized.Height)
 	if err != nil {
diff --git a/storage/locks.go b/storage/locks.go
@@ -29,6 +29,8 @@ 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"
 )
 
 // Locks returns a list of all named locks used by the storage layer.
@@ -42,6 +44,8 @@ func Locks() []string {
 		LockInsertCollection,
 		LockBootstrapping,
 		LockInsertChunkDataPack,
+		LockInsertSafetyData,
+		LockInsertLivenessData,
 	}
 }
 
@@ -65,6 +69,9 @@ func makeLockPolicy() lockctx.Policy {
 	return lockctx.NewDAGPolicyBuilder().
 		Add(LockInsertBlock, LockFinalizeBlock).
 		Add(LockFinalizeBlock, LockBootstrapping).
+		Add(LockBootstrapping, LockInsertSafetyData).
+		Add(LockInsertSafetyData, LockInsertLivenessData).
+		Add(LockInsertOrFinalizeClusterBlock, LockInsertSafetyData).
 		Add(LockInsertOwnReceipt, LockInsertChunkDataPack).
 		Build()
 }
diff --git a/storage/operation/views.go b/storage/operation/views.go
@@ -1,6 +1,10 @@
 package operation
 
 import (
+	"fmt"
+
+	"github.com/jordanschalm/lockctx"
+
 	"github.com/onflow/flow-go/consensus/hotstuff"
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/storage"
@@ -9,11 +13,14 @@ import (
 // UpsertSafetyData inserts or updates the given safety data for this node.
 // Intended for consensus participants only (consensus and collector nodes).
 // Here, `chainID` specifies which consensus instance specifically the node participates in.
-// CAUTION: OVERWRITES existing data (potential for data corruption).
 //
 // No errors are expected during normal operation.
-func UpsertSafetyData(w storage.Writer, chainID flow.ChainID, safetyData *hotstuff.SafetyData) error {
-	return UpsertByKey(w, MakePrefix(codeSafetyData, chainID), safetyData)
+func UpsertSafetyData(lctx lockctx.Proof, rw storage.ReaderBatchWriter, chainID flow.ChainID, safetyData *hotstuff.SafetyData) error {
+	if !lctx.HoldsLock(storage.LockInsertSafetyData) {
+		return fmt.Errorf("missing required lock: storage.LockInsertSafetyData")
+	}
+
+	return UpsertByKey(rw.Writer(), MakePrefix(codeSafetyData, chainID), safetyData)
 }
 
 // RetrieveSafetyData retrieves the safety data for this node.
@@ -28,11 +35,14 @@ func RetrieveSafetyData(r storage.Reader, chainID flow.ChainID, safetyData *hots
 // UpsertLivenessData inserts or updates the given liveness data for this node.
 // Intended for consensus participants only (consensus and collector nodes).
 // Here, `chainID` specifies which consensus instance specifically the node participates in.
-// CAUTION: OVERWRITES existing data (potential for data corruption).
 //
 // No errors are expected during normal operation.
-func UpsertLivenessData(w storage.Writer, chainID flow.ChainID, livenessData *hotstuff.LivenessData) error {
-	return UpsertByKey(w, MakePrefix(codeLivenessData, chainID), livenessData)
+func UpsertLivenessData(lctx lockctx.Proof, rw storage.ReaderBatchWriter, chainID flow.ChainID, livenessData *hotstuff.LivenessData) error {
+	if !lctx.HoldsLock(storage.LockInsertLivenessData) {
+		return fmt.Errorf("missing required lock: storage.LockInsertLivenessData")
+	}
+
+	return UpsertByKey(rw.Writer(), MakePrefix(codeLivenessData, chainID), livenessData)
 }
 
 // RetrieveSafetyData retrieves the safety data for this node.

Original file line number	Diff line number	Diff line change
`@@ -621,7 +621,7 @@ func main() {`
`621`	`621`	`notifier.AddFollowerConsumer(followerDistributor)`
`622`	`622`
`623`	`623`	`// initialize the persister`
`624`		`- persist, err := persister.New(node.ProtocolDB, node.RootChainID)`
	`624`	`+ persist, err := persister.New(node.ProtocolDB, node.RootChainID, node.StorageLockMgr)`
`625`	`625`	`if err != nil {`
`626`	`626`	`return nil, err`
`627`	`627`	`}`
Original file line number	Diff line number	Diff line change
`@@ -303,7 +303,7 @@ func (fnb *FlowNodeBuilder) EnqueuePingService() {`
`303`	`303`	`var hotstuffViewFunc func() (uint64, error)`
`304`	`304`	`// Setup consensus nodes to report their HotStuff view`
`305`	`305`	`if fnb.BaseConfig.NodeRole == flow.RoleConsensus.String() {`
`306`		`- hotstuffReader, err := persister.NewReader(node.ProtocolDB, node.RootChainID)`
	`306`	`+ hotstuffReader, err := persister.NewReader(node.ProtocolDB, node.RootChainID, node.StorageLockMgr)`
`307`	`307`	`if err != nil {`
`308`	`308`	`return nil, err`
`309`	`309`	`}`