Fix race condition in DropAll (#656)

- Fix a race condition in DropAll, caused if concurrent writes are going on.
- Use `blockWrites` during DB.Close as well, so users don't see channel panics when doing writes after DB.Close.
- Port tests over for Stream framework from Dgraph.

Changes:
* Starting to add a test for stream
* Add a test for Stream. Allow repeated calls to Orchestrate on the same Stream object.
* Fix various race conditions with concurrent writes during DropAll.
* Use blockWrites during Close as well, to avoid the chan panics users see when doing writes after Close.
* Add license for stream_test.go
* 100 chars

backup.go

@@ -37,6 +37,7 @@ import (
 // This can be used to backup the data in a database at a given point in time.
 func (db *DB) Backup(w io.Writer, since uint64) (uint64, error) {
 	stream := db.NewStream()
+	stream.LogPrefix = "DB.Backup"
 	stream.KeyToList = func(key []byte, itr *Iterator) (*pb.KVList, error) {
 		list := &pb.KVList{}
 		for ; itr.Valid(); itr.Next() {
@@ -105,7 +106,7 @@ func (db *DB) Backup(w io.Writer, since uint64) (uint64, error) {
 		return nil
 	}
 
-	if err := stream.Orchestrate(context.Background(), 8, "DB.Backup"); err != nil {
+	if err := stream.Orchestrate(context.Background()); err != nil {
 		return 0, err
 	}
 	return maxVersion, nil

db.go

@@ -319,6 +319,8 @@ func Open(opt Options) (db *DB, err error) {
 // cause panic.
 func (db *DB) Close() (err error) {
 	db.elog.Printf("Closing database")
+	atomic.StoreInt32(&db.blockWrites, 1)
+
 	// Stop value GC first.
 	db.closers.valueGC.SignalAndWait()
 
@@ -573,11 +575,6 @@ func (db *DB) writeRequests(reqs []*request) error {
 			r.Wg.Done()
 		}
 	}
-	if atomic.LoadInt32(&db.blockWrites) > 0 {
-		done(ErrBlockedWrites)
-		return ErrBlockedWrites
-	}
-
 	db.elog.Printf("writeRequests called. Writing to value log")
 
 	err := db.vlog.write(reqs)
@@ -1243,23 +1240,30 @@ func (db *DB) Flatten(workers int) error {
 // any reads while DropAll is going on, otherwise they may result in panics. Ideally, both reads and
 // writes are paused before running DropAll, and resumed after it is finished.
 func (db *DB) DropAll() error {
+	if db.opt.ReadOnly {
+		panic("Attempting to drop data in read-only mode.")
+	}
 	Infof("DropAll called. Blocking writes...")
 	// Stop accepting new writes.
 	atomic.StoreInt32(&db.blockWrites, 1)
 
-	// Wait for writeCh to reach size of zero. This is not ideal, but a very
-	// simple way to allow writeCh to flush out, before we proceed.
-	tick := time.NewTicker(100 * time.Millisecond)
-	for range tick.C {
-		if len(db.writeCh) == 0 {
-			break
-		}
-	}
-	tick.Stop()
-	Infof("All previous writes done. Stopping compactions...")
+	// Make all pending writes finish. The following will also close writeCh.
+	db.closers.writes.SignalAndWait()
+	Infof("Writes flushed. Stopping compactions now...")
 
+	// Stop all compactions.
 	db.stopCompactions()
-	defer db.startCompactions()
+	defer func() {
+		Infof("Resuming writes")
+		db.startCompactions()
+
+		db.writeCh = make(chan *request, kvWriteChCapacity)
+		db.closers.writes = y.NewCloser(1)
+		go db.doWrites(db.closers.writes)
+
+		// Resume writes.
+		atomic.StoreInt32(&db.blockWrites, 0)
+	}()
 	Infof("Compactions stopped. Dropping all SSTables...")
 
 	// Remove inmemory tables. Calling DecrRef for safety. Not sure if they're absolutely needed.
@@ -1281,11 +1285,6 @@ func (db *DB) DropAll() error {
 		return err
 	}
 	db.vhead = valuePointer{} // Zero it out.
-	Infof("Deleted %d value log files. Resuming operations...\n", num)
-
-	// Resume writes.
-	atomic.StoreInt32(&db.blockWrites, 0)
-
-	Infof("DropAll done")
+	Infof("Deleted %d value log files. DropAll done.\n", num)
 	return nil
 }

errors.go

@@ -97,9 +97,9 @@ var (
 
 	// ErrTruncateNeeded is returned when the value log gets corrupt, and requires truncation of
 	// corrupt data to allow Badger to run properly.
-	ErrTruncateNeeded = errors.New("Value log truncate required to run DB. This might result in data loss.")
+	ErrTruncateNeeded = errors.New("Value log truncate required to run DB. This might result in data loss")
 
 	// ErrBlockedWrites is returned if the user called DropAll. During the process of dropping all
 	// data from Badger, we stop accepting new writes, by returning this error.
-	ErrBlockedWrites = errors.New("Writes are blocked possibly due to DropAll")
+	ErrBlockedWrites = errors.New("Writes are blocked, possibly due to DropAll or Close")
 )

managed_db_test.go

@@ -134,6 +134,59 @@ func TestDropAll(t *testing.T) {
 	db2.Close()
 }
 
+func TestDropReadOnly(t *testing.T) {
+	dir, err := ioutil.TempDir("", "badger")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+	opts := getTestOptions(dir)
+	opts.ValueLogFileSize = 5 << 20
+	db, err := Open(opts)
+	require.NoError(t, err)
+	N := uint64(1000)
+	populate := func(db *DB) {
+		writer := db.NewWriteBatch()
+		for i := uint64(0); i < N; i++ {
+			require.NoError(t, writer.Set([]byte(key("key", int(i))), val(true), 0))
+		}
+		require.NoError(t, writer.Flush())
+	}
+
+	populate(db)
+	require.Equal(t, int(N), numKeys(db))
+	require.NoError(t, db.Close())
+
+	opts.ReadOnly = true
+	db2, err := Open(opts)
+	require.NoError(t, err)
+	require.Panics(t, func() { db2.DropAll() })
+}
+
+func TestWriteAfterClose(t *testing.T) {
+	dir, err := ioutil.TempDir(".", "badger-test")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+	opts := getTestOptions(dir)
+	opts.ValueLogFileSize = 5 << 20
+	db, err := Open(opts)
+	require.NoError(t, err)
+	N := uint64(1000)
+	populate := func(db *DB) {
+		writer := db.NewWriteBatch()
+		for i := uint64(0); i < N; i++ {
+			require.NoError(t, writer.Set([]byte(key("key", int(i))), val(true), 0))
+		}
+		require.NoError(t, writer.Flush())
+	}
+
+	populate(db)
+	require.Equal(t, int(N), numKeys(db))
+	require.NoError(t, db.Close())
+	err = db.Update(func(txn *Txn) error {
+		return txn.Set([]byte("a"), []byte("b"))
+	})
+	require.Equal(t, ErrBlockedWrites, err)
+}
+
 func TestDropAllRace(t *testing.T) {
 	dir, err := ioutil.TempDir("", "badger")
 	require.NoError(t, err)

stream.go

@@ -30,15 +30,20 @@ import (
 const pageSize = 4 << 20 // 4MB
 
 // Stream provides a framework to concurrently iterate over a snapshot of Badger, pick up
-// key-values, batch them up and call Send. Because, Stream does concurrent iteration over many
-// smaller key ranges, it does NOT send keys in a strictly lexicographical order.
+// 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
+// order, use Iterator.
 type Stream struct {
 	// Prefix to only iterate over certain range of keys. If set to nil (default), Stream would
 	// iterate over the entire DB.
 	Prefix []byte
 
-	readTs uint64
-	db     *DB
+	// Number of goroutines to use for iterating over key ranges. Defaults to 16.
+	NumGo int
+
+	// Badger would produce log entries in Infof to indicate the progress of Stream. LogPrefix can
+	// be used to help differentiate them from other activities. Default is "Badger.Stream".
+	LogPrefix string
 
 	// ChooseKey is invoked each time a new key is encountered. Note that this is not called
 	// on every version of the value, only the first encountered version (i.e. the highest version
@@ -61,6 +66,8 @@ type Stream struct {
 	// single goroutine, i.e. logic within Send method can expect single threaded execution.
 	Send func(*pb.KVList) error
 
+	readTs  uint64
+	db      *DB
 	rangeCh chan keyRange
 	kvChan  chan *pb.KVList
 }
@@ -194,7 +201,7 @@ func (st *Stream) produceKVs(ctx context.Context) error {
 	}
 }
 
-func (st *Stream) streamKVs(ctx context.Context, logPrefix string) error {
+func (st *Stream) streamKVs(ctx context.Context) error {
 	var count int
 	var bytesSent uint64
 	t := time.NewTicker(time.Second)
@@ -223,7 +230,7 @@ func (st *Stream) streamKVs(ctx context.Context, logPrefix string) error {
 			return err
 		}
 		Infof("%s Created batch of size: %s in %s.\n",
-			logPrefix, humanize.Bytes(sz), time.Since(t))
+			st.LogPrefix, humanize.Bytes(sz), time.Since(t))
 		return nil
 	}
 
@@ -241,8 +248,8 @@ outer:
 				continue
 			}
 			speed := bytesSent / durSec
-			Infof("%s Time elapsed: %s, bytes sent: %s, speed: %s/sec\n",
-				logPrefix, y.FixedDuration(dur), humanize.Bytes(bytesSent), humanize.Bytes(speed))
+			Infof("%s Time elapsed: %s, bytes sent: %s, speed: %s/sec\n", st.LogPrefix,
+				y.FixedDuration(dur), humanize.Bytes(bytesSent), humanize.Bytes(speed))
 
 		case kvs, ok := <-st.kvChan:
 			if !ok {
@@ -256,23 +263,23 @@ outer:
 		}
 	}
 
-	Infof("%s Sent %d keys\n", logPrefix, count)
+	Infof("%s Sent %d keys\n", st.LogPrefix, count)
 	return nil
 }
 
-// Orchestrate runs Stream. It picks up ranges from the SSTables, then runs numGo number of
-// goroutines to iterate over these ranges and batch up KVs in lists. It then runs a single
+// Orchestrate runs Stream. It picks up ranges from the SSTables, then runs NumGo number of
+// goroutines to iterate over these ranges and batch up KVs in lists. It concurrently runs a single
 // goroutine to pick these lists, batch them up further and send to Output.Send. Orchestrate also
-// spits logs out to Infof, using the logPrefix string provided.  Note that all calls to
-// Output.Send are serial. In case any of these steps encounter an error, Orchestrate would stop
-// execution and return that error. Orchestrate should only be called once on the same Stream
-// object.
-func (st *Stream) Orchestrate(ctx context.Context, numGo int, logPrefix string) error {
+// spits logs out to Infof, using provided LogPrefix. Note that all calls to Output.Send
+// are serial. In case any of these steps encounter an error, Orchestrate would stop execution and
+// return that error. Orchestrate can be called multiple times, but in serial order.
+func (st *Stream) Orchestrate(ctx context.Context) error {
 	st.rangeCh = make(chan keyRange, 3) // Contains keys for posting lists.
 
 	// kvChan should only have a small capacity to ensure that we don't buffer up too much data if
-	// sending is slow. So, setting this to 3.
-	st.kvChan = make(chan *pb.KVList, 3)
+	// sending is slow. Page size is set to 4MB, which is used to lazily cap the size of each
+	// KVList. To get around 64MB buffer, we can set the channel size to 16.
+	st.kvChan = make(chan *pb.KVList, 16)
 
 	if st.KeyToList == nil {
 		st.KeyToList = st.ToList
@@ -283,7 +290,7 @@ func (st *Stream) Orchestrate(ctx context.Context, numGo int, logPrefix string)
 
 	errCh := make(chan error, 1) // Stores error by consumeKeys.
 	var wg sync.WaitGroup
-	for i := 0; i < numGo; i++ {
+	for i := 0; i < st.NumGo; i++ {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
@@ -301,7 +308,7 @@ func (st *Stream) Orchestrate(ctx context.Context, numGo int, logPrefix string)
 	kvErr := make(chan error, 1)
 	go func() {
 		// Picks up KV lists from kvChan, and sends them to Output.
-		kvErr <- st.streamKVs(ctx, logPrefix)
+		kvErr <- st.streamKVs(ctx)
 	}()
 	wg.Wait()        // Wait for produceKVs to be over.
 	close(st.kvChan) // Now we can close kvChan.
@@ -319,12 +326,16 @@ func (st *Stream) Orchestrate(ctx context.Context, numGo int, logPrefix string)
 	return nil
 }
 
+func (db *DB) newStream() *Stream {
+	return &Stream{db: db, NumGo: 16, LogPrefix: "Badger.Stream"}
+}
+
 // NewStream creates a new Stream.
 func (db *DB) NewStream() *Stream {
 	if db.opt.managedTxns {
 		panic("This API can not be called in managed mode.")
 	}
-	return &Stream{db: db}
+	return db.newStream()
 }
 
 // NewStreamAt creates a new Stream at a particular timestamp. Should only be used with managed DB.
@@ -332,5 +343,7 @@ func (db *DB) NewStreamAt(readTs uint64) *Stream {
 	if !db.opt.managedTxns {
 		panic("This API can only be called in managed mode.")
 	}
-	return &Stream{db: db, readTs: readTs}
+	stream := db.newStream()
+	stream.readTs = readTs
+	return stream
 }

stream_test.go

@@ -0,0 +1,169 @@
+/*
+ * Copyright 2018 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package badger
+
+import (
+	"context"
+	"fmt"
+	"io/ioutil"
+	"math"
+	"os"
+	"strconv"
+	"strings"
+	"testing"
+
+	bpb "github.com/dgraph-io/badger/pb"
+	"github.com/dgraph-io/badger/y"
+	"github.com/stretchr/testify/require"
+)
+
+func openManaged(dir string) (*DB, error) {
+	opt := DefaultOptions
+	opt.Dir = dir
+	opt.ValueDir = dir
+
+	return OpenManaged(opt)
+}
+
+func keyWithPrefix(prefix string, k int) []byte {
+	return []byte(fmt.Sprintf("%s-%d", prefix, k))
+}
+
+func keyToInt(k []byte) (string, int) {
+	splits := strings.Split(string(k), "-")
+	key, err := strconv.Atoi(splits[1])
+	y.Check(err)
+	return splits[0], key
+}
+
+func value(k int) []byte {
+	return []byte(fmt.Sprintf("%08d", k))
+}
+
+type collector struct {
+	kv []*bpb.KV
+}
+
+func (c *collector) Send(list *bpb.KVList) error {
+	c.kv = append(c.kv, list.Kv...)
+	return nil
+}
+
+var ctxb = context.Background()
+
+func TestStream(t *testing.T) {
+	dir, err := ioutil.TempDir("", "badger")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	db, err := openManaged(dir)
+	require.NoError(t, err)
+
+	var count int
+	for _, prefix := range []string{"p0", "p1", "p2"} {
+		txn := db.NewTransactionAt(math.MaxUint64, true)
+		for i := 1; i <= 100; i++ {
+			require.NoError(t, txn.Set(keyWithPrefix(prefix, i), value(i)))
+			count++
+		}
+		require.NoError(t, txn.CommitAt(5, nil))
+	}
+
+	stream := db.NewStreamAt(math.MaxUint64)
+	stream.LogPrefix = "Testing"
+	c := &collector{}
+	stream.Send = func(list *bpb.KVList) error {
+		return c.Send(list)
+	}
+
+	// Test case 1. Retrieve everything.
+	err = stream.Orchestrate(ctxb)
+	require.NoError(t, err)
+	require.Equal(t, 300, len(c.kv), "Expected 300. 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, 100, count, "Count mismatch for pred: %s", pred)
+	}
+
+	// Test case 2. Retrieve only 1 predicate.
+	stream.Prefix = []byte("p1")
+	c.kv = c.kv[:0]
+	err = stream.Orchestrate(ctxb)
+	require.NoError(t, err)
+	require.Equal(t, 100, len(c.kv), "Expected 100. 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, 1, len(m))
+	for pred, count := range m {
+		require.Equal(t, 100, count, "Count mismatch for pred: %s", pred)
+	}
+
+	// Test case 3. Retrieve select keys within the predicate.
+	c.kv = c.kv[:0]
+	stream.ChooseKey = func(item *Item) bool {
+		_, k := keyToInt(item.Key())
+		return k%2 == 0
+	}
+	err = stream.Orchestrate(ctxb)
+	require.NoError(t, err)
+	require.Equal(t, 50, len(c.kv), "Expected 50. 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, 1, len(m))
+	for pred, count := range m {
+		require.Equal(t, 50, count, "Count mismatch for pred: %s", pred)
+	}
+
+	// Test case 4. Retrieve select keys from all predicates.
+	c.kv = c.kv[:0]
+	stream.Prefix = []byte{}
+	err = stream.Orchestrate(ctxb)
+	require.NoError(t, err)
+	require.Equal(t, 150, len(c.kv), "Expected 150. 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, 50, count, "Count mismatch for pred: %s", pred)
+	}
+}