remove consistency check in InsertProtocolKVStore

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,20 +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).
-//
-// 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.
+// ProtocolStateVersion).
 //
-// Expected error returns 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_storage.go

@@ -20,13 +20,8 @@ type ProtocolKVStore interface {
 	// implementations of [protocol.KVStoreReader] should be able to successfully encode their state into a
 	// data blob. If the encoding fails, an error is returned.
 	//
-	// 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 the [storage.LockInsertBlock] lock and hold it until the database
-	// write has been committed.
-	//
-	// Expected error returns during normal operations:
-	// - [storage.ErrAlreadyExists] if a KV store with the given ID has already been stored
-	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

state/protocol/protocol_state/state/protocol_state.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)
 		})
 	}

storage/operation/protocol_kv_store.go

@@ -11,23 +11,11 @@ import (
 )
 
 // InsertProtocolKVStore inserts a protocol KV store by protocol kv store ID.
-// Error returns:
-//   - [storage.ErrAlreadyExists] if the key already exists in the database.
-func InsertProtocolKVStore(lctx lockctx.Proof, rw storage.ReaderBatchWriter, protocolKVStoreID flow.Identifier, kvStore *flow.PSKeyValueStoreData) error {
-	if !lctx.HoldsLock(storage.LockInsertBlock) {
-		return fmt.Errorf("missing required lock: %s", storage.LockInsertBlock)
-	}
-
-	key := MakePrefix(codeProtocolKVStore, protocolKVStoreID)
-	exists, err := KeyExists(rw.GlobalReader(), key)
-	if err != nil {
-		return fmt.Errorf("could not check if kv-store snapshot %x exists: %w", protocolKVStoreID[:], irrecoverable.NewException(err))
-	}
-	if exists {
-		return fmt.Errorf("a kv-store snapshot with id %x already exists: %w", protocolKVStoreID[:], storage.ErrAlreadyExists)
-	}
-
-	return UpsertByKey(rw.Writer(), key, kvStore)
+// 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.

storage/operation/protocol_kv_store_test.go

@@ -24,10 +24,8 @@ func TestInsertProtocolKVStore(t *testing.T) {
 		}
 
 		kvStoreStateID := unittest.IdentifierFixture()
-		err := unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.InsertProtocolKVStore(lctx, rw, kvStoreStateID, expected)
-			})
+		err := db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+			return operation.InsertProtocolKVStore(rw, kvStoreStateID, expected)
 		})
 		require.NoError(t, err)
 
@@ -53,47 +51,6 @@ func TestInsertProtocolKVStore(t *testing.T) {
 	})
 }
 
-// TestInsertProtocolKVStore_ErrAlreadyExists tests that InsertProtocolKVStore returns ErrAlreadyExists
-// when attempting to insert a protocol KV store that already exists.
-func TestInsertProtocolKVStore_ErrAlreadyExists(t *testing.T) {
-	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
-		lockManager := storage.NewTestingLockManager()
-		expected := &flow.PSKeyValueStoreData{
-			Version: 2,
-			Data:    unittest.RandomBytes(32),
-		}
-
-		kvStoreStateID := unittest.IdentifierFixture()
-
-		// First insertion should succeed
-		err := unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.InsertProtocolKVStore(lctx, rw, kvStoreStateID, expected)
-			})
-		})
-		require.NoError(t, err)
-
-		// Second insertion with same ID should fail with ErrAlreadyExists
-		differentData := &flow.PSKeyValueStoreData{
-			Version: 3,
-			Data:    unittest.RandomBytes(32),
-		}
-		err = unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.InsertProtocolKVStore(lctx, rw, kvStoreStateID, differentData)
-			})
-		})
-		require.Error(t, err)
-		require.ErrorIs(t, err, storage.ErrAlreadyExists)
-
-		// Verify original data is still there and unchanged
-		var actual flow.PSKeyValueStoreData
-		err = operation.RetrieveProtocolKVStore(db.Reader(), kvStoreStateID, &actual)
-		require.NoError(t, err)
-		assert.Equal(t, expected, &actual)
-	})
-}
-
 // TestIndexProtocolKVStore_ErrAlreadyExists tests that IndexProtocolKVStore returns ErrAlreadyExists
 // when attempting to index a protocol KV store for a block ID that already has an index.
 func TestIndexProtocolKVStore_ErrAlreadyExists(t *testing.T) {
@@ -108,10 +65,8 @@ func TestIndexProtocolKVStore_ErrAlreadyExists(t *testing.T) {
 		blockID := unittest.IdentifierFixture()
 
 		// Insert the protocol KV store first
-		err := unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.InsertProtocolKVStore(lctx, rw, kvStoreStateID, expected)
-			})
+		err := db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+			return operation.InsertProtocolKVStore(rw, kvStoreStateID, expected)
 		})
 		require.NoError(t, err)
 
