Repository: dgraph-io/badger
Pull request #1524 (2020-09-17 14:38:46)

Refactor estimatedSize of table
Users reported mismatch between size of data on disk and size reported by dgraph. Due to compression, the real size of data on disk is much smaller than what we estimate from index.EstimatedSize.

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


CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/dgraph-io/badger?pullRequest=1524) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/dgraph-io/badger?pullRequest=1524) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/dgraph-io/badger?pullRequest=1524) it.</sub>

Patch:
diff --git a/db.go b/db.go
index 463f7ee..53b39ae 100644
--- a/db.go
+++ b/db.go
@@ -193,43 +193,42 @@ func (db *DB) replayFunction() func(Entry, valuePointer) error {
 	}
 }
 
-// Open returns a new DB object.
-func Open(opt Options) (db *DB, err error) {
+func checkAndSetOptions(opt *Options) error {
 	// It's okay to have zero compactors which will disable all compactions but
 	// we cannot have just one compactor otherwise we will end up with all data
 	// on level 2.
 	if opt.NumCompactors == 1 {
-		return nil, errors.New("Cannot have 1 compactor. Need at least 2")
+		return errors.New("Cannot have 1 compactor. Need at least 2")
 	}
 	if opt.InMemory && (opt.Dir != "" || opt.ValueDir != "") {
-		return nil, errors.New("Cannot use badger in Disk-less mode with Dir or ValueDir set")
+		return errors.New("Cannot use badger in Disk-less mode with Dir or ValueDir set")
 	}
 	opt.maxBatchSize = (15 * opt.MaxTableSize) / 100
 	opt.maxBatchCount = opt.maxBatchSize / int64(skl.MaxNodeSize)
 
 	// We are limiting opt.ValueThreshold to maxValueThreshold for now.
 	if opt.ValueThreshold > maxValueThreshold {
-		return nil, errors.Errorf("Invalid ValueThreshold, must be less or equal to %d",
+		return errors.Errorf("Invalid ValueThreshold, must be less or equal to %d",
 			maxValueThreshold)
 	}
 
 	// If ValueThreshold is greater than opt.maxBatchSize, we won't be able to push any data using
 	// the transaction APIs. Transaction batches entries into batches of size opt.maxBatchSize.
 	if int64(opt.ValueThreshold) > opt.maxBatchSize {
-		return nil, errors.Errorf("Valuethreshold greater than max batch size of %d. Either "+
+		return errors.Errorf("Valuethreshold greater than max batch size of %d. Either "+
 			"reduce opt.ValueThreshold or increase opt.MaxTableSize.", opt.maxBatchSize)
 	}
 	if !(opt.ValueLogFileSize <= 2<<30 && opt.ValueLogFileSize >= 1<<20) {
-		return nil, ErrValueLogSize
+		return ErrValueLogSize
 	}
 	if !(opt.ValueLogLoadingMode == options.FileIO ||
 		opt.ValueLogLoadingMode == options.MemoryMap) {
-		return nil, ErrInvalidLoadingMode
+		return ErrInvalidLoadingMode
 	}
 
 	// Return error if badger is built without cgo and compression is set to ZSTD.
 	if opt.Compression == options.ZSTD && !y.CgoEnabled {
-		return nil, y.ErrZstdCgo
+		return y.ErrZstdCgo
 	}
 	// Keep L0 in memory if either KeepL0InMemory is set or if InMemory is set.
 	opt.KeepL0InMemory = opt.KeepL0InMemory || opt.InMemory
@@ -244,6 +243,20 @@ func Open(opt Options) (db *DB, err error) {
 		// Do not perform compaction in read only mode.
 		opt.CompactL0OnClose = false
 	}
+
+	needCache := (opt.Compression != options.None) || (len(opt.EncryptionKey) > 0)
+	if needCache && opt.BlockCacheSize == 0 {
+		opt.Warningf("BlockCacheSize should be set " +
+			"since compression/encryption are enabled")
+	}
+	return nil
+}
+
+// Open returns a new DB object.
+func Open(opt Options) (db *DB, err error) {
+	if err := checkAndSetOptions(&opt); err != nil {
+		return nil, err
+	}
 	var dirLockGuard, valueDirLockGuard *directoryLockGuard
 
 	// Create directories and acquire lock on it only if badger is not running in InMemory mode.
@@ -1864,3 +1877,32 @@ func (db *DB) StreamDB(outOptions Options) error {
 func (db *DB) Opts() Options {
 	return db.opt
 }
+
+// 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/iterator.go b/iterator.go
index 11d0c27..2df066e 100644
--- a/iterator.go
+++ b/iterator.go
@@ -41,20 +41,21 @@ const (
 // Item is returned during iteration. Both the Key() and Value() output is only valid until
 // iterator.Next() is called.
 type Item struct {
-	status    prefetchStatus
-	err       error
-	wg        sync.WaitGroup
-	db        *DB
 	key       []byte
 	vptr      []byte
-	meta      byte // We need to store meta to know about bitValuePointer.
-	userMeta  byte
-	expiresAt uint64
 	val       []byte
-	slice     *y.Slice // Used only during prefetching.
-	next      *Item
 	version   uint64
-	txn       *Txn
+	expiresAt uint64
+
+	slice *y.Slice // Used only during prefetching.
+	next  *Item
+	txn   *Txn
+
+	err      error
+	wg       sync.WaitGroup
+	status   prefetchStatus
+	meta     byte // We need to store meta to know about bitValuePointer.
+	userMeta byte
 }
 
 // String returns a string representation of Item
@@ -167,10 +168,11 @@ func (item *Item) yieldItemValue() ([]byte, func(), error) {
 
 		var vp valuePointer
 		vp.Decode(item.vptr)
-		result, cb, err := item.db.vlog.Read(vp, item.slice)
+		db := item.txn.db
+		result, cb, err := db.vlog.Read(vp, item.slice)
 		if err != ErrRetry {
 			if err != nil {
-				item.db.opt.Logger.Errorf(`Unable to read: Key: %v, Version : %v,
+				db.opt.Logger.Errorf(`Unable to read: Key: %v, Version : %v,
 				meta: %v, userMeta: %v`, key, item.version, item.meta, item.userMeta)
 			}
 			return result, cb, err
@@ -194,7 +196,7 @@ func (item *Item) yieldItemValue() ([]byte, func(), error) {
 		// change the key user sees before and after this call. Also, this move
 		// logic is internal logic and should not impact the external behavior
 		// of the retrieval.
-		vs, err := item.db.get(key)
+		vs, err := item.txn.db.get(key)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -226,7 +228,7 @@ func (item *Item) prefetchValue() {
 	if val == nil {
 		return
 	}
-	if item.db.opt.ValueLogLoadingMode == options.MemoryMap {
+	if item.txn.db.opt.ValueLogLoadingMode == options.MemoryMap {
 		buf := item.slice.Resize(len(val))
 		copy(buf, val)
 		item.val = buf
@@ -329,20 +331,21 @@ func (l *list) pop() *Item {
 // should work for most applications. Consider using that as a starting point
 // before customizing it for your own needs.
 type IteratorOptions struct {
-	// Indicates whether we should prefetch values during iteration and store them.
-	PrefetchValues bool
-	// How many KV pairs to prefetch while iterating. Valid only if PrefetchValues is true.
+	// PrefetchSize is the number of KV pairs to prefetch while iterating.
+	// Valid only if PrefetchValues is true.
 	PrefetchSize int
-	Reverse      bool // Direction of iteration. False is forward, true is backward.
-	AllVersions  bool // Fetch all valid versions of the same key.
+	// PrefetchValues Indicates whether we should prefetch values during
+	// iteration and store them.
+	PrefetchValues bool
+	Reverse        bool // Direction of iteration. False is forward, true is backward.
+	AllVersions    bool // Fetch all valid versions of the same key.
+	InternalAccess bool // Used to allow internal access to badger keys.
 
-	// The following option is used to narrow down the SSTables that iterator picks up. If
-	// Prefix is specified, only tables which could have this prefix are picked based on their range
-	// of keys.
-	Prefix      []byte // Only iterate over this given prefix.
+	// The following option is used to narrow down the SSTables that iterator
+	// picks up. If Prefix is specified, only tables which could have this
+	// prefix are picked based on their range of keys.
 	prefixIsKey bool   // If set, use the prefix for bloom filter lookup.
-
-	InternalAccess bool // Used to allow internal access to badger keys.
+	Prefix      []byte // Only iterate over this given prefix.
 }
 
 func (opt *IteratorOptions) compareToPrefix(key []byte) int {
@@ -505,7 +508,7 @@ func (txn *Txn) NewKeyIterator(key []byte, opt IteratorOptions) *Iterator {
 func (it *Iterator) newItem() *Item {
 	item := it.waste.pop()
 	if item == nil {
-		item = &Item{slice: new(y.Slice), db: it.txn.db, txn: it.txn}
+		item = &Item{slice: new(y.Slice), txn: it.txn}
 	}
 	return item
 }
diff --git a/structs.go b/structs.go
index 469cdc4..38355ba 100644
--- a/structs.go
+++ b/structs.go
@@ -140,13 +140,13 @@ func (h *header) DecodeFrom(reader *hashReader) (int, error) {
 type Entry struct {
 	Key       []byte
 	Value     []byte
-	UserMeta  byte
 	ExpiresAt uint64 // time.Unix
-	meta      byte
 	version   uint64
+	offset    uint32 // offset is an internal field.
+	UserMeta  byte
+	meta      byte
 
 	// Fields maintained internally.
-	offset   uint32
 	skipVlog bool
 	hlen     int // Length of the header.
 }
diff --git a/table/table.go b/table/table.go
index 76125cc..0f2ee70 100644
--- a/table/table.go
+++ b/table/table.go
@@ -446,7 +446,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.hasBloomFilter = len(index.BloomFilter) > 0
 	t.noOfBlocks = len(index.Offsets)
 
diff --git a/txn.go b/txn.go
index 8ff06d2..21b0db5 100644
--- a/txn.go
+++ b/txn.go
@@ -249,9 +249,11 @@ func (o *oracle) doneCommit(cts uint64) {
 type Txn struct {
 	readTs   uint64
 	commitTs uint64
+	size     int64
+	count    int64
+	db       *DB
 
-	update bool     // update is used to conditionally keep track of reads.
-	reads  []uint64 // contains fingerprints of keys read.
+	reads []uint64 // contains fingerprints of keys read.
 	// contains fingerprints of keys written. This is used for conflict detection.
 	conflictKeys map[uint64]struct{}
 	readsLock    sync.Mutex // guards the reads slice. See addReadKey.
@@ -259,13 +261,10 @@ type Txn struct {
 	pendingWrites   map[string]*Entry // cache stores any writes done by txn.
 	duplicateWrites []*Entry          // Used in managed mode to store duplicate entries.
 
-	db        *DB
-	discarded bool
-	doneRead  bool
-
-	size         int64
-	count        int64
 	numIterators int32
+	discarded    bool
+	doneRead     bool
+	update       bool // update is used to conditionally keep track of reads.
 }
 
 type pendingWritesIterator struct {
@@ -483,7 +482,6 @@ func (txn *Txn) Get(key []byte) (item *Item, rerr error) {
 	item.version = vs.Version
 	item.meta = vs.Meta
 	item.userMeta = vs.UserMeta
-	item.db = txn.db
 	item.vptr = y.SafeCopy(item.vptr, vs.Value)
 	item.txn = txn
 	item.expiresAt = vs.ExpiresAt
diff --git a/y/iterator.go b/y/iterator.go
index 7c9b211..ef032ef 100644
--- a/y/iterator.go
+++ b/y/iterator.go
@@ -46,10 +46,6 @@ func sizeVarint(x uint64) (n int) {
 // EncodedSize is the size of the ValueStruct when encoded
 func (v *ValueStruct) EncodedSize() uint32 {
 	sz := len(v.Value) + 2 // meta, usermeta.
-	if v.ExpiresAt == 0 {
-		return uint32(sz + 1)
-	}
-
 	enc := sizeVarint(v.ExpiresAt)
 	return uint32(sz + enc)
 }


Test patch:
diff --git a/db2_test.go b/db2_test.go
index 2a19991..7b165e0 100644
--- a/db2_test.go
+++ b/db2_test.go
@@ -858,3 +858,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 b9f6c98..88b7bd9 100644
--- a/table/table_test.go
+++ b/table/table_test.go
@@ -923,13 +923,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"
diff --git a/y/y_test.go b/y/y_test.go
index d1b9631..151d711 100644
--- a/y/y_test.go
+++ b/y/y_test.go
@@ -2,6 +2,7 @@ package y
 
 import (
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"math/rand"
@@ -254,3 +255,24 @@ func TestPagebufferReader4(t *testing.T) {
 	require.Equal(t, err, io.EOF, "should return EOF")
 	require.Equal(t, n, 0)
 }
+
+func TestSizeVarintForZero(t *testing.T) {
+	siz := sizeVarint(0)
+	require.Equal(t, 1, siz)
+}
+
+func TestEncodedSize(t *testing.T) {
+	valBufSize := uint32(rand.Int31n(1e5))
+	expiry := rand.Uint64()
+	expiryVarintBuf := make([]byte, 64)
+	expVarintSize := uint32(binary.PutUvarint(expiryVarintBuf, expiry))
+	valBuf := make([]byte, valBufSize)
+	_, _ = rand.Read(valBuf)
+
+	valStruct := &ValueStruct{
+		Value:     valBuf,
+		ExpiresAt: expiry,
+	}
+
+	require.Equal(t, valBufSize+uint32(2)+expVarintSize, valStruct.EncodedSize())
+}