update safety data operation

Author: zhangchiqing

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,

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
 			}

consensus/hotstuff/persister/persister.go

@@ -3,6 +3,7 @@ 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 +16,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 +28,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 +65,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 +88,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)
+		})
 	})
 }

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
 	}

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,

state/cluster/badger/state.go

@@ -36,6 +36,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)
@@ -86,12 +94,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)
 		}

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 {

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()
 }

storage/operation/views.go

@@ -1,6 +1,9 @@
 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 +12,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 +34,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.