/
rtree.go
227 lines (205 loc) · 5.96 KB
/
rtree.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
// Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
package rtree
import (
"bytes"
"github.com/google/btree"
backuppb "github.com/pingcap/kvproto/pkg/brpb"
"github.com/pingcap/log"
"github.com/tikv/migration/br/pkg/logutil"
)
// Range represents a backup response.
type Range struct {
StartKey []byte
EndKey []byte
Files []*backuppb.File
}
// BytesAndKeys returns total bytes and keys in a range.
func (rg *Range) BytesAndKeys() (bytes, keys uint64) {
for _, f := range rg.Files {
bytes += f.TotalBytes
keys += f.TotalKvs
}
return
}
// Intersect returns intersect range in the tree.
func (rg *Range) Intersect(
start, end []byte,
) (subStart, subEnd []byte, isIntersect bool) {
// empty mean the max end key
if len(rg.EndKey) != 0 && bytes.Compare(start, rg.EndKey) >= 0 {
isIntersect = false
return
}
if len(end) != 0 && bytes.Compare(end, rg.StartKey) <= 0 {
isIntersect = false
return
}
isIntersect = true
if bytes.Compare(start, rg.StartKey) >= 0 {
subStart = start
} else {
subStart = rg.StartKey
}
switch {
case len(end) == 0:
subEnd = rg.EndKey
case len(rg.EndKey) == 0:
subEnd = end
case bytes.Compare(end, rg.EndKey) < 0:
subEnd = end
default:
subEnd = rg.EndKey
}
return
}
// Contains check if the range contains the given key, [start, end).
func (rg *Range) Contains(key []byte) bool {
start, end := rg.StartKey, rg.EndKey
return bytes.Compare(key, start) >= 0 &&
(len(end) == 0 || bytes.Compare(key, end) < 0)
}
// Less impls btree.Item.
func (rg *Range) Less(than btree.Item) bool {
// rg.StartKey < than.StartKey
ta := than.(*Range)
return bytes.Compare(rg.StartKey, ta.StartKey) < 0
}
var _ btree.Item = &Range{}
// RangeTree is sorted tree for Ranges.
// All the ranges it stored do not overlap.
type RangeTree struct {
*btree.BTree
hasValidRange bool
}
// NewRangeTree returns an empty range tree.
func NewRangeTree() RangeTree {
return RangeTree{
BTree: btree.New(32),
hasValidRange: false,
}
}
// Find is a helper function to find an item that contains the range start
// key.
func (rangeTree *RangeTree) Find(rg *Range) *Range {
var ret *Range
rangeTree.DescendLessOrEqual(rg, func(i btree.Item) bool {
ret = i.(*Range)
return false
})
if ret == nil || !ret.Contains(rg.StartKey) {
return nil
}
return ret
}
// getOverlaps gets the ranges which are overlapped with the specified range range.
func (rangeTree *RangeTree) getOverlaps(rg *Range) []*Range {
// note that find() gets the last item that is less or equal than the range.
// in the case: |_______a_______|_____b_____|___c___|
// new range is |______d______|
// find() will return Range of range_a
// and both startKey of range_a and range_b are less than endKey of range_d,
// thus they are regarded as overlapped ranges.
found := rangeTree.Find(rg)
if found == nil {
found = rg
}
var overlaps []*Range
rangeTree.AscendGreaterOrEqual(found, func(i btree.Item) bool {
over := i.(*Range)
if len(rg.EndKey) > 0 && bytes.Compare(rg.EndKey, over.StartKey) <= 0 {
return false
}
overlaps = append(overlaps, over)
return true
})
return overlaps
}
// Update inserts range into tree and delete overlapping ranges.
func (rangeTree *RangeTree) Update(rg Range) {
overlaps := rangeTree.getOverlaps(&rg)
// Range has backuped, overwrite overlapping range.
for _, item := range overlaps {
log.Info("delete overlapping range",
logutil.Key("startKey", item.StartKey),
logutil.Key("endKey", item.EndKey))
rangeTree.Delete(item)
}
rangeTree.ReplaceOrInsert(&rg)
rangeTree.hasValidRange = true
}
// Put forms a range and inserts it into tree.
func (rangeTree *RangeTree) Put(
startKey, endKey []byte, files []*backuppb.File,
) {
rg := Range{
StartKey: startKey,
EndKey: endKey,
Files: files,
}
rangeTree.Update(rg)
}
// InsertRange inserts ranges into the range tree.
// It returns a non-nil range if there are soe overlapped ranges.
func (rangeTree *RangeTree) InsertRange(rg Range) *Range {
out := rangeTree.ReplaceOrInsert(&rg)
if out == nil {
return nil
}
return out.(*Range)
}
// GetSortedRanges collects and returns sorted ranges.
func (rangeTree *RangeTree) GetSortedRanges() []Range {
sortedRanges := make([]Range, 0, rangeTree.Len())
rangeTree.Ascend(func(rg btree.Item) bool {
if rg == nil {
return false
}
sortedRanges = append(sortedRanges, *rg.(*Range))
return true
})
return sortedRanges
}
// GetIncompleteRange returns missing range covered by startKey and endKey.
func (rangeTree *RangeTree) GetIncompleteRange(
startKey, endKey []byte,
) []Range {
if len(startKey) != 0 && bytes.Equal(startKey, endKey) {
return []Range{}
}
incomplete := make([]Range, 0, 64)
requsetRange := Range{StartKey: startKey, EndKey: endKey}
lastEndKey := startKey
pviot := &Range{StartKey: startKey}
if first := rangeTree.Find(pviot); first != nil {
pviot.StartKey = first.StartKey
}
rangeTree.AscendGreaterOrEqual(pviot, func(i btree.Item) bool {
rg := i.(*Range)
if bytes.Compare(lastEndKey, rg.StartKey) < 0 {
start, end, isIntersect :=
requsetRange.Intersect(lastEndKey, rg.StartKey)
if isIntersect {
// There is a gap between the last item and the current item.
incomplete =
append(incomplete, Range{StartKey: start, EndKey: end})
}
}
lastEndKey = rg.EndKey
return len(endKey) == 0 || bytes.Compare(rg.EndKey, endKey) < 0
})
// Check whether we need append the last range
if !bytes.Equal(lastEndKey, endKey) && len(lastEndKey) != 0 &&
(len(endKey) == 0 || bytes.Compare(lastEndKey, endKey) < 0) {
start, end, isIntersect := requsetRange.Intersect(lastEndKey, endKey)
if isIntersect {
incomplete =
append(incomplete, Range{StartKey: start, EndKey: end})
}
}
// is no valid range in Btree, and startKey/endKey are both empty, incomplete is empty too.
// This is not the result we need, append the whole input range.
if !rangeTree.hasValidRange && len(incomplete) == 0 {
incomplete = append(incomplete, Range{StartKey: startKey, EndKey: endKey})
}
return incomplete
}