chainstore lifecycle: close via Close() instead of context.

EpiK-Protocol · Jan 7, 2021 · 912fd8a · 912fd8a
1 parent 87924ba
commit 912fd8a
Show file tree

Hide file tree

Showing 13 changed files with 59 additions and 25 deletions.
diff --git a/chain/gen/gen.go b/chain/gen/gen.go
@@ -248,7 +248,7 @@ func NewGeneratorWithSectors(numSectors int) (*ChainGen, error) {
 		return nil, xerrors.Errorf("make genesis block failed: %w", err)
 	}
 
-	cs := store.NewChainStore(context.Background(), bs, bs, ds, sys, j)
+	cs := store.NewChainStore(bs, bs, ds, sys, j)
 
 	genfb := &types.FullBlock{Header: genb.Genesis}
 	gents := store.NewFullTipSet([]*types.FullBlock{genfb})

diff --git a/chain/gen/genesis/genesis.go b/chain/gen/genesis/genesis.go
@@ -543,7 +543,7 @@ func MakeGenesisBlock(ctx context.Context, j journal.Journal, bs bstore.Blocksto
 	}
 
 	// temp chainstore
-	cs := store.NewChainStore(context.Background(), bs, bs, datastore.NewMapDatastore(), sys, j)
+	cs := store.NewChainStore(bs, bs, datastore.NewMapDatastore(), sys, j)
 
 	// Verify PreSealed Data
 	stateroot, err = VerifyPreSealedData(ctx, cs, stateroot, template, keyIDs)

