From f8f7b00691d8bcac1db796e2b5e5a9d74c85e903 Mon Sep 17 00:00:00 2001
From: Ibrahim Jarif <ibrahim@dgraph.io>
Date: Fri, 9 Aug 2019 13:34:34 +0530
Subject: [PATCH]  Fix transaction too big issue in restore  (#957)

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 the table size is too small). This commit fixes the
issue. We flush the entries once the batch reaches its max capacity.

(cherry picked from commit f19b4c13b86137edb87f3859d73fadeadf92d919)
---
 backup.go | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/backup.go b/backup.go
index 2569b3100..e4bdc8861 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 = make([]*Entry, 0, l.db.opt.maxBatchCount)
+	l.entriesSize = 0
 	return nil
 }