Repository: dgraph-io/badger (go)

Issue:
WriteBatch.Flush panics if called twice
I understand that flushing twice doesn't make sense, since a batch can only be flushed once, but the second call should error with a helpful message instead of panicking.

```
$ go version
go version devel +152ffca82f Mon Jun 29 08:06:32 2020 +0000 linux/amd64
$ go list -m github.com/dgraph-io/badger/v2
github.com/dgraph-io/badger/v2 v2.0.3
$ cat main.go
package main

import (
	"io/ioutil"
	"os"

	"github.com/dgraph-io/badger/v2"
)

func main() {
	dir, err := ioutil.TempDir("", "badger-test")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(dir)
	db, err := badger.Open(badger.DefaultOptions(dir))
	if err != nil {
		panic(err)
	}
	defer db.Close()

	wb := db.NewWriteBatch()
	wb.Flush()
	wb.Flush()
}
$ go run main.go
badger 2020/06/30 18:36:15 INFO: All 0 tables opened in 0s
badger 2020/06/30 18:36:15 INFO: Got compaction priority: {level:0 score:1.73 dropPrefix:[]}
panic: send on closed channel

goroutine 1 [running]:
github.com/dgraph-io/badger/v2/y.(*Throttle).Do(0xc0000200c0, 0x6e7260, 0xc0000200c0)
	/home/mvdan/go/pkg/mod/github.com/dgraph-io/badger/v2@v2.0.3/y/y.go:260 +0xc9
github.com/dgraph-io/badger/v2.(*WriteBatch).commit(0xc000020180, 0x0, 0x0)
	/home/mvdan/go/pkg/mod/github.com/dgraph-io/badger/v2@v2.0.3/batch.go:141 +0x45
github.com/dgraph-io/badger/v2.(*WriteBatch).Flush(0xc000020180, 0x0, 0x0)
	/home/mvdan/go/pkg/mod/github.com/dgraph-io/badger/v2@v2.0.3/batch.go:155 +0x45
main.main()
	/home/mvdan/src/test/main.go:24 +0x1fa
exit status 2
```

Pull request:
Throttle.Do does not panic after Finish
Fixes #1394

I'm not sure about an error at `Do`. Maybe it will be ok to just return `nil`.

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