diff --git a/chain/store/index_test.go b/chain/store/index_test.go
@@ -31,7 +31,8 @@ func TestIndexSeeks(t *testing.T) {
 	ctx := context.TODO()
 
 	nbs := blockstore.NewTemporarySync()
-	cs := store.NewChainStore(context.Background(), nbs, nbs, syncds.MutexWrap(datastore.NewMapDatastore()), nil, nil)
+	cs := store.NewChainStore(nbs, nbs, syncds.MutexWrap(datastore.NewMapDatastore()), nil, nil)
+	defer cs.Close() //nolint:errcheck
 
 	_, err = cs.Import(bytes.NewReader(gencar))
 	if err != nil {

diff --git a/chain/store/store.go b/chain/store/store.go
@@ -134,24 +134,30 @@ type ChainStore struct {
 
 	evtTypes [1]journal.EventType
 	journal  journal.Journal
+
+	cancelFn context.CancelFunc
+	wg       sync.WaitGroup
 }
 
 // localbs is guaranteed to fail Get* if requested block isn't stored locally
-func NewChainStore(ctx context.Context, bs bstore.Blockstore, localbs bstore.Blockstore, ds dstore.Batching, vmcalls vm.SyscallBuilder, j journal.Journal) *ChainStore {
-	c, _ := lru.NewARC(DefaultMsgMetaCacheSize)
-	tsc, _ := lru.NewARC(DefaultTipSetCacheSize)
+func NewChainStore(bs bstore.Blockstore, localbs bstore.Blockstore, ds dstore.Batching, vmcalls vm.SyscallBuilder, j journal.Journal) *ChainStore {
+	mmCache, _ := lru.NewARC(DefaultMsgMetaCacheSize)
+	tsCache, _ := lru.NewARC(DefaultTipSetCacheSize)
 	if j == nil {
 		j = journal.NilJournal()
 	}
+
+	ctx, cancel := context.WithCancel(context.Background())
 	cs := &ChainStore{
 		bs:       bs,
 		localbs:  localbs,
 		ds:       ds,
 		bestTips: pubsub.New(64),
 		tipsets:  make(map[abi.ChainEpoch][]cid.Cid),
-		mmCache:  c,
-		tsCache:  tsc,
+		mmCache:  mmCache,
+		tsCache:  tsCache,
 		vmcalls:  vmcalls,
+		cancelFn: cancel,
 		journal:  j,
 	}
 
@@ -192,7 +198,7 @@ func NewChainStore(ctx context.Context, bs bstore.Blockstore, localbs bstore.Blo
 
 	hcmetric := func(rev, app []*types.TipSet) error {
 		for _, r := range app {
-			stats.Record(ctx, metrics.ChainNodeHeight.M(int64(r.Height())))
+			stats.Record(context.Background(), metrics.ChainNodeHeight.M(int64(r.Height())))
 		}
 		return nil
 	}
@@ -203,6 +209,12 @@ func NewChainStore(ctx context.Context, bs bstore.Blockstore, localbs bstore.Blo
 	return cs
 }
 
+func (cs *ChainStore) Close() error {
+	cs.cancelFn()
+	cs.wg.Wait()
+	return nil
+}
+
 func (cs *ChainStore) Load() error {
 	head, err := cs.ds.Get(chainHeadKey)
 	if err == dstore.ErrNotFound {
@@ -405,7 +417,9 @@ func (cs *ChainStore) reorgWorker(ctx context.Context, initialNotifees []ReorgNo
 	notifees := make([]ReorgNotifee, len(initialNotifees))
 	copy(notifees, initialNotifees)
 
+	cs.wg.Add(1)
 	go func() {
+		defer cs.wg.Done()
 		defer log.Warn("reorgWorker quit")
 
 		for {

diff --git a/chain/store/store_test.go b/chain/store/store_test.go
@@ -70,7 +70,8 @@ func BenchmarkGetRandomness(b *testing.B) {
 		b.Fatal(err)
 	}
 
-	cs := store.NewChainStore(context.Background(), bs, bs, mds, nil, nil)
+	cs := store.NewChainStore(bs, bs, mds, nil, nil)
+	defer cs.Close() //nolint:errcheck
 
 	b.ResetTimer()
 
@@ -104,7 +105,8 @@ func TestChainExportImport(t *testing.T) {
 	}
 
 	nbs := blockstore.NewTemporary()
-	cs := store.NewChainStore(context.Background(), nbs, nbs, datastore.NewMapDatastore(), nil, nil)
+	cs := store.NewChainStore(nbs, nbs, datastore.NewMapDatastore(), nil, nil)
+	defer cs.Close() //nolint:errcheck
 
 	root, err := cs.Import(buf)
 	if err != nil {
@@ -138,7 +140,9 @@ func TestChainExportImportFull(t *testing.T) {
 	}
 
 	nbs := blockstore.NewTemporary()
-	cs := store.NewChainStore(context.Background(), nbs, nbs, datastore.NewMapDatastore(), nil, nil)
+	cs := store.NewChainStore(nbs, nbs, datastore.NewMapDatastore(), nil, nil)
+	defer cs.Close() //nolint:errcheck
+
 	root, err := cs.Import(buf)
 	if err != nil {
 		t.Fatal(err)

diff --git a/cmd/epik-bench/import.go b/cmd/epik-bench/import.go
@@ -262,7 +262,9 @@ var importBenchCmd = &cli.Command{
 		}
 
 		metadataDs := datastore.NewMapDatastore()
-		cs := store.NewChainStore(context.Background(), bs, bs, metadataDs, vm.Syscalls(verifier), nil)
+		cs := store.NewChainStore(bs, bs, metadataDs, vm.Syscalls(verifier), nil)
+		defer cs.Close() //nolint:errcheck
+
 		stm := stmgr.NewStateManager(cs)
 
 		startTime := time.Now()

diff --git a/cmd/epik-shed/balances.go b/cmd/epik-shed/balances.go
@@ -184,7 +184,8 @@ var chainBalanceStateCmd = &cli.Command{
 			return err
 		}
 
-		cs := store.NewChainStore(context.Background(), bs, bs, mds, vm.Syscalls(ffiwrapper.ProofVerifier), nil)
+		cs := store.NewChainStore(bs, bs, mds, vm.Syscalls(ffiwrapper.ProofVerifier), nil)
+		defer cs.Close() //nolint:errcheck
 
 		cst := cbor.NewCborStore(bs)
 		store := adt.WrapStore(ctx, cst)
@@ -406,6 +407,7 @@ func printAccountInfos(infos []accountInfo, minerInfo bool) {
 		}
 
 		cs := store.NewChainStore(bs, bs, mds, vm.Syscalls(ffiwrapper.ProofVerifier), nil)
+		defer cs.Close() //nolint:errcheck
 
 		cst := cbor.NewCborStore(bs)
 		store := adt.WrapStore(ctx, cst)

diff --git a/cmd/epik-shed/export.go b/cmd/epik-shed/export.go
@@ -90,7 +90,9 @@ var exportChainCmd = &cli.Command{
 			return err
 		}
 
-		cs := store.NewChainStore(context.Background(), bs, bs, mds, nil, nil)
+		cs := store.NewChainStore(bs, bs, mds, nil, nil)
+		defer cs.Close() //nolint:errcheck
+
 		if err := cs.Load(); err != nil {
 			return err
 		}

diff --git a/cmd/epik-shed/genesis-verify.go b/cmd/epik-shed/genesis-verify.go
@@ -52,7 +52,8 @@ var genesisVerifyCmd = &cli.Command{
 		}
 		bs := blockstore.NewBlockstore(datastore.NewMapDatastore())
 
-		cs := store.NewChainStore(context.Background(), bs, bs, datastore.NewMapDatastore(), nil, nil)
+		cs := store.NewChainStore(bs, bs, datastore.NewMapDatastore(), nil, nil)
+		defer cs.Close() //nolint:errcheck
 
 		cf := cctx.Args().Get(0)
 		f, err := os.Open(cf)

diff --git a/cmd/epik-shed/pruning.go b/cmd/epik-shed/pruning.go
@@ -169,7 +169,9 @@ var stateTreePruneCmd = &cli.Command{
 			return nil
 		}
 
-		cs := store.NewChainStore(context.Background(), bs, bs, mds, vm.Syscalls(ffiwrapper.ProofVerifier), nil)
+		cs := store.NewChainStore(bs, bs, mds, vm.Syscalls(ffiwrapper.ProofVerifier), nil)
+		defer cs.Close() //nolint:errcheck
+
 		if err := cs.Load(); err != nil {
 			return fmt.Errorf("loading chainstore: %w", err)
 		}

diff --git a/cmd/epik/daemon.go b/cmd/epik/daemon.go
@@ -422,9 +422,8 @@ func ImportChain(r repo.Repo, fname string, snapshot bool) (err error) {
 		return xerrors.Errorf("failed to open journal: %w", err)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	cst := store.NewChainStore(ctx, bs, bs, mds, vm.Syscalls(ffiwrapper.ProofVerifier), j)
+	cst := store.NewChainStore(bs, bs, mds, vm.Syscalls(ffiwrapper.ProofVerifier), j)
+	defer cst.Close()
 
 	log.Infof("importing chain from %s...", fname)
 
@@ -469,7 +468,7 @@ func ImportChain(r repo.Repo, fname string, snapshot bool) (err error) {
 	}
 
 	log.Infof("accepting %s as new head", ts.Cids())
-	if err := cst.ForceHeadSilent(ctx, ts); err != nil {
+	if err := cst.ForceHeadSilent(context.Background(), ts); err != nil {
 		return err
 	}
 

diff --git a/conformance/driver.go b/conformance/driver.go
@@ -87,10 +87,12 @@ func (d *Driver) ExecuteTipset(bs blockstore.Blockstore, ds ds.Batching, preroot
 		syscalls = vm.Syscalls(ffiwrapper.ProofVerifier)
 		vmRand   = NewFixedRand()
 
-		cs = store.NewChainStore(context.Background(), bs, bs, ds, syscalls, nil)
+		cs = store.NewChainStore(bs, bs, ds, syscalls, nil)
 		sm = stmgr.NewStateManager(cs)
 	)
 
+	defer cs.Close() //nolint:errcheck
+
 	blocks := make([]store.BlockMessages, 0, len(tipset.Blocks))
 	for _, b := range tipset.Blocks {
 		sb := store.BlockMessages{

diff --git a/node/modules/chain.go b/node/modules/chain.go
@@ -111,14 +111,19 @@ func SetupFallbackBlockstore(cbs dtypes.ChainBlockstore, rem dtypes.ChainBitswap
 	return nil
 }
 
-func ChainStore(lc fx.Lifecycle, mctx helpers.MetricsCtx, bs dtypes.ChainBlockstore, lbs dtypes.ChainRawBlockstore, ds dtypes.MetadataDS, syscalls vm.SyscallBuilder, j journal.Journal) *store.ChainStore {
-	ctx := helpers.LifecycleCtx(mctx, lc)
-	chain := store.NewChainStore(ctx, bs, lbs, ds, syscalls, j)
+func ChainStore(lc fx.Lifecycle, bs dtypes.ChainBlockstore, lbs dtypes.ChainRawBlockstore, ds dtypes.MetadataDS, syscalls vm.SyscallBuilder, j journal.Journal) *store.ChainStore {
+	chain := store.NewChainStore(bs, lbs, ds, syscalls, j)
 
 	if err := chain.Load(); err != nil {
 		log.Warnf("loading chain state from disk: %s", err)
 	}
 
+	lc.Append(fx.Hook{
+		OnStop: func(_ context.Context) error {
+			return chain.Close()
+		},
+	})
+
 	return chain
 }