BREAKING CHANGE: Marshal pb.list in Stream.Backup (#807)

* BREAKING CHANGE: Marshal pb.list in Stream.Backup

WARN: This is a breaking change as it changes format of data to be
written by backup tool on disk. Backups taken using code before this
commit cannot be restored using code from this commit onwards.

Backup tool uses Stream to read DB entries and writing it to a
file. It implements Stream.Send to write entry one by one. This
commit modifies Stream.Send to write entries in batches by marshaling
pb.list instead of pb.KV. This improves backup speed by ~4-5x.

* Use buffered writer for writing in backup tool

backup.go

@@ -105,11 +105,8 @@ func (stream *Stream) Backup(w io.Writer, since uint64) (uint64, error) {
 			if maxVersion < kv.Version {
 				maxVersion = kv.Version
 			}
-			if err := writeTo(kv, w); err != nil {
-				return err
-			}
 		}
-		return nil
+		return writeTo(list, w)
 	}
 
 	if err := stream.Orchestrate(context.Background()); err != nil {
@@ -118,11 +115,11 @@ func (stream *Stream) Backup(w io.Writer, since uint64) (uint64, error) {
 	return maxVersion, nil
 }
 
-func writeTo(entry *pb.KV, w io.Writer) error {
-	if err := binary.Write(w, binary.LittleEndian, uint64(entry.Size())); err != nil {
+func writeTo(list *pb.KVList, w io.Writer) error {
+	if err := binary.Write(w, binary.LittleEndian, uint64(list.Size())); err != nil {
 		return err
 	}
-	buf, err := entry.Marshal()
+	buf, err := list.Marshal()
 	if err != nil {
 		return err
 	}
@@ -213,17 +210,22 @@ func (db *DB) Load(r io.Reader, maxPendingWrites int) error {
 		if _, err = io.ReadFull(br, unmarshalBuf[:sz]); err != nil {
 			return err
 		}
-		kv := &pb.KV{}
-		if err = kv.Unmarshal(unmarshalBuf[:sz]); err != nil {
-			return err
-		}
-		if err := ldr.set(kv); err != nil {
+
+		list := &pb.KVList{}
+		if err := list.Unmarshal(unmarshalBuf[:sz]); err != nil {
 			return err
 		}
-		// Update nextTxnTs, memtable stores this timestamp in badger head
-		// when flushed.
-		if kv.Version >= db.orc.nextTxnTs {
-			db.orc.nextTxnTs = kv.Version + 1
+
+		for _, kv := range list.Kv {
+			if err := ldr.set(kv); err != nil {
+				return err
+			}
+
+			// Update nextTxnTs, memtable stores this
+			// timestamp in badger head when flushed.
+			if kv.Version >= db.orc.nextTxnTs {
+				db.orc.nextTxnTs = kv.Version + 1
+			}
 		}
 	}
 

badger/cmd/backup.go

@@ -17,6 +17,7 @@
 package cmd
 
 import (
+	"bufio"
 	"os"
 
 	"github.com/dgraph-io/badger"
@@ -64,9 +65,19 @@ func doBackup(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
-	defer f.Close()
 
-	// Run Backup
-	_, err = db.Backup(f, 0)
-	return err
+	bw := bufio.NewWriterSize(f, 64<<20)
+	if _, err = db.Backup(bw, 0); err != nil {
+		return err
+	}
+
+	if err = bw.Flush(); err != nil {
+		return err
+	}
+
+	if err = f.Sync(); err != nil {
+		return err
+	}
+
+	return f.Close()
 }