Patch:
diff --git a/badger/cmd/write_bench.go b/badger/cmd/write_bench.go
index 5bb61a4..7045fcc 100644
--- a/badger/cmd/write_bench.go
+++ b/badger/cmd/write_bench.go
@@ -41,19 +41,20 @@ var writeBenchCmd = &cobra.Command{
 	Use:   "write",
 	Short: "Writes random data to Badger to benchmark write speed.",
 	Long: `
-  This command writes random data to Badger to benchmark write speed. Useful for testing and
-  performance analysis.
-  `,
+	 This command writes random data to Badger to benchmark write speed. Useful for testing and
+	 performance analysis.
+	 `,
 	RunE: writeBench,
 }
 
 var (
-	keySz    int
-	valSz    int
-	numKeys  float64
-	force    bool
-	sorted   bool
-	showLogs bool
+	keySz      int
+	valSz      int
+	numKeys    float64
+	syncWrites bool
+	force      bool
+	sorted     bool
+	showLogs   bool
 
 	sizeWritten    uint64
 	entriesWritten uint64
@@ -96,27 +97,29 @@ func init() {
 	writeBenchCmd.Flags().IntVarP(&valSz, "val-size", "v", 128, "Size of value")
 	writeBenchCmd.Flags().Float64VarP(&numKeys, "keys-mil", "m", 10.0,
 		"Number of keys to add in millions")
+	writeBenchCmd.Flags().BoolVar(&syncWrites, "sync", true,
+		"If true, sync writes to disk.")
 	writeBenchCmd.Flags().BoolVarP(&force, "force-compact", "f", true,
 		"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,
+	writeBenchCmd.Flags().Int64VarP(&maxCacheSize, "max-cache", "C", 0, "Max size of cache")
+	writeBenchCmd.Flags().BoolVarP(&keepBlockIdxInCache, "keep-bidx", "b", false,
 		"Keep block indices in cache")
-	writeBenchCmd.Flags().BoolVarP(&keepBlocksInCache, "keep-blocks", "B", true,
+	writeBenchCmd.Flags().BoolVarP(&keepBlocksInCache, "keep-blocks", "B", false,
 		"Keep blocks in cache")
-	writeBenchCmd.Flags().Int64VarP(&maxBfCacheSize, "max-bf-cache", "c", 500<<20,
+	writeBenchCmd.Flags().Int64VarP(&maxBfCacheSize, "max-bf-cache", "c", 0,
 		"Maximum Bloom Filter Cache Size")
-	writeBenchCmd.Flags().Uint32Var(&vlogMaxEntries, "vlog-maxe", 10000, "Value log Max Entries")
+	writeBenchCmd.Flags().Uint32Var(&vlogMaxEntries, "vlog-maxe", 1000000, "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,
+	writeBenchCmd.Flags().BoolVar(&loadBloomsOnOpen, "load-blooms", true,
 		"Load Bloom filter on DB open.")
-	writeBenchCmd.Flags().BoolVar(&detectConflicts, "conficts", false,
+	writeBenchCmd.Flags().BoolVar(&detectConflicts, "conficts", true,
 		"If true, it badger will detect the conflicts")
 	writeBenchCmd.Flags().BoolVar(&compression, "compression", false,
 		"If true, badger will use ZSTD mode")
@@ -243,7 +246,7 @@ func writeBench(cmd *cobra.Command, args []string) error {
 	opt := badger.DefaultOptions(sstDir).
 		WithValueDir(vlogDir).
 		WithTruncate(truncate).
-		WithSyncWrites(false).
+		WithSyncWrites(syncWrites).
 		WithCompactL0OnClose(force).
 		WithValueThreshold(valueThreshold).
 		WithNumVersionsToKeep(numVersions).
diff --git a/batch.go b/batch.go
index f74c2f5..4ec54ca 100644
--- a/batch.go
+++ b/batch.go
@@ -34,6 +34,7 @@ type WriteBatch struct {
 
 	isManaged bool
 	commitTs  uint64
+	finished  bool
 }
 
 // NewWriteBatch creates a new WriteBatch. This provides a way to conveniently do a lot of writes,
@@ -72,6 +73,9 @@ func (wb *WriteBatch) SetMaxPendingTxns(max int) {
 //
 // Note that any committed writes would still go through despite calling Cancel.
 func (wb *WriteBatch) Cancel() {
+	wb.Lock()
+	defer wb.Unlock()
+	wb.finished = true
 	if err := wb.throttle.Finish(); err != nil {
 		wb.db.opt.Errorf("WatchBatch.Cancel error while finishing: %v", err)
 	}
@@ -180,6 +184,9 @@ func (wb *WriteBatch) commit() error {
 	if wb.err != nil {
 		return wb.err
 	}
+	if wb.finished {
+		return y.ErrCommitAfterFinish
+	}
 	if err := wb.throttle.Do(); err != nil {
 		wb.err = err
 		return wb.err
@@ -194,8 +201,12 @@ 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()
+	err := wb.commit()
+	if err != nil {
+		wb.Unlock()
+		return err
+	}
+	wb.finished = true
 	wb.txn.Discard()
 	wb.Unlock()
 
diff --git a/y/y.go b/y/y.go
index 5e6dd5b..93b47c7 100644
--- a/y/y.go
+++ b/y/y.go
@@ -41,6 +41,10 @@ var (
 	// compression algorithm is being used for compression. ZSTD cannot work
 	// without CGO.
 	ErrZstdCgo = errors.New("zstd compression requires building badger with cgo enabled")
+
+	// ErrCommitAfterFinish indicates that write batch commit was called after
+	// finish
+	ErrCommitAfterFinish = errors.New("Batch commit not permitted after finish")
 )
 
 const (


Test patch:
diff --git a/batch_test.go b/batch_test.go
index 8ff14a5..4a970d4 100644
--- a/batch_test.go
+++ b/batch_test.go
@@ -21,6 +21,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/dgraph-io/badger/v2/y"
+
 	"github.com/stretchr/testify/require"
 )
 
@@ -123,3 +125,22 @@ func TestEmptyWriteBatch(t *testing.T) {
 		})
 	})
 }
+
+// This test ensures we don't panic during flush.
+// See issue: https://github.com/dgraph-io/badger/issues/1394
+func TestFlushPanic(t *testing.T) {
+	t.Run("flush after flush", func(t *testing.T) {
+		runBadgerTest(t, nil, func(t *testing.T, db *DB) {
+			wb := db.NewWriteBatch()
+			wb.Flush()
+			require.Error(t, y.ErrCommitAfterFinish, wb.Flush())
+		})
+	})
+	t.Run("flush after cancel", func(t *testing.T) {
+		runBadgerTest(t, nil, func(t *testing.T, db *DB) {
+			wb := db.NewWriteBatch()
+			wb.Cancel()
+			require.Error(t, y.ErrCommitAfterFinish, wb.Flush())
+		})
+	})
+}