@@ -141,29 +96,6 @@ func TestIndexProtocolKVStore_ErrAlreadyExists(t *testing.T) {
 	})
 }
 
-// TestInsertProtocolKVStore_MissingLock tests that InsertProtocolKVStore requires LockInsertBlock.
-func TestInsertProtocolKVStore_MissingLock(t *testing.T) {
-	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
-		lockManager := storage.NewTestingLockManager()
-		expected := &flow.PSKeyValueStoreData{
-			Version: 2,
-			Data:    unittest.RandomBytes(32),
-		}
-
-		kvStoreStateID := unittest.IdentifierFixture()
-
-		// Attempt to insert without holding the required lock
-		lctx := lockManager.NewContext()
-		defer lctx.Release()
-
-		err := db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-			return operation.InsertProtocolKVStore(lctx, rw, kvStoreStateID, expected)
-		})
-		require.Error(t, err)
-		require.Contains(t, err.Error(), storage.LockInsertBlock)
-	})
-}
-
 // TestIndexProtocolKVStore_MissingLock tests that IndexProtocolKVStore requires LockInsertBlock.
 func TestIndexProtocolKVStore_MissingLock(t *testing.T) {
 	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
@@ -183,28 +115,6 @@ func TestIndexProtocolKVStore_MissingLock(t *testing.T) {
 	})
 }
 
