impl bucket set informer

Signed-off-by: bufferflies <1045931706@qq.com>
tikv · Mar 4, 2022 · ea4c1d7 · ea4c1d7
1 parent 05ebcdb
commit ea4c1d7
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 7 deletions.
diff --git a/go.mod b/go.mod
@@ -28,7 +28,7 @@ require (
 	github.com/pingcap/errcode v0.3.0
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20200702092429-9f69995143ce
-	github.com/pingcap/kvproto v0.0.0-20220228094105-9bb22e5a97fc
+	github.com/pingcap/kvproto v0.0.0-20220304032058-ccd676426a27
 	github.com/pingcap/log v0.0.0-20210906054005-afc726e70354
 	github.com/pingcap/sysutil v0.0.0-20211208032423-041a72e5860d
 	github.com/pingcap/tidb-dashboard v0.0.0-20220117082709-e8076b5c79ba

diff --git a/go.sum b/go.sum
@@ -402,6 +402,8 @@ github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO
 github.com/pingcap/kvproto v0.0.0-20200411081810-b85805c9476c/go.mod h1:IOdRDPLyda8GX2hE/jO7gqaCV/PNFh8BZQCQZXfIOqI=
 github.com/pingcap/kvproto v0.0.0-20220228094105-9bb22e5a97fc h1:s4ObV2nLm0n8Y1f71qNbJUseGe6r96p8Uq+XGQHb7Kc=
 github.com/pingcap/kvproto v0.0.0-20220228094105-9bb22e5a97fc/go.mod h1:IOdRDPLyda8GX2hE/jO7gqaCV/PNFh8BZQCQZXfIOqI=
+github.com/pingcap/kvproto v0.0.0-20220304032058-ccd676426a27 h1:+Ax2NXyAFIITrzgSPWBo3SBZtX/D60VeELCG0B0hqiY=
+github.com/pingcap/kvproto v0.0.0-20220304032058-ccd676426a27/go.mod h1:IOdRDPLyda8GX2hE/jO7gqaCV/PNFh8BZQCQZXfIOqI=
 github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20200511115504-543df19646ad/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=

diff --git a/pkg/mock/mockcluster/mockcluster.go b/pkg/mock/mockcluster/mockcluster.go
@@ -40,7 +40,7 @@ import (
 const (
 	defaultStoreCapacity = 100 * (1 << 30) // 100GiB
 	defaultRegionSize    = 96 * (1 << 20)  // 96MiB
-	mb                   = (1 << 20)       // 1MiB
+	mb                   = 1 << 20         // 1MiB
 )
 
 // Cluster is used to mock a cluster for test purpose.

diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go
@@ -771,6 +771,16 @@ func (c *RaftCluster) GetRegions() []*core.RegionInfo {
 	return c.core.GetRegions()
 }
 
+// GetBucketByRegionID returns the bucket that the region belongs to.
+func (c *RaftCluster) GetBucketByRegionID(regionID uint64) *metapb.Buckets {
+	return c.core.GetBucketByRegionID(regionID)
+}
+
+// GetStoreRegions returns the bucket that key located.
+func (c *RaftCluster) GetBucketByKey(key []byte) *metapb.Buckets {
+	return c.core.GetBucketByKey(key)
+}
+
 // GetRegionCount returns total count of regions
 func (c *RaftCluster) GetRegionCount() int {
 	return c.core.GetRegionCount()

diff --git a/server/core/basic_cluster.go b/server/core/basic_cluster.go
@@ -44,12 +44,19 @@ func NewBasicCluster() *BasicCluster {
 }
 
 // GetBuckets returns the buckets of the given buckets.
-func (bc *BasicCluster) GetBuckets(regionID uint64) *metapb.Buckets {
+func (bc *BasicCluster) GetBucketByRegionID(regionID uint64) *metapb.Buckets {
 	bc.RLock()
 	defer bc.RUnlock()
 	return bc.Buckets.GetByRegionID(regionID)
 }
 
+// GetBucketByKey returns the bucket of the given key.
+func (bc *BasicCluster) GetBucketByKey(startKey []byte) *metapb.Buckets {
+	bc.RLock()
+	defer bc.RUnlock()
+	return bc.Buckets.GetKey(startKey)
+}
+
 // GetStores returns all Stores in the cluster.
 func (bc *BasicCluster) GetStores() []*StoreInfo {
 	bc.RLock()

diff --git a/server/core/bucket.go b/server/core/bucket.go
@@ -106,9 +106,20 @@ func (b *BucketsInfo) GetByRegionID(regionID uint64) *metapb.Buckets {
 	return b.buckets[regionID]
 }
 
-func (b *BucketsInfo) find(bucket *metapb.Buckets) *regionKeyItem {
+// GetKey returns the key of the bucket.
+func (b *BucketsInfo) GetKey(key []byte) *metapb.Buckets {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+	item := b.find(key)
+	if item != nil {
+		return b.buckets[item.ID]
+	}
+	return nil
+}
+
+func (b *BucketsInfo) find(startKey []byte) *regionKeyItem {
 	item := &regionKeyItem{
-		startKey: bucket.Keys[0],
+		startKey: startKey,
 	}
 	var result *regionKeyItem
 	b.mu.RLock()
@@ -124,8 +135,7 @@ func (b *BucketsInfo) find(bucket *metapb.Buckets) *regionKeyItem {
 }
 
 func (b *BucketsInfo) scanRange(startKey []byte, f func(item *regionKeyItem) bool) {
-	bucket := &metapb.Buckets{Keys: [][]byte{startKey}}
-	startItem := b.find(bucket)
+	startItem := b.find(startKey)
 	if startItem == nil {
 		startItem = &regionKeyItem{
 			startKey: startKey,

diff --git a/server/schedule/cluster.go b/server/schedule/cluster.go
@@ -25,6 +25,7 @@ type Cluster interface {
 	core.RegionSetInformer
 	core.StoreSetInformer
 	core.StoreSetController
+	core.BucketSetInformer
 
 	statistics.RegionStatInformer
 	statistics.StoreStatInformer