lightning: fix encode kvs size greater than 4.0g caused pebble panic

pkg/lightning/backend/kv/sql2kv.go

@@ -408,6 +408,14 @@ func (kvcodec *tableKVEncoder) Encode(
 	return kvPairs(pairs), nil
 }
 
+func (kvs kvPairs) Size() uint64 {
+	size := uint64(0)
+	for _, kv := range kvs {
+		size += uint64(len(kv.Key) + len(kv.Val))
+	}
+	return size
+}
+
 func (kvs kvPairs) ClassifyAndAppend(
 	data *Rows,
 	dataChecksum *verification.KVChecksum,

pkg/lightning/backend/kv/types.go

@@ -35,6 +35,9 @@ type Row interface {
 		indices *Rows,
 		indexChecksum *verification.KVChecksum,
 	)
+
+	// Size represents the total kv size of this Row.
+	Size() uint64
 }
 
 // Rows represents a collection of encoded rows.

pkg/lightning/backend/tidb/tidb.go

@@ -94,6 +94,10 @@ func NewTiDBBackend(db *sql.DB, onDuplicate string) backend.Backend {
 	return backend.MakeBackend(&tidbBackend{db: db, onDuplicate: onDuplicate})
 }
 
+func (row tidbRow) Size() uint64 {
+	return uint64(len(row))
+}
+
 func (row tidbRow) ClassifyAndAppend(data *kv.Rows, checksum *verification.KVChecksum, _ *kv.Rows, _ *verification.KVChecksum) {
 	rows := (*data).(tidbRows)
 	// Cannot do `rows := data.(*tidbRows); *rows = append(*rows, row)`.

pkg/lightning/restore/restore.go

@@ -2601,6 +2601,7 @@ func (cr *chunkRestore) encodeLoop(
 		canDeliver := false
 		kvPacket := make([]deliveredKVs, 0, maxKvPairsCnt)
 		var newOffset, rowID int64
+		var kvSize uint64
 	outLoop:
 		for !canDeliver {
 			readDurStart := time.Now()
@@ -2636,6 +2637,14 @@ func (cr *chunkRestore) encodeLoop(
 				return
 			}
 			kvPacket = append(kvPacket, deliveredKVs{kvs: kvs, columns: columnNames, offset: newOffset, rowID: rowID})
+			kvSize += kvs.Size()
+			// pebble cannot allow > 4.0G kv in one batch.
+			// we will meet pebble panic when import sql file and each kv has the size larger than 4G / maxKvPairsCnt.
+			// so add this check.
+			if kvSize > minDeliverBytes {
+				canDeliver = true
+				kvSize = 0
+			}
 			if len(kvPacket) >= maxKvPairsCnt || newOffset == cr.chunk.Chunk.EndOffset {
 				canDeliver = true
 			}