Repository: dgraph-io/badger
Pull request #1531 (2020-09-21 13:36:34)

release/v2.2007 - feat(db): Add db.MaxVersion API (#1526)
This PR adds a new `db.MaxVersion()` API that returns the
maximum version across all keys in the DB.

(cherry picked from commit 36af85fc7a483e414c49d0c4fc9e701b0bebf78f)

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



Patch:
diff --git a/db.go b/db.go
index 91e2811..820f599 100644
--- a/db.go
+++ b/db.go
@@ -1818,3 +1818,53 @@ 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.KeyToList = func(key []byte, itr *Iterator) (*pb.KVList, error) {
+		maxVs := uint64(0)
+		for ; itr.Valid(); itr.Next() {
+			item := itr.Item()
+			if item.IsDeletedOrExpired() {
+				break
+			}
+			if !bytes.Equal(key, item.Key()) {
+				// Break out on the first encounter with another key.
+				break
+			}
+
+			if item.Version() > maxVs {
+				maxVs = item.Version()
+			}
+			if db.opt.NumVersionsToKeep == 1 {
+				break
+			}
+
+			if item.DiscardEarlierVersions() {
+				break
+			}
+		}
+		mu.Lock()
+		if maxVs > maxVersion {
+			maxVersion = maxVs
+		}
+		mu.Unlock()
+		return nil, nil
+	}
+	stream.Send = nil
+	if err := stream.Orchestrate(context.Background()); err != nil {
+		return 0, err
+	}
+	return maxVersion, nil
+
+}


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())
+	})
+}