Merge pull request #7967 from onflow/leo/refactor-index-protocol-kv-store

[Storage] Refactor index protocol kv store

Author: zhangchiqing

state/protocol/badger/state.go

@@ -232,7 +232,7 @@ func bootstrapProtocolState(
 	// The sealing segment contains a protocol state entry for every block in the segment, including the root block.
 	for protocolStateID, stateEntry := range segment.ProtocolStateEntries {
 		// Store the protocol KV Store entry
-		err := protocolKVStoreSnapshots.BatchStore(lctx, rw, protocolStateID, &stateEntry.KVStore)
+		err := protocolKVStoreSnapshots.BatchStore(rw, protocolStateID, &stateEntry.KVStore)
 		if err != nil {
 			return fmt.Errorf("could not store protocol state kvstore: %w", err)
 		}

state/protocol/protocol_state/kvstore/kvstore_storage.go

@@ -3,8 +3,6 @@ package kvstore
 import (
 	"fmt"
 
-	"github.com/jordanschalm/lockctx"
-
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/state/protocol"
 	"github.com/onflow/flow-go/state/protocol/protocol_state"
@@ -37,16 +35,15 @@ func NewProtocolKVStore(protocolStateSnapshots storage.ProtocolKVStore) *Protoco
 // data blob. If the encoding fails, an error is returned.
 // BatchStore is idempotent, i.e. it accepts repeated calls with the same pairs of (stateID, kvStore).
 // Here, the ID is expected to be a collision-resistant hash of the snapshot (including the
-// ProtocolStateVersion). Hence, for the same ID (key), BatchStore will reject changing the data (value).
+// ProtocolStateVersion).
 //
-// Expected errors during normal operations:
-// - storage.ErrDataMismatch if a _different_ KV store for the given stateID has already been persisted
-func (p *ProtocolKVStore) BatchStore(lctx lockctx.Proof, rw storage.ReaderBatchWriter, stateID flow.Identifier, kvStore protocol.KVStoreReader) error {
+// No error is exepcted during normal operations
+func (p *ProtocolKVStore) BatchStore(rw storage.ReaderBatchWriter, stateID flow.Identifier, kvStore protocol.KVStoreReader) error {
 	version, data, err := kvStore.VersionedEncode()
 	if err != nil {
 		return fmt.Errorf("failed to VersionedEncode protocol state: %w", err)
 	}
-	return p.ProtocolKVStore.BatchStore(lctx, rw, stateID, &flow.PSKeyValueStoreData{
+	return p.ProtocolKVStore.BatchStore(rw, stateID, &flow.PSKeyValueStoreData{
 		Version: version,
 		Data:    data,
 	})

state/protocol/protocol_state/kvstore/kvstore_storage_test.go

@@ -29,7 +29,6 @@ func TestProtocolKVStore_StoreTx(t *testing.T) {
 	// On the happy path, where the input `kvState` encodes its state successfully, the wrapped store
 	// should be called to persist the version-encoded snapshot.
 	t.Run("happy path", func(t *testing.T) {
-		lockManager := storage.NewTestingLockManager()
 		expectedVersion := uint64(13)
 		encData := unittest.RandomBytes(117)
 		versionedSnapshot := &flow.PSKeyValueStoreData{
@@ -38,32 +37,26 @@ func TestProtocolKVStore_StoreTx(t *testing.T) {
 		}
 		kvState.On("VersionedEncode").Return(expectedVersion, encData, nil).Once()
 
-		err := unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			rw := storagemock.NewReaderBatchWriter(t)
-			llStorage.On("BatchStore", lctx, rw, kvStateID, versionedSnapshot).Return(nil).Once()
+		rw := storagemock.NewReaderBatchWriter(t)
+		llStorage.On("BatchStore", rw, kvStateID, versionedSnapshot).Return(nil).Once()
 
-			// TODO: potentially update - we might be bringing back a functor here, because we acquire a lock  as explained in slack thread https://flow-foundation.slack.com/archives/C071612SJJE/p1754600182033289?thread_ts=1752912083.194619&cid=C071612SJJE
-			// Calling `BatchStore` should return the output of the wrapped low-level storage, which is a deferred database
-			// update. Conceptually, it is possible that `ProtocolKVStore` wraps the deferred database operation in faulty
-			// code, such that it cannot be executed. Therefore, we execute the top-level deferred database update below
-			// and verify that the deferred database operation returned by the lower-level is actually reached.
-			return store.BatchStore(lctx, rw, kvStateID, kvState)
-		})
+		// Calling `BatchStore` should return the output of the wrapped low-level storage, which is a deferred database
+		// update. Conceptually, it is possible that `ProtocolKVStore` wraps the deferred database operation in faulty
+		// code, such that it cannot be executed. Therefore, we execute the top-level deferred database update below
+		// and verify that the deferred database operation returned by the lower-level is actually reached.
+		err := store.BatchStore(rw, kvStateID, kvState)
 		require.NoError(t, err)
 	})
 
 	// On the unhappy path, i.e. when the encoding of input `kvState` failed, `ProtocolKVStore` should produce
 	// a deferred database update that always returns the encoding error.
 	t.Run("encoding fails", func(t *testing.T) {
-		lockManager := storage.NewTestingLockManager()
 		encodingError := errors.New("encoding error")
 
 		kvState.On("VersionedEncode").Return(uint64(0), nil, encodingError).Once()
 
-		err := unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			rw := storagemock.NewReaderBatchWriter(t)
-			return store.BatchStore(lctx, rw, kvStateID, kvState)
-		})
+		rw := storagemock.NewReaderBatchWriter(t)
+		err := store.BatchStore(rw, kvStateID, kvState)
 		require.ErrorIs(t, err, encodingError)
 	})
 }

state/protocol/protocol_state/kvstore_storage.go

@@ -19,13 +19,9 @@ type ProtocolKVStore interface {
 	// BatchStore adds the KV-store snapshot in the database using the given ID as key. Per convention, all
 	// implementations of [protocol.KVStoreReader] should be able to successfully encode their state into a
 	// data blob. If the encoding fails, an error is returned.
-	// BatchStore is idempotent, i.e. it accepts repeated calls with the same pairs of (stateID, kvStore).
-	// Here, the ID is expected to be a collision-resistant hash of the snapshot (including the
-	// ProtocolStateVersion). Hence, for the same ID (key), BatchStore will reject changing the data (value).
 	//
-	// Expected errors during normal operations:
-	// - storage.ErrDataMismatch if a _different_ KV store for the given stateID has already been persisted
-	BatchStore(lctx lockctx.Proof, rw storage.ReaderBatchWriter, stateID flow.Identifier, kvStore protocol.KVStoreReader) error
+	// No error is expected during normal operations
+	BatchStore(rw storage.ReaderBatchWriter, stateID flow.Identifier, kvStore protocol.KVStoreReader) error
 
 	// BatchIndex writes the blockID->stateID index to the input write batch.
 	// In a nutshell, we want to maintain a map from `blockID` to `stateID`, where `blockID` references the
@@ -34,12 +30,15 @@ type ProtocolKVStore interface {
 	//   - Consider block B, whose ingestion might potentially lead to an updated KV store. For example,
 	//     the KV store changes if we seal some execution results emitting specific service events.
 	//   - For the key `blockID`, we use the identity of block B which _proposes_ this updated KV store.
-	//   - CAUTION: The updated state requires confirmation by a QC and will only become active at the
+	//   - IMPORTANT: The updated state requires confirmation by a QC and will only become active at the
 	//     child block, _after_ validating the QC.
 	//
-	// It requires the caller to acquire storage.LockInsertBlock lock
+	// 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.
+	//
 	// Expected errors of the returned anonymous function:
-	//   - storage.ErrAlreadyExists if a KV store for the given blockID has already been indexed.
+	//   - [storage.ErrAlreadyExists] if a KV store for the given blockID has already been indexed.
 	BatchIndex(lctx lockctx.Proof, rw storage.ReaderBatchWriter, blockID flow.Identifier, stateID flow.Identifier) error
 
 	// ByID retrieves the KV store snapshot with the given ID.

state/protocol/protocol_state/mock/protocol_kv_store.go

@@ -109,7 +109,7 @@ func (s *StateMutatorSuite) testEvolveState(seals []*flow.Seal, expectedResultin
 	// expect calls to prepare a deferred update for indexing and storing the resulting state:
 	// as state has not changed, we expect the parent blocks protocol state ID
 	if stateChangeExpected {
-		s.protocolKVStoreDB.On("BatchStore", mock.Anything, rw, expectedResultingStateID, &s.evolvingState).Return(nil).Once()
+		s.protocolKVStoreDB.On("BatchStore", rw, expectedResultingStateID, &s.evolvingState).Return(nil).Once()
 	}
 
 	deferredDBOps := deferred.NewDeferredBlockPersist()
@@ -553,7 +553,7 @@ func (s *StateMutatorSuite) Test_EncodeFailed() {
 
 	rw := storagemock.NewReaderBatchWriter(s.T())
 	s.protocolKVStoreDB.On("BatchIndex", mock.Anything, mock.Anything, s.candidate.ID(), expectedResultingStateID).Return(nil).Once()
-	s.protocolKVStoreDB.On("BatchStore", mock.Anything, mock.Anything, expectedResultingStateID, &s.evolvingState).Return(exception).Once()
+	s.protocolKVStoreDB.On("BatchStore", mock.Anything, expectedResultingStateID, &s.evolvingState).Return(exception).Once()
 
 	deferredDBOps := deferred.NewDeferredBlockPersist()
 	_, err := s.mutableState.EvolveState(deferredDBOps, s.candidate.ParentID, s.candidate.View, []*flow.Seal{})

state/protocol/protocol_state/state/mutable_protocol_state_test.go

@@ -349,15 +349,11 @@ func (s *MutableProtocolState) build(
 
 	if parentStateID != resultingStateID {
 		deferredDBOps.AddNextOperation(func(lctx lockctx.Proof, blockID flow.Identifier, rw storage.ReaderBatchWriter) error {
-			err := s.kvStoreSnapshots.BatchStore(lctx, rw, resultingStateID, evolvingState)
+			// no need to hold any lock, because the resultingStateID is a full content hash of the value
+			err := s.kvStoreSnapshots.BatchStore(rw, resultingStateID, evolvingState)
 			if err == nil {
 				return nil
 			}
-			// The only error that `ProtocolKVStore.BatchStore` might return is `storage.ErrDataMismatch`.
-			// Repeated requests to store the same state for the same id should be no-ops. It should be noted
-			// that the `resultingStateID` is a collision-resistant hash of the encoded state (including the
-			// state's version). Hence, mismatching data for the same id indicates a security-critical bug
-			// or state corruption, making continuation impossible.
 			return irrecoverable.NewExceptionf("unexpected error while trying to store new protocol state: %w", err)
 		})
 	}

state/protocol/protocol_state/state/protocol_state.go

@@ -1,38 +1,58 @@
 package operation
 
 import (
+	"fmt"
+
+	"github.com/jordanschalm/lockctx"
+
 	"github.com/onflow/flow-go/model/flow"
+	"github.com/onflow/flow-go/module/irrecoverable"
 	"github.com/onflow/flow-go/storage"
 )
 
-// InsertProtocolKVStore inserts a protocol KV store by ID.
-// Error returns:
-//   - storage.ErrAlreadyExists if the key already exists in the database.
-//   - generic error in case of unexpected failure from the database layer or encoding failure.
-func InsertProtocolKVStore(w storage.Writer, protocolKVStoreID flow.Identifier, kvStore *flow.PSKeyValueStoreData) error {
-	return UpsertByKey(w, MakePrefix(codeProtocolKVStore, protocolKVStoreID), kvStore)
+// 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 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 {
+	return UpsertByKey(rw.Writer(), MakePrefix(codeProtocolKVStore, protocolKVStoreID), kvStore)
 }
 
 // RetrieveProtocolKVStore retrieves a protocol KV store 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
+// Expected error returns during normal operations:
+//   - [storage.ErrNotFound] if no protocol KV with the given ID store exists
 func RetrieveProtocolKVStore(r storage.Reader, protocolKVStoreID flow.Identifier, kvStore *flow.PSKeyValueStoreData) error {
 	return RetrieveByKey(r, MakePrefix(codeProtocolKVStore, protocolKVStoreID), kvStore)
 }
 
 // IndexProtocolKVStore indexes a protocol KV store by block ID.
-// Error returns:
-//   - storage.ErrAlreadyExists if the key already exists in the database.
-//   - generic error in case of unexpected failure from the database layer
-func IndexProtocolKVStore(w storage.Writer, blockID flow.Identifier, protocolKVStoreID flow.Identifier) error {
-	return UpsertByKey(w, MakePrefix(codeProtocolKVStoreByBlockID, blockID), protocolKVStoreID)
+//
+// 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.
+//
+// Expected error returns during normal operations:
+//   - [storage.ErrAlreadyExists] if a KV store for the given blockID has already been indexed
+func IndexProtocolKVStore(lctx lockctx.Proof, rw storage.ReaderBatchWriter, blockID flow.Identifier, protocolKVStoreID flow.Identifier) error {
+	if !lctx.HoldsLock(storage.LockInsertBlock) {
+		return fmt.Errorf("missing required lock: %s", storage.LockInsertBlock)
+	}
+
+	key := MakePrefix(codeProtocolKVStoreByBlockID, blockID)
+	exists, err := KeyExists(rw.GlobalReader(), key)
+	if err != nil {
+		return fmt.Errorf("could not check if kv-store snapshot with block id (%x) exists: %w", blockID[:], irrecoverable.NewException(err))
+	}
+	if exists {
+		return fmt.Errorf("a kv-store snapshot for block id (%x) already exists: %w", blockID[:], storage.ErrAlreadyExists)
+	}
+
+	return UpsertByKey(rw.Writer(), key, protocolKVStoreID)
 }
 
 // LookupProtocolKVStore finds protocol KV store 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
+// Expected error returns during normal operations:
+//   - [storage.ErrNotFound] if the given ID does not correspond to any known block
 func LookupProtocolKVStore(r storage.Reader, blockID flow.Identifier, protocolKVStoreID *flow.Identifier) error {
 	return RetrieveByKey(r, MakePrefix(codeProtocolKVStoreByBlockID, blockID), protocolKVStoreID)
 }