@@ -22,15 +22,19 @@ import (
2222 "encoding/binary"
2323 "fmt"
2424 "path"
25+ "slices"
2526 "strconv"
2627 "strings"
2728
2829 "github.com/gogo/protobuf/proto"
2930 "github.com/pingcap/errors"
3031 "github.com/pingcap/failpoint"
32+ backuppb "github.com/pingcap/kvproto/pkg/brpb"
3133 "github.com/pingcap/log"
3234 berrors "github.com/pingcap/tidb/br/pkg/errors"
3335 "github.com/pingcap/tidb/br/pkg/logutil"
36+ "github.com/pingcap/tidb/br/pkg/restore/split"
37+ restoreutils "github.com/pingcap/tidb/br/pkg/restore/utils"
3438 "github.com/pingcap/tidb/br/pkg/storage"
3539 "github.com/pingcap/tidb/br/pkg/utils"
3640 "github.com/pingcap/tidb/pkg/domain"
@@ -389,3 +393,158 @@ func HasRestoreIDColumn(dom *domain.Domain) bool {
389393 }
390394 return false
391395}
396+
397+ type regionScanner struct {
398+ regionClient split.SplitClient
399+ regionCache []* split.RegionInfo
400+ cacheSize int
401+ }
402+
403+ func NewRegionScanner (regionClient split.SplitClient , cacheSize int ) * regionScanner {
404+ return & regionScanner {
405+ regionClient : regionClient ,
406+ cacheSize : cacheSize ,
407+ }
408+ }
409+
410+ func (scanner * regionScanner ) locateRegionFromRemote (ctx context.Context , key []byte ) (* split.RegionInfo , error ) {
411+ regionInfos , err := split .ScanRegionsWithRetry (ctx , scanner .regionClient , key , []byte ("" ), scanner .cacheSize )
412+ if err != nil {
413+ return nil , errors .Trace (err )
414+ }
415+ scanner .regionCache = regionInfos
416+ return scanner .regionCache [0 ], nil
417+ }
418+
419+ func (scanner * regionScanner ) locateRegionFromCache (ctx context.Context , key []byte ) (* split.RegionInfo , error ) {
420+ if len (scanner .regionCache ) == 0 {
421+ return scanner .locateRegionFromRemote (ctx , key )
422+ }
423+ if bytes .Compare (key , scanner .regionCache [len (scanner .regionCache )- 1 ].Region .EndKey ) >= 0 {
424+ return scanner .locateRegionFromRemote (ctx , key )
425+ }
426+ i , ok := slices .BinarySearchFunc (scanner .regionCache , key , func (regionInfo * split.RegionInfo , k []byte ) int {
427+ startCmpRet := bytes .Compare (regionInfo .Region .StartKey , k )
428+ if startCmpRet <= 0 && (len (regionInfo .Region .EndKey ) == 0 || bytes .Compare (regionInfo .Region .EndKey , k ) > 0 ) {
429+ return 0
430+ }
431+ return startCmpRet
432+ })
433+ if ! ok {
434+ return scanner .locateRegionFromRemote (ctx , key )
435+ }
436+ scanner .regionCache = scanner .regionCache [i :]
437+ return scanner .regionCache [0 ], nil
438+ }
439+
440+ func (scanner * regionScanner ) IsKeyRangeInOneRegion (ctx context.Context , startKey , endKey []byte ) (bool , error ) {
441+ regionInfo , err := scanner .locateRegionFromCache (ctx , startKey )
442+ if err != nil {
443+ return false , errors .Trace (err )
444+ }
445+ return len (regionInfo .Region .EndKey ) == 0 || bytes .Compare (endKey , regionInfo .Region .EndKey ) < 0 , nil
446+ }
447+
448+ type BackupFileSetWithKeyRange struct {
449+ backupFileSet BackupFileSet
450+ startKey []byte
451+ endKey []byte
452+ }
453+
454+ func GroupOverlappedBackupFileSetsIter (ctx context.Context , regionClient split.SplitClient , backupFileSets []BackupFileSet , fn func (BatchBackupFileSet )) error {
455+ backupFileSetWithKeyRanges := make ([]* BackupFileSetWithKeyRange , 0 , len (backupFileSets ))
456+ for _ , backupFileSet := range backupFileSets {
457+ startKey , endKey , err := getKeyRangeForBackupFileSet (backupFileSet )
458+ if err != nil {
459+ return errors .Trace (err )
460+ }
461+ backupFileSetWithKeyRanges = append (backupFileSetWithKeyRanges , & BackupFileSetWithKeyRange {
462+ backupFileSet : backupFileSet ,
463+ startKey : startKey ,
464+ endKey : endKey ,
465+ })
466+ }
467+ slices .SortFunc (backupFileSetWithKeyRanges , func (a , b * BackupFileSetWithKeyRange ) int {
468+ startKeyCmp := bytes .Compare (a .startKey , b .startKey )
469+ if startKeyCmp == 0 {
470+ return bytes .Compare (a .endKey , b .endKey )
471+ }
472+ return startKeyCmp
473+ })
474+ regionScanner := NewRegionScanner (regionClient , 64 )
475+ var thisBackupFileSet * BackupFileSet = nil
476+ thisBatchBackupFileSet := make ([]BackupFileSet , 0 )
477+ lastEndKey := []byte {}
478+ for _ , file := range backupFileSetWithKeyRanges {
479+ if bytes .Compare (lastEndKey , file .startKey ) < 0 {
480+ // the next file is not overlapped with this backup file set anymore, so add the set
481+ // into the batch set.
482+ if thisBackupFileSet != nil {
483+ thisBatchBackupFileSet = append (thisBatchBackupFileSet , * thisBackupFileSet )
484+ thisBackupFileSet = nil
485+ }
486+ // create new this backup file set
487+ thisBackupFileSet = & BackupFileSet {
488+ TableID : file .backupFileSet .TableID ,
489+ SSTFiles : make ([]* backuppb.File , 0 ),
490+ RewriteRules : file .backupFileSet .RewriteRules ,
491+ }
492+ thisBackupFileSet .SSTFiles = append (thisBackupFileSet .SSTFiles , file .backupFileSet .SSTFiles ... )
493+ // check whether [lastEndKey, file.startKey] is in the one region
494+ inOneRegion , err := regionScanner .IsKeyRangeInOneRegion (ctx , lastEndKey , file .startKey )
495+ if err != nil {
496+ return errors .Trace (err )
497+ }
498+ if ! inOneRegion && len (thisBatchBackupFileSet ) > 0 {
499+ // not in the same region, so this batch backup file set can be output
500+ fn (thisBatchBackupFileSet )
501+ thisBatchBackupFileSet = make ([]BackupFileSet , 0 )
502+ }
503+ lastEndKey = file .endKey
504+ } else {
505+ // the next file is overlapped with this backup file set, so add the file
506+ // into the set.
507+ thisBackupFileSet .SSTFiles = append (thisBackupFileSet .SSTFiles , file .backupFileSet .SSTFiles ... )
508+ if thisBackupFileSet .TableID != file .backupFileSet .TableID || ! thisBackupFileSet .RewriteRules .Equal (file .backupFileSet .RewriteRules ) {
509+ log .Error ("the overlapped SST must have the same table id and rewrite rules" ,
510+ zap .Int64 ("set table id" , thisBackupFileSet .TableID ),
511+ zap .Int64 ("file table id" , file .backupFileSet .TableID ),
512+ zap .Reflect ("set rewrite rule" , thisBackupFileSet .RewriteRules ),
513+ zap .Reflect ("file rewrite rule" , file .backupFileSet .RewriteRules ),
514+ )
515+ return errors .Errorf ("the overlapped SST must have the same table id(%d<>%d) and rewrite rules" ,
516+ thisBackupFileSet .TableID , file .backupFileSet .TableID )
517+ }
518+ // update lastEndKey if file.endKey is larger
519+ if bytes .Compare (lastEndKey , file .endKey ) < 0 {
520+ lastEndKey = file .endKey
521+ }
522+ }
523+ }
524+ // add the set into the batch set.
525+ if thisBackupFileSet != nil {
526+ thisBatchBackupFileSet = append (thisBatchBackupFileSet , * thisBackupFileSet )
527+ }
528+ // output the last batch backup file set
529+ if len (thisBatchBackupFileSet ) > 0 {
530+ fn (thisBatchBackupFileSet )
531+ }
532+ return nil
533+ }
534+
535+ func getKeyRangeForBackupFileSet (backupFileSet BackupFileSet ) ([]byte , []byte , error ) {
536+ var startKey , endKey []byte
537+ for _ , f := range backupFileSet .SSTFiles {
538+ start , end , err := restoreutils .GetRewriteRawKeys (f , backupFileSet .RewriteRules )
539+ if err != nil {
540+ return nil , nil , errors .Trace (err )
541+ }
542+ if len (startKey ) == 0 || bytes .Compare (start , startKey ) < 0 {
543+ startKey = start
544+ }
545+ if len (endKey ) == 0 || bytes .Compare (endKey , end ) < 0 {
546+ endKey = end
547+ }
548+ }
549+ return startKey , endKey , nil
550+ }
0 commit comments