Fix transaction too big issue in restore

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
 }