*: optimize heartbeat process with concurrent runner - part 1

pkg/cluster/cluster.go

@@ -28,6 +28,7 @@
 	GetLabelStats() *statistics.LabelStatistics
 	GetCoordinator() *schedule.Coordinator
 	GetRuleManager() *placement.RuleManager
+	GetBasicCluster() *core.BasicCluster
 }
 
 // HandleStatsAsync handles the flow asynchronously.
@@ -55,8 +56,17 @@
 }
 
 // Collect collects the cluster information.
-func Collect(c Cluster, region *core.RegionInfo, stores []*core.StoreInfo, hasRegionStats bool) {
+func Collect(c Cluster, region *core.RegionInfo, hasRegionStats bool) {
 	if hasRegionStats {
-		c.GetRegionStats().Observe(region, stores)
+		// get region again from root tree. make sure the observed region is the latest.
+		bc := c.GetBasicCluster()
+		if bc == nil {
+			return
+		}
+		region = bc.GetRegion(region.GetID())
+		if region == nil {
+			return
+		}
+		c.GetRegionStats().Observe(region, c.GetBasicCluster().GetRegionStores(region))
 	}
 }

pkg/core/context.go

@@ -0,0 +1,40 @@
+// Copyright 2024 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 core
+
+import (
+	"context"
+
+	"github.com/tikv/pd/pkg/ratelimit"
+)
+
+// MetaProcessContext is a context for meta process.
+type MetaProcessContext struct {
+	context.Context
+	Tracer     RegionHeartbeatProcessTracer
+	TaskRunner ratelimit.Runner
+	Limiter    *ratelimit.ConcurrencyLimiter
+}
+
+// NewMetaProcessContext creates a new MetaProcessContext.
+// used in tests, can be changed if no need to test concurrency.
+func ContextTODO() *MetaProcessContext {
+	return &MetaProcessContext{
+		Context:    context.TODO(),
+		Tracer:     NewNoopHeartbeatProcessTracer(),
+		TaskRunner: ratelimit.NewSyncRunner(),
+		// Limit default is nil
+	}
+}

pkg/core/region.go

@@ -16,6 +16,7 @@
 
 import (
 	"bytes"
+	"context"
 	"encoding/hex"
 	"fmt"
 	"math"
@@ -35,6 +36,7 @@
 	"github.com/pingcap/kvproto/pkg/replication_modepb"
 	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/errs"
+	"github.com/tikv/pd/pkg/ratelimit"
 	"github.com/tikv/pd/pkg/utils/logutil"
 	"github.com/tikv/pd/pkg/utils/syncutil"
 	"github.com/tikv/pd/pkg/utils/typeutil"
@@ -711,20 +713,50 @@
 
 // RegionGuideFunc is a function that determines which follow-up operations need to be performed based on the origin
 // and new region information.
-type RegionGuideFunc func(region, origin *RegionInfo) (saveKV, saveCache, needSync bool)
+type RegionGuideFunc func(ctx *MetaProcessContext, region, origin *RegionInfo) (saveKV, saveCache, needSync bool)
 
 // GenerateRegionGuideFunc is used to generate a RegionGuideFunc. Control the log output by specifying the log function.
 // nil means do not print the log.
 func GenerateRegionGuideFunc(enableLog bool) RegionGuideFunc {
 	noLog := func(string, ...zap.Field) {}
-	debug, info := noLog, noLog
+	d, i := noLog, noLog
 	if enableLog {
-		debug = log.Debug
-		info = log.Info
+		d = log.Debug
+		i = log.Info
 	}
 	// Save to storage if meta is updated.
 	// Save to cache if meta or leader is updated, or contains any down/pending peer.
-	return func(region, origin *RegionInfo) (saveKV, saveCache, needSync bool) {
+	return func(ctx *MetaProcessContext, region, origin *RegionInfo) (saveKV, saveCache, needSync bool) {
+		taskRunner := ctx.TaskRunner
+		limiter := ctx.Limiter
+		// print log asynchronously
+		debug, info := d, i
+		if taskRunner != nil {
+			debug = func(msg string, fields ...zap.Field) {
+				taskRunner.RunTask(
+					ctx.Context,
+					ratelimit.TaskOpts{
+						TaskName: "Log",
+						Limit:    limiter,
+					},
+					func(_ context.Context) {
+						d(msg, fields...)
+					},
+				)
+			}
+			info = func(msg string, fields ...zap.Field) {
+				taskRunner.RunTask(
+					ctx.Context,
+					ratelimit.TaskOpts{
+						TaskName: "Log",
+						Limit:    limiter,
+					},
+					func(_ context.Context) {
+						i(msg, fields...)
+					},
+				)
+			}
+		}
 		if origin == nil {
 			if log.GetLevel() <= zap.DebugLevel {
 				debug("insert new region",
@@ -789,7 +821,7 @@
 			}
 			if !SortedPeersStatsEqual(region.GetDownPeers(), origin.GetDownPeers()) {
 				if log.GetLevel() <= zap.DebugLevel {
-					debug("down-peers changed", zap.Uint64("region-id", region.GetID()))
+					debug("down-peers changed", zap.Uint64("region-id", region.GetID()), zap.Reflect("before", origin.GetDownPeers()), zap.Reflect("after", region.GetDownPeers()))
 				}
 				saveCache, needSync = true, true
 				return
@@ -912,7 +944,7 @@
 	if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) {
 		ols = r.tree.overlaps(&regionItem{RegionInfo: region})
 	}
-	err := check(region, origin, ols)
+	err := check(region, origin, convertItemsToRegions(ols))
 	if err != nil {
 		log.Debug("region is stale", zap.Stringer("origin", origin.GetMeta()), errs.ZapError(err))
 		// return the state region to delete.
@@ -933,48 +965,59 @@
 }
 
 // PreCheckPutRegion checks if the region is valid to put.
-func (r *RegionsInfo) PreCheckPutRegion(region *RegionInfo, trace RegionHeartbeatProcessTracer) (*RegionInfo, []*regionItem, error) {
-	origin, overlaps := r.GetRelevantRegions(region, trace)
+func (r *RegionsInfo) PreCheckPutRegion(region *RegionInfo) (*RegionInfo, []*RegionInfo, error) {
+	origin, overlaps := r.GetRelevantRegions(region)
 	err := check(region, origin, overlaps)
 	return origin, overlaps, err
 }
 
+func convertItemsToRegions(items []*regionItem) []*RegionInfo {
+	regions := make([]*RegionInfo, 0, len(items))
+	for _, item := range items {
+		regions = append(regions, item.RegionInfo)
+	}
+	return regions
+}
+
 // AtomicCheckAndPutRegion checks if the region is valid to put, if valid then put.
-func (r *RegionsInfo) AtomicCheckAndPutRegion(region *RegionInfo, trace RegionHeartbeatProcessTracer) ([]*RegionInfo, error) {
+func (r *RegionsInfo) AtomicCheckAndPutRegion(ctx *MetaProcessContext, region *RegionInfo) ([]*RegionInfo, error) {
+	tracer := ctx.Tracer
 	r.t.Lock()
 	var ols []*regionItem
 	origin := r.getRegionLocked(region.GetID())
 	if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) {
 		ols = r.tree.overlaps(&regionItem{RegionInfo: region})
 	}
-	trace.OnCheckOverlapsFinished()
-	err := check(region, origin, ols)
+	tracer.OnCheckOverlapsFinished()
+	err := check(region, origin, convertItemsToRegions(ols))
 	if err != nil {
 		r.t.Unlock()
-		trace.OnValidateRegionFinished()
+		tracer.OnValidateRegionFinished()
 		return nil, err
 	}
-	trace.OnValidateRegionFinished()
+	tracer.OnValidateRegionFinished()
 	origin, overlaps, rangeChanged := r.setRegionLocked(region, true, ols...)
 	r.t.Unlock()
-	trace.OnSetRegionFinished()
+	tracer.OnSetRegionFinished()
 	r.UpdateSubTree(region, origin, overlaps, rangeChanged)
-	trace.OnUpdateSubTreeFinished()
+	tracer.OnUpdateSubTreeFinished()
 	return overlaps, nil
 }
 
 // GetRelevantRegions returns the relevant regions for a given region.
-func (r *RegionsInfo) GetRelevantRegions(region *RegionInfo, _ RegionHeartbeatProcessTracer) (origin *RegionInfo, overlaps []*regionItem) {
+func (r *RegionsInfo) GetRelevantRegions(region *RegionInfo) (origin *RegionInfo, overlaps []*RegionInfo) {
 	r.t.RLock()
 	defer r.t.RUnlock()
 	origin = r.getRegionLocked(region.GetID())
 	if origin == nil || !bytes.Equal(origin.GetStartKey(), region.GetStartKey()) || !bytes.Equal(origin.GetEndKey(), region.GetEndKey()) {
-		overlaps = r.tree.overlaps(&regionItem{RegionInfo: region})
+		for _, item := range r.tree.overlaps(&regionItem{RegionInfo: region}) {
+			overlaps = append(overlaps, item.RegionInfo)
+		}
 	}
 	return
 }
 
-func check(region, origin *RegionInfo, overlaps []*regionItem) error {
+func check(region, origin *RegionInfo, overlaps []*RegionInfo) error {
 	for _, item := range overlaps {
 		// PD ignores stale regions' heartbeats, unless it is recreated recently by unsafe recover operation.
 		if region.GetRegionEpoch().GetVersion() < item.GetRegionEpoch().GetVersion() && !region.isRegionRecreated() {
@@ -1043,7 +1086,6 @@
 		item = &regionItem{RegionInfo: region}
 		r.regions[region.GetID()] = item
 	}
-
 	var overlaps []*RegionInfo
 	if rangeChanged {
 		overlaps = r.tree.update(item, withOverlaps, ol...)

pkg/core/region_test.go

@@ -363,7 +363,7 @@ func TestNeedSync(t *testing.T) {
 	for _, testCase := range testCases {
 		regionA := region.Clone(testCase.optionsA...)
 		regionB := region.Clone(testCase.optionsB...)
-		_, _, needSync := RegionGuide(regionA, regionB)
+		_, _, needSync := RegionGuide(ContextTODO(), regionA, regionB)
 		re.Equal(testCase.needSync, needSync)
 	}
 }
@@ -459,9 +459,9 @@ func TestSetRegionConcurrence(t *testing.T) {
 	regions := NewRegionsInfo()
 	region := NewTestRegionInfo(1, 1, []byte("a"), []byte("b"))
 	go func() {
-		regions.AtomicCheckAndPutRegion(region, NewNoopHeartbeatProcessTracer())
+		regions.AtomicCheckAndPutRegion(ContextTODO(), region)
 	}()
-	regions.AtomicCheckAndPutRegion(region, NewNoopHeartbeatProcessTracer())
+	regions.AtomicCheckAndPutRegion(ContextTODO(), region)
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/core/UpdateSubTree"))
 }
 

pkg/core/region_tree.go

@@ -35,6 +35,11 @@ func (r *regionItem) GetStartKey() []byte {
 	return r.meta.StartKey
 }
 
+// GetID returns the ID of the region.
+func (r *regionItem) GetID() uint64 {
+	return r.meta.GetId()
+}
+
 // GetEndKey returns the end key of the region.
 func (r *regionItem) GetEndKey() []byte {
 	return r.meta.EndKey

pkg/core/region_tree_test.go

@@ -158,6 +158,8 @@ func TestRegionTree(t *testing.T) {
 
 	updateNewItem(tree, regionA)
 	updateNewItem(tree, regionC)
+	re.Nil(tree.overlaps(newRegionItem([]byte("b"), []byte("c"))))
+	re.Equal(regionC, tree.overlaps(newRegionItem([]byte("a"), []byte("cc")))[1].RegionInfo)
 	re.Nil(tree.search([]byte{}))
 	re.Equal(regionA, tree.search([]byte("a")))
 	re.Nil(tree.search([]byte("b")))