Repository: dgraph-io/badger
Pull request #1520 (2020-09-15 07:59:21)

release/v2.2007 - fix(flatten): Add --num_versions flag (#1518)
Add --num_versions flag to set NumVersionsToKeep option in badger.

(cherry picked from commit c68fbc1ed6fc54b49624a7cb3ee02548519acce5)

<!-- 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/1520)
<!-- Reviewable:end -->



Patch:
diff --git a/badger/cmd/flatten.go b/badger/cmd/flatten.go
index 1882cb9..e1259f7 100644
--- a/badger/cmd/flatten.go
+++ b/badger/cmd/flatten.go
@@ -17,6 +17,9 @@
 package cmd
 
 import (
+	"fmt"
+	"math"
+
 	"github.com/dgraph-io/badger/v2"
 	"github.com/spf13/cobra"
 )
@@ -37,13 +40,23 @@ func init() {
 	flattenCmd.Flags().IntVarP(&numWorkers, "num-workers", "w", 1,
 		"Number of concurrent compactors to run. More compactors would use more"+
 			" server resources to potentially achieve faster compactions.")
+	flattenCmd.Flags().IntVarP(&numVersions, "num_versions", "", 1,
+		"Option to configure the maximum number of versions per key. "+
+			"Values <= 0 will be considered to have the max number of versions.")
 }
 
 func flatten(cmd *cobra.Command, args []string) error {
-	db, err := badger.Open(badger.DefaultOptions(sstDir).
+	if numVersions <= 0 {
+		// Keep all versions.
+		numVersions = math.MaxInt32
+	}
+	opt := badger.DefaultOptions(sstDir).
 		WithValueDir(vlogDir).
 		WithTruncate(truncate).
-		WithNumCompactors(0))
+		WithNumVersionsToKeep(numVersions).
+		WithNumCompactors(0)
+	fmt.Printf("Opening badger with options = %+v\n", opt)
+	db, err := badger.Open(opt)
 	if err != nil {
 		return err
 	}
diff --git a/db.go b/db.go
index 91e2811..9a3be1e 100644
--- a/db.go
+++ b/db.go
@@ -1818,3 +1818,32 @@ func (db *DB) StreamDB(outOptions Options) error {
 	}
 	return nil
 }
+
+// MaxVersion returns the maximum commited version across all keys in the DB. It
+// uses the stream framework to find the maximum version.
+func (db *DB) MaxVersion() (uint64, error) {
+	maxVersion := uint64(0)
+	var mu sync.Mutex
+	var stream *Stream
+	if db.opt.managedTxns {
+		stream = db.NewStreamAt(math.MaxUint64)
+	} else {
+		stream = db.NewStream()
+	}
+
+	stream.ChooseKey = func(item *Item) bool {
+		mu.Lock()
+		if item.Version() > maxVersion {
+			maxVersion = item.Version()
+		}
+		mu.Unlock()
+		return false
+	}
+	stream.KeyToList = nil
+	stream.Send = nil
+	if err := stream.Orchestrate(context.Background()); err != nil {
+		return 0, err
+	}
+	return maxVersion, nil
+
+}
diff --git a/table/table.go b/table/table.go
index fa49df1..f8841b6 100644
--- a/table/table.go
+++ b/table/table.go
@@ -381,7 +381,13 @@ func (t *Table) initIndex() (*pb.BlockOffset, error) {
 		return nil, err
 	}
 
-	t.estimatedSize = index.EstimatedSize
+	if t.opt.Compression == options.None {
+		t.estimatedSize = index.EstimatedSize
+	} else {
+		// Due to compression the real size on disk is much
+		// smaller than what we estimate from index.EstimatedSize.
+		t.estimatedSize = uint64(t.tableSize)
+	}
 	t.noOfBlocks = len(index.Offsets)
 
 	// No cache


Test patch:
diff --git a/db2_test.go b/db2_test.go
index 645971b..493e54f 100644
--- a/db2_test.go
+++ b/db2_test.go
@@ -861,3 +861,72 @@ func TestIsClosed(t *testing.T) {
 	})
 
 }
