Repository: dgraph-io/badger
Pull request #1421 (2020-07-13 14:36:33)

Handler error in writeBatch.Flush
The `writeBatch.Flush` would skip the error returned by `wb.Commit()` call. We shouldn't ignore the error.

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



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9f4200b..3ecb531 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,11 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
+## [20.07.0-rc2] - 2020-07-14
+
+### Fixed
+- Add a limit to the size of the batches sent over a stream. (#1412)
+
 ## [20.07.0-rc1] - 2020-07-11
 
 ### Fixed
diff --git a/badger/cmd/write_bench.go b/badger/cmd/write_bench.go
index fb669ea..cb15b7d 100644
--- a/badger/cmd/write_bench.go
+++ b/badger/cmd/write_bench.go
@@ -29,6 +29,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/dgraph-io/badger/v2"
+	"github.com/dgraph-io/badger/v2/options"
 	"github.com/dgraph-io/badger/v2/pb"
 	"github.com/dgraph-io/badger/v2/y"
 )
@@ -53,6 +54,21 @@ var (
 
 	sizeWritten    uint64
 	entriesWritten uint64
+
+	valueThreshold      int
+	numVersions         int
+	maxCacheSize        int64
+	keepBlockIdxInCache bool
+	keepBlocksInCache   bool
+	maxBfCacheSize      int64
+	vlogMaxEntries      uint32
+	loadBloomsOnOpen    bool
+	detectConflicts     bool
+	compression         bool
+
+	gcPeriod       string
+	gcDiscardRatio float64
+	gcSuccess      uint64
 )
 
 const (
@@ -69,6 +85,28 @@ func init() {
 		"Force compact level 0 on close.")
 	writeBenchCmd.Flags().BoolVarP(&sorted, "sorted", "s", false, "Write keys in sorted order.")
 	writeBenchCmd.Flags().BoolVarP(&showLogs, "logs", "l", false, "Show Badger logs.")
+	writeBenchCmd.Flags().IntVarP(&valueThreshold, "value-th", "t", 1<<10, "Value threshold")
+	writeBenchCmd.Flags().IntVarP(&numVersions, "num-version", "n", 1, "Number of versions to keep")
+	writeBenchCmd.Flags().Int64VarP(&maxCacheSize, "max-cache", "C", 1<<30, "Max size of cache")
+	writeBenchCmd.Flags().BoolVarP(&keepBlockIdxInCache, "keep-bidx", "b", true,
+		"Keep block indices in cache")
+	writeBenchCmd.Flags().BoolVarP(&keepBlocksInCache, "keep-blocks", "B", true,
+		"Keep blocks in cache")
+	writeBenchCmd.Flags().Int64VarP(&maxBfCacheSize, "max-bf-cache", "c", 500<<20,
+		"Maximum Bloom Filter Cache Size")
+	writeBenchCmd.Flags().Uint32Var(&vlogMaxEntries, "vlog-maxe", 10000, "Value log Max Entries")
+	writeBenchCmd.Flags().StringVarP(&encryptionKey, "encryption-key", "e", "",
+		"If it is true, badger will encrypt all the data stored on the disk.")
+	writeBenchCmd.Flags().StringVar(&loadingMode, "loading-mode", "mmap",
+		"Mode for accessing SSTables")
+	writeBenchCmd.Flags().BoolVar(&loadBloomsOnOpen, "load-blooms", false,
+		"Load Bloom filter on DB open.")
+	writeBenchCmd.Flags().BoolVar(&detectConflicts, "conficts", false,
+		"If true, it badger will detect the conflicts")
+	writeBenchCmd.Flags().BoolVar(&compression, "compression", false,
+		"If true, badger will use ZSTD mode")
+	writeBenchCmd.Flags().StringVarP(&gcPeriod, "gc-every", "g", "5m", "GC Period.")
+	writeBenchCmd.Flags().Float64VarP(&gcDiscardRatio, "gc-ratio", "r", 0.5, "GC discard ratio.")
 }
 
 func writeRandom(db *badger.DB, num uint64) error {
@@ -157,11 +195,30 @@ func writeSorted(db *badger.DB, num uint64) error {
 }
 
 func writeBench(cmd *cobra.Command, args []string) error {
+	var cmode options.CompressionType
+	if compression {
+		cmode = options.ZSTD
+	} else {
+		cmode = options.None
+	}
+	mode := getLoadingMode(loadingMode)
 	opt := badger.DefaultOptions(sstDir).
 		WithValueDir(vlogDir).
 		WithTruncate(truncate).
 		WithSyncWrites(false).
-		WithCompactL0OnClose(force)
+		WithCompactL0OnClose(force).
+		WithValueThreshold(valueThreshold).
+		WithNumVersionsToKeep(numVersions).
+		WithMaxCacheSize(maxCacheSize).
+		WithKeepBlockIndicesInCache(keepBlockIdxInCache).
+		WithKeepBlocksInCache(keepBlocksInCache).
+		WithMaxBfCacheSize(maxBfCacheSize).
+		WithValueLogMaxEntries(vlogMaxEntries).
+		WithTableLoadingMode(mode).
+		WithEncryptionKey([]byte(encryptionKey)).
+		WithLoadBloomsOnOpen(loadBloomsOnOpen).
+		WithDetectConflicts(detectConflicts).
+		WithCompression(cmode)
 
 	if !showLogs {
 		opt = opt.WithLogger(nil)
@@ -183,8 +240,9 @@ func writeBench(cmd *cobra.Command, args []string) error {
 
 	startTime = time.Now()
 	num := uint64(numKeys * mil)
-	c := y.NewCloser(1)
+	c := y.NewCloser(2)
 	go reportStats(c)
+	go runGC(c, db)
 
 	if sorted {
 		err = writeSorted(db, num)
@@ -212,8 +270,28 @@ func reportStats(c *y.Closer) {
 			bytesRate := sz / uint64(dur.Seconds())
 			entriesRate := entries / uint64(dur.Seconds())
 			fmt.Printf("Time elapsed: %s, bytes written: %s, speed: %s/sec, "+
-				"entries written: %d, speed: %d/sec\n", y.FixedDuration(time.Since(startTime)),
-				humanize.Bytes(sz), humanize.Bytes(bytesRate), entries, entriesRate)
+				"entries written: %d, speed: %d/sec, gcSuccess: %d\n", y.FixedDuration(time.Since(startTime)),
+				humanize.Bytes(sz), humanize.Bytes(bytesRate), entries, entriesRate, gcSuccess)
+		}
+	}
+}
+
+func runGC(c *y.Closer, db *badger.DB) {
+	defer c.Done()
+	period, err := time.ParseDuration(gcPeriod)
+	y.Check(err)
+	t := time.NewTicker(period)
+	defer t.Stop()
+	for {
+		select {
+		case <-c.HasBeenClosed():
+			return
+		case <-t.C:
+			if err := db.RunValueLogGC(gcDiscardRatio); err == nil {
+				atomic.AddUint64(&gcSuccess, 1)
+			} else {
+				log.Printf("[GC] Failed due to following err %v", err)
+			}
 		}
 	}
 }
diff --git a/batch.go b/batch.go
index ff94e86..f74c2f5 100644
--- a/batch.go
+++ b/batch.go
@@ -181,7 +181,8 @@ func (wb *WriteBatch) commit() error {
 		return wb.err
 	}
 	if err := wb.throttle.Do(); err != nil {
-		return err
+		wb.err = err
+		return wb.err
 	}
 	wb.txn.CommitWith(wb.callback)
 	wb.txn = wb.db.newTransaction(true, wb.isManaged)
@@ -193,11 +194,15 @@ func (wb *WriteBatch) commit() error {
 // returns any error stored by WriteBatch.
 func (wb *WriteBatch) Flush() error {
 	wb.Lock()
+	// commit will set the wb.err so no need to check the error here.
 	_ = wb.commit()
 	wb.txn.Discard()
 	wb.Unlock()
 
 	if err := wb.throttle.Finish(); err != nil {
+		if wb.err != nil {
+			return errors.Errorf("wb.err: %s err: %s", wb.err, err)
+		}
 		return err
 	}
 
diff --git a/stream.go b/stream.go
index 30814e7..eb7341e 100644
--- a/stream.go
+++ b/stream.go
@@ -32,6 +32,12 @@ import (
 
 const pageSize = 4 << 20 // 4MB
 
+// maxStreamSize is the maximum allowed size of a stream batch. This is a soft limit
+// as a single list that is still over the limit will have to be sent as is since it
+// cannot be split further. This limit prevents the framework from creating batches
+// so big that sending them causes issues (e.g running into the max size gRPC limit).
+var maxStreamSize = uint64(100 << 20) // 100MB
+
 // Stream provides a framework to concurrently iterate over a snapshot of Badger, pick up
 // key-values, batch them up and call Send. Stream does concurrent iteration over many smaller key
 // ranges. It does NOT send keys in lexicographical sorted order. To get keys in sorted
@@ -248,9 +254,29 @@ func (st *Stream) streamKVs(ctx context.Context) error {
 	defer t.Stop()
 	now := time.Now()
 
+	sendBatch := func(batch *pb.KVList) error {
+		sz := uint64(proto.Size(batch))
+		bytesSent += sz
+		count += len(batch.Kv)
+		t := time.Now()
+		if err := st.Send(batch); err != nil {
+			return err
+		}
+		st.db.opt.Infof("%s Created batch of size: %s in %s.\n",
+			st.LogPrefix, humanize.Bytes(sz), time.Since(t))
+		return nil
+	}
+
 	slurp := func(batch *pb.KVList) error {
 	loop:
 		for {
+			// Send the batch immediately if it already exceeds the maximum allowed size.
+			// If the size of the batch exceeds maxStreamSize, break from the loop to
+			// avoid creating a batch that is so big that certain limits are reached.
+			sz := uint64(proto.Size(batch))
+			if sz > maxStreamSize {
+				break loop
+			}
 			select {
 			case kvs, ok := <-st.kvChan:
 				if !ok {
@@ -262,16 +288,7 @@ func (st *Stream) streamKVs(ctx context.Context) error {
 				break loop
 			}
 		}
-		sz := uint64(proto.Size(batch))
-		bytesSent += sz
-		count += len(batch.Kv)
-		t := time.Now()
-		if err := st.Send(batch); err != nil {
-			return err
-		}
-		st.db.opt.Infof("%s Created batch of size: %s in %s.\n",
-			st.LogPrefix, humanize.Bytes(sz), time.Since(t))
-		return nil
+		return sendBatch(batch)
 	}
 
 outer:
@@ -297,6 +314,8 @@ outer:
 			}
 			y.AssertTrue(kvs != nil)
 			batch = kvs
+
+			// Otherwise, slurp more keys into this batch.
 			if err := slurp(batch); err != nil {
 				return err
 			}


Test patch:
diff --git a/stream_test.go b/stream_test.go
index f60e3de..f3c4b1d 100644
--- a/stream_test.go
+++ b/stream_test.go
@@ -77,9 +77,7 @@ func TestStream(t *testing.T) {
 	stream := db.NewStreamAt(math.MaxUint64)
 	stream.LogPrefix = "Testing"
 	c := &collector{}
-	stream.Send = func(list *bpb.KVList) error {
-		return c.Send(list)
-	}
+	stream.Send = c.Send
 
 	// Test case 1. Retrieve everything.
 	err = stream.Orchestrate(ctxb)
@@ -186,9 +184,7 @@ func TestStreamWithThreadId(t *testing.T) {
 		return stream.ToList(key, itr)
 	}
 	c := &collector{}
-	stream.Send = func(list *bpb.KVList) error {
-		return c.Send(list)
-	}
+	stream.Send = c.Send
 
 	err = stream.Orchestrate(ctxb)
 	require.NoError(t, err)
@@ -207,3 +203,54 @@ func TestStreamWithThreadId(t *testing.T) {
 	}
 	require.NoError(t, db.Close())
 }
+
+func TestBigStream(t *testing.T) {
+	// Set the maxStreamSize to 1MB for the duration of the test so that the it can use a smaller
+	// dataset than it would otherwise need.
+	originalMaxStreamSize := maxStreamSize
+	maxStreamSize = 1 << 20
+	defer func() {
+		maxStreamSize = originalMaxStreamSize
+	}()
+
+	testSize := int(1e6)
+	dir, err := ioutil.TempDir("", "badger-big-test")
+	require.NoError(t, err)
+	defer removeDir(dir)
+
+	db, err := OpenManaged(DefaultOptions(dir))
+	require.NoError(t, err)
+
+	var count int
+	wb := db.NewWriteBatchAt(5)
+	for _, prefix := range []string{"p0", "p1", "p2"} {
+		for i := 1; i <= testSize; i++ {
+			require.NoError(t, wb.SetEntry(NewEntry(keyWithPrefix(prefix, i), value(i))))
+			count++
+		}
+	}
+	require.NoError(t, wb.Flush())
+
+	stream := db.NewStreamAt(math.MaxUint64)
+	stream.LogPrefix = "Testing"
+	c := &collector{}
+	stream.Send = c.Send
+
+	// Test case 1. Retrieve everything.
+	err = stream.Orchestrate(ctxb)
+	require.NoError(t, err)
+	require.Equal(t, 3*testSize, len(c.kv), "Expected 30000. Got: %d", len(c.kv))
+
+	m := make(map[string]int)
+	for _, kv := range c.kv {
+		prefix, ki := keyToInt(kv.Key)
+		expected := value(ki)
+		require.Equal(t, expected, kv.Value)
+		m[prefix]++
+	}
+	require.Equal(t, 3, len(m))
+	for pred, count := range m {
+		require.Equal(t, testSize, count, "Count mismatch for pred: %s", pred)
+	}
+	require.NoError(t, db.Close())
+}