Pass context to NewKVStore (#8125)

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/db/db.go

@@ -3,11 +3,13 @@
 package db
 
 import (
+	"context"
+
 	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/beacon-chain/db/kv"
 )
 
 // NewDB initializes a new DB.
-func NewDB(dirPath string, stateSummaryCache *cache.StateSummaryCache) (Database, error) {
-	return kv.NewKVStore(dirPath, stateSummaryCache)
+func NewDB(ctx context.Context, dirPath string, stateSummaryCache *cache.StateSummaryCache) (Database, error) {
+	return kv.NewKVStore(ctx, dirPath, stateSummaryCache)
 }

beacon-chain/db/kv/backup_test.go

@@ -39,7 +39,7 @@ func TestStore_Backup(t *testing.T) {
 	// our NewKVStore function expects when opening a database.
 	require.NoError(t, os.Rename(oldFilePath, newFilePath))
 
-	backedDB, err := NewKVStore(backupsPath, nil)
+	backedDB, err := NewKVStore(ctx, backupsPath, nil)
 	require.NoError(t, err, "Failed to instantiate DB")
 	t.Cleanup(func() {
 		require.NoError(t, backedDB.Close(), "Failed to close database")

beacon-chain/db/kv/kv.go

@@ -3,6 +3,7 @@
 package kv
 
 import (
+	"context"
 	"os"
 	"path"
 	"time"
@@ -45,12 +46,13 @@ type Store struct {
 	blockCache          *ristretto.Cache
 	validatorIndexCache *ristretto.Cache
 	stateSummaryCache   *cache.StateSummaryCache
+	ctx                 context.Context
 }
 
 // NewKVStore initializes a new boltDB key-value store at the directory
 // path specified, creates the kv-buckets based on the schema, and stores
 // an open connection db object as a property of the Store struct.
-func NewKVStore(dirPath string, stateSummaryCache *cache.StateSummaryCache) (*Store, error) {
+func NewKVStore(ctx context.Context, dirPath string, stateSummaryCache *cache.StateSummaryCache) (*Store, error) {
 	hasDir, err := fileutil.HasDir(dirPath)
 	if err != nil {
 		return nil, err
@@ -93,6 +95,7 @@ func NewKVStore(dirPath string, stateSummaryCache *cache.StateSummaryCache) (*St
 		blockCache:          blockCache,
 		validatorIndexCache: validatorCache,
 		stateSummaryCache:   stateSummaryCache,
+		ctx:                 ctx,
 	}
 
 	if err := kv.db.Update(func(tx *bolt.Tx) error {

beacon-chain/db/kv/kv_test.go

@@ -1,6 +1,7 @@
 package kv
 
 import (
+	"context"
 	"testing"
 
 	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
@@ -9,7 +10,7 @@ import (
 
 // setupDB instantiates and returns a Store instance.
 func setupDB(t testing.TB) *Store {
-	db, err := NewKVStore(t.TempDir(), cache.NewStateSummaryCache())
+	db, err := NewKVStore(context.Background(), t.TempDir(), cache.NewStateSummaryCache())
 	require.NoError(t, err, "Failed to instantiate DB")
 	t.Cleanup(func() {
 		require.NoError(t, db.Close(), "Failed to close database")

beacon-chain/db/restore_test.go

@@ -22,7 +22,7 @@ func TestRestore(t *testing.T) {
 	logHook := logTest.NewGlobal()
 	ctx := context.Background()
 
-	backupDb, err := kv.NewKVStore(t.TempDir(), cache.NewStateSummaryCache())
+	backupDb, err := kv.NewKVStore(context.Background(), t.TempDir(), cache.NewStateSummaryCache())
 	defer func() {
 		require.NoError(t, backupDb.Close())
 	}()
@@ -58,7 +58,7 @@ func TestRestore(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, 1, len(files))
 	assert.Equal(t, kv.DatabaseFileName, files[0].Name())
-	restoredDb, err := kv.NewKVStore(path.Join(restoreDir, kv.BeaconNodeDbDirName), nil)
+	restoredDb, err := kv.NewKVStore(context.Background(), path.Join(restoreDir, kv.BeaconNodeDbDirName), nil)
 	defer func() {
 		require.NoError(t, restoredDb.Close())
 	}()

beacon-chain/db/testing/setup_db.go

@@ -3,6 +3,7 @@
 package testing
 
 import (
+	"context"
 	"testing"
 
 	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
@@ -13,7 +14,7 @@ import (
 // SetupDB instantiates and returns database backed by key value store.
 func SetupDB(t testing.TB) (db.Database, *cache.StateSummaryCache) {
 	sc := cache.NewStateSummaryCache()
-	s, err := kv.NewKVStore(t.TempDir(), sc)
+	s, err := kv.NewKVStore(context.Background(), t.TempDir(), sc)
 	if err != nil {
 		t.Fatal(err)
 	}

beacon-chain/node/node.go

@@ -297,7 +297,7 @@ func (b *BeaconNode) startDB(cliCtx *cli.Context) error {
 
 	log.WithField("database-path", dbPath).Info("Checking DB")
 
-	d, err := db.NewDB(dbPath, b.stateSummaryCache)
+	d, err := db.NewDB(b.ctx, dbPath, b.stateSummaryCache)
 	if err != nil {
 		return err
 	}
@@ -319,7 +319,7 @@ func (b *BeaconNode) startDB(cliCtx *cli.Context) error {
 		if err := d.ClearDB(); err != nil {
 			return errors.Wrap(err, "could not clear database")
 		}
-		d, err = db.NewDB(dbPath, b.stateSummaryCache)
+		d, err = db.NewDB(b.ctx, dbPath, b.stateSummaryCache)
 		if err != nil {
 			return errors.Wrap(err, "could not create new database")
 		}

fuzz/block_fuzz.go

@@ -60,7 +60,7 @@ func init() {
 
 	var err error
 
-	db1, err = db.NewDB(dbPath, ssc)
+	db1, err = db.NewDB(context.Background(), dbPath, ssc)
 	if err != nil {
 		panic(err)
 	}

tools/blocktree/main.go

@@ -38,7 +38,7 @@ type node struct {
 
 func main() {
 	flag.Parse()
-	db, err := db.NewDB(*datadir, cache.NewStateSummaryCache())
+	db, err := db.NewDB(context.Background(), *datadir, cache.NewStateSummaryCache())
 	if err != nil {
 		panic(err)
 	}

tools/extractor/main.go

@@ -24,7 +24,7 @@ func main() {
 	defer resetCfg()
 	flag.Parse()
 	fmt.Println("Starting process...")
-	d, err := db.NewDB(*datadir, cache.NewStateSummaryCache())
+	d, err := db.NewDB(context.Background(), *datadir, cache.NewStateSummaryCache())
 	if err != nil {
 		panic(err)
 	}

tools/interop/export-genesis/main.go

@@ -21,7 +21,7 @@ func main() {
 
 	fmt.Printf("Reading db at %s and writing ssz output to %s.\n", os.Args[1], os.Args[2])
 
-	d, err := db.NewDB(os.Args[1], cache.NewStateSummaryCache())
+	d, err := db.NewDB(context.Background(), os.Args[1], cache.NewStateSummaryCache())
 	if err != nil {
 		panic(err)
 	}