From 0297be4b9a70f65d7645b60af55948a48d4dd888 Mon Sep 17 00:00:00 2001
From: Jakub Sztandera <kubuxu@protocol.ai>
Date: Wed, 28 Oct 2020 23:24:45 +0100
Subject: [PATCH] Add `lotus-shed datastore rewrite` command

Signed-off-by: Jakub Sztandera <kubuxu@protocol.ai>
---
 cmd/lotus-shed/datastore.go | 65 +++++++++++++++++++++++++++++++++++++
 node/repo/fsrepo_ds.go      | 23 +++++++------
 2 files changed, 78 insertions(+), 10 deletions(-)

diff --git a/cmd/lotus-shed/datastore.go b/cmd/lotus-shed/datastore.go
index c6bac6815..854845346 100644
--- a/cmd/lotus-shed/datastore.go
+++ b/cmd/lotus-shed/datastore.go
@@ -1,17 +1,22 @@
 package main
 
 import (
+	"bufio"
 	"encoding/json"
 	"fmt"
+	"io"
 	"os"
 	"strings"
 
 	"github.com/docker/go-units"
 	"github.com/ipfs/go-datastore"
 	dsq "github.com/ipfs/go-datastore/query"
+	badgerds "github.com/ipfs/go-ds-badger2"
 	logging "github.com/ipfs/go-log"
+	"github.com/mitchellh/go-homedir"
 	"github.com/polydawn/refmt/cbor"
 	"github.com/urfave/cli/v2"
+	"go.uber.org/multierr"
 	"golang.org/x/xerrors"
 
 	"github.com/filecoin-project/lotus/lib/backupds"
@@ -25,6 +30,7 @@ var datastoreCmd = &cli.Command{
 		datastoreBackupCmd,
 		datastoreListCmd,
 		datastoreGetCmd,
+		datastoreRewriteCmd,
 	},
 }
 
@@ -288,3 +294,62 @@ func printVal(enc string, val []byte) error {
 
 	return nil
 }
+
+var datastoreRewriteCmd = &cli.Command{
+	Name:        "rewrite",
+	Description: "rewrites badger datastore to compact it and possibly change params",
+	ArgsUsage:   "source destination",
+	Action: func(cctx *cli.Context) error {
+		if cctx.NArg() != 2 {
+			return xerrors.Errorf("expected 2 arguments, got %d", cctx.NArg())
+		}
+		fromPath, err := homedir.Expand(cctx.Args().Get(0))
+		if err != nil {
+			return xerrors.Errorf("cannot get fromPath: %w", err)
+		}
+		toPath, err := homedir.Expand(cctx.Args().Get(1))
+		if err != nil {
+			return xerrors.Errorf("cannot get toPath: %w", err)
+		}
+
+		opts := repo.ChainBadgerOptions()
+		opts.Options = opts.Options.WithSyncWrites(false)
+		to, err := badgerds.NewDatastore(toPath, &opts)
+		if err != nil {
+			return xerrors.Errorf("opennig 'to' datastore: %w", err)
+		}
+
+		opts.Options = opts.Options.WithReadOnly(false)
+		from, err := badgerds.NewDatastore(fromPath, &opts)
+		if err != nil {
+			return xerrors.Errorf("opennig 'from' datastore: %w", err)
+		}
+
+		pr, pw := io.Pipe()
+		errCh := make(chan error)
+		go func() {
+			bw := bufio.NewWriterSize(pw, 64<<20)
+			_, err := from.DB.Backup(bw, 0)
+			bw.Flush()
+			pw.CloseWithError(err)
+			errCh <- err
+		}()
+		go func() {
+			err := to.DB.Load(pr, 256)
+			errCh <- err
+		}()
+
+		err = <-errCh
+		if err != nil {
+			select {
+			case nerr := <-errCh:
+				err = multierr.Append(err, nerr)
+			default:
+			}
+			return err
+		}
+
+		err = <-errCh
+		return err
+	},
+}
diff --git a/node/repo/fsrepo_ds.go b/node/repo/fsrepo_ds.go
index 28d737210..e7746cb8e 100644
--- a/node/repo/fsrepo_ds.go
+++ b/node/repo/fsrepo_ds.go
@@ -4,18 +4,27 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/ipfs/go-datastore"
+	dgbadger "github.com/dgraph-io/badger/v2"
+	ldbopts "github.com/syndtr/goleveldb/leveldb/opt"
 	"golang.org/x/xerrors"
 
-	dgbadger "github.com/dgraph-io/badger/v2"
+	"github.com/ipfs/go-datastore"
 	badger "github.com/ipfs/go-ds-badger2"
 	levelds "github.com/ipfs/go-ds-leveldb"
 	measure "github.com/ipfs/go-ds-measure"
-	ldbopts "github.com/syndtr/goleveldb/leveldb/opt"
 )
 
 type dsCtor func(path string, readonly bool) (datastore.Batching, error)
 
+func ChainBadgerOptions() badger.Options {
+	opts := badger.DefaultOptions
+	opts.GcInterval = 0 // disable GC for chain datastore
+
+	opts.Options = dgbadger.DefaultOptions("").WithTruncate(true).
+		WithValueThreshold(128)
+	return opts
+}
+
 var fsDatastores = map[string]dsCtor{
 	"chain":    chainBadgerDs,
 	"metadata": levelDs,
@@ -27,13 +36,8 @@ var fsDatastores = map[string]dsCtor{
 }
 
 func chainBadgerDs(path string, readonly bool) (datastore.Batching, error) {
-	opts := badger.DefaultOptions
-	opts.GcInterval = 0 // disable GC for chain datastore
+	opts := ChainBadgerOptions()
 	opts.ReadOnly = readonly
-
-	opts.Options = dgbadger.DefaultOptions("").WithTruncate(true).
-		WithValueThreshold(128)
-
 	return badger.NewDatastore(path, &opts)
 }
 
@@ -43,7 +47,6 @@ func badgerDs(path string, readonly bool) (datastore.Batching, error) {
 
 	opts.Options = dgbadger.DefaultOptions("").WithTruncate(true).
 		WithValueThreshold(1 << 10)
-
 	return badger.NewDatastore(path, &opts)
 }