Repository: dgraph-io/badger
Pull request #1416 (2020-07-12 16:36:31)

Add changes.sh script
The changes.sh script can be used to generate PR descriptions for badger
update PRs on dgraph.

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



Patch:
diff --git a/.github/ISSUE_TEMPLATE b/.github/ISSUE_TEMPLATE
index 893e1e0..7b0145e 100644
--- a/.github/ISSUE_TEMPLATE
+++ b/.github/ISSUE_TEMPLATE
@@ -1,43 +1,1 @@
-<!-- Please answer these questions before submitting your issue. Thanks! -->
-
-### What version of Go are you using (`go version`)?
-
-<pre>
-$ go version
-
-</pre>
-
-### What operating system are you using?
-
-
-
-### What version of Badger are you using?
-
-
-
-### Does this issue reproduce with the latest master?
-
-
-
-### Steps to Reproduce the issue
-<!--
-If possible, provide a recipe for reproducing the error. A complete runnable program is great.
--->
-
-
-### What Badger options were set? 
-<!--
-Provide the list of Badger options that were set.
--->
-
-
-
-### What did you do?
-
-
-
-### What did you expect to see?
-
-
-
-### What did you see instead?
+**GitHub Issues are deprecated. Use [Discuss Issues](https://discuss.dgraph.io/c/issues/badger/37) for reporting issues about this repository.**
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
new file mode 100644
index 0000000..5f1aa26
--- /dev/null
+++ b/.github/workflows/main.yml
@@ -0,0 +1,24 @@
+# This is a basic workflow to help you get started with Actions
+
+name: Issue Closer
+
+# Controls when the action will run. Triggers the workflow on push or pull request
+# events but only for the master branch
+on:
+  issues:
+    types: [ opened, reopened ]
+
+# A workflow run is made up of one or more jobs that can run sequentially or in parallel
+jobs:
+  # This workflow contains a single job called "build"
+  build:
+    # The type of runner that the job will run on
+    runs-on: ubuntu-latest
+
+    # Steps represent a sequence of tasks that will be executed as part of the job
+    steps:
+      - name: Close Issue
+        uses: peter-evans/close-issue@v1.0.1
+        with:
+          comment: |
+           **Use [Discuss Issues](https://discuss.dgraph.io/c/issues/badger/37) for reporting issues about this repository.**
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/README.md b/README.md
index 2a22091..5e00e82 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,8 @@ in pure Go. It is the underlying database for [Dgraph](https://dgraph.io), a
 fast, distributed graph database. It's meant to be a performant alternative to
 non-Go-based key-value stores like RocksDB.
 
+**Use [Discuss Issues](https://discuss.dgraph.io/c/issues/badger/37) for reporting issues about this repository.**
+
 ## Project Status [March 24, 2020]
 
 Badger is stable and is being used to serve data sets worth hundreds of
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/changes.sh b/changes.sh
new file mode 100755
index 0000000..e7cede9
--- /dev/null
+++ b/changes.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+set -e
+GHORG=${GHORG:-dgraph-io}
+GHREPO=${GHREPO:-badger}
+cat <<EOF
+This description was generated using this script:
+\`\`\`sh
+`cat $0`
+\`\`\`
+Invoked as:
+
+    `echo GHORG=${GHORG} GHREPO=${GHREPO} $(basename $0) ${@:1}`
+
+EOF
+git log --oneline --reverse ${@:1} \
+  | sed -E "s/^(\S{7}\s)//g" \
+  | sed -E "s/([\s|\(| ])#([0-9]+)/\1${GHORG}\/${GHREPO}#\2/g" \
diff --git a/levels.go b/levels.go
index 71f0a36..ecacdba 100644
--- a/levels.go
+++ b/levels.go
@@ -663,6 +663,9 @@ nextTable:
 			return tbl, errors.Wrapf(err, "Unable to open table: %q", fd.Name())
 		}
 		if builder.Empty() {
+			// Cleanup builder resources:
+			builder.Finish()
+			builder.Close()
 			continue
 		}
 		numBuilds++
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())
+}