diff --git a/pkg/rangetree/range_tree.go b/pkg/rangetree/range_tree.go new file mode 100644 index 00000000000..d647e1e78d0 --- /dev/null +++ b/pkg/rangetree/range_tree.go @@ -0,0 +1,164 @@ +// Copyright 2022 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package rangetree + +import ( + "bytes" + + "github.com/tikv/pd/pkg/btree" +) + +// RangeItem is one key range tree item. +type RangeItem interface { + btree.Item + GetStartKey() []byte + GetEndKey() []byte +} + +// DebrisFactory is the factory that generates some debris when updating items. +type DebrisFactory func(startKey, EndKey []byte, item RangeItem) []RangeItem + +// RangeTree is the tree contains RangeItems. +type RangeTree struct { + tree *btree.BTree + factory DebrisFactory +} + +// NewRangeTree is the constructor of the range tree. +func NewRangeTree(degree int, factory DebrisFactory) *RangeTree { + return &RangeTree{ + tree: btree.New(degree), + factory: factory, + } +} + +// Update insert the item and delete overlaps. +func (r *RangeTree) Update(item RangeItem) []RangeItem { + overlaps := r.GetOverlaps(item) + for _, old := range overlaps { + r.tree.Delete(old) + children := r.factory(item.GetStartKey(), item.GetEndKey(), old) + for _, child := range children { + if bytes.Compare(child.GetStartKey(), child.GetEndKey()) < 0 { + r.tree.ReplaceOrInsert(child) + } + } + } + r.tree.ReplaceOrInsert(item) + return overlaps +} + +// GetOverlaps returns the range items that has some intersections with the given items. +func (r *RangeTree) GetOverlaps(item RangeItem) []RangeItem { + // note that Find() gets the last item that is less or equal than the item. + // in the case: |_______a_______|_____b_____|___c___| + // new item is |______d______| + // Find() will return RangeItem of item_a + // and both startKey of item_a and item_b are less than endKey of item_d, + // thus they are regarded as overlapped items. + result := r.Find(item) + if result == nil { + result = item + } + + var overlaps []RangeItem + r.tree.AscendGreaterOrEqual(result, func(i btree.Item) bool { + over := i.(RangeItem) + if len(item.GetEndKey()) > 0 && bytes.Compare(item.GetEndKey(), over.GetStartKey()) <= 0 { + return false + } + overlaps = append(overlaps, over) + return true + }) + return overlaps +} + +// Find returns the range item contains the start key. +func (r *RangeTree) Find(item RangeItem) RangeItem { + var result RangeItem + r.tree.DescendLessOrEqual(item, func(i btree.Item) bool { + result = i.(RangeItem) + return false + }) + + if result == nil || !contains(result, item.GetStartKey()) { + return nil + } + + return result +} + +func contains(item RangeItem, key []byte) bool { + start, end := item.GetStartKey(), item.GetEndKey() + return bytes.Compare(key, start) >= 0 && (len(end) == 0 || bytes.Compare(key, end) < 0) +} + +// Remove removes the given item and return the deleted item. +func (r *RangeTree) Remove(item RangeItem) RangeItem { + if r := r.tree.Delete(item); r != nil { + return r.(RangeItem) + } + return nil +} + +// Len returns the count of the range tree. +func (r *RangeTree) Len() int { + return r.tree.Len() +} + +// ScanRange scan the start item util the result of the function is false. +func (r *RangeTree) ScanRange(start RangeItem, f func(_ RangeItem) bool) { + // Find if there is one item with key range [s, d), s < startKey < d + startItem := r.Find(start) + if startItem == nil { + startItem = start + } + r.tree.AscendGreaterOrEqual(startItem, func(item btree.Item) bool { + return f(item.(RangeItem)) + }) +} + +// GetAdjacentItem returns the adjacent range item. +func (r *RangeTree) GetAdjacentItem(item RangeItem) (prev RangeItem, next RangeItem) { + r.tree.AscendGreaterOrEqual(item, func(i btree.Item) bool { + if bytes.Equal(item.GetStartKey(), i.(RangeItem).GetStartKey()) { + return true + } + next = i.(RangeItem) + return false + }) + r.tree.DescendLessOrEqual(item, func(i btree.Item) bool { + if bytes.Equal(item.GetStartKey(), i.(RangeItem).GetStartKey()) { + return true + } + prev = i.(RangeItem) + return false + }) + return prev, next +} + +// GetAt returns the given index item. +func (r *RangeTree) GetAt(index int) RangeItem { + return r.tree.GetAt(index).(RangeItem) +} + +// GetWithIndex returns index and item for the given item. +func (r *RangeTree) GetWithIndex(item RangeItem) (RangeItem, int) { + rst, index := r.tree.GetWithIndex(item) + if rst == nil { + return nil, index + } + return rst.(RangeItem), index +} diff --git a/pkg/rangetree/range_tree_test.go b/pkg/rangetree/range_tree_test.go new file mode 100644 index 00000000000..d1e9cd79de5 --- /dev/null +++ b/pkg/rangetree/range_tree_test.go @@ -0,0 +1,145 @@ +// Copyright 2022 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package rangetree + +import ( + "bytes" + "testing" + + . "github.com/pingcap/check" + "github.com/tikv/pd/pkg/btree" +) + +func Test(t *testing.T) { + TestingT(t) +} + +var _ = Suite(&testRangeTreeSuite{}) + +type testRangeTreeSuite struct { +} + +type simpleBucketItem struct { + startKey []byte + endKey []byte +} + +func newSimpleBucketItem(startKey, endKey []byte) *simpleBucketItem { + return &simpleBucketItem{ + startKey: startKey, + endKey: endKey, + } +} + +// Less returns true if the start key of the item is less than the start key of the argument. +func (s *simpleBucketItem) Less(than btree.Item) bool { + return bytes.Compare(s.GetStartKey(), than.(RangeItem).GetStartKey()) < 0 +} + +// StartKey returns the start key of the item. +func (s *simpleBucketItem) GetStartKey() []byte { + return s.startKey +} + +// EndKey returns the end key of the item. +func (s *simpleBucketItem) GetEndKey() []byte { + return s.endKey +} + +func minKey(a, b []byte) []byte { + if bytes.Compare(a, b) < 0 { + return a + } + return b +} + +func maxKey(a, b []byte) []byte { + if bytes.Compare(a, b) > 0 { + return a + } + return b +} + +// Debris returns the debris of the item. +// details: https://leetcode.cn/problems/interval-list-intersections/ +func bucketDebrisFactory(startKey, endKey []byte, item RangeItem) []RangeItem { + var res []RangeItem + + left := maxKey(startKey, item.GetStartKey()) + right := minKey(endKey, item.GetEndKey()) + // they have no intersection if they are neighbour like |010 - 100| and |100 - 200|. + if bytes.Compare(left, right) >= 0 { + return nil + } + // the left has oen intersection like |010 - 100| and |020 - 100|. + if !bytes.Equal(item.GetStartKey(), left) { + res = append(res, newSimpleBucketItem(item.GetStartKey(), left)) + } + // the right has oen intersection like |010 - 100| and |010 - 099|. + if !bytes.Equal(right, item.GetEndKey()) { + res = append(res, newSimpleBucketItem(right, item.GetEndKey())) + } + return res +} + +func (bs *testRangeTreeSuite) TestRingPutItem(c *C) { + bucketTree := NewRangeTree(2, bucketDebrisFactory) + bucketTree.Update(newSimpleBucketItem([]byte("002"), []byte("100"))) + c.Assert(bucketTree.Len(), Equals, 1) + bucketTree.Update(newSimpleBucketItem([]byte("100"), []byte("200"))) + c.Assert(bucketTree.Len(), Equals, 2) + + // init key range: [002,100], [100,200] + c.Assert(bucketTree.GetOverlaps(newSimpleBucketItem([]byte("000"), []byte("002"))), HasLen, 0) + c.Assert(bucketTree.GetOverlaps(newSimpleBucketItem([]byte("000"), []byte("009"))), HasLen, 1) + c.Assert(bucketTree.GetOverlaps(newSimpleBucketItem([]byte("010"), []byte("090"))), HasLen, 1) + c.Assert(bucketTree.GetOverlaps(newSimpleBucketItem([]byte("010"), []byte("110"))), HasLen, 2) + c.Assert(bucketTree.GetOverlaps(newSimpleBucketItem([]byte("200"), []byte("300"))), HasLen, 0) + + // test1： insert one key range, the old overlaps will retain like split buckets. + // key range: [002,010],[010,090],[090,100],[100,200] + bucketTree.Update(newSimpleBucketItem([]byte("010"), []byte("090"))) + c.Assert(bucketTree.Len(), Equals, 4) + c.Assert(bucketTree.GetOverlaps(newSimpleBucketItem([]byte("010"), []byte("090"))), HasLen, 1) + + // test2: insert one key range, the old overlaps will retain like merge . + // key range: [001,080], [080,090],[090,100],[100,200] + bucketTree.Update(newSimpleBucketItem([]byte("001"), []byte("080"))) + c.Assert(bucketTree.Len(), Equals, 4) + c.Assert(bucketTree.GetOverlaps(newSimpleBucketItem([]byte("010"), []byte("090"))), HasLen, 2) + + // test2: insert one keyrange, the old overlaps will retain like merge . + // key range: [001,120],[120,200] + bucketTree.Update(newSimpleBucketItem([]byte("001"), []byte("120"))) + c.Assert(bucketTree.Len(), Equals, 2) + c.Assert(bucketTree.GetOverlaps(newSimpleBucketItem([]byte("010"), []byte("090"))), HasLen, 1) +} + +func (bs *testRangeTreeSuite) TestDebris(c *C) { + ringItem := newSimpleBucketItem([]byte("010"), []byte("090")) + var overlaps []RangeItem + overlaps = bucketDebrisFactory([]byte("000"), []byte("100"), ringItem) + c.Assert(overlaps, HasLen, 0) + overlaps = bucketDebrisFactory([]byte("000"), []byte("080"), ringItem) + c.Assert(overlaps, HasLen, 1) + overlaps = bucketDebrisFactory([]byte("020"), []byte("080"), ringItem) + c.Assert(overlaps, HasLen, 2) + overlaps = bucketDebrisFactory([]byte("010"), []byte("090"), ringItem) + c.Assert(overlaps, HasLen, 0) + overlaps = bucketDebrisFactory([]byte("010"), []byte("100"), ringItem) + c.Assert(overlaps, HasLen, 0) + overlaps = bucketDebrisFactory([]byte("100"), []byte("200"), ringItem) + c.Assert(overlaps, HasLen, 0) +} diff --git a/server/core/region_tree.go b/server/core/region_tree.go index c0f26d41ba0..4bb1dc9d0b9 100644 --- a/server/core/region_tree.go +++ b/server/core/region_tree.go @@ -23,19 +23,30 @@ import ( "github.com/tikv/pd/pkg/btree" "github.com/tikv/pd/pkg/errs" "github.com/tikv/pd/pkg/logutil" + "github.com/tikv/pd/pkg/rangetree" "go.uber.org/zap" ) -var _ btree.Item = ®ionItem{} +var _ rangetree.RangeItem = ®ionItem{} type regionItem struct { region *RegionInfo } +// GetStartKey returns the start key of the region. +func (r *regionItem) GetStartKey() []byte { + return r.region.GetStartKey() +} + +// GetEndKey returns the end key of the region. +func (r *regionItem) GetEndKey() []byte { + return r.region.GetEndKey() +} + // Less returns true if the region start key is less than the other. func (r *regionItem) Less(other btree.Item) bool { left := r.region.GetStartKey() - right := other.(*regionItem).region.GetStartKey() + right := other.(rangetree.RangeItem).GetStartKey() return bytes.Compare(left, right) < 0 } @@ -49,7 +60,7 @@ const ( ) type regionTree struct { - tree *btree.BTree + tree *rangetree.RangeTree // Statistics totalSize int64 totalWriteBytesRate float64 @@ -57,8 +68,11 @@ type regionTree struct { } func newRegionTree() *regionTree { + factory := func(_, _ []byte, _ rangetree.RangeItem) []rangetree.RangeItem { + return nil + } return ®ionTree{ - tree: btree.New(defaultBTreeDegree), + tree: rangetree.NewRangeTree(defaultBTreeDegree, factory), totalSize: 0, totalWriteBytesRate: 0, totalWriteKeysRate: 0, @@ -75,27 +89,11 @@ func (t *regionTree) length() int { // getOverlaps gets the regions which are overlapped with the specified region range. func (t *regionTree) getOverlaps(region *RegionInfo) []*RegionInfo { item := ®ionItem{region: region} - - // note that find() gets the last item that is less or equal than the region. - // in the case: |_______a_______|_____b_____|___c___| - // new region is |______d______| - // find() will return regionItem of region_a - // and both startKey of region_a and region_b are less than endKey of region_d, - // thus they are regarded as overlapped regions. - result := t.find(region) - if result == nil { - result = item + result := t.tree.GetOverlaps(item) + overlaps := make([]*RegionInfo, len(result)) + for i, r := range result { + overlaps[i] = r.(*regionItem).region } - - var overlaps []*RegionInfo - t.tree.AscendGreaterOrEqual(result, func(i btree.Item) bool { - over := i.(*regionItem) - if len(region.GetEndKey()) > 0 && bytes.Compare(region.GetEndKey(), over.region.GetStartKey()) <= 0 { - return false - } - overlaps = append(overlaps, over.region) - return true - }) return overlaps } @@ -109,21 +107,22 @@ func (t *regionTree) update(item *regionItem) []*RegionInfo { t.totalWriteBytesRate += regionWriteBytesRate t.totalWriteKeysRate += regionWriteKeysRate - overlaps := t.getOverlaps(region) - for _, old := range overlaps { + overlaps := t.tree.Update(item) + result := make([]*RegionInfo, len(overlaps)) + for i, overlap := range overlaps { + old := overlap.(*regionItem).region + result[i] = old log.Debug("overlapping region", zap.Uint64("region-id", old.GetID()), logutil.ZapRedactStringer("delete-region", RegionToHexMeta(old.GetMeta())), logutil.ZapRedactStringer("update-region", RegionToHexMeta(region.GetMeta()))) - t.tree.Delete(®ionItem{old}) t.totalSize -= old.approximateSize regionWriteBytesRate, regionWriteKeysRate = old.GetWriteRate() t.totalWriteBytesRate -= regionWriteBytesRate t.totalWriteKeysRate -= regionWriteKeysRate } - t.tree.ReplaceOrInsert(item) - return overlaps + return result } // updateStat is used to update statistics when regionItem.region is directly replaced. @@ -146,8 +145,9 @@ func (t *regionTree) remove(region *RegionInfo) { if t.length() == 0 { return } - result := t.find(region) - if result == nil || result.region.GetID() != region.GetID() { + item := ®ionItem{region: region} + result := t.tree.Find(item) + if result == nil || result.(*regionItem).region.GetID() != region.GetID() { return } @@ -155,7 +155,7 @@ func (t *regionTree) remove(region *RegionInfo) { regionWriteBytesRate, regionWriteKeysRate := region.GetWriteRate() t.totalWriteBytesRate -= regionWriteBytesRate t.totalWriteKeysRate -= regionWriteKeysRate - t.tree.Delete(result) + t.tree.Remove(result) } // search returns a region that contains the key. @@ -188,19 +188,14 @@ func (t *regionTree) searchPrev(regionKey []byte) *RegionInfo { // find is a helper function to find an item that contains the regions start // key. func (t *regionTree) find(region *RegionInfo) *regionItem { - item := ®ionItem{region: region} - - var result *regionItem - t.tree.DescendLessOrEqual(item, func(i btree.Item) bool { - result = i.(*regionItem) - return false - }) - - if result == nil || !result.Contains(region.GetStartKey()) { + item := t.tree.Find(®ionItem{region: region}) + if item == nil { return nil } - - return result + if result, ok := item.(*regionItem); ok && result.Contains(region.GetStartKey()) { + return result + } + return nil } // scanRage scans from the first region containing or behind the start key @@ -208,13 +203,11 @@ func (t *regionTree) find(region *RegionInfo) *regionItem { func (t *regionTree) scanRange(startKey []byte, f func(*RegionInfo) bool) { region := &RegionInfo{meta: &metapb.Region{StartKey: startKey}} // find if there is a region with key range [s, d), s < startKey < d - startItem := t.find(region) - if startItem == nil { - startItem = ®ionItem{region: &RegionInfo{meta: &metapb.Region{StartKey: startKey}}} + fn := func(item rangetree.RangeItem) bool { + r := item.(*regionItem) + return f(r.region) } - t.tree.AscendGreaterOrEqual(startItem, func(item btree.Item) bool { - return f(item.(*regionItem).region) - }) + t.tree.ScanRange(®ionItem{region: region}, fn) } func (t *regionTree) scanRanges() []*RegionInfo { @@ -231,21 +224,14 @@ func (t *regionTree) scanRanges() []*RegionInfo { func (t *regionTree) getAdjacentRegions(region *RegionInfo) (*regionItem, *regionItem) { item := ®ionItem{region: &RegionInfo{meta: &metapb.Region{StartKey: region.GetStartKey()}}} + prevItem, nextItem := t.tree.GetAdjacentItem(item) var prev, next *regionItem - t.tree.AscendGreaterOrEqual(item, func(i btree.Item) bool { - if bytes.Equal(item.region.GetStartKey(), i.(*regionItem).region.GetStartKey()) { - return true - } - next = i.(*regionItem) - return false - }) - t.tree.DescendLessOrEqual(item, func(i btree.Item) bool { - if bytes.Equal(item.region.GetStartKey(), i.(*regionItem).region.GetStartKey()) { - return true - } - prev = i.(*regionItem) - return false - }) + if prevItem != nil { + prev = prevItem.(*regionItem) + } + if nextItem != nil { + next = nextItem.(*regionItem) + } return prev, next }