Merge ManagedDB back into DB (#568)

* Remove the extra ManagedDB struct, which was causing maintenance headaches, due to the way Go handles polymorphism. Instead just expose the ManagedTxns option, and use it directly.

* Bring back the OpenManaged function for convenience.

db.go

@@ -242,7 +242,7 @@ func Open(opt Options) (db *DB, err error) {
 	}()
 
 	orc := &oracle{
-		isManaged:  opt.managedTxns,
+		isManaged:  opt.ManagedTxns,
 		nextCommit: 1,
 		commits:    make(map[uint64]uint64),
 		readMark:   y.WaterMark{},
@@ -1069,7 +1069,13 @@ func (seq *Sequence) updateLease() error {
 // available, in the database. Sequence can be used to get a list of monotonically increasing
 // integers. Multiple sequences can be created by providing different keys. Bandwidth sets the
 // size of the lease, determining how many Next() requests can be served from memory.
+//
+// GetSequence is not supported on ManagedDB. Calling this would result in a panic.
 func (db *DB) GetSequence(key []byte, bandwidth uint64) (*Sequence, error) {
+	if db.opt.ManagedTxns {
+		panic("Cannot use GetSequence with ManagedTxns=true.")
+	}
+
 	switch {
 	case len(key) == 0:
 		return nil, ErrEmptyKey

db_test.go

@@ -290,7 +290,8 @@ func TestForceCompactL0(t *testing.T) {
 
 	opts := getTestOptions(dir)
 	opts.ValueLogFileSize = 15 << 20
-	db, err := OpenManaged(opts)
+	opts.ManagedTxns = true
+	db, err := Open(opts)
 	require.NoError(t, err)
 
 	data := func(i int) []byte {
@@ -310,7 +311,8 @@ func TestForceCompactL0(t *testing.T) {
 	}
 	db.Close()
 
-	db, err = OpenManaged(opts)
+	opts.ManagedTxns = true
+	db, err = Open(opts)
 	defer db.Close()
 	require.Equal(t, len(db.lc.levels[0].tables), 0)
 }

managed_db.go

@@ -26,36 +26,14 @@ import (
 	"github.com/pkg/errors"
 )
 
-// ManagedDB allows end users to manage the transactions themselves. Transaction
-// start and commit timestamps are set by end-user.
+// OpenManaged returns a new DB, which allows more control over setting
+// transaction timestamps, by setting ManagedTxns=true.
 //
 // This is only useful for databases built on top of Badger (like Dgraph), and
 // can be ignored by most users.
-//
-// WARNING: This is an experimental feature and may be changed significantly in
-// a future release. So please proceed with caution.
-type ManagedDB struct {
-	*DB
-}
-
-// OpenManaged returns a new ManagedDB, which allows more control over setting
-// transaction timestamps.
-//
-// This is only useful for databases built on top of Badger (like Dgraph), and
-// can be ignored by most users.
-func OpenManaged(opts Options) (*ManagedDB, error) {
-	opts.managedTxns = true
-	db, err := Open(opts)
-	if err != nil {
-		return nil, err
-	}
-	return &ManagedDB{db}, nil
-}
-
-// NewTransaction overrides DB.NewTransaction() and panics when invoked. Use
-// NewTransactionAt() instead.
-func (db *ManagedDB) NewTransaction(update bool) {
-	panic("Cannot use NewTransaction() for ManagedDB. Use NewTransactionAt() instead.")
+func OpenManaged(opts Options) (*DB, error) {
+	opts.ManagedTxns = true
+	return Open(opts)
 }
 
 // NewTransactionAt follows the same logic as DB.NewTransaction(), but uses the
@@ -63,35 +41,35 @@ func (db *ManagedDB) NewTransaction(update bool) {
 //
 // This is only useful for databases built on top of Badger (like Dgraph), and
 // can be ignored by most users.
-func (db *ManagedDB) NewTransactionAt(readTs uint64, update bool) *Txn {
-	txn := db.DB.NewTransaction(update)
+func (db *DB) NewTransactionAt(readTs uint64, update bool) *Txn {
+	if !db.opt.ManagedTxns {
+		panic("Cannot use NewTransactionAt with ManagedTxns=false. Use NewTransaction instead.")
+	}
+	txn := db.NewTransaction(update)
 	txn.readTs = readTs
 	return txn
 }
 
 // CommitAt commits the transaction, following the same logic as Commit(), but
-// at the given commit timestamp. This will panic if not used with ManagedDB.
+// at the given commit timestamp. This will panic if not used with managed transactions.
 //
 // This is only useful for databases built on top of Badger (like Dgraph), and
 // can be ignored by most users.
 func (txn *Txn) CommitAt(commitTs uint64, callback func(error)) error {
-	if !txn.db.opt.managedTxns {
-		return ErrManagedTxn
+	if !txn.db.opt.ManagedTxns {
+		panic("Cannot use CommitAt with ManagedTxns=false. Use Commit instead.")
 	}
 	txn.commitTs = commitTs
 	return txn.Commit(callback)
 }
 
-// GetSequence is not supported on ManagedDB. Calling this would result
-// in a panic.
-func (db *ManagedDB) GetSequence(_ []byte, _ uint64) (*Sequence, error) {
-	panic("Cannot use GetSequence for ManagedDB.")
-}
-
 // SetDiscardTs sets a timestamp at or below which, any invalid or deleted
 // versions can be discarded from the LSM tree, and thence from the value log to
-// reclaim disk space.
-func (db *ManagedDB) SetDiscardTs(ts uint64) {
+// reclaim disk space. Can only be used with managed transactions.
+func (db *DB) SetDiscardTs(ts uint64) {
+	if !db.opt.ManagedTxns {
+		panic("Cannot use SetDiscardTs with ManagedTxns=false.")
+	}
 	db.orc.setDiscardTs(ts)
 }
 
@@ -106,7 +84,12 @@ var errDone = errors.New("Done deleting keys")
 // - Iterate over the KVs in Level 0, and run deletes on them via transactions.
 // - The deletions are done at the same timestamp as the latest version of the
 // key. Thus, we could write the keys back at the same timestamp as before.
-func (db *ManagedDB) DropAll() error {
+//
+// DropAll is only available with managed transactions.
+func (db *DB) DropAll() error {
+	if !db.opt.ManagedTxns {
+		panic("DropAll is only available with ManagedTxns=true.")
+	}
 	// Stop accepting new writes.
 	atomic.StoreInt32(&db.blockWrites, 1)
 

managed_db_test.go

@@ -24,8 +24,8 @@ func val(large bool) []byte {
 	return buf
 }
 
-func numKeys(mdb *ManagedDB, readTs uint64) int {
-	txn := mdb.NewTransactionAt(readTs, false)
+func numKeys(db *DB, readTs uint64) int {
+	txn := db.NewTransactionAt(readTs, false)
 	defer txn.Discard()
 
 	itr := txn.NewIterator(DefaultIteratorOptions)
@@ -43,11 +43,12 @@ func TestDropAll(t *testing.T) {
 	require.NoError(t, err)
 	defer os.RemoveAll(dir)
 	opts := getTestOptions(dir)
-	mdb, err := OpenManaged(opts)
+	opts.ManagedTxns = true
+	db, err := Open(opts)
 	require.NoError(t, err)
 
 	N := uint64(10000)
-	populate := func(db *ManagedDB, start uint64) {
+	populate := func(db *DB, start uint64) {
 		var wg sync.WaitGroup
 		for i := start; i < start+N; i++ {
 			wg.Add(1)
@@ -61,22 +62,23 @@ func TestDropAll(t *testing.T) {
 		wg.Wait()
 	}
 
-	populate(mdb, 1)
-	require.Equal(t, int(N), numKeys(mdb, math.MaxUint64))
+	populate(db, 1)
+	require.Equal(t, int(N), numKeys(db, math.MaxUint64))
 
-	require.NoError(t, mdb.DropAll())
-	require.Equal(t, 0, numKeys(mdb, math.MaxUint64))
+	require.NoError(t, db.DropAll())
+	require.Equal(t, 0, numKeys(db, math.MaxUint64))
 
 	// Check that we can still write to mdb, and using the same timestamps as before.
-	populate(mdb, 1)
-	require.Equal(t, int(N), numKeys(mdb, math.MaxUint64))
-	mdb.Close()
+	populate(db, 1)
+	require.Equal(t, int(N), numKeys(db, math.MaxUint64))
+	db.Close()
 
 	// Ensure that value log is correctly replayed, that we are preserving badgerHead.
-	mdb2, err := OpenManaged(opts)
+	opts.ManagedTxns = true
+	db2, err := Open(opts)
 	require.NoError(t, err)
-	require.Equal(t, int(N), numKeys(mdb2, math.MaxUint64))
-	mdb2.Close()
+	require.Equal(t, int(N), numKeys(db2, math.MaxUint64))
+	db2.Close()
 }
 
 func TestDropAllRace(t *testing.T) {
@@ -84,7 +86,8 @@ func TestDropAllRace(t *testing.T) {
 	require.NoError(t, err)
 	defer os.RemoveAll(dir)
 	opts := getTestOptions(dir)
-	mdb, err := OpenManaged(opts)
+	opts.ManagedTxns = true
+	db, err := Open(opts)
 	require.NoError(t, err)
 
 	N := 10000
@@ -100,7 +103,7 @@ func TestDropAllRace(t *testing.T) {
 			select {
 			case <-ticker.C:
 				i++
-				txn := mdb.NewTransactionAt(math.MaxUint64, true)
+				txn := db.NewTransactionAt(math.MaxUint64, true)
 				require.NoError(t, txn.Set([]byte(key("key", i)), val(false)))
 				_ = txn.CommitAt(uint64(i), nil)
 			case <-closer.HasBeenClosed():
@@ -112,7 +115,7 @@ func TestDropAllRace(t *testing.T) {
 	var wg sync.WaitGroup
 	for i := 1; i <= N; i++ {
 		wg.Add(1)
-		txn := mdb.NewTransactionAt(math.MaxUint64, true)
+		txn := db.NewTransactionAt(math.MaxUint64, true)
 		require.NoError(t, txn.Set([]byte(key("key", i)), val(false)))
 		require.NoError(t, txn.CommitAt(uint64(i), func(err error) {
 			require.NoError(t, err)
@@ -121,14 +124,14 @@ func TestDropAllRace(t *testing.T) {
 	}
 	wg.Wait()
 
-	before := numKeys(mdb, math.MaxUint64)
+	before := numKeys(db, math.MaxUint64)
 	require.True(t, before > N)
 
-	require.NoError(t, mdb.DropAll())
+	require.NoError(t, db.DropAll())
 	closer.SignalAndWait()
 
-	after := numKeys(mdb, math.MaxUint64)
+	after := numKeys(db, math.MaxUint64)
 	t.Logf("Before: %d. After dropall: %d\n", before, after)
 	require.True(t, after < before)
-	mdb.Close()
+	db.Close()
 }

options.go

@@ -83,9 +83,10 @@ type Options struct {
 	// Number of compaction workers to run concurrently.
 	NumCompactors int
 
-	// Transaction start and commit timestamps are manaVgedTxns by end-user. This
-	// is a private option used by ManagedDB.
-	managedTxns bool
+	// Transaction start and commit timestamps are managed by end-user.
+	// This is only useful for databases built on top of Badger (like Dgraph).
+	// Not recommended for most users.
+	ManagedTxns bool
 
 	// 4. Flags for testing purposes
 	// ------------------------------

transaction.go

@@ -474,8 +474,8 @@ func (txn *Txn) Discard() {
 // If error is nil, the transaction is successfully committed. In case of a non-nil error, the LSM
 // tree won't be updated, so there's no need for any rollback.
 func (txn *Txn) Commit(callback func(error)) error {
-	if txn.commitTs == 0 && txn.db.opt.managedTxns {
-		return ErrManagedTxn
+	if txn.commitTs == 0 && txn.db.opt.ManagedTxns {
+		panic("Commit cannot be called with ManagedTxns=true. Use CommitAt.")
 	}
 	if txn.discarded {
 		return ErrDiscardedTxn
@@ -577,11 +577,14 @@ func (db *DB) NewTransaction(update bool) *Txn {
 
 // View executes a function creating and managing a read-only transaction for the user. Error
 // returned by the function is relayed by the View method.
+// If View is used with managed transactions, it would assume a read timestamp of MaxUint64.
 func (db *DB) View(fn func(txn *Txn) error) error {
-	if db.opt.managedTxns {
-		return ErrManagedTxn
+	var txn *Txn
+	if db.opt.ManagedTxns {
+		txn = db.NewTransactionAt(math.MaxUint64, false)
+	} else {
+		txn = db.NewTransaction(false)
 	}
-	txn := db.NewTransaction(false)
 	defer txn.Discard()
 
 	return fn(txn)
@@ -589,9 +592,10 @@ func (db *DB) View(fn func(txn *Txn) error) error {
 
 // Update executes a function, creating and managing a read-write transaction
 // for the user. Error returned by the function is relayed by the Update method.
+// Update cannot be used with managed transactions.
 func (db *DB) Update(fn func(txn *Txn) error) error {
-	if db.opt.managedTxns {
-		return ErrManagedTxn
+	if db.opt.ManagedTxns {
+		panic("Update can only be used with ManagedTxns=false.")
 	}
 	txn := db.NewTransaction(true)
 	defer txn.Discard()

transaction_test.go

@@ -51,7 +51,7 @@ func TestTxnSimple(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, []byte("val=8"), val)
 
-		require.Equal(t, ErrManagedTxn, txn.CommitAt(100, nil))
+		require.Panics(t, func() { txn.CommitAt(100, nil) })
 		require.NoError(t, txn.Commit(nil))
 	})
 }
@@ -709,9 +709,10 @@ func TestManagedDB(t *testing.T) {
 	defer os.RemoveAll(dir)
 
 	opt := getTestOptions(dir)
-	kv, err := OpenManaged(opt)
+	opt.ManagedTxns = true
+	db, err := Open(opt)
 	require.NoError(t, err)
-	defer kv.Close()
+	defer db.Close()
 
 	key := func(i int) []byte {
 		return []byte(fmt.Sprintf("key-%02d", i))
@@ -721,25 +722,23 @@ func TestManagedDB(t *testing.T) {
 		return []byte(fmt.Sprintf("val-%d", i))
 	}
 
-	// Don't allow these APIs in ManagedDB
-	require.Panics(t, func() { kv.NewTransaction(false) })
+	require.Panics(t, func() {
+		db.Update(func(tx *Txn) error { return nil })
+	})
 
-	err = kv.Update(func(tx *Txn) error { return nil })
-	require.Equal(t, ErrManagedTxn, err)
-
-	err = kv.View(func(tx *Txn) error { return nil })
-	require.Equal(t, ErrManagedTxn, err)
+	err = db.View(func(tx *Txn) error { return nil })
+	require.NoError(t, err)
 
 	// Write data at t=3.
-	txn := kv.NewTransactionAt(3, true)
+	txn := db.NewTransactionAt(3, true)
 	for i := 0; i <= 3; i++ {
 		require.NoError(t, txn.Set(key(i), val(i)))
 	}
-	require.Equal(t, ErrManagedTxn, txn.Commit(nil))
+	require.Panics(t, func() { txn.Commit(nil) })
 	require.NoError(t, txn.CommitAt(3, nil))
 
 	// Read data at t=2.
-	txn = kv.NewTransactionAt(2, false)
+	txn = db.NewTransactionAt(2, false)
 	for i := 0; i <= 3; i++ {
 		_, err := txn.Get(key(i))
 		require.Equal(t, ErrKeyNotFound, err)
@@ -747,7 +746,7 @@ func TestManagedDB(t *testing.T) {
 	txn.Discard()
 
 	// Read data at t=3.
-	txn = kv.NewTransactionAt(3, false)
+	txn = db.NewTransactionAt(3, false)
 	for i := 0; i <= 3; i++ {
 		item, err := txn.Get(key(i))
 		require.NoError(t, err)
@@ -759,7 +758,7 @@ func TestManagedDB(t *testing.T) {
 	txn.Discard()
 
 	// Write data at t=7.
-	txn = kv.NewTransactionAt(6, true)
+	txn = db.NewTransactionAt(6, true)
 	for i := 0; i <= 7; i++ {
 		_, err := txn.Get(key(i))
 		if err == nil {
@@ -770,7 +769,7 @@ func TestManagedDB(t *testing.T) {
 	require.NoError(t, txn.CommitAt(7, nil))
 
 	// Read data at t=9.
-	txn = kv.NewTransactionAt(9, false)
+	txn = db.NewTransactionAt(9, false)
 	for i := 0; i <= 9; i++ {
 		item, err := txn.Get(key(i))
 		if i <= 7 {

value.go

@@ -467,7 +467,7 @@ func (vlog *valueLog) rewrite(f *logFile, tr trace.Trace) error {
 	return nil
 }
 
-func (vlog *valueLog) deleteMoveKeysFor(fid uint32, tr trace.Trace) {
+func (vlog *valueLog) deleteMoveKeysFor(fid uint32, tr trace.Trace) error {
 	db := vlog.kv
 	var result []*Entry
 	var count, pointers uint64
@@ -498,7 +498,7 @@ func (vlog *valueLog) deleteMoveKeysFor(fid uint32, tr trace.Trace) {
 	if err != nil {
 		tr.LazyPrintf("Got error while iterating move keys: %v", err)
 		tr.SetError()
-		return
+		return err
 	}
 	tr.LazyPrintf("Num total move keys: %d. Num pointers: %d", count, pointers)
 	tr.LazyPrintf("Number of invalid move keys found: %d", len(result))
@@ -516,12 +516,12 @@ func (vlog *valueLog) deleteMoveKeysFor(fid uint32, tr trace.Trace) {
 			}
 			tr.LazyPrintf("Error while doing batchSet: %v", err)
 			tr.SetError()
-			return
+			return err
 		}
 		i += batchSize
 	}
 	tr.LazyPrintf("Move keys deletion done.")
-	return
+	return nil
 }
 
 func (vlog *valueLog) incrIteratorCount() {
@@ -1215,8 +1215,7 @@ func (vlog *valueLog) runGC(discardRatio float64, head valuePointer) error {
 			tried[lf.fid] = true
 			err = vlog.doRunGC(lf, discardRatio, tr)
 			if err == nil {
-				vlog.deleteMoveKeysFor(lf.fid, tr)
-				return nil
+				return vlog.deleteMoveKeysFor(lf.fid, tr)
 			}
 		}
 		return err

value_test.go

@@ -21,9 +21,11 @@ import (
 	"io/ioutil"
 	"math/rand"
 	"os"
+	"sync"
 	"testing"
 
 	"github.com/dgraph-io/badger/y"
+	humanize "github.com/dustin/go-humanize"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/trace"
 )
@@ -84,6 +86,65 @@ func TestValueBasic(t *testing.T) {
 
 }
 
+func TestValueGCManaged(t *testing.T) {
+	dir, err := ioutil.TempDir("", "badger")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+
+	N := 10000
+	opt := getTestOptions(dir)
+	opt.ValueLogMaxEntries = uint32(N / 10)
+	opt.ManagedTxns = true
+	db, err := Open(opt)
+	require.NoError(t, err)
+	defer db.Close()
+
+	var ts uint64
+	newTs := func() uint64 {
+		ts += 1
+		return ts
+	}
+
+	sz := 64 << 10
+	var wg sync.WaitGroup
+	for i := 0; i < N; i++ {
+		v := make([]byte, sz)
+		rand.Read(v[:rand.Intn(sz)])
+
+		wg.Add(1)
+		txn := db.NewTransactionAt(newTs(), true)
+		require.NoError(t, txn.Set([]byte(fmt.Sprintf("key%d", i)), v))
+		require.NoError(t, txn.CommitAt(newTs(), func(err error) {
+			wg.Done()
+			require.NoError(t, err)
+		}))
+	}
+
+	for i := 0; i < N; i++ {
+		wg.Add(1)
+		txn := db.NewTransactionAt(newTs(), true)
+		require.NoError(t, txn.Delete([]byte(fmt.Sprintf("key%d", i))))
+		require.NoError(t, txn.CommitAt(newTs(), func(err error) {
+			wg.Done()
+			require.NoError(t, err)
+		}))
+	}
+	wg.Wait()
+	files, err := ioutil.ReadDir(dir)
+	require.NoError(t, err)
+	for _, fi := range files {
+		t.Logf("File: %s. Size: %s\n", fi.Name(), humanize.Bytes(uint64(fi.Size())))
+	}
+
+	for i := 0; i < 100; i++ {
+		// Try at max 100 times to GC even a single value log file.
+		if err := db.RunValueLogGC(0.0001); err == nil {
+			return // Done
+		}
+	}
+	require.Fail(t, "Unable to GC even a single value log file.")
+}
+
 func TestValueGC(t *testing.T) {
 	dir, err := ioutil.TempDir("", "badger")
 	require.NoError(t, err)