Repository: dgraph-io/badger
Pull request #1465 (2020-08-18 11:42:19)

Fix(cleanup): Avoid truncating in value.Open on error
The vlog.Open() function can return an error denoting that the open was
unsuccessful but we have `db.cleanup` which will be called to stop all
the running goroutines in case badger couldn't start. The db.cleanup
function calls vlog.Close() which will truncate the maxFid vlog file
based on the vlog.writableLogOffset. The vlog.writableLogOffset was not
being updated in case open failed. As a result of this, we will end up
truncating the vlog file to length 0 if open fails.

This PR fixes this by using `vlog.stopDiscardStatFlush()` instead of `vlog.close`.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/dgraph-io/badger/1465)
<!-- Reviewable:end -->



Patch:
diff --git a/badger/cmd/stream.go b/badger/cmd/stream.go
new file mode 100644
index 0000000..1683cc5
--- /dev/null
+++ b/badger/cmd/stream.go
@@ -0,0 +1,97 @@
+/*
+ * Copyright 2020 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package cmd
+
+import (
+	"io"
+	"math"
+	"os"
+
+	"github.com/dgraph-io/badger/v2"
+	"github.com/dgraph-io/badger/v2/options"
+	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
+)
+
+var streamCmd = &cobra.Command{
+	Use:   "stream",
+	Short: "Stream DB into another DB with different options",
+	Long: `
+This command streams the contents of this DB into another DB with the given options.
+`,
+	RunE: stream,
+}
+
+var outDir string
+var compressionType uint32
+
+func init() {
+	// TODO: Add more options.
+	RootCmd.AddCommand(streamCmd)
+	streamCmd.Flags().StringVarP(&outDir, "out", "o", "",
+		"Path to output DB. The directory should be empty.")
+	streamCmd.Flags().BoolVarP(&truncate, "truncate", "", false, "Option to truncate the DBs")
+	streamCmd.Flags().BoolVarP(&readOnly, "read_only", "", true,
+		"Option to open input DB in read-only mode")
+	streamCmd.Flags().IntVarP(&numVersions, "num_versions", "", 0,
+		"Option to configure the maximum number of versions per key. "+
+			"Values <= 0 will be considered to have the max number of versions.")
+	streamCmd.Flags().Uint32VarP(&compressionType, "compression", "", 0,
+		"Option to configure the compression type in output DB. "+
+			"0 to disable, 1 for Snappy, and 2 for ZSTD.")
+}
+
+func stream(cmd *cobra.Command, args []string) error {
+	// Check that outDir doesn't exist or is empty.
+	if _, err := os.Stat(outDir); err == nil {
+		f, err := os.Open(outDir)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+
+		_, err = f.Readdirnames(1)
+		if err != io.EOF {
+			return errors.Errorf("cannot run stream tool on non-empty output directory %s", outDir)
+		}
+	}
+
+	// Options for input DB.
+	if numVersions <= 0 {
+		numVersions = math.MaxInt32
+	}
+	inOpt := badger.DefaultOptions(sstDir).
+		WithReadOnly(readOnly).
+		WithTruncate(truncate).
+		WithValueThreshold(1 << 10 /* 1KB */).
+		WithNumVersionsToKeep(numVersions)
+
+	// Options for output DB.
+	if compressionType < 0 || compressionType > 2 {
+		return errors.Errorf(
+			"compression value must be one of 0 (disabled), 1 (Snappy), or 2 (ZSTD)")
+	}
+	outOpt := inOpt.WithDir(outDir).WithValueDir(outDir).
+		WithCompression(options.CompressionType(compressionType)).WithReadOnly(false)
+
+	inDB, err := badger.OpenManaged(inOpt)
+	if err != nil {
+		return errors.Wrapf(err, "cannot open DB at %s", sstDir)
+	}
+	defer inDB.Close()
+	return inDB.StreamDB(outOpt)
+}
diff --git a/db.go b/db.go
index 586e352..d434994 100644
--- a/db.go
+++ b/db.go
@@ -21,6 +21,7 @@ import (
 	"context"
 	"encoding/binary"
 	"expvar"
+	"fmt"
 	"math"
 	"os"
 	"path/filepath"
@@ -447,7 +448,10 @@ func (db *DB) cleanup() {
 	}
 
 	db.orc.Stop()
-	db.vlog.Close()
+
+	// Do not use vlog.Close() here. vlog.Close truncates the files. We don't
+	// want to truncate files unless the user has specified the truncate flag.
+	db.vlog.stopFlushDiscardStats()
 }
 
 // BlockCacheMetrics returns the metrics for the underlying block cache.
@@ -1785,3 +1789,34 @@ func createDirs(opt Options) error {
 	}
 	return nil
 }
