pkg: add new range tree to store key range struct (#4913)

ref #4726

Signed-off-by: bufferflies <[email protected]>

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
+}

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)
+}