From ff8c0af8c82c022b86fd7f1043460711c31a6c43 Mon Sep 17 00:00:00 2001
From: Jakub Sztandera <kubuxu@protocol.ai>
Date: Tue, 22 Sep 2020 00:54:11 +0200
Subject: [PATCH] Add only-import option

Signed-off-by: Jakub Sztandera <kubuxu@protocol.ai>
---
 cmd/lotus-bench/import.go | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/cmd/lotus-bench/import.go b/cmd/lotus-bench/import.go
index c3554f939..6af1bfc02 100644
--- a/cmd/lotus-bench/import.go
+++ b/cmd/lotus-bench/import.go
@@ -90,6 +90,9 @@ var importBenchCmd = &cli.Command{
 		&cli.Int64Flag{
 			Name: "start-at",
 		},
+		&cli.BoolFlag{
+			Name: "only-import",
+		},
 	},
 	Action: func(cctx *cli.Context) error {
 		vm.BatchSealVerifyParallelism = cctx.Int("batch-seal-verify-threads")
@@ -119,6 +122,13 @@ var importBenchCmd = &cli.Command{
 			tdir = tmp
 		}
 
+		bdgOpt := badger.DefaultOptions
+		bdgOpt.GcInterval = 0
+		bdgOpt.Options = bdg.DefaultOptions("")
+		bdgOpt.Options.SyncWrites = false
+		bdgOpt.Options.Truncate = true
+		bdgOpt.Options.DetectConflicts = false
+
 		var bds datastore.Batching
 		if false {
 			cache := 512
@@ -142,17 +152,12 @@ var importBenchCmd = &cli.Command{
 				Logger: log,
 			})
 		} else {
-			bdgOpt := badger.DefaultOptions
-			bdgOpt.GcInterval = 0
-			bdgOpt.Options = bdg.DefaultOptions("")
-			bdgOpt.Options.SyncWrites = false
-			bdgOpt.Options.Truncate = true
-			bdgOpt.Options.DetectConflicts = false
 			bds, err = badger.NewDatastore(tdir, &bdgOpt)
 		}
 		if err != nil {
 			return err
 		}
+		defer bds.Close()
 
 		if cctx.Bool("only-gc") {
 			log.Info("calling CollectGarbage on main ds")
@@ -176,6 +181,7 @@ var importBenchCmd = &cli.Command{
 			if err != nil {
 				return xerrors.Errorf("opening syscall-cache datastore: %w", err)
 			}
+			defer scds.Close()
 
 			if cctx.Bool("only-gc") {
 				log.Info("calling CollectGarbage on syscall ds")
@@ -223,6 +229,10 @@ var importBenchCmd = &cli.Command{
 			}
 		}
 
+		if cctx.Bool("only-import") {
+			return nil
+		}
+
 		gb, err := cs.GetTipsetByHeight(context.TODO(), 0, head, true)
 		if err != nil {
 			return err