feat(Skiplist): Introduce a way to hand over skiplists to Badger (#1696)

In Dgraph, we already use Raft write-ahead log. Also, when we commit transactions, we update tens of thousands of keys in one go. To optimize this write path, this PR introduces a way to directly hand over Skiplist to Badger, short circuiting Badger's Value Log and WAL.

This feature allows Dgraph to generate Skiplists while processing mutations and just hand them over to Badger during commits. It also accepts a callback which can be run when Skiplist is written to disk. This is useful for determining when to create a snapshot in Dgraph.

backup_test.go

@@ -446,7 +446,7 @@ func TestBackupLoadIncremental(t *testing.T) {
 				if err := txn.SetEntry(entry); err != nil {
 					return err
 				}
-				updates[i] = bitDiscardEarlierVersions
+				updates[i] = BitDiscardEarlierVersions
 			}
 			return nil
 		})

db.go

@@ -763,16 +763,9 @@ var requestPool = sync.Pool{
 }
 
 func (db *DB) writeToLSM(b *request) error {
-	// We should check the length of b.Prts and b.Entries only when badger is not
-	// running in InMemory mode. In InMemory mode, we don't write anything to the
-	// value log and that's why the length of b.Ptrs will always be zero.
-	if !db.opt.InMemory && len(b.Ptrs) != len(b.Entries) {
-		return errors.Errorf("Ptrs and Entries don't match: %+v", b)
-	}
-
 	for i, entry := range b.Entries {
 		var err error
-		if entry.skipVlogAndSetThreshold(db.valueThreshold()) {
+		if db.opt.managedTxns || entry.skipVlogAndSetThreshold(db.valueThreshold()) {
 			// Will include deletion / tombstone case.
 			err = db.mt.Put(entry.Key,
 				y.ValueStruct{
@@ -818,10 +811,13 @@ func (db *DB) writeRequests(reqs []*request) error {
 		}
 	}
 	db.opt.Debugf("writeRequests called. Writing to value log")
-	err := db.vlog.write(reqs)
-	if err != nil {
-		done(err)
-		return err
+	if !db.opt.managedTxns {
+		// Don't do value log writes in managed mode.
+		err := db.vlog.write(reqs)
+		if err != nil {
+			done(err)
+			return err
+		}
 	}
 
 	db.opt.Debugf("Sending updates to subscribers")
@@ -834,6 +830,7 @@ func (db *DB) writeRequests(reqs []*request) error {
 		}
 		count += len(b.Entries)
 		var i uint64
+		var err error
 		for err = db.ensureRoomForWrite(); err == errNoRoom; err = db.ensureRoomForWrite() {
 			i++
 			if i%100 == 0 {
@@ -1010,16 +1007,61 @@ func (db *DB) ensureRoomForWrite() error {
 	}
 }
 
+func (db *DB) HandoverSkiplist(skl *skl.Skiplist, callback func()) error {
+	if !db.opt.managedTxns {
+		panic("Handover Skiplist is only available in managed mode.")
+	}
+	db.lock.Lock()
+	defer db.lock.Unlock()
+
+	// If we have some data in db.mt, we should push that first, so the ordering of writes is
+	// maintained.
+	if !db.mt.sl.Empty() {
+		sz := db.mt.sl.MemSize()
+		db.opt.Infof("Handover found %d B data in current memtable. Pushing to flushChan.", sz)
+		var err error
+		select {
+		case db.flushChan <- flushTask{mt: db.mt}:
+			db.imm = append(db.imm, db.mt)
+			db.mt, err = db.newMemTable()
+			if err != nil {
+				return y.Wrapf(err, "cannot push current memtable")
+			}
+		default:
+			return errNoRoom
+		}
+	}
+
+	mt := &memTable{sl: skl}
+	select {
+	case db.flushChan <- flushTask{mt: mt, cb: callback}:
+		db.imm = append(db.imm, mt)
+		return nil
+	default:
+		return errNoRoom
+	}
+}
+
 func arenaSize(opt Options) int64 {
 	return opt.MemTableSize + opt.maxBatchSize + opt.maxBatchCount*int64(skl.MaxNodeSize)
 }
 
+func (db *DB) NewSkiplist() *skl.Skiplist {
+	return skl.NewSkiplist(arenaSize(db.opt))
+}
+
 // buildL0Table builds a new table from the memtable.
 func buildL0Table(ft flushTask, bopts table.Options) *table.Builder {
-	iter := ft.mt.sl.NewIterator()
+	var iter y.Iterator
+	if ft.itr != nil {
+		iter = ft.itr
+	} else {
+		iter = ft.mt.sl.NewUniIterator(false)
+	}
 	defer iter.Close()
+
 	b := table.NewTableBuilder(bopts)
-	for iter.SeekToFirst(); iter.Valid(); iter.Next() {
+	for iter.Rewind(); iter.Valid(); iter.Next() {
 		if len(ft.dropPrefixes) > 0 && hasAnyPrefixes(iter.Key(), ft.dropPrefixes) {
 			continue
 		}
@@ -1035,16 +1077,14 @@ func buildL0Table(ft flushTask, bopts table.Options) *table.Builder {
 
 type flushTask struct {
 	mt           *memTable
+	cb           func()
+	itr          y.Iterator
 	dropPrefixes [][]byte
 }
 
 // handleFlushTask must be run serially.
 func (db *DB) handleFlushTask(ft flushTask) error {
-	// There can be a scenario, when empty memtable is flushed.
-	if ft.mt.sl.Empty() {
-		return nil
-	}
-
+	// ft.mt could be nil with ft.itr being the valid field.
 	bopts := buildTableOptions(db)
 	builder := buildL0Table(ft, bopts)
 	defer builder.Close()
@@ -1080,11 +1120,52 @@ func (db *DB) handleFlushTask(ft flushTask) error {
 func (db *DB) flushMemtable(lc *z.Closer) error {
 	defer lc.Done()
 
+	var sz int64
+	var itrs []y.Iterator
+	var mts []*memTable
+	var cbs []func()
+	slurp := func() {
+		for {
+			select {
+			case more := <-db.flushChan:
+				if more.mt == nil {
+					return
+				}
+				sl := more.mt.sl
+				itrs = append(itrs, sl.NewUniIterator(false))
+				mts = append(mts, more.mt)
+				cbs = append(cbs, more.cb)
+
+				sz += sl.MemSize()
+				if sz > db.opt.MemTableSize {
+					return
+				}
+			default:
+				return
+			}
+		}
+	}
+
 	for ft := range db.flushChan {
 		if ft.mt == nil {
 			// We close db.flushChan now, instead of sending a nil ft.mt.
 			continue
 		}
+		sz = ft.mt.sl.MemSize()
+		// Reset of itrs, mts etc. is being done below.
+		y.AssertTrue(len(itrs) == 0 && len(mts) == 0 && len(cbs) == 0)
+		itrs = append(itrs, ft.mt.sl.NewUniIterator(false))
+		mts = append(mts, ft.mt)
+		cbs = append(cbs, ft.cb)
+
+		// Pick more memtables, so we can really fill up the L0 table.
+		slurp()
+
+		// db.opt.Infof("Picked %d memtables. Size: %d\n", len(itrs), sz)
+		ft.mt = nil
+		ft.itr = table.NewMergeIterator(itrs, false)
+		ft.cb = nil
+
 		for {
 			err := db.handleFlushTask(ft)
 			if err == nil {
@@ -1095,17 +1176,26 @@ func (db *DB) flushMemtable(lc *z.Closer) error {
 				// which would arrive here would match db.imm[0], because we acquire a
 				// lock over DB when pushing to flushChan.
 				// TODO: This logic is dirty AF. Any change and this could easily break.
-				y.AssertTrue(ft.mt == db.imm[0])
-				db.imm = db.imm[1:]
-				ft.mt.DecrRef() // Return memory.
+				for _, mt := range mts {
+					y.AssertTrue(mt == db.imm[0])
+					db.imm = db.imm[1:]
+					mt.DecrRef() // Return memory.
+				}
 				db.lock.Unlock()
 
+				for _, cb := range cbs {
+					if cb != nil {
+						cb()
+					}
+				}
 				break
 			}
 			// Encountered error. Retry indefinitely.
 			db.opt.Errorf("Failure while flushing memtable to disk: %v. Retrying...\n", err)
 			time.Sleep(time.Second)
 		}
+		// Reset everything.
+		itrs, mts, cbs, sz = itrs[:0], mts[:0], cbs[:0], 0
 	}
 	return nil
 }

db_test.go

@@ -2088,7 +2088,7 @@ func TestVerifyChecksum(t *testing.T) {
 			y.Check2(rand.Read(value))
 			st := 0
 
-			buf := z.NewBuffer(10 << 20, "test")
+			buf := z.NewBuffer(10<<20, "test")
 			defer buf.Release()
 			for i := 0; i < 1000; i++ {
 				key := make([]byte, 8)

iterator.go

@@ -146,7 +146,7 @@ func (item *Item) IsDeletedOrExpired() bool {
 // DiscardEarlierVersions returns whether the item was created with the
 // option to discard earlier versions of a key when multiple are available.
 func (item *Item) DiscardEarlierVersions() bool {
-	return item.meta&bitDiscardEarlierVersions > 0
+	return item.meta&BitDiscardEarlierVersions > 0
 }
 
 func (item *Item) yieldItemValue() ([]byte, func(), error) {

levels.go

@@ -716,7 +716,7 @@ func (s *levelsController) subcompact(it y.Iterator, kr keyRange, cd compactDef,
 				}
 				lastKey = y.SafeCopy(lastKey, it.Key())
 				numVersions = 0
-				firstKeyHasDiscardSet = it.Value().Meta&bitDiscardEarlierVersions > 0
+				firstKeyHasDiscardSet = it.Value().Meta&BitDiscardEarlierVersions > 0
 
 				if len(tableKr.left) == 0 {
 					tableKr.left = y.SafeCopy(tableKr.left, it.Key())
@@ -753,7 +753,7 @@ func (s *levelsController) subcompact(it y.Iterator, kr keyRange, cd compactDef,
 				// - The `discardEarlierVersions` bit is set OR
 				// - We've already processed `NumVersionsToKeep` number of versions
 				// (including the current item being processed)
-				lastValidVersion := vs.Meta&bitDiscardEarlierVersions > 0 ||
+				lastValidVersion := vs.Meta&BitDiscardEarlierVersions > 0 ||
 					numVersions == s.kv.opt.NumVersionsToKeep
 
 				if isExpired || lastValidVersion {

levels_test.go

@@ -707,11 +707,11 @@ func TestDiscardFirstVersion(t *testing.T) {
 
 	runBadgerTest(t, &opt, func(t *testing.T, db *DB) {
 		l0 := []keyValVersion{{"foo", "bar", 1, 0}}
-		l01 := []keyValVersion{{"foo", "bar", 2, bitDiscardEarlierVersions}}
+		l01 := []keyValVersion{{"foo", "bar", 2, BitDiscardEarlierVersions}}
 		l02 := []keyValVersion{{"foo", "bar", 3, 0}}
 		l03 := []keyValVersion{{"foo", "bar", 4, 0}}
 		l04 := []keyValVersion{{"foo", "bar", 9, 0}}
-		l05 := []keyValVersion{{"foo", "bar", 10, bitDiscardEarlierVersions}}
+		l05 := []keyValVersion{{"foo", "bar", 10, BitDiscardEarlierVersions}}
 
 		// Level 0 has all the tables.
 		createAndOpen(db, l0, 0)
@@ -742,11 +742,11 @@ func TestDiscardFirstVersion(t *testing.T) {
 		// - Version 1 is below DiscardTS and below the first "bitDiscardEarlierVersions"
 		//   marker so IT WILL BE REMOVED.
 		ExpectedKeys := []keyValVersion{
-			{"foo", "bar", 10, bitDiscardEarlierVersions},
+			{"foo", "bar", 10, BitDiscardEarlierVersions},
 			{"foo", "bar", 9, 0},
 			{"foo", "bar", 4, 0},
 			{"foo", "bar", 3, 0},
-			{"foo", "bar", 2, bitDiscardEarlierVersions}}
+			{"foo", "bar", 2, BitDiscardEarlierVersions}}
 
 		getAllAndCheck(t, db, ExpectedKeys)
 	})
@@ -1060,15 +1060,15 @@ func TestSameLevel(t *testing.T) {
 	opt.LmaxCompaction = true
 	runBadgerTest(t, &opt, func(t *testing.T, db *DB) {
 		l6 := []keyValVersion{
-			{"A", "bar", 4, bitDiscardEarlierVersions}, {"A", "bar", 3, 0},
+			{"A", "bar", 4, BitDiscardEarlierVersions}, {"A", "bar", 3, 0},
 			{"A", "bar", 2, 0}, {"Afoo", "baz", 2, 0},
 		}
 		l61 := []keyValVersion{
-			{"B", "bar", 4, bitDiscardEarlierVersions}, {"B", "bar", 3, 0},
+			{"B", "bar", 4, BitDiscardEarlierVersions}, {"B", "bar", 3, 0},
 			{"B", "bar", 2, 0}, {"Bfoo", "baz", 2, 0},
 		}
 		l62 := []keyValVersion{
-			{"C", "bar", 4, bitDiscardEarlierVersions}, {"C", "bar", 3, 0},
+			{"C", "bar", 4, BitDiscardEarlierVersions}, {"C", "bar", 3, 0},
 			{"C", "bar", 2, 0}, {"Cfoo", "baz", 2, 0},
 		}
 		createAndOpen(db, l6, 6)
@@ -1077,11 +1077,11 @@ func TestSameLevel(t *testing.T) {
 		require.NoError(t, db.lc.validate())
 
 		getAllAndCheck(t, db, []keyValVersion{
-			{"A", "bar", 4, bitDiscardEarlierVersions}, {"A", "bar", 3, 0},
+			{"A", "bar", 4, BitDiscardEarlierVersions}, {"A", "bar", 3, 0},
 			{"A", "bar", 2, 0}, {"Afoo", "baz", 2, 0},
-			{"B", "bar", 4, bitDiscardEarlierVersions}, {"B", "bar", 3, 0},
+			{"B", "bar", 4, BitDiscardEarlierVersions}, {"B", "bar", 3, 0},
 			{"B", "bar", 2, 0}, {"Bfoo", "baz", 2, 0},
-			{"C", "bar", 4, bitDiscardEarlierVersions}, {"C", "bar", 3, 0},
+			{"C", "bar", 4, BitDiscardEarlierVersions}, {"C", "bar", 3, 0},
 			{"C", "bar", 2, 0}, {"Cfoo", "baz", 2, 0},
 		})
 
@@ -1097,11 +1097,11 @@ func TestSameLevel(t *testing.T) {
 		db.SetDiscardTs(3)
 		require.NoError(t, db.lc.runCompactDef(-1, 6, cdef))
 		getAllAndCheck(t, db, []keyValVersion{
-			{"A", "bar", 4, bitDiscardEarlierVersions}, {"A", "bar", 3, 0},
+			{"A", "bar", 4, BitDiscardEarlierVersions}, {"A", "bar", 3, 0},
 			{"A", "bar", 2, 0}, {"Afoo", "baz", 2, 0},
-			{"B", "bar", 4, bitDiscardEarlierVersions}, {"B", "bar", 3, 0},
+			{"B", "bar", 4, BitDiscardEarlierVersions}, {"B", "bar", 3, 0},
 			{"B", "bar", 2, 0}, {"Bfoo", "baz", 2, 0},
-			{"C", "bar", 4, bitDiscardEarlierVersions}, {"C", "bar", 3, 0},
+			{"C", "bar", 4, BitDiscardEarlierVersions}, {"C", "bar", 3, 0},
 			{"C", "bar", 2, 0}, {"Cfoo", "baz", 2, 0},
 		})
 
@@ -1118,9 +1118,9 @@ func TestSameLevel(t *testing.T) {
 		cdef.t.baseLevel = 1
 		require.NoError(t, db.lc.runCompactDef(-1, 6, cdef))
 		getAllAndCheck(t, db, []keyValVersion{
-			{"A", "bar", 4, bitDiscardEarlierVersions}, {"Afoo", "baz", 2, 0},
-			{"B", "bar", 4, bitDiscardEarlierVersions}, {"Bfoo", "baz", 2, 0},
-			{"C", "bar", 4, bitDiscardEarlierVersions}, {"Cfoo", "baz", 2, 0}})
+			{"A", "bar", 4, BitDiscardEarlierVersions}, {"Afoo", "baz", 2, 0},
+			{"B", "bar", 4, BitDiscardEarlierVersions}, {"Bfoo", "baz", 2, 0},
+			{"C", "bar", 4, BitDiscardEarlierVersions}, {"Cfoo", "baz", 2, 0}})
 		require.NoError(t, db.lc.validate())
 	})
 }
@@ -1186,7 +1186,7 @@ func TestStaleDataCleanup(t *testing.T) {
 			for i := count; i > 0; i-- {
 				var meta byte
 				if i == 0 {
-					meta = bitDiscardEarlierVersions
+					meta = BitDiscardEarlierVersions
 				}
 				b.AddStaleKey(y.KeyWithTs(key, i), y.ValueStruct{Meta: meta, Value: val}, 0)
 			}

managed_db_test.go

@@ -771,6 +771,54 @@ func TestWriteBatchDuplicate(t *testing.T) {
 	})
 }
 
+func TestWriteViaSkip(t *testing.T) {
+	key := func(i int) []byte {
+		return []byte(fmt.Sprintf("%10d", i))
+	}
+	val := func(i int) []byte {
+		return []byte(fmt.Sprintf("%128d", i))
+	}
+	opt := DefaultOptions("")
+	opt.managedTxns = true
+	runBadgerTest(t, &opt, func(t *testing.T, db *DB) {
+		s := db.NewSkiplist()
+		for i := 0; i < 100; i++ {
+			s.Put(y.KeyWithTs(key(i), math.MaxUint64), y.ValueStruct{Value: val(i)})
+		}
+		{
+			// Update key timestamps by directly changing them in the skiplist.
+			itr := s.NewUniIterator(false)
+			defer itr.Close()
+			itr.Rewind()
+			for itr.Valid() {
+				y.SetKeyTs(itr.Key(), 101)
+				itr.Next()
+			}
+		}
+
+		// Hand over skiplist to Badger.
+		require.NoError(t, db.HandoverSkiplist(s, nil))
+
+		// Read the data back.
+		txn := db.NewTransactionAt(101, false)
+		defer txn.Discard()
+		itr := txn.NewIterator(DefaultIteratorOptions)
+		defer itr.Close()
+
+		i := 0
+		for itr.Rewind(); itr.Valid(); itr.Next() {
+			item := itr.Item()
+			require.Equal(t, string(key(i)), string(item.Key()))
+			require.Equal(t, item.Version(), uint64(101))
+			valcopy, err := item.ValueCopy(nil)
+			require.NoError(t, err)
+			require.Equal(t, val(i), valcopy)
+			i++
+		}
+		require.Equal(t, 100, i)
+	})
+}
+
 func TestZeroDiscardStats(t *testing.T) {
 	N := uint64(10000)
 	populate := func(t *testing.T, db *DB) {

merge.go

@@ -116,7 +116,7 @@ func (op *MergeOperator) compact() error {
 		{
 			Key:   y.KeyWithTs(op.key, version),
 			Value: val,
-			meta:  bitDiscardEarlierVersions,
+			meta:  BitDiscardEarlierVersions,
 		},
 	}
 	// Write value back to the DB. It is important that we do not set the bitMergeEntry bit

options.go

@@ -144,7 +144,7 @@ func DefaultOptions(path string) Options {
 		NumCompactors:           4, // Run at least 2 compactors. Zero-th compactor prioritizes L0.
 		NumLevelZeroTables:      5,
 		NumLevelZeroTablesStall: 15,
-		NumMemtables:            5,
+		NumMemtables:            15,
 		BloomFalsePositive:      0.01,
 		BlockSize:               4 * 1024,
 		SyncWrites:              false,

skl/skl.go

@@ -584,3 +584,7 @@ func (b *Builder) Add(k []byte, v y.ValueStruct) {
 		b.prev[i] = nodeOffset
 	}
 }
+
+func (b *Builder) Skiplist() *Skiplist {
+	return b.s
+}

structs.go

@@ -206,7 +206,7 @@ func (e *Entry) WithMeta(meta byte) *Entry {
 // have a higher setting for NumVersionsToKeep (in Dgraph, we set it to infinity), you can use this
 // method to indicate that all the older versions can be discarded and removed during compactions.
 func (e *Entry) WithDiscard() *Entry {
-	e.meta = bitDiscardEarlierVersions
+	e.meta = BitDiscardEarlierVersions
 	return e
 }
 

txn.go

@@ -358,9 +358,30 @@ func exceedsSize(prefix string, max int64, key []byte) error {
 		prefix, len(key), max, prefix, hex.Dump(key[:1<<10]))
 }
 
-func (txn *Txn) modify(e *Entry) error {
-	const maxKeySize = 65000
+const maxKeySize = 65000
+const maxValSize = 1 << 20
 
+func ValidEntry(db *DB, key, val []byte) error {
+	switch {
+	case len(key) == 0:
+		return ErrEmptyKey
+	case bytes.HasPrefix(key, badgerPrefix):
+		return ErrInvalidKey
+	case len(key) > maxKeySize:
+		// Key length can't be more than uint16, as determined by table::header.  To
+		// keep things safe and allow badger move prefix and a timestamp suffix, let's
+		// cut it down to 65000, instead of using 65536.
+		return exceedsSize("Key", maxKeySize, key)
+	case int64(len(val)) > maxValSize:
+		return exceedsSize("Value", maxValSize, val)
+	}
+	if err := db.isBanned(key); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (txn *Txn) modify(e *Entry) error {
 	switch {
 	case !txn.update:
 		return ErrReadOnlyTxn
@@ -384,7 +405,6 @@ func (txn *Txn) modify(e *Entry) error {
 	if err := txn.db.isBanned(e.Key); err != nil {
 		return err
 	}
-
 	if err := txn.checkSize(e); err != nil {
 		return err
 	}

value.go

@@ -47,7 +47,7 @@ var maxVlogFileSize uint32 = math.MaxUint32
 const (
 	bitDelete                 byte = 1 << 0 // Set if the key has been deleted.
 	bitValuePointer           byte = 1 << 1 // Set if the value is NOT stored directly next to key.
-	bitDiscardEarlierVersions byte = 1 << 2 // Set if earlier versions can be discarded.
+	BitDiscardEarlierVersions byte = 1 << 2 // Set if earlier versions can be discarded.
 	// Set if item shouldn't be discarded via compactions (used by merge operator)
 	bitMergeEntry byte = 1 << 3
 	// The MSB 2 bits are for transactions.

value_test.go

@@ -136,6 +136,8 @@ func TestValueBasic(t *testing.T) {
 }
 
 func TestValueGCManaged(t *testing.T) {
+	t.Skipf("Value Log is not used in managed mode.")
+
 	dir, err := ioutil.TempDir("", "badger-test")
 	require.NoError(t, err)
 	defer removeDir(dir)

y/y.go

@@ -121,6 +121,11 @@ func Copy(a []byte) []byte {
 	return b
 }
 
+func SetKeyTs(key []byte, ts uint64) {
+	start := len(key) - 8
+	binary.BigEndian.PutUint64(key[start:], math.MaxUint64-ts)
+}
+
 // KeyWithTs generates a new key by appending ts to key.
 func KeyWithTs(key []byte, ts uint64) []byte {
 	out := make([]byte, len(key)+8)