[Storage] Refactor index execution result

cmd/access/node_builder/access_node_builder.go

@@ -2236,6 +2236,8 @@ func (builder *FlowAccessNodeBuilder) Build() (cmd.Node, error) {
 				node.EngineRegistry,
 				node.State,
 				node.Me,
+				node.StorageLockMgr,
+				node.ProtocolDB,
 				node.Storage.Blocks,
 				node.Storage.Results,
 				node.Storage.Receipts,

cmd/util/cmd/execution-state-extract/execution_state_extract_test.go

@@ -67,7 +67,7 @@ func TestExtractExecutionState(t *testing.T) {
 			blockID := unittest.IdentifierFixture()
 			stateCommitment := unittest.StateCommitmentFixture()
 
-			err := unittest.WithLock(t, lockManager, storage.LockInsertOwnReceipt, func(lctx lockctx.Context) error {
+			err := unittest.WithLock(t, lockManager, storage.LockIndexStateCommitment, func(lctx lockctx.Context) error {
 				return storageDB.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
 					// Store the state commitment for the block ID
 					return operation.IndexStateCommitment(lctx, rw, blockID, stateCommitment)
@@ -139,7 +139,7 @@ func TestExtractExecutionState(t *testing.T) {
 				// generate random block and map it to state commitment
 				blockID := unittest.IdentifierFixture()
 
-				err = unittest.WithLock(t, lockManager, storage.LockInsertOwnReceipt, func(lctx lockctx.Context) error {
+				err = unittest.WithLock(t, lockManager, storage.LockIndexStateCommitment, func(lctx lockctx.Context) error {
 					return storageDB.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
 						return operation.IndexStateCommitment(lctx, rw, blockID, flow.StateCommitment(stateCommitment))
 					})

cmd/util/cmd/root.go

@@ -39,7 +39,6 @@ import (
 	read_execution_state "github.com/onflow/flow-go/cmd/util/cmd/read-execution-state"
 	read_hotstuff "github.com/onflow/flow-go/cmd/util/cmd/read-hotstuff/cmd"
 	read_protocol_state "github.com/onflow/flow-go/cmd/util/cmd/read-protocol-state/cmd"
-	index_er "github.com/onflow/flow-go/cmd/util/cmd/reindex/cmd"
 	rollback_executed_height "github.com/onflow/flow-go/cmd/util/cmd/rollback-executed-height/cmd"
 	run_script "github.com/onflow/flow-go/cmd/util/cmd/run-script"
 	"github.com/onflow/flow-go/cmd/util/cmd/snapshot"
@@ -111,7 +110,6 @@ func addCommands() {
 	rootCmd.AddCommand(leaders.Cmd)
 	rootCmd.AddCommand(epochs.RootCmd)
 	rootCmd.AddCommand(edbs.RootCmd)
-	rootCmd.AddCommand(index_er.RootCmd)
 	rootCmd.AddCommand(rollback_executed_height.Cmd)
 	rootCmd.AddCommand(read_execution_state.Cmd)
 	rootCmd.AddCommand(snapshot.Cmd)

engine/access/access_test.go

@@ -544,6 +544,7 @@ func (suite *Suite) TestGetBlockByIDAndHeight() {
 
 func (suite *Suite) TestGetExecutionResultByBlockID() {
 	suite.RunTest(func(handler *rpc.Handler, db storage.DB, all *store.All) {
+		lockManager := storage.NewTestingLockManager()
 
 		// test block1 get by ID
 		nonexistingID := unittest.IdentifierFixture()
@@ -553,8 +554,16 @@ func (suite *Suite) TestGetExecutionResultByBlockID() {
 			unittest.WithExecutionResultBlockID(blockID),
 			unittest.WithServiceEvents(3))
 
-		require.NoError(suite.T(), all.Results.Store(er))
-		require.NoError(suite.T(), all.Results.Index(blockID, er.ID()))
+		require.NoError(suite.T(), storage.WithLock(lockManager, storage.LockIndexExecutionResult, func(lctx lockctx.Context) error {
+			return db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+				err := all.Results.BatchStore(er, rw)
+				if err != nil {
+					return err
+				}
+				// requires LockIndexExecutionResult
+				return all.Results.BatchIndex(lctx, rw, blockID, er.ID())
+			})
+		}))
 
 		assertResp := func(
 			resp *accessproto.ExecutionResultForBlockIDResponse,
@@ -626,6 +635,7 @@ func (suite *Suite) TestGetExecutionResultByBlockID() {
 // is reported as sealed
 func (suite *Suite) TestGetSealedTransaction() {
 	unittest.RunWithPebbleDB(suite.T(), func(pdb *pebble.DB) {
+		lockManager := storage.NewTestingLockManager()
 		db := pebbleimpl.ToDB(pdb)
 		all := store.InitAll(metrics.NewNoopCollector(), db)
 		enIdentities := unittest.IdentityListFixture(2, unittest.WithRole(flow.RoleExecution))
@@ -740,6 +750,8 @@ func (suite *Suite) TestGetSealedTransaction() {
 			suite.net,
 			suite.state,
 			suite.me,
+			lockManager,
+			db,
 			all.Blocks,
 			all.Results,
 			all.Receipts,
@@ -816,6 +828,7 @@ func (suite *Suite) TestGetSealedTransaction() {
 // transaction ID, block ID, and collection ID.
 func (suite *Suite) TestGetTransactionResult() {
 	unittest.RunWithPebbleDB(suite.T(), func(pdb *pebble.DB) {
+		lockManager := storage.NewTestingLockManager()
 		db := pebbleimpl.ToDB(pdb)
 		all := store.InitAll(metrics.NewNoopCollector(), db)
 		originID := unittest.IdentifierFixture()
@@ -963,6 +976,8 @@ func (suite *Suite) TestGetTransactionResult() {
 			suite.net,
 			suite.state,
 			suite.me,
+			lockManager,
+			db,
 			all.Blocks,
 			all.Results,
 			all.Receipts,
@@ -1135,6 +1150,7 @@ func (suite *Suite) TestGetTransactionResult() {
 // the correct block id
 func (suite *Suite) TestExecuteScript() {
 	unittest.RunWithPebbleDB(suite.T(), func(pdb *pebble.DB) {
+		lockManager := storage.NewTestingLockManager()
 		db := pebbleimpl.ToDB(pdb)
 		all := store.InitAll(metrics.NewNoopCollector(), db)
 		identities := unittest.IdentityListFixture(2, unittest.WithRole(flow.RoleExecution))
@@ -1229,6 +1245,8 @@ func (suite *Suite) TestExecuteScript() {
 			suite.net,
 			suite.state,
 			suite.me,
+			lockManager,
+			db,
 			all.Blocks,
 			all.Results,
 			all.Receipts,

engine/access/ingestion/engine.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/jordanschalm/lockctx"
 	"github.com/rs/zerolog"
 
 	"github.com/onflow/flow-go/consensus/hotstuff/model"
@@ -78,6 +79,8 @@ type Engine struct {
 
 	// storage
 	// FIX: remove direct DB access by substituting indexer module
+	db                storage.DB
+	lockManager       storage.LockManager
 	blocks            storage.Blocks
 	executionReceipts storage.ExecutionReceipts
 	maxReceiptHeight  uint64
@@ -101,6 +104,8 @@ func New(
 	net network.EngineRegistry,
 	state protocol.State,
 	me module.Local,
+	lockManager storage.LockManager,
+	db storage.DB,
 	blocks storage.Blocks,
 	executionResults storage.ExecutionResults,
 	executionReceipts storage.ExecutionReceipts,
@@ -133,6 +138,8 @@ func New(
 		log:                      log.With().Str("engine", "ingestion").Logger(),
 		state:                    state,
 		me:                       me,
+		lockManager:              lockManager,
+		db:                       db,
 		blocks:                   blocks,
 		executionResults:         executionResults,
 		executionReceipts:        executionReceipts,
@@ -372,19 +379,32 @@ func (e *Engine) processFinalizedBlock(block *flow.Block) error {
 	// TODO: substitute an indexer module as layer between engine and storage
 
 	// index the block storage with each of the collection guarantee
-	err := e.blocks.IndexBlockContainingCollectionGuarantees(block.ID(), flow.GetIDs(block.Payload.Guarantees))
+	err := storage.WithLocks(e.lockManager, []string{
+		storage.LockIndexCollectionsByBlock,
+		storage.LockIndexExecutionResult,
+	}, func(lctx lockctx.Context) error {
+		return e.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+			// requires [storage.LockIndexCollectionsByBlock] lock
+			err := e.blocks.BatchIndexBlockContainingCollectionGuarantees(lctx, rw, block.ID(), flow.GetIDs(block.Payload.Guarantees))
+			if err != nil {
+				return fmt.Errorf("could not index block for collections: %w", err)
+			}
+
+			// loop through seals and index ID -> result ID
+			for _, seal := range block.Payload.Seals {
+				// requires [storage.LockIndexExecutionResult] lock
+				err := e.executionResults.BatchIndex(lctx, rw, seal.BlockID, seal.ResultID)
+				if err != nil {
+					return fmt.Errorf("could not index block for execution result: %w", err)
+				}
+			}
+			return nil
+		})
+	})
 	if err != nil {
 		return fmt.Errorf("could not index block for collections: %w", err)
 	}
 
-	// loop through seals and index ID -> result ID
-	for _, seal := range block.Payload.Seals {
-		err := e.executionResults.Index(seal.BlockID, seal.ResultID)
-		if err != nil {
-			return fmt.Errorf("could not index block for execution result: %w", err)
-		}
-	}
-
 	e.collectionSyncer.RequestCollectionsForBlock(block.Height, block.Payload.Guarantees)
 	e.collectionExecutedMetric.BlockFinalized(block)
 

engine/access/ingestion/engine_test.go

@@ -131,6 +131,7 @@ func (s *Suite) SetupTest() {
 	s.receipts = new(storagemock.ExecutionReceipts)
 	s.transactions = new(storagemock.Transactions)
 	s.results = new(storagemock.ExecutionResults)
+	s.results.On("BatchIndex", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
 	collectionsToMarkFinalized := stdmap.NewTimes(100)
 	collectionsToMarkExecuted := stdmap.NewTimes(100)
 	blocksToMarkExecuted := stdmap.NewTimes(100)
@@ -213,6 +214,8 @@ func (s *Suite) initEngineAndSyncer(ctx irrecoverable.SignalerContext) (*Engine,
 		s.net,
 		s.proto.state,
 		s.me,
+		s.lockManager,
+		s.db,
 		s.blocks,
 		s.results,
 		s.receipts,
@@ -294,7 +297,7 @@ func (s *Suite) TestOnFinalizedBlockSingle() {
 	}
 
 	// expect that the block storage is indexed with each of the collection guarantee
-	s.blocks.On("IndexBlockContainingCollectionGuarantees", block.ID(), []flow.Identifier(flow.GetIDs(block.Payload.Guarantees))).Return(nil).Once()
+	s.blocks.On("BatchIndexBlockContainingCollectionGuarantees", mock.Anything, mock.Anything, block.ID(), []flow.Identifier(flow.GetIDs(block.Payload.Guarantees))).Return(nil).Once()
 	for _, seal := range block.Payload.Seals {
 		s.results.On("Index", seal.BlockID, seal.ResultID).Return(nil).Once()
 	}
@@ -321,7 +324,7 @@ func (s *Suite) TestOnFinalizedBlockSingle() {
 	// assert that the block was retrieved and all collections were requested
 	s.headers.AssertExpectations(s.T())
 	s.request.AssertNumberOfCalls(s.T(), "EntityByID", len(block.Payload.Guarantees))
-	s.results.AssertNumberOfCalls(s.T(), "Index", len(block.Payload.Seals))
+	s.results.AssertNumberOfCalls(s.T(), "BatchIndex", len(block.Payload.Seals))
 }
 
 // TestOnFinalizedBlockSeveralBlocksAhead checks OnFinalizedBlock with a block several blocks newer than the last block processed
@@ -370,7 +373,7 @@ func (s *Suite) TestOnFinalizedBlockSeveralBlocksAhead() {
 
 	// expected all new blocks after last block processed
 	for _, block := range blocks {
-		s.blocks.On("IndexBlockContainingCollectionGuarantees", block.ID(), []flow.Identifier(flow.GetIDs(block.Payload.Guarantees))).Return(nil).Once()
+		s.blocks.On("BatchIndexBlockContainingCollectionGuarantees", mock.Anything, mock.Anything, block.ID(), []flow.Identifier(flow.GetIDs(block.Payload.Guarantees))).Return(nil).Once()
 
 		for _, cg := range block.Payload.Guarantees {
 			s.request.On("EntityByID", cg.CollectionID, mock.Anything).Return().Run(func(args mock.Arguments) {
@@ -398,9 +401,9 @@ func (s *Suite) TestOnFinalizedBlockSeveralBlocksAhead() {
 	}
 
 	s.headers.AssertExpectations(s.T())
-	s.blocks.AssertNumberOfCalls(s.T(), "IndexBlockContainingCollectionGuarantees", newBlocksCount)
+	s.blocks.AssertNumberOfCalls(s.T(), "BatchIndexBlockContainingCollectionGuarantees", newBlocksCount)
 	s.request.AssertNumberOfCalls(s.T(), "EntityByID", expectedEntityByIDCalls)
-	s.results.AssertNumberOfCalls(s.T(), "Index", expectedIndexCalls)
+	s.results.AssertNumberOfCalls(s.T(), "BatchIndex", expectedIndexCalls)
 }
 
 // TestOnCollection checks that when a Collection is received, it is persisted

engine/access/ingestion2/collection_syncer.go

@@ -189,12 +189,12 @@ func (s *CollectionSyncer) StartWorkerLoop(ctx irrecoverable.SignalerContext, re
 
 			// Create a lock context for indexing
 			lctx := s.lockManager.NewContext()
+			defer lctx.Release()
 			err := lctx.AcquireLock(storage.LockInsertCollection)
 			if err != nil {
 				ctx.Throw(fmt.Errorf("could not acquire lock for collection indexing: %w", err))
 				return
 			}
-			defer lctx.Release()
 
 			err = indexer.IndexCollection(lctx, collection, s.collections, s.logger, s.collectionExecutedMetric)
 			if err != nil {