pkg: add new range tree to store key range struct

server/statistics/buckets/bucket_tree.go

@@ -0,0 +1,97 @@
+// 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 buckets
+
+import (
+	"bytes"
+
+	"github.com/tikv/pd/pkg/btree"
+)
+
+// BucketTree is a buffer, the key range must be continuous.
+type BucketTree struct {
+	tree *btree.BTree
+}
+
+// NewBucketTree creates a new bucket tree.
+func NewBucketTree(degree int) *BucketTree {
+	return &BucketTree{
+		tree: btree.New(degree),
+	}
+}
+
+// BucketItem is bucket tree item.
+type BucketItem interface {
+	Less(than btree.Item) bool
+	StartKey() []byte
+	EndKey() []byte
+	// Debris returns the debris after replacing the key range.
+	Debris(startKey, endKey []byte) []BucketItem
+	String() string
+}
+
+// Len returns the length of the bucket tree.
+func (r *BucketTree) Len() int {
+	return r.tree.Len()
+}
+
+// GetRange returns the items that belong the key range.
+// cache key range:  |001-----100|100-----200|
+// request key range:   |005-----120|
+// return items:     |001-----100|100-----200|
+func (r *BucketTree) GetRange(item BucketItem) []BucketItem {
+	var res []BucketItem
+
+	var first BucketItem
+	r.tree.DescendLessOrEqual(item, func(i btree.Item) bool {
+		first = i.(BucketItem)
+		return false
+	})
+
+	// find the first item that contains the start key. if not found,
+	// it will use the param.
+	if first == nil || !(bytes.Compare(first.StartKey(), item.StartKey()) <= 0 &&
+		bytes.Compare(item.StartKey(), first.EndKey()) < 0) {
+		first = item
+	}
+
+	// find the next item util the item greater than end key.
+	r.tree.AscendGreaterOrEqual(first, func(i btree.Item) bool {
+		ringItem := i.(BucketItem)
+		if len(item.EndKey()) > 0 && bytes.Compare(ringItem.StartKey(), item.EndKey()) >= 0 {
+			return false
+		}
+		res = append(res, ringItem)
+		return true
+	})
+	return res
+}
+
+// Put puts a new item into the bucket tree.
+func (r *BucketTree) Put(item BucketItem) {
+	overlaps := r.GetRange(item)
+	for _, overlap := range overlaps {
+		r.tree.Delete(overlap)
+		others := overlap.Debris(item.StartKey(), item.EndKey())
+		for _, other := range others {
+			if bytes.Equal(other.StartKey(), other.EndKey()) {
+				r.tree.Delete(other)
+			} else {
+				r.tree.ReplaceOrInsert(other)
+			}
+		}
+	}
+	r.tree.ReplaceOrInsert(item)
+}

server/statistics/buckets/bucket_treee_test.go

@@ -0,0 +1,151 @@
+// 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 buckets
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	. "github.com/pingcap/check"
+	"github.com/tikv/pd/pkg/btree"
+)
+
+func Test(t *testing.T) {
+	TestingT(t)
+}
+
+var _ = Suite(&testBucketSuite{})
+
+type testBucketSuite struct {
+}
+
+type simpleBucketItem struct {
+	startKey []byte
+	endKey   []byte
+}
+
+func newSimpleBucketItem(startKey, endKey []byte) *simpleBucketItem {
+	return &simpleBucketItem{
+		startKey: startKey,
+		endKey:   endKey,
+	}
+}
+
+// String implements String.
+func (s *simpleBucketItem) String() string {
+	return fmt.Sprintf("key-range: [%s, %s]", s.startKey, s.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.StartKey(), than.(BucketItem).StartKey()) < 0
+}
+
+// Debris returns the debris of the item.
+// details: https://leetcode.cn/problems/interval-list-intersections/
+func (s simpleBucketItem) Debris(startKey, endKey []byte) []BucketItem {
+	var res []BucketItem
+
+	left := maxKey(startKey, s.startKey)
+	right := minKey(endKey, s.endKey)
+	// 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(s.startKey, left) {
+		res = append(res, newSimpleBucketItem(s.startKey, left))
+	}
+	// the right has oen intersection like |010 - 100| and |010 - 099|.
+	if !bytes.Equal(right, s.endKey) {
+		res = append(res, newSimpleBucketItem(right, s.endKey))
+	}
+	return res
+}
+
+// EndKey returns the end key of the item.
+func (s *simpleBucketItem) EndKey() []byte {
+	return s.endKey
+}
+
+// StartKey returns the start key of the item.
+func (s *simpleBucketItem) StartKey() []byte {
+	return s.startKey
+}
+
+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
+}
+
+func (bs *testBucketSuite) TestRingPutItem(c *C) {
+	bucketTree := NewBucketTree(2)
+	bucketTree.Put(newSimpleBucketItem([]byte("002"), []byte("100")))
+	c.Assert(bucketTree.tree.Len(), Equals, 1)
+	bucketTree.Put(newSimpleBucketItem([]byte("100"), []byte("200")))
+	c.Assert(bucketTree.tree.Len(), Equals, 2)
+
+	// init key range: [002,100], [100,200]
+	c.Assert(bucketTree.GetRange(newSimpleBucketItem([]byte("000"), []byte("002"))), HasLen, 0)
+	c.Assert(bucketTree.GetRange(newSimpleBucketItem([]byte("000"), []byte("009"))), HasLen, 1)
+	c.Assert(bucketTree.GetRange(newSimpleBucketItem([]byte("010"), []byte("090"))), HasLen, 1)
+	c.Assert(bucketTree.GetRange(newSimpleBucketItem([]byte("010"), []byte("110"))), HasLen, 2)
+	c.Assert(bucketTree.GetRange(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.Put(newSimpleBucketItem([]byte("010"), []byte("090")))
+	c.Assert(bucketTree.tree.Len(), Equals, 4)
+	c.Assert(bucketTree.GetRange(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.Put(newSimpleBucketItem([]byte("001"), []byte("080")))
+	c.Assert(bucketTree.tree.Len(), Equals, 4)
+	c.Assert(bucketTree.GetRange(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.Put(newSimpleBucketItem([]byte("001"), []byte("120")))
+	c.Assert(bucketTree.tree.Len(), Equals, 2)
+	c.Assert(bucketTree.GetRange(newSimpleBucketItem([]byte("010"), []byte("090"))), HasLen, 1)
+}
+
+func (bs *testBucketSuite) TestDebris(c *C) {
+	ringItem := newSimpleBucketItem([]byte("010"), []byte("090"))
+	var overlaps []BucketItem
+	overlaps = ringItem.Debris([]byte("000"), []byte("100"))
+	c.Assert(overlaps, HasLen, 0)
+	overlaps = ringItem.Debris([]byte("000"), []byte("080"))
+	c.Assert(overlaps, HasLen, 1)
+	overlaps = ringItem.Debris([]byte("020"), []byte("080"))
+	c.Assert(overlaps, HasLen, 2)
+	overlaps = ringItem.Debris([]byte("010"), []byte("090"))
+	c.Assert(overlaps, HasLen, 0)
+	overlaps = ringItem.Debris([]byte("010"), []byte("100"))
+	c.Assert(overlaps, HasLen, 0)
+	overlaps = ringItem.Debris([]byte("100"), []byte("200"))
+	c.Assert(overlaps, HasLen, 0)
+}