[Access] Refactor storage collections for access node

cmd/access/node_builder/access_node_builder.go

@@ -11,8 +11,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/cockroachdb/pebble"
-	"github.com/dgraph-io/badger/v2"
 	"github.com/ipfs/boxo/bitswap"
 	"github.com/ipfs/go-cid"
 	"github.com/ipfs/go-datastore"
@@ -117,8 +115,6 @@ import (
 	"github.com/onflow/flow-go/state/protocol/blocktimer"
 	"github.com/onflow/flow-go/storage"
 	bstorage "github.com/onflow/flow-go/storage/badger"
-	"github.com/onflow/flow-go/storage/operation/badgerimpl"
-	"github.com/onflow/flow-go/storage/operation/pebbleimpl"
 	pstorage "github.com/onflow/flow-go/storage/pebble"
 	"github.com/onflow/flow-go/storage/store"
 	"github.com/onflow/flow-go/utils/grpcutils"
@@ -351,6 +347,8 @@ type FlowAccessNodeBuilder struct {
 	events                         storage.Events
 	lightTransactionResults        storage.LightTransactionResults
 	transactionResultErrorMessages storage.TransactionResultErrorMessages
+	transactions                   storage.Transactions
+	collections                    storage.Collections
 
 	// The sync engine participants provider is the libp2p peer store for the access node
 	// which is not available until after the network has started.
@@ -579,6 +577,12 @@ func (builder *FlowAccessNodeBuilder) BuildExecutionSyncComponents() *FlowAccess
 		AdminCommand("read-execution-data", func(config *cmd.NodeConfig) commands.AdminCommand {
 			return stateSyncCommands.NewReadExecutionDataCommand(builder.ExecutionDataStore)
 		}).
+		Module("transactions and collections storage", func(node *cmd.NodeConfig) error {
+			transactions := store.NewTransactions(node.Metrics.Cache, node.ProtocolDB)
+			builder.collections = store.NewCollections(node.ProtocolDB, transactions)
+			builder.transactions = transactions
+			return nil
+		}).
 		Module("execution data datastore and blobstore", func(node *cmd.NodeConfig) error {
 			datastoreDir := filepath.Join(builder.executionDataDir, "blobstore")
 			err := os.MkdirAll(datastoreDir, 0700)
@@ -618,29 +622,15 @@ func (builder *FlowAccessNodeBuilder) BuildExecutionSyncComponents() *FlowAccess
 		Module("processed block height consumer progress", func(node *cmd.NodeConfig) error {
 			// Note: progress is stored in the datastore's DB since that is where the jobqueue
 			// writes execution data to.
-			var db storage.DB
-			edmdb := builder.ExecutionDatastoreManager.DB()
-
-			if bdb, ok := edmdb.(*badger.DB); ok {
-				db = badgerimpl.ToDB(bdb)
-			} else if pdb, ok := edmdb.(*pebble.DB); ok {
-				db = pebbleimpl.ToDB(pdb)
-			} else {
-				return fmt.Errorf("unsupported execution data DB type: %T", edmdb)
-			}
+			db := builder.ExecutionDatastoreManager.DB()
 
 			processedBlockHeight = store.NewConsumerProgress(db, module.ConsumeProgressExecutionDataRequesterBlockHeight)
 			return nil
 		}).
 		Module("processed notifications consumer progress", func(node *cmd.NodeConfig) error {
 			// Note: progress is stored in the datastore's DB since that is where the jobqueue
 			// writes execution data to.
-			var db storage.DB
-			if executionDataDBMode == execution_data.ExecutionDataDBModeBadger {
-				db = badgerimpl.ToDB(builder.ExecutionDatastoreManager.DB().(*badger.DB))
-			} else {
-				db = pebbleimpl.ToDB(builder.ExecutionDatastoreManager.DB().(*pebble.DB))
-			}
+			db := builder.ExecutionDatastoreManager.DB()
 			processedNotifications = store.NewConsumerProgress(db, module.ConsumeProgressExecutionDataRequesterNotification)
 			return nil
 		}).
@@ -876,7 +866,7 @@ func (builder *FlowAccessNodeBuilder) BuildExecutionSyncComponents() *FlowAccess
 			}).
 			Module("indexed block height consumer progress", func(node *cmd.NodeConfig) error {
 				// Note: progress is stored in the MAIN db since that is where indexed execution data is stored.
-				indexedBlockHeight = store.NewConsumerProgress(badgerimpl.ToDB(builder.DB), module.ConsumeProgressExecutionDataIndexerBlockHeight)
+				indexedBlockHeight = store.NewConsumerProgress(builder.ProtocolDB, module.ConsumeProgressExecutionDataIndexerBlockHeight)
 				return nil
 			}).
 			Module("transaction results storage", func(node *cmd.NodeConfig) error {
@@ -974,10 +964,10 @@ func (builder *FlowAccessNodeBuilder) BuildExecutionSyncComponents() *FlowAccess
 					builder.ProtocolDB,
 					builder.Storage.RegisterIndex,
 					builder.Storage.Headers,
-					builder.events,
-					builder.Storage.Collections,
-					builder.Storage.Transactions,
-					builder.lightTransactionResults,
+					notNil(builder.events),
+					notNil(builder.collections),
+					notNil(builder.transactions),
+					notNil(builder.lightTransactionResults),
 					builder.RootChainID.Chain(),
 					indexerDerivedChainData,
 					builder.collectionExecutedMetric,
@@ -1622,7 +1612,7 @@ func (builder *FlowAccessNodeBuilder) Initialize() error {
 	builder.EnqueueNetworkInit()
 
 	builder.AdminCommand("get-transactions", func(conf *cmd.NodeConfig) commands.AdminCommand {
-		return storageCommands.NewGetTransactionsCommand(conf.State, conf.Storage.Payloads, conf.Storage.Collections)
+		return storageCommands.NewGetTransactionsCommand(conf.State, conf.Storage.Payloads, notNil(builder.collections))
 	})
 
 	// if this is an access node that supports public followers, enqueue the public network
@@ -1867,13 +1857,13 @@ func (builder *FlowAccessNodeBuilder) Build() (cmd.Node, error) {
 			return nil
 		}).
 		Module("processed finalized block height consumer progress", func(node *cmd.NodeConfig) error {
-			processedFinalizedBlockHeight = store.NewConsumerProgress(badgerimpl.ToDB(builder.DB), module.ConsumeProgressIngestionEngineBlockHeight)
+			processedFinalizedBlockHeight = store.NewConsumerProgress(builder.ProtocolDB, module.ConsumeProgressIngestionEngineBlockHeight)
 			return nil
 		}).
 		Module("processed last full block height monotonic consumer progress", func(node *cmd.NodeConfig) error {
 			rootBlockHeight := node.State.Params().FinalizedRoot().Height
 
-			progress, err := store.NewConsumerProgress(badgerimpl.ToDB(builder.DB), module.ConsumeProgressLastFullBlockHeight).Initialize(rootBlockHeight)
+			progress, err := store.NewConsumerProgress(builder.ProtocolDB, module.ConsumeProgressLastFullBlockHeight).Initialize(rootBlockHeight)
 			if err != nil {
 				return err
 			}
@@ -2043,8 +2033,8 @@ func (builder *FlowAccessNodeBuilder) Build() (cmd.Node, error) {
 				HistoricalAccessNodes: builder.HistoricalAccessRPCs,
 				Blocks:                node.Storage.Blocks,
 				Headers:               node.Storage.Headers,
-				Collections:           node.Storage.Collections,
-				Transactions:          node.Storage.Transactions,
+				Collections:           notNil(builder.collections),
+				Transactions:          notNil(builder.transactions),
 				ExecutionReceipts:     node.Storage.Receipts,
 				ExecutionResults:      node.Storage.Results,
 				TxResultErrorMessages: builder.transactionResultErrorMessages,
@@ -2146,8 +2136,8 @@ func (builder *FlowAccessNodeBuilder) Build() (cmd.Node, error) {
 				builder.RequestEng,
 				node.Storage.Blocks,
 				node.Storage.Headers,
-				node.Storage.Collections,
-				node.Storage.Transactions,
+				notNil(builder.collections),
+				notNil(builder.transactions),
 				node.Storage.Results,
 				node.Storage.Receipts,
 				builder.collectionExecutedMetric,
@@ -2180,7 +2170,7 @@ func (builder *FlowAccessNodeBuilder) Build() (cmd.Node, error) {
 	if builder.storeTxResultErrorMessages {
 		builder.Module("processed error messages block height consumer progress", func(node *cmd.NodeConfig) error {
 			processedTxErrorMessagesBlockHeight = store.NewConsumerProgress(
-				badgerimpl.ToDB(builder.DB),
+				builder.ProtocolDB,
 				module.ConsumeProgressEngineTxErrorMessagesBlockHeight,
 			)
 			return nil
@@ -2393,3 +2383,15 @@ func (builder *FlowAccessNodeBuilder) initPublicLibp2pNode(networkKey crypto.Pri
 
 	return libp2pNode, nil
 }
+
+// notNil ensures that the input is not nil and returns it
+// the usage is to ensure the dependencies are initialized before initializing a module.
+// for instance, the IngestionEngine depends on storage.Collections, which is initialized in a
+// different function, so we need to ensure that the storage.Collections is initialized before
+// creating the IngestionEngine.
+func notNil[T any](dep T) T {
+	if any(dep) == nil {
+		panic("dependency is nil")
+	}
+	return dep
+}

cmd/execution_builder.go

@@ -144,6 +144,7 @@ type ExecutionNode struct {
 	receipts      storageerr.ExecutionReceipts
 	myReceipts    storageerr.MyExecutionReceipts
 	commits       storageerr.Commits
+	collections   storageerr.Collections
 
 	chunkDataPackDB        *pebble.DB
 	chunkDataPacks         storageerr.ChunkDataPacks
@@ -194,9 +195,6 @@ func (builder *ExecutionNodeBuilder) LoadComponentsAndModules() {
 		AdminCommand("set-uploader-enabled", func(config *NodeConfig) commands.AdminCommand {
 			return uploaderCommands.NewToggleUploaderCommand(exeNode.blockDataUploader)
 		}).
-		AdminCommand("get-transactions", func(conf *NodeConfig) commands.AdminCommand {
-			return storageCommands.NewGetTransactionsCommand(conf.State, conf.Storage.Payloads, conf.Storage.Collections)
-		}).
 		AdminCommand("protocol-snapshot", func(conf *NodeConfig) commands.AdminCommand {
 			return storageCommands.NewProtocolSnapshotCommand(
 				conf.Logger,
@@ -206,6 +204,7 @@ func (builder *ExecutionNodeBuilder) LoadComponentsAndModules() {
 				exeNode.exeConf.triedir,
 			)
 		}).
+		Module("load collections", exeNode.LoadCollections).
 		Module("mutable follower state", exeNode.LoadMutableFollowerState).
 		Module("system specs", exeNode.LoadSystemSpecs).
 		Module("execution metrics", exeNode.LoadExecutionMetrics).
@@ -218,6 +217,9 @@ func (builder *ExecutionNodeBuilder) LoadComponentsAndModules() {
 		Module("blobservice peer manager dependencies", exeNode.LoadBlobservicePeerManagerDependencies).
 		Module("bootstrap", exeNode.LoadBootstrapper).
 		Module("register store", exeNode.LoadRegisterStore).
+		AdminCommand("get-transactions", func(conf *NodeConfig) commands.AdminCommand {
+			return storageCommands.NewGetTransactionsCommand(conf.State, conf.Storage.Payloads, exeNode.collections)
+		}).
 		Component("execution state ledger", exeNode.LoadExecutionStateLedger).
 
 		// TODO: Modules should be able to depends on components
@@ -259,6 +261,13 @@ func (builder *ExecutionNodeBuilder) LoadComponentsAndModules() {
 		Component("grpc server", exeNode.LoadGrpcServer)
 }
 
+func (exeNode *ExecutionNode) LoadCollections(node *NodeConfig) error {
+	db := badgerimpl.ToDB(node.DB)
+	transactions := store.NewTransactions(node.Metrics.Cache, db)
+	exeNode.collections = store.NewCollections(db, transactions)
+	return nil
+}
+
 func (exeNode *ExecutionNode) LoadMutableFollowerState(node *NodeConfig) error {
 	// For now, we only support state implementations from package badger.
 	// If we ever support different implementations, the following can be replaced by a type-aware factory
@@ -467,7 +476,7 @@ func (exeNode *ExecutionNode) LoadGCPBlockDataUploader(
 		asyncUploader,
 		node.Storage.Blocks,
 		exeNode.commits,
-		node.Storage.Collections,
+		exeNode.collections,
 		exeNode.events,
 		exeNode.results,
 		exeNode.txResults,
@@ -745,7 +754,7 @@ func (exeNode *ExecutionNode) LoadExecutionState(
 		return nil
 	})
 	chunkDataPacks := store.NewChunkDataPacks(node.Metrics.Cache,
-		pebbleimpl.ToDB(chunkDataPackDB), node.Storage.Collections, exeNode.exeConf.chunkDataPackCacheSize)
+		pebbleimpl.ToDB(chunkDataPackDB), exeNode.collections, exeNode.exeConf.chunkDataPackCacheSize)
 
 	getLatestFinalized := func() (uint64, error) {
 		final, err := node.State.Final().Head()
@@ -763,7 +772,6 @@ func (exeNode *ExecutionNode) LoadExecutionState(
 		exeNode.commits,
 		node.Storage.Blocks,
 		node.Storage.Headers,
-		node.Storage.Collections,
 		chunkDataPacks,
 		exeNode.results,
 		exeNode.myReceipts,
@@ -1107,7 +1115,7 @@ func (exeNode *ExecutionNode) LoadIngestionEngine(
 		colFetcher,
 		node.Storage.Headers,
 		node.Storage.Blocks,
-		node.Storage.Collections,
+		exeNode.collections,
 		exeNode.executionState,
 		node.State,
 		exeNode.collector,

cmd/observer/node_builder/observer_builder.go

@@ -13,8 +13,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/cockroachdb/pebble"
-	"github.com/dgraph-io/badger/v2"
 	"github.com/ipfs/boxo/bitswap"
 	"github.com/ipfs/go-cid"
 	"github.com/ipfs/go-datastore"
@@ -109,7 +107,6 @@ import (
 	"github.com/onflow/flow-go/storage"
 	bstorage "github.com/onflow/flow-go/storage/badger"
 	"github.com/onflow/flow-go/storage/operation/badgerimpl"
-	"github.com/onflow/flow-go/storage/operation/pebbleimpl"
 	pstorage "github.com/onflow/flow-go/storage/pebble"
 	"github.com/onflow/flow-go/storage/store"
 	"github.com/onflow/flow-go/utils/grpcutils"
@@ -1129,25 +1126,16 @@ func (builder *ObserverServiceBuilder) BuildExecutionSyncComponents() *ObserverS
 		Module("processed block height consumer progress", func(node *cmd.NodeConfig) error {
 			// Note: progress is stored in the datastore's DB since that is where the jobqueue
 			// writes execution data to.
-			var db storage.DB
-			if executionDataDBMode == execution_data.ExecutionDataDBModeBadger {
-				db = badgerimpl.ToDB(builder.ExecutionDatastoreManager.DB().(*badger.DB))
-			} else {
-				db = pebbleimpl.ToDB(builder.ExecutionDatastoreManager.DB().(*pebble.DB))
-			}
+			db := builder.ExecutionDatastoreManager.DB()
 
 			processedBlockHeight = store.NewConsumerProgress(db, module.ConsumeProgressExecutionDataRequesterBlockHeight)
 			return nil
 		}).
 		Module("processed notifications consumer progress", func(node *cmd.NodeConfig) error {
 			// Note: progress is stored in the datastore's DB since that is where the jobqueue
 			// writes execution data to.
-			var db storage.DB
-			if executionDataDBMode == execution_data.ExecutionDataDBModeBadger {
-				db = badgerimpl.ToDB(builder.ExecutionDatastoreManager.DB().(*badger.DB))
-			} else {
-				db = pebbleimpl.ToDB(builder.ExecutionDatastoreManager.DB().(*pebble.DB))
-			}
+			db := builder.ExecutionDatastoreManager.DB()
+
 			processedNotifications = store.NewConsumerProgress(db, module.ConsumeProgressExecutionDataRequesterNotification)
 			return nil
 		}).

cmd/util/cmd/rollback-executed-height/cmd/rollback_executed_height_test.go

@@ -46,8 +46,6 @@ func TestReExecuteBlock(t *testing.T) {
 			myReceipts := store.NewMyExecutionReceipts(metrics, db, receipts)
 			events := store.NewEvents(metrics, db)
 			serviceEvents := store.NewServiceEvents(metrics, db)
-			transactions := bstorage.NewTransactions(metrics, bdb)
-			collections := bstorage.NewCollections(bdb, transactions)
 
 			err = headers.Store(genesis)
 			require.NoError(t, err)
@@ -62,7 +60,6 @@ func TestReExecuteBlock(t *testing.T) {
 				commits,
 				nil,
 				headers,
-				collections,
 				chunkDataPacks,
 				results,
 				myReceipts,
@@ -200,7 +197,6 @@ func TestReExecuteBlockWithDifferentResult(t *testing.T) {
 				commits,
 				nil,
 				headers,
-				collections,
 				chunkDataPacks,
 				results,
 				myReceipts,

engine/execution/state/state.go

@@ -93,7 +93,6 @@ type state struct {
 	commits            storage.Commits
 	blocks             storage.Blocks
 	headers            storage.Headers
-	collections        storage.Collections
 	chunkDataPacks     storage.ChunkDataPacks
 	results            storage.ExecutionResults
 	myReceipts         storage.MyExecutionReceipts
@@ -115,7 +114,6 @@ func NewExecutionState(
 	commits storage.Commits,
 	blocks storage.Blocks,
 	headers storage.Headers,
-	collections storage.Collections,
 	chunkDataPacks storage.ChunkDataPacks,
 	results storage.ExecutionResults,
 	myReceipts storage.MyExecutionReceipts,
@@ -134,7 +132,6 @@ func NewExecutionState(
 		commits:             commits,
 		blocks:              blocks,
 		headers:             headers,
-		collections:         collections,
 		chunkDataPacks:      chunkDataPacks,
 		results:             results,
 		myReceipts:          myReceipts,

engine/execution/state/state_storehouse_test.go

@@ -54,7 +54,6 @@ func prepareStorehouseTest(f func(t *testing.T, es state.ExecutionState, l *ledg
 			stateCommitments.On("BatchStore", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 			headers := storage.NewHeaders(t)
 			blocks := storage.NewBlocks(t)
-			collections := storage.NewCollections(t)
 			events := storage.NewEvents(t)
 			events.On("BatchStore", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 			serviceEvents := storage.NewServiceEvents(t)
@@ -93,7 +92,7 @@ func prepareStorehouseTest(f func(t *testing.T, es state.ExecutionState, l *ledg
 				}
 
 				es := state.NewExecutionState(
-					ls, stateCommitments, blocks, headers, collections, chunkDataPacks, results, myReceipts, events, serviceEvents, txResults, badgerimpl.ToDB(badgerDB),
+					ls, stateCommitments, blocks, headers, chunkDataPacks, results, myReceipts, events, serviceEvents, txResults, badgerimpl.ToDB(badgerDB),
 					getLatestFinalized,
 					trace.NewNoopTracer(),
 					rs,

engine/execution/state/state_test.go

@@ -40,7 +40,6 @@ func prepareTest(f func(t *testing.T, es state.ExecutionState, l *ledger.Ledger,
 			stateCommitments := storage.NewCommits(t)
 			headers := storage.NewHeaders(t)
 			blocks := storage.NewBlocks(t)
-			collections := storage.NewCollections(t)
 			events := storage.NewEvents(t)
 			serviceEvents := storage.NewServiceEvents(t)
 			txResults := storage.NewTransactionResults(t)
@@ -54,7 +53,7 @@ func prepareTest(f func(t *testing.T, es state.ExecutionState, l *ledger.Ledger,
 
 			db := badgerimpl.ToDB(badgerDB)
 			es := state.NewExecutionState(
-				ls, stateCommitments, blocks, headers, collections, chunkDataPacks, results, myReceipts, events, serviceEvents, txResults, db, getLatestFinalized, trace.NewNoopTracer(),
+				ls, stateCommitments, blocks, headers, chunkDataPacks, results, myReceipts, events, serviceEvents, txResults, db, getLatestFinalized, trace.NewNoopTracer(),
 				nil,
 				false,
 			)

engine/testutil/nodes.go

@@ -640,7 +640,7 @@ func ExecutionNode(t *testing.T, hub *stub.Hub, identity bootstrap.NodeInfo, ide
 		return final.Height, nil
 	}
 	execState := executionState.NewExecutionState(
-		ls, commitsStorage, node.Blocks, node.Headers, collectionsStorage, chunkDataPackStorage, results, myReceipts, eventsStorage, serviceEventsStorage, txResultStorage, db, getLatestFinalized, node.Tracer,
+		ls, commitsStorage, node.Blocks, node.Headers, chunkDataPackStorage, results, myReceipts, eventsStorage, serviceEventsStorage, txResultStorage, db, getLatestFinalized, node.Tracer,
 		// TODO: test with register store
 		registerStore,
 		storehouseEnabled,