context for canceling

dymensionxyz · Sep 16, 2024 · 454ae23 · 454ae23
1 parent 7807dd5
commit 454ae23
Show file tree

Hide file tree

Showing 9 changed files with 19 additions and 16 deletions.
diff --git a/da/grpc/mockserv/cmd/main.go b/da/grpc/mockserv/cmd/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"log"
 	"net"
@@ -18,7 +19,7 @@ func main() {
 	flag.StringVar(&conf.Host, "host", "0.0.0.0", "listening address")
 	flag.Parse()
 
-	kv := store.NewDefaultKVStore(".", "db", "dymint")
+	kv := store.NewDefaultKVStore(context.TODO(), ".", "db", "dymint")
 	lis, err := net.Listen("tcp", conf.Host+":"+strconv.Itoa(conf.Port))
 	if err != nil {
 		log.Panic(err)

diff --git a/indexers/txindex/kv/kv_bench_test.go b/indexers/txindex/kv/kv_bench_test.go
@@ -20,7 +20,7 @@ func BenchmarkTxSearch(b *testing.B) {
 		b.Errorf("create temporary directory: %s", err)
 	}
 
-	db := store.NewDefaultKVStore(dbDir, "db", "benchmark_tx_search_test")
+	db := store.NewDefaultKVStore(context.TODO(), dbDir, "db", "benchmark_tx_search_test")
 	if err != nil {
 		b.Errorf("create database: %s", err)
 	}

diff --git a/indexers/txindex/kv/kv_test.go b/indexers/txindex/kv/kv_test.go
@@ -335,7 +335,7 @@ func benchmarkTxIndex(txsCount int64, b *testing.B) {
 		require.NoError(b, err)
 	}()
 
-	store := store.NewDefaultKVStore(dir, "db", "tx_index")
+	store := store.NewDefaultKVStore(context.TODO(), dir, "db", "tx_index")
 	require.NoError(b, err)
 	indexer := NewTxIndex(store)
 

diff --git a/node/node.go b/node/node.go
@@ -106,7 +106,7 @@ func NewNode(
 		dstore = datastore.NewMapDatastore()
 	} else {
 		// TODO(omritoptx): Move dymint to const
-		baseKV = store.NewKVStore(conf.RootDir, conf.DBPath, "dymint", conf.DBConfig.SyncWrites)
+		baseKV = store.NewKVStore(ctx, conf.RootDir, conf.DBPath, "dymint", conf.DBConfig.SyncWrites)
 		path := filepath.Join(store.Rootify(conf.RootDir, conf.DBPath), "blocksync")
 		var err error
 		dstore, err = leveldb.NewDatastore(path, &leveldb.Options{})

diff --git a/settlement/grpc/mockserv/mockserv.go b/settlement/grpc/mockserv/mockserv.go
@@ -67,7 +67,7 @@ func (s *server) SetBatch(ctx context.Context, in *slmock.SLSetBatchRequest) (*s
 func GetServer(conf settlement.GrpcConfig) *grpc.Server {
 	srv := grpc.NewServer()
 
-	slstore := store.NewDefaultKVStore(".", "db", "settlement")
+	slstore := store.NewDefaultKVStore(context.TODO(), ".", "db", "settlement")
 	kv := store.NewPrefixKV(slstore, settlementKVPrefix)
 
 	mockImpl := &server{kv: kv}

diff --git a/settlement/local/local.go b/settlement/local/local.go
@@ -108,7 +108,7 @@ func initConfig(conf settlement.Config) (slstore store.KV, proposer string, err
 			proposer = hex.EncodeToString(pubKeybytes)
 		}
 	} else {
-		slstore = store.NewDefaultKVStore(conf.KeyringHomeDir, "data", kvStoreDBName)
+		slstore = store.NewDefaultKVStore(context.TODO(), conf.KeyringHomeDir, "data", kvStoreDBName)
 		if conf.ProposerPubKey != "" {
 			proposer = conf.ProposerPubKey
 		} else {

diff --git a/store/badger.go b/store/badger.go
@@ -14,7 +14,7 @@ import (
 
 const (
 	gcTimeout    = 1 * time.Minute
-	discardRatio = 0.5
+	discardRatio = 0.5 // Recommended by badger. Indicates that a file will be rewritten if half the space can be discarded.
 )
 
 var (
@@ -38,7 +38,7 @@ func NewDefaultInMemoryKVStore() KV {
 	}
 }
 
-func NewKVStore(rootDir, dbPath, dbName string, syncWrites bool) KV {
+func NewKVStore(ctx context.Context, rootDir, dbPath, dbName string, syncWrites bool) KV {
 	path := filepath.Join(Rootify(rootDir, dbPath), dbName)
 	opts := memoryEfficientBadgerConfig(path, syncWrites)
 	db, err := badger.Open(*opts)
@@ -48,14 +48,15 @@ func NewKVStore(rootDir, dbPath, dbName string, syncWrites bool) KV {
 	b := &BadgerKV{
 		db: db,
 	}
-	go b.gc(gcTimeout, discardRatio)
-
+	if ctx != context.TODO() {
+		go b.gc(ctx, gcTimeout, discardRatio)
+	}
 	return b
 }
 
 // NewDefaultKVStore creates instance of default key-value store.
-func NewDefaultKVStore(rootDir, dbPath, dbName string) KV {
-	return NewKVStore(rootDir, dbPath, dbName, true)
+func NewDefaultKVStore(ctx context.Context, rootDir, dbPath, dbName string) KV {
+	return NewKVStore(ctx, rootDir, dbPath, dbName, true)
 }
 
 // Rootify is helper function to make config creation independent of root dir
@@ -71,11 +72,10 @@ func (b *BadgerKV) Close() error {
 	return b.db.Close()
 }
 
-func (b *BadgerKV) gc(period time.Duration, discardRatio float64) {
+func (b *BadgerKV) gc(ctx context.Context, period time.Duration, discardRatio float64) {
 	gcTimeout := time.NewTimer(period)
 	defer gcTimeout.Stop()
 
-	ctx := context.Background()
 	for {
 		select {
 		case <-ctx.Done():

diff --git a/store/store_test.go b/store/store_test.go
@@ -1,6 +1,7 @@
 package store_test
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -49,7 +50,7 @@ func TestStoreLoad(t *testing.T) {
 		}
 	}()
 
-	for _, kv := range []store.KV{store.NewDefaultInMemoryKVStore(), store.NewDefaultKVStore(tmpDir, "db", "test")} {
+	for _, kv := range []store.KV{store.NewDefaultInMemoryKVStore(), store.NewDefaultKVStore(context.TODO(), tmpDir, "db", "test")} {
 		for _, c := range cases {
 			t.Run(c.name, func(t *testing.T) {
 				assert := assert.New(t)

diff --git a/test/loadtime/cmd/report/main.go b/test/loadtime/cmd/report/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"encoding/csv"
 	"flag"
 	"fmt"
@@ -39,7 +40,7 @@ func (b *BlockStore) Base() uint64 {
 func getStore(directory string) *store.PrefixKV {
 	dataDirectory := directory[strings.LastIndex(directory, "/")+1:]
 	baseDirectory := directory[:len(directory)-len(dataDirectory)]
-	baseKV := store.NewDefaultKVStore(baseDirectory, dataDirectory, blockStoreDBName)
+	baseKV := store.NewDefaultKVStore(context.TODO(), baseDirectory, dataDirectory, blockStoreDBName)
 	mainKV := store.NewPrefixKV(baseKV, mainPrefix[:])
 	return mainKV
 }