From 30b27d8b4a08c1c69f7a9ab9da9d1aa563afd4d4 Mon Sep 17 00:00:00 2001
From: Ibrahim Jarif <ibrahim@dgraph.io>
Date: Thu, 1 Aug 2019 18:44:44 +0530
Subject: [PATCH 1/2] Fix transaction too big issue in restore

When restore happens, we bunch entries together and perform batch
write. The existing implementation would batch 1000 entries. The batch
of 1000 entries can overflow the transactional limits in certain
circumstances (eg: when table size is too small). This commit fixes the
issue. We flush the entries once the batch reaches its max capacity.
---
 backup.go | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/backup.go b/backup.go
index 2569b3100..7133cfbf4 100644
--- a/backup.go
+++ b/backup.go
@@ -129,9 +129,10 @@ func writeTo(list *pb.KVList, w io.Writer) error {
 
 // KVLoader is used to write KVList objects in to badger. It can be used to restore a backup.
 type KVLoader struct {
-	db       *DB
-	throttle *y.Throttle
-	entries  []*Entry
+	db          *DB
+	throttle    *y.Throttle
+	entries     []*Entry
+	entriesSize int64
 }
 
 // NewKVLoader returns a new instance of KVLoader.
@@ -139,6 +140,7 @@ func (db *DB) NewKVLoader(maxPendingWrites int) *KVLoader {
 	return &KVLoader{
 		db:       db,
 		throttle: y.NewThrottle(maxPendingWrites),
+		entries:  make([]*Entry, 0, db.opt.maxBatchCount),
 	}
 }
 
@@ -151,17 +153,23 @@ func (l *KVLoader) Set(kv *pb.KV) error {
 	if len(kv.Meta) > 0 {
 		meta = kv.Meta[0]
 	}
-
-	l.entries = append(l.entries, &Entry{
+	e := &Entry{
 		Key:       y.KeyWithTs(kv.Key, kv.Version),
 		Value:     kv.Value,
 		UserMeta:  userMeta,
 		ExpiresAt: kv.ExpiresAt,
 		meta:      meta,
-	})
-	if len(l.entries) >= 1000 {
-		return l.send()
 	}
+	estimatedSize := int64(e.estimateSize(l.db.opt.ValueThreshold))
+	// Flush entries if inserting the next entry would overflow the transactional limits.
+	if int64(len(l.entries))+1 >= l.db.opt.maxBatchCount ||
+		l.entriesSize+estimatedSize >= l.db.opt.maxBatchSize {
+		if err := l.send(); err != nil {
+			return err
+		}
+	}
+	l.entries = append(l.entries, e)
+	l.entriesSize += estimatedSize
 	return nil
 }
 
@@ -175,7 +183,8 @@ func (l *KVLoader) send() error {
 		return err
 	}
 
-	l.entries = make([]*Entry, 0, 1000)
+	l.entries = l.entries[:0]
+	l.entriesSize = 0
 	return nil
 }
 

From 1dca2e26aac554fc1ba983b544186d02bf8a033a Mon Sep 17 00:00:00 2001
From: Ibrahim Jarif <ibrahim@dgraph.io>
Date: Sat, 3 Aug 2019 00:11:57 +0530
Subject: [PATCH 2/2] Do not reuse l.entries slice

---
 backup.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/backup.go b/backup.go
index 7133cfbf4..e4bdc8861 100644
--- a/backup.go
+++ b/backup.go
@@ -183,7 +183,7 @@ func (l *KVLoader) send() error {
 		return err
 	}
 
-	l.entries = l.entries[:0]
+	l.entries = make([]*Entry, 0, l.db.opt.maxBatchCount)
 	l.entriesSize = 0
 	return nil
 }