+
+// Stream the contents of this DB to a new DB with options outOptions that will be
+// created in outDir.
+func (db *DB) StreamDB(outOptions Options) error {
+	outDir := outOptions.Dir
+
+	// Open output DB.
+	outDB, err := OpenManaged(outOptions)
+	if err != nil {
+		return errors.Wrapf(err, "cannot open out DB at %s", outDir)
+	}
+	defer outDB.Close()
+	writer := outDB.NewStreamWriter()
+	if err := writer.Prepare(); err != nil {
+		errors.Wrapf(err, "cannot create stream writer in out DB at %s", outDir)
+	}
+
+	// Stream contents of DB to the output DB.
+	stream := db.NewStreamAt(math.MaxUint64)
+	stream.LogPrefix = fmt.Sprintf("Streaming DB to new DB at %s", outDir)
+	stream.Send = func(kvs *pb.KVList) error {
+		return writer.Write(kvs)
+	}
+	if err := stream.Orchestrate(context.Background()); err != nil {
+		return errors.Wrapf(err, "cannot stream DB to out DB at %s", outDir)
+	}
+	if err := writer.Flush(); err != nil {
+		return errors.Wrapf(err, "cannot flush writer")
+	}
+	return nil
+}
diff --git a/value.go b/value.go
index fd29f5a..695240c 100644
--- a/value.go
+++ b/value.go
@@ -1047,6 +1047,8 @@ func (vlog *valueLog) replayLog(lf *logFile, offset uint32, replayFn logEntry) e
 	// End offset is different from file size. So, we should truncate the file
 	// to that size.
 	if !vlog.opt.Truncate {
+		vlog.db.opt.Warningf("Truncate Needed. File %s size: %d Endoffset: %d",
+			lf.fd.Name(), fi.Size(), endOffset)
 		return ErrTruncateNeeded
 	}
 
@@ -1063,6 +1065,7 @@ func (vlog *valueLog) replayLog(lf *logFile, offset uint32, replayFn logEntry) e
 		return lf.bootstrap()
 	}
 
+	vlog.db.opt.Infof("Truncating vlog file %s to offset: %d", lf.fd.Name(), endOffset)
 	if err := lf.fd.Truncate(int64(endOffset)); err != nil {
 		return errFile(err, lf.path, fmt.Sprintf(
 			"Truncation needed at offset %d. Can be done manually as well.", endOffset))
@@ -1228,6 +1231,12 @@ func (lf *logFile) init() error {
 	return nil
 }
 
+func (vlog *valueLog) stopFlushDiscardStats() {
+	if vlog.lfDiscardStats != nil {
+		vlog.lfDiscardStats.closer.Signal()
+	}
+}
+
 func (vlog *valueLog) Close() error {
 	if vlog == nil || vlog.db == nil || vlog.db.opt.InMemory {
 		return nil
@@ -1245,6 +1254,9 @@ func (vlog *valueLog) Close() error {
 		}
 
 		maxFid := vlog.maxFid
+		// TODO(ibrahim) - Do we need the following truncations on non-windows
+		// platforms? We expand the file only on windows and the vlog.woffset()
+		// should point to end of file on all other platforms.
 		if !vlog.opt.ReadOnly && id == maxFid {
 			// truncate writable log file to correct offset.
 			if truncErr := f.fd.Truncate(


Test patch:
diff --git a/db_test.go b/db_test.go
index febf6d1..c3aa8dc 100644
--- a/db_test.go
+++ b/db_test.go
@@ -376,6 +376,49 @@ func TestForceCompactL0(t *testing.T) {
 	require.NoError(t, db.Close())
 }
 
+func TestStreamDB(t *testing.T) {
+	dir, err := ioutil.TempDir("", "badger-test")
+	require.NoError(t, err)
+	defer removeDir(dir)
+	opts := getTestOptions(dir).WithCompression(options.ZSTD)
+
+	db, err := OpenManaged(opts)
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, db.Close())
+	}()
+
+	writer := db.NewManagedWriteBatch()
+	for i := 0; i < 100; i++ {
+		key := []byte(fmt.Sprintf("key%d", i))
+		val := []byte(fmt.Sprintf("val%d", i))
+		require.NoError(t, writer.SetEntryAt(NewEntry(key, val).WithMeta(0x00), 1))
+	}
+	require.NoError(t, writer.Flush())
+
+	outDir, err := ioutil.TempDir("", "badger-test")
+	require.NoError(t, err)
+	outOpt := getTestOptions(outDir)
+	require.NoError(t, db.StreamDB(outOpt))
+
+	outDB, err := OpenManaged(outOpt)
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, outDB.Close())
+	}()
+
+	for i := 0; i < 100; i++ {
+		key := []byte(fmt.Sprintf("key%d", i))
+		val := []byte(fmt.Sprintf("val%d", i))
+		txn := outDB.NewTransactionAt(1, false)
+		item, err := txn.Get(key)
+		require.NoError(t, err)
+		require.EqualValues(t, val, getItemValue(t, item))
+		require.Equal(t, byte(0x00), item.UserMeta())
+		txn.Discard()
+	}
+}
+
 func dirSize(path string) (int64, error) {
 	var size int64
 	err := filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {