Merge pull request #8005 from onflow/leo/refactor-index-result

[Storage] Refactor index execution result

Author: zhangchiqing

cmd/access/node_builder/access_node_builder.go

@@ -2267,6 +2267,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/reindex/cmd/results.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)

cmd/util/cmd/reindex/cmd/root.go

@@ -545,6 +545,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()
@@ -554,8 +555,15 @@ 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
+				}
+				return all.Results.BatchIndex(lctx, rw, blockID, er.ID()) // requires storage.LockIndexExecutionResult
+			})
+		}))
 
 		assertResp := func(
 			resp *accessproto.ExecutionResultForBlockIDResponse,
@@ -627,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))
@@ -751,6 +760,8 @@ func (suite *Suite) TestGetSealedTransaction() {
 			suite.net,
 			suite.state,
 			suite.me,
+			lockManager,
+			db,
 			all.Blocks,
 			all.Results,
 			all.Receipts,
@@ -841,6 +852,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()
@@ -1010,6 +1022,8 @@ func (suite *Suite) TestGetTransactionResult() {
 			suite.net,
 			suite.state,
 			suite.me,
+			lockManager,
+			db,
 			all.Blocks,
 			all.Results,
 			all.Receipts,
@@ -1172,6 +1186,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))
@@ -1276,6 +1291,8 @@ func (suite *Suite) TestExecuteScript() {
 			suite.net,
 			suite.state,
 			suite.me,
+			lockManager,
+			db,
 			all.Blocks,
 			all.Results,
 			all.Receipts,

cmd/util/cmd/reindex/main.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/jordanschalm/lockctx"
 	"github.com/rs/zerolog"
 
 	"github.com/onflow/flow-go/consensus/hotstuff/model"
@@ -58,6 +59,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
@@ -83,6 +86,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,
@@ -116,6 +121,8 @@ func New(
 		log:                      log.With().Str("engine", "ingestion").Logger(),
 		state:                    state,
 		me:                       me,
+		lockManager:              lockManager,
+		db:                       db,
 		blocks:                   blocks,
 		executionResults:         executionResults,
 		executionReceipts:        executionReceipts,
@@ -355,19 +362,29 @@ 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, storage.LockGroupAccessFinalizingBlock, func(lctx lockctx.Context) error {
+		return e.db.WithReaderBatchWriter(func(rw storage.ReaderBatchWriter) error {
+			// requires [storage.LockIndexBlockByPayloadGuarantees] 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)
-		}
-	}
-
 	err = e.collectionSyncer.RequestCollectionsForBlock(block.Height, block.Payload.Guarantees)
 	if err != nil {
 		return fmt.Errorf("could not request collections for block: %w", err)

cmd/util/cmd/root.go

@@ -125,6 +125,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)
@@ -208,6 +209,8 @@ func (s *Suite) initEngineAndSyncer() (*Engine, *collections.Syncer, *collection
 		s.net,
 		s.proto.state,
 		s.me,
+		s.lockManager,
+		s.db,
 		s.blocks,
 		s.results,
 		s.receipts,
@@ -297,7 +300,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()
 	}
@@ -324,7 +327,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
@@ -384,7 +387,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) {
@@ -412,9 +415,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)
 }
 
 // TestExecutionReceiptsAreIndexed checks that execution receipts are properly indexed
@@ -517,9 +520,9 @@ func (s *Suite) TestCollectionSyncing() {
 
 	// setup finalized block indexer mocks
 	guaranteeIDs := []flow.Identifier(flow.GetIDs(block.Payload.Guarantees))
-	s.blocks.On("IndexBlockContainingCollectionGuarantees", block.ID(), guaranteeIDs).Return(nil).Once()
+	s.blocks.On("BatchIndexBlockContainingCollectionGuarantees", mock.Anything, mock.Anything, block.ID(), guaranteeIDs).Return(nil).Once()
 	for _, seal := range payload.Seals {
-		s.results.On("Index", seal.BlockID, seal.ResultID).Return(nil).Once()
+		s.results.On("BatchIndex", mock.Anything, mock.Anything, seal.BlockID, seal.ResultID).Return(nil).Once()
 	}
 
 	// initialize the engine using the initial finalized block