github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/restore/range.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package restore
     4  
     5  import (
     6  	"bytes"
     7  	"sort"
     8  	"sync"
     9  
    10  	"github.com/pingcap/errors"
    11  	"github.com/pingcap/kvproto/pkg/import_sstpb"
    12  	"github.com/pingcap/kvproto/pkg/metapb"
    13  	"github.com/pingcap/log"
    14  	"github.com/pingcap/tidb/tablecodec"
    15  	"go.uber.org/zap"
    16  
    17  	berrors "github.com/pingcap/br/pkg/errors"
    18  	"github.com/pingcap/br/pkg/logutil"
    19  	"github.com/pingcap/br/pkg/rtree"
    20  )
    21  
    22  // Range record start and end key for localStoreDir.DB
    23  // so we can write it to tikv in streaming
    24  type Range struct {
    25  	Start []byte
    26  	End   []byte
    27  }
    28  
    29  type syncdRanges struct {
    30  	sync.Mutex
    31  	ranges []Range
    32  }
    33  
    34  func (r *syncdRanges) add(g Range) {
    35  	r.Lock()
    36  	r.ranges = append(r.ranges, g)
    37  	r.Unlock()
    38  }
    39  
    40  func (r *syncdRanges) take() []Range {
    41  	r.Lock()
    42  	rg := r.ranges
    43  	r.ranges = []Range{}
    44  	r.Unlock()
    45  	if len(rg) > 0 {
    46  		sort.Slice(rg, func(i, j int) bool {
    47  			return bytes.Compare(rg[i].Start, rg[j].Start) < 0
    48  		})
    49  	}
    50  	return rg
    51  }
    52  
    53  func newSyncdRanges() *syncdRanges {
    54  	return &syncdRanges{
    55  		ranges: make([]Range, 0, 128),
    56  	}
    57  }
    58  
    59  // SortRanges checks if the range overlapped and sort them.
    60  func SortRanges(ranges []rtree.Range, rewriteRules *RewriteRules) ([]rtree.Range, error) {
    61  	rangeTree := rtree.NewRangeTree()
    62  	for _, rg := range ranges {
    63  		if rewriteRules != nil {
    64  			startID := tablecodec.DecodeTableID(rg.StartKey)
    65  			endID := tablecodec.DecodeTableID(rg.EndKey)
    66  			var rule *import_sstpb.RewriteRule
    67  			if startID == endID {
    68  				rg.StartKey, rule = replacePrefix(rg.StartKey, rewriteRules)
    69  				if rule == nil {
    70  					log.Warn("cannot find rewrite rule", logutil.Key("key", rg.StartKey))
    71  				} else {
    72  					log.Debug(
    73  						"rewrite start key",
    74  						logutil.Key("key", rg.StartKey), logutil.RewriteRule(rule))
    75  				}
    76  				rg.EndKey, rule = replacePrefix(rg.EndKey, rewriteRules)
    77  				if rule == nil {
    78  					log.Warn("cannot find rewrite rule", logutil.Key("key", rg.EndKey))
    79  				} else {
    80  					log.Debug(
    81  						"rewrite end key",
    82  						logutil.Key("key", rg.EndKey),
    83  						logutil.RewriteRule(rule))
    84  				}
    85  			} else {
    86  				log.Warn("table id does not match",
    87  					logutil.Key("startKey", rg.StartKey),
    88  					logutil.Key("endKey", rg.EndKey),
    89  					zap.Int64("startID", startID),
    90  					zap.Int64("endID", endID))
    91  				return nil, errors.Annotate(berrors.ErrRestoreTableIDMismatch, "table id mismatch")
    92  			}
    93  		}
    94  		if out := rangeTree.InsertRange(rg); out != nil {
    95  			log.Error("insert ranges overlapped",
    96  				logutil.Key("startKeyOut", out.StartKey),
    97  				logutil.Key("endKeyOut", out.EndKey),
    98  				logutil.Key("startKeyIn", rg.StartKey),
    99  				logutil.Key("endKeyIn", rg.EndKey))
   100  			return nil, errors.Annotatef(berrors.ErrRestoreInvalidRange, "ranges overlapped")
   101  		}
   102  	}
   103  	sortedRanges := rangeTree.GetSortedRanges()
   104  	return sortedRanges, nil
   105  }
   106  
   107  // RegionInfo includes a region and the leader of the region.
   108  type RegionInfo struct {
   109  	Region *metapb.Region
   110  	Leader *metapb.Peer
   111  }
   112  
   113  // ContainsInterior returns whether the region contains the given key, and also
   114  // that the key does not fall on the boundary (start key) of the region.
   115  func (region *RegionInfo) ContainsInterior(key []byte) bool {
   116  	return bytes.Compare(key, region.Region.GetStartKey()) > 0 &&
   117  		(len(region.Region.GetEndKey()) == 0 ||
   118  			bytes.Compare(key, region.Region.GetEndKey()) < 0)
   119  }
   120  
   121  // RewriteRules contains rules for rewriting keys of tables.
   122  type RewriteRules struct {
   123  	Data []*import_sstpb.RewriteRule
   124  }
   125  
   126  // Append append its argument to this rewrite rules.
   127  func (r *RewriteRules) Append(other RewriteRules) {
   128  	r.Data = append(r.Data, other.Data...)
   129  }
   130  
   131  // EmptyRewriteRule make a new, empty rewrite rule.
   132  func EmptyRewriteRule() *RewriteRules {
   133  	return &RewriteRules{
   134  		Data: []*import_sstpb.RewriteRule{},
   135  	}
   136  }