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.

backup.go

@@ -129,16 +129,18 @@ 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.
 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
 }