-// TestInsertProtocolKVStore_WrongLock tests that InsertProtocolKVStore fails when holding wrong locks.
-func TestInsertProtocolKVStore_WrongLock(t *testing.T) {
-	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
-		lockManager := storage.NewTestingLockManager()
-		expected := &flow.PSKeyValueStoreData{
-			Version: 2,
-			Data:    unittest.RandomBytes(32),
-		}
-
-		kvStoreStateID := unittest.IdentifierFixture()
-
-		// Test with LockFinalizeBlock (wrong lock)
-		err := unittest.WithLock(t, lockManager, storage.LockFinalizeBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return operation.InsertProtocolKVStore(lctx, rw, kvStoreStateID, expected)
-			})
-		})
-		require.Error(t, err)
-		require.Contains(t, err.Error(), storage.LockInsertBlock)
-	})
-}
-
 // TestIndexProtocolKVStore_WrongLock tests that IndexProtocolKVStore fails when holding wrong locks.
 func TestIndexProtocolKVStore_WrongLock(t *testing.T) {
 	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {

storage/protocol_kv_store.go

@@ -17,15 +17,10 @@ type ProtocolKVStore interface {
 	// BatchStore persists the KV-store snapshot in the database using the given ID as key.
 	// 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).
 	//
-	// 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.ErrDataMismatch] if a _different_ KV store for the given stateID has already been persisted
-	BatchStore(lctx lockctx.Proof, rw ReaderBatchWriter, stateID flow.Identifier, data *flow.PSKeyValueStoreData) error
+	// No error is expected during normal operations.
+	BatchStore(rw ReaderBatchWriter, stateID flow.Identifier, data *flow.PSKeyValueStoreData) error
 
 	// BatchIndex appends the following operation to the provided write batch:
 	// we extend the map from `blockID` to `stateID`, where `blockID` references the
@@ -45,7 +40,7 @@ type ProtocolKVStore interface {
 	// database write has been committed.
 	//
 	// Expected error returns during normal operations:
-	// - storage.ErrDataMismatch if a _different_ KV store for the given stateID has already been persisted
+	// - [storage.ErrAlreadyExists] if a KV store for the given blockID has already been indexed
 	BatchIndex(lctx lockctx.Proof, rw ReaderBatchWriter, blockID flow.Identifier, stateID flow.Identifier) error
 
 	// ByID retrieves the KV store snapshot with the given ID.

storage/store/protocol_kv_store.go

@@ -70,7 +70,7 @@ func NewProtocolKVStore(collector module.CacheMetrics,
 		db: db,
 		cache: newCache(collector, metrics.ResourceProtocolKVStore,
 			withLimit[flow.Identifier, *flow.PSKeyValueStoreData](kvStoreCacheSize),
-			withStoreWithLock(operation.InsertProtocolKVStore),
+			withStore(operation.InsertProtocolKVStore),
 			withRetrieve(retrieveByStateID)),
 		byBlockIdCache: newCache(collector, metrics.ResourceProtocolKVStoreByBlockID,
 			withLimit[flow.Identifier, flow.Identifier](kvStoreByBlockIDCacheSize),
@@ -84,15 +84,10 @@ func NewProtocolKVStore(collector module.CacheMetrics,
 // Here, the ID is expected to be a collision-resistant hash of the snapshot (including the
 // ProtocolStateVersion). Hence, for the same ID, BatchStore will reject changing the data.
 //
-// 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 the [storage.LockInsertBlock] lock and hold it until the database
-// write has been committed.
-//
-// Expected error returns during normal operations:
-// - [storage.ErrAlreadyExist] if a KV store with the given ID has already been stored
-func (s *ProtocolKVStore) BatchStore(lctx lockctx.Proof, rw storage.ReaderBatchWriter, stateID flow.Identifier, data *flow.PSKeyValueStoreData) error {
+// No error is expected during normal operations.
+func (s *ProtocolKVStore) BatchStore(rw storage.ReaderBatchWriter, stateID flow.Identifier, data *flow.PSKeyValueStoreData) error {
 	return s.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-		return s.cache.PutWithLockTx(lctx, rw, stateID, data)
+		return s.cache.PutTx(rw, stateID, data)
 	})
 }
 

storage/store/protocol_kv_store_test.go

@@ -31,8 +31,10 @@ func TestKeyValueStoreStorage(t *testing.T) {
 		// store protocol state and auxiliary info
 		err := unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
 			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				err := store.BatchStore(lctx, rw, stateID, expected)
-				require.NoError(t, err)
+				err := store.BatchStore(rw, stateID, expected)
+				if err != nil {
+					return err
+				}
 				return store.BatchIndex(lctx, rw, blockID, stateID)
 			})
 		})
@@ -50,13 +52,11 @@ func TestKeyValueStoreStorage(t *testing.T) {
 	})
 }
 
-// TestProtocolKVStore_StoreTx tests that StoreTx handles storage request correctly, when a snapshot with
-// the given id has already been  stored:
-//   - if the KV-store snapshot is exactly the same as the one already stored (incl. the version),  `BatchStore` should return without an error
-//   - if we request to store a _different_  KV-store snapshot, an `storage.ErrDataMismatch` should be returned.
+// TestProtocolKVStore_StoreTx tests that StoreTx handles storage request correctly.
+// Since BatchStore is now idempotent and doesn't return errors for duplicate data,
+// we test that it can be called multiple times without issues.
 func TestProtocolKVStore_StoreTx(t *testing.T) {
 	dbtest.RunWithDB(t, func(t *testing.T, db storage.DB) {
-		lockManager := storage.NewTestingLockManager()
 		metrics := metrics.NewNoopCollector()
 		store := NewProtocolKVStore(metrics, db, DefaultProtocolKVStoreCacheSize, DefaultProtocolKVStoreByBlockIDCacheSize)
 
@@ -67,46 +67,21 @@ func TestProtocolKVStore_StoreTx(t *testing.T) {
 		}
 
 		// Store initial data
-		err := unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return store.BatchStore(lctx, rw, stateID, expected)
-			})
+		err := db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+			return store.BatchStore(rw, stateID, expected)
 		})
 		require.NoError(t, err)
 
-		// Store same data again - should error
-		err = unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return store.BatchStore(lctx, rw, stateID, expected)
-			})
+		// Store same data again - should succeed (idempotent)
+		err = db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+			return store.BatchStore(rw, stateID, expected)
 		})
-		require.ErrorIs(t, err, storage.ErrAlreadyExists)
-
-		// Attempt to store different data with the same stateID
-		dataDifferent := &flow.PSKeyValueStoreData{
-			Version: 2,
-			Data:    unittest.RandomBytes(32),
-		}
-
-		err = unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return store.BatchStore(lctx, rw, stateID, dataDifferent)
-			})
-		})
-		require.ErrorIs(t, err, storage.ErrAlreadyExists)
-
-		// Attempt to store different version with the same stateID
-		versionDifferent := &flow.PSKeyValueStoreData{
-			Version: 3,
-			Data:    expected.Data,
-		}
+		require.NoError(t, err)
 
-		err = unittest.WithLock(t, lockManager, storage.LockInsertBlock, func(lctx lockctx.Context) error {
-			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
-				return store.BatchStore(lctx, rw, stateID, versionDifferent)
-			})
-		})
-		require.ErrorIs(t, err, storage.ErrAlreadyExists)
+		// Verify the data can still be retrieved
+		actual, err := store.ByID(stateID)
+		require.NoError(t, err)
+		assert.Equal(t, expected, actual)
 	})
 }