Replace TxnWriter with WriteBatch

posting/index.go

@@ -567,9 +567,10 @@ func (r *rebuilder) Run(ctx context.Context) error {
 	// We set it to 1 in case there are no keys found and NewStreamAt is called with ts=0.
 	var counter uint64 = 1
 
-	// TODO(Aman): Replace TxnWriter with WriteBatch. While we do that we should ensure that
-	// WriteBatch has a mechanism for throttling. Also, find other places where TxnWriter
-	// could be replaced with WriteBatch in the code
+	// WriteBatch can not be used here because it doesn't have an API to allow writing
+	// multiple versions. We wish to store same keys with diff version/timestamp to
+	// ensure that we get all of them back when doing roll-up. WriteBatch can only be
+	// used when we want to write all txns at the same timestamp.
 	tmpWriter := NewTxnWriter(tmpDB)
 	stream := pstore.NewStreamAt(r.startTs)
 	stream.LogPrefix = fmt.Sprintf("Rebuilding index for predicate %s (1/2):", r.attr)
@@ -644,7 +645,7 @@ func (r *rebuilder) Run(ctx context.Context) error {
 			r.attr, time.Since(start))
 	}()
 
-	writer := NewTxnWriter(pstore)
+	batchWriter := pstore.NewWriteBatchAt(r.startTs)
 	tmpStream := tmpDB.NewStreamAt(counter)
 	tmpStream.LogPrefix = fmt.Sprintf("Rebuilding index for predicate %s (2/2):", r.attr)
 	tmpStream.KeyToList = func(key []byte, itr *badger.Iterator) (*bpb.KVList, error) {
@@ -670,7 +671,8 @@ func (r *rebuilder) Run(ctx context.Context) error {
 
 			// We choose to write the PL at r.startTs, so it won't be read by txns,
 			// which occurred before this schema mutation.
-			if err := writer.SetAt(kv.Key, kv.Value, BitCompletePosting, r.startTs); err != nil {
+			e := &badger.Entry{Key: kv.Key, Value: kv.Value, UserMeta: BitCompletePosting}
+			if err := batchWriter.SetEntry(e); err != nil {
 				return errors.Wrap(err, "error in writing index to pstore")
 			}
 		}
@@ -682,7 +684,7 @@ func (r *rebuilder) Run(ctx context.Context) error {
 		return err
 	}
 	glog.V(1).Infof("Rebuilding index for predicate %s: Flushing all writes.\n", r.attr)
-	return writer.Flush()
+	return batchWriter.Flush()
 }
 
 // IndexRebuild holds the info needed to initiate a rebuilt of the indices.

posting/writer_test.go

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2019 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 posting
+
+import (
+	"fmt"
+	"io/ioutil"
+	"math"
+	"testing"
+
+	"github.com/dgraph-io/badger/v2"
+	"github.com/dgraph-io/badger/v2/options"
+	"github.com/dgraph-io/dgraph/x"
+)
+
+type kv struct {
+	key   []byte
+	value []byte
+}
+
+var tmpIndexDir, err = ioutil.TempDir("", "dgraph_index_")
+
+var dbOpts = badger.DefaultOptions(tmpIndexDir).
+	WithSyncWrites(false).
+	WithNumVersionsToKeep(math.MaxInt64).
+	WithLogger(&x.ToGlog{}).
+	WithCompression(options.None).
+	WithEventLogging(false).
+	WithLogRotatesToFlush(10).
+	WithMaxCacheSize(50) // TODO(Aman): Disable cache altogether
+
+var db, err2 = badger.OpenManaged(dbOpts)
+
+func createKVList() []kv {
+	var KVList = []kv{}
+	for i := 0; i < 50000; i++ {
+		n := kv{key: []byte(string(i)), value: []byte("Check Value")}
+		KVList = append(KVList, n)
+	}
+	return KVList
+}
+
+func BenchmarkTxnWriter(b *testing.B) {
+	KVList := createKVList()
+	w := NewTxnWriter(db)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for _, typ := range KVList {
+			k := typ.key
+			v := typ.value
+			x.Check(w.SetAt(k, v, BitSchemaPosting, 1))
+		}
+		if err := w.Flush(); err != nil {
+			fmt.Printf("Got error while flushing txnwriter: %v\n", err)
+		}
+	}
+
+}
+
+func BenchmarkWriteBatch(b *testing.B) {
+	KVList := createKVList()
+	batch := db.NewWriteBatchAt(1)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for _, typ := range KVList {
+			k := typ.key
+			v := typ.value
+			x.Check(batch.Set(k, v))
+		}
+		if err := batch.Flush(); err != nil {
+			fmt.Printf("Got error while flushing batch: %v\n", err)
+		}
+	}
+
+}

worker/draft.go

@@ -639,6 +639,7 @@ func (n *node) processApplyCh() {
 	}
 }
 
+// TODO(Anurag): Are we using pkey? Remove if unused.
 func (n *node) commitOrAbort(pkey string, delta *pb.OracleDelta) error {
 	// First let's commit all mutations to disk.
 	writer := posting.NewTxnWriter(pstore)