+
+func TestMaxVersion(t *testing.T) {
+	N := uint64(10000)
+	key := func(i int) []byte {
+		return []byte(fmt.Sprintf("%d%10d", i, i))
+	}
+	t.Run("normal", func(t *testing.T) {
+		runBadgerTest(t, nil, func(t *testing.T, db *DB) {
+			// This will create commits from 1 to N.
+			for i := 0; i < int(N); i++ {
+				txnSet(t, db, key(i), nil, 0)
+			}
+			ver, err := db.MaxVersion()
+			require.NoError(t, err)
+			require.Equal(t, N, ver)
+		})
+	})
+	t.Run("multiple versions", func(t *testing.T) {
+		dir, err := ioutil.TempDir("", "badger-test")
+		require.NoError(t, err)
+		defer removeDir(dir)
+
+		opt := getTestOptions(dir)
+		opt.NumVersionsToKeep = 100
+		db, err := OpenManaged(opt)
+		require.NoError(t, err)
+
+		wb := db.NewManagedWriteBatch()
+		defer wb.Cancel()
+
+		k := make([]byte, 100)
+		rand.Read(k)
+		// Create multiple version of the same key.
+		for i := uint64(1); i <= N; i++ {
+			wb.SetEntryAt(&Entry{Key: k}, i)
+		}
+		require.NoError(t, wb.Flush())
+
+		ver, err := db.MaxVersion()
+		require.NoError(t, err)
+		require.Equal(t, N, ver)
+
+		require.NoError(t, db.Close())
+	})
+	t.Run("Managed mode", func(t *testing.T) {
+		dir, err := ioutil.TempDir("", "badger-test")
+		require.NoError(t, err)
+		defer removeDir(dir)
+
+		opt := getTestOptions(dir)
+		db, err := OpenManaged(opt)
+		require.NoError(t, err)
+
+		wb := db.NewManagedWriteBatch()
+		defer wb.Cancel()
+
+		// This will create commits from 1 to N.
+		for i := uint64(1); i <= N; i++ {
+			wb.SetEntryAt(&Entry{Key: []byte(fmt.Sprintf("%d", i))}, i)
+		}
+		require.NoError(t, wb.Flush())
+
+		ver, err := db.MaxVersion()
+		require.NoError(t, err)
+		require.Equal(t, N, ver)
+
+		require.NoError(t, db.Close())
+	})
+}
diff --git a/table/table_test.go b/table/table_test.go
index ab84427..bac7641 100644
--- a/table/table_test.go
+++ b/table/table_test.go
@@ -930,13 +930,30 @@ func TestMain(m *testing.M) {
 }
 
 func TestOpenKVSize(t *testing.T) {
-	opts := getTestTableOptions()
-	table, err := OpenTable(buildTestTable(t, "foo", 1, opts), opts)
-	require.NoError(t, err)
+	t.Run("compression", func(t *testing.T) {
+		// When compression is on
+		opts := getTestTableOptions()
+		opts.Compression = options.ZSTD
+		table, err := OpenTable(buildTestTable(t, "foo", 1000, opts), opts)
+		defer table.Close()
+		require.NoError(t, err)
+
+		// The estimated size is same as table size in case compression is enabled.
+		require.Equal(t, uint64(table.tableSize), table.EstimatedSize())
+	})
 
-	// The following values might change if the table/header structure is changed.
-	var entrySize uint64 = 15 /* DiffKey len */ + 4 /* Header Size */ + 4 /* Encoded vp */
-	require.Equal(t, entrySize, table.EstimatedSize())
+	t.Run("no compressin", func(t *testing.T) {
+		// When compression is off
+		opts := getTestTableOptions()
+		opts.Compression = options.None
+		table, err := OpenTable(buildTestTable(t, "foo", 1, opts), opts)
+		require.NoError(t, err)
+		defer table.Close()
+
+		stat, err := table.fd.Stat()
+		require.NoError(t, err)
+		require.Less(t, table.EstimatedSize(), uint64(stat.Size()))
+	})
 }
 
 // Run this test with command "go test -race -run TestDoesNotHaveRace"