Add storage locking for index operations

Author: peterargue

module/state_synchronization/indexer/indexer_core.go

@@ -171,6 +171,12 @@ func (c *IndexerCore) IndexBlockData(data *execution_data.BlockExecutionDataEnti
 			return fmt.Errorf("could not collect scheduled transaction data: %w", err)
 		}
 
+		lctx := c.lockManager.NewContext()
+		defer lctx.Release()
+		if err = lctx.AcquireLock(storage.LockIndexScheduledTransaction); err != nil {
+			return fmt.Errorf("could not acquire lock for indexing scheduled transactions: %w", err)
+		}
+
 		err = c.protocolDB.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
 			err := c.events.BatchStore(data.BlockID, []flow.EventsList{events}, rw)
 			if err != nil {
@@ -183,7 +189,7 @@ func (c *IndexerCore) IndexBlockData(data *execution_data.BlockExecutionDataEnti
 			}
 
 			for txID, scheduledTxID := range scheduledTransactionData {
-				err = c.scheduledTransactions.BatchIndex(data.BlockID, txID, scheduledTxID, rw)
+				err = c.scheduledTransactions.BatchIndex(lctx, data.BlockID, txID, scheduledTxID, rw)
 				if err != nil {
 					return fmt.Errorf("could not index scheduled transaction (%d) %s at height %d: %w", scheduledTxID, txID, header.Height, err)
 				}

module/state_synchronization/indexer/indexer_core_test.go

@@ -245,7 +245,7 @@ func TestExecutionState_IndexBlockData(t *testing.T) {
 			test.collections.On("StoreAndIndexByTransaction", mock.Anything, collection).Return(&flow.LightCollection{}, nil)
 		}
 		for txID, scheduledTxID := range tf.ExpectedScheduledTransactions {
-			test.scheduledTransactions.On("BatchIndex", blockID, txID, scheduledTxID, mock.Anything).Return(nil)
+			test.scheduledTransactions.On("BatchIndex", mock.Anything, blockID, txID, scheduledTxID, mock.Anything).Return(nil)
 		}
 
 		err := test.indexer.IndexBlockData(tf.ExecutionDataEntity())

storage/locks.go

@@ -27,8 +27,10 @@ const (
 	LockInsertCollection = "lock_insert_collection"
 	// LockBootstrapping protects data that is *exclusively* written during bootstrapping.
 	LockBootstrapping = "lock_bootstrapping"
-	// LockInsertChunkDataPack protects the insertion of chunk data packs (not yet used anywhere
+	// LockInsertChunkDataPack protects the insertion of chunk data packs (not yet used anywhere)
 	LockInsertChunkDataPack = "lock_insert_chunk_data_pack"
+	// LockIndexScheduledTransaction protects the indexing of scheduled transactions.
+	LockIndexScheduledTransaction = "lock_index_scheduled_transaction"
 )
 
 // Locks returns a list of all named locks used by the storage layer.
@@ -42,6 +44,7 @@ func Locks() []string {
 		LockInsertCollection,
 		LockBootstrapping,
 		LockInsertChunkDataPack,
+		LockIndexScheduledTransaction,
 	}
 }
 

storage/mock/scheduled_transactions.go

@@ -1,6 +1,9 @@
 package operation
 
 import (
+	"fmt"
+
+	"github.com/jordanschalm/lockctx"
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/storage"
 )
@@ -22,16 +25,44 @@ func RetrieveBlockIDByScheduledTransactionID(r storage.Reader, txID flow.Identif
 	return RetrieveByKey(r, MakePrefix(codeBlockIDByScheduledTransactionID, txID), blockID)
 }
 
-// BatchIndexScheduledTransactionID indexes the scheduled transaction's transaction ID by its scheduled transaction ID.
+// IndexScheduledTransactionID indexes the scheduled transaction's transaction ID by its scheduled transaction ID.
 //
-// No errors are expected during normal operation.
-func BatchIndexScheduledTransactionID(w storage.Writer, scheduledTxID uint64, txID flow.Identifier) error {
-	return UpsertByKey(w, MakePrefix(codeTransactionIDByScheduledTransactionID, scheduledTxID), txID)
+// Expected error returns during normal operation:
+//   - [storage.ErrAlreadyExists]: if the scheduled transaction ID is already indexed
+func IndexScheduledTransactionID(lctx lockctx.Proof, rw storage.ReaderBatchWriter, scheduledTxID uint64, txID flow.Identifier) error {
+	if !lctx.HoldsLock(storage.LockIndexScheduledTransaction) {
+		return fmt.Errorf("missing lock: %v", storage.LockIndexScheduledTransaction)
+	}
+
+	key := MakePrefix(codeTransactionIDByScheduledTransactionID, scheduledTxID)
+	exists, err := KeyExists(rw.GlobalReader(), key)
+	if err != nil {
+		return err
+	}
+	if exists {
+		return fmt.Errorf("scheduled transaction ID already indexed: %w", storage.ErrAlreadyExists)
+	}
+
+	return UpsertByKey(rw.Writer(), key, txID)
 }
 
-// BatchIndexScheduledTransactionBlockID indexes the scheduled transaction's block ID by its transaction ID.
+// IndexScheduledTransactionBlockID indexes the scheduled transaction's block ID by its transaction ID.
 //
-// No errors are expected during normal operation.
-func BatchIndexScheduledTransactionBlockID(w storage.Writer, txID flow.Identifier, blockID flow.Identifier) error {
-	return UpsertByKey(w, MakePrefix(codeBlockIDByScheduledTransactionID, txID), blockID)
+// Expected error returns during normal operation:
+//   - [storage.ErrAlreadyExists]: if the scheduled transaction block ID is already indexed
+func IndexScheduledTransactionBlockID(lctx lockctx.Proof, rw storage.ReaderBatchWriter, txID flow.Identifier, blockID flow.Identifier) error {
+	if !lctx.HoldsLock(storage.LockIndexScheduledTransaction) {
+		return fmt.Errorf("missing lock: %v", storage.LockIndexScheduledTransaction)
+	}
+
+	key := MakePrefix(codeBlockIDByScheduledTransactionID, txID)
+	exists, err := KeyExists(rw.GlobalReader(), key)
+	if err != nil {
+		return err
+	}
+	if exists {
+		return fmt.Errorf("scheduled transaction block ID already indexed: %w", storage.ErrAlreadyExists)
+	}
+
+	return UpsertByKey(rw.Writer(), key, blockID)
 }

storage/operation/scheduled_transactions.go

@@ -1,6 +1,7 @@
 package storage
 
 import (
+	"github.com/jordanschalm/lockctx"
 	"github.com/onflow/flow-go/model/flow"
 )
 
@@ -30,7 +31,8 @@ type ScheduledTransactions interface {
 	// BatchIndex indexes the scheduled transaction by its block ID, transaction ID, and scheduled transaction ID.
 	// `scheduledTxID` is the uint64 id field returned by the system smart contract.
 	// `txID` is be the TransactionBody.ID of the scheduled transaction.
+	// Requires the lock: [storage.LockIndexScheduledTransaction]
 	//
 	// No errors are expected during normal operation.
-	BatchIndex(blockID flow.Identifier, txID flow.Identifier, scheduledTxID uint64, batch ReaderBatchWriter) error
+	BatchIndex(lctx lockctx.Proof, blockID flow.Identifier, txID flow.Identifier, scheduledTxID uint64, batch ReaderBatchWriter) error
 }

storage/scheduled_transactions.go

@@ -3,6 +3,7 @@ package store
 import (
 	"fmt"
 
+	"github.com/jordanschalm/lockctx"
 	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/module"
 	"github.com/onflow/flow-go/module/metrics"
@@ -57,19 +58,19 @@ func NewScheduledTransactions(collector module.CacheMetrics, db storage.DB, cach
 }
 
 // BatchIndex indexes the scheduled transaction by its block ID, transaction ID, and scheduled transaction ID.
-// `scheduledTxID` is the uint64 id field returned by the system smart contract.
 // `txID` is be the TransactionBody.ID of the scheduled transaction.
+// `scheduledTxID` is the uint64 id field returned by the system smart contract.
+// Requires the lock: [storage.LockIndexScheduledTransaction]
 //
-// No errors are expected during normal operation.
-func (st *ScheduledTransactions) BatchIndex(blockID flow.Identifier, txID flow.Identifier, scheduledTxID uint64, batch storage.ReaderBatchWriter) error {
-	writer := batch.Writer()
-
-	err := operation.BatchIndexScheduledTransactionID(writer, scheduledTxID, txID)
+// Expected error returns during normal operation:
+//   - [storage.ErrAlreadyExists]: if the scheduled transaction is already indexed
+func (st *ScheduledTransactions) BatchIndex(lctx lockctx.Proof, blockID flow.Identifier, txID flow.Identifier, scheduledTxID uint64, batch storage.ReaderBatchWriter) error {
+	err := operation.IndexScheduledTransactionID(lctx, batch, scheduledTxID, txID)
 	if err != nil {
 		return fmt.Errorf("failed to batch index scheduled transaction: %w", err)
 	}
 
-	err = operation.BatchIndexScheduledTransactionBlockID(writer, txID, blockID)
+	err = operation.IndexScheduledTransactionBlockID(lctx, batch, txID, blockID)
 	if err != nil {
 		return fmt.Errorf("failed to batch index scheduled transaction block ID: %w", err)
 	}

storage/store/scheduled_transactions.go

@@ -38,12 +38,19 @@ func runTest(t *testing.T, g *fixtures.GeneratorSuite, db storage.DB, store *Sch
 		data[txID] = g.Random().Uint64()
 	}
 
+	lockManager := storage.NewTestingLockManager()
+	lctx := lockManager.NewContext()
+	defer lctx.Release()
+	if err := lctx.AcquireLock(storage.LockIndexScheduledTransaction); err != nil {
+		t.Fatalf("could not acquire lock for indexing scheduled transactions: %v", err)
+	}
+
 	// index data within a batch
 	batch := db.NewBatch()
 	defer batch.Close()
 
 	for txID, scheduledTxID := range data {
-		err := store.BatchIndex(blockID, txID, scheduledTxID, batch)
+		err := store.BatchIndex(lctx, blockID, txID, scheduledTxID, batch)
 		require.NoError(t, err)
 	}