*: optimize heartbeat process with asyun runner

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

pkg/cluster/cluster.go

@@ -28,6 +28,7 @@ type Cluster interface {
 	GetLabelStats() *statistics.LabelStatistics
 	GetCoordinator() *schedule.Coordinator
 	GetRuleManager() *placement.RuleManager
+	GetBasicCluster() *core.BasicCluster
 }
 
 // HandleStatsAsync handles the flow asynchronously.
@@ -55,8 +56,10 @@ func HandleOverlaps(c Cluster, overlaps []*core.RegionInfo) {
 }
 
 // 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.
+		region = c.GetBasicCluster().GetRegion(region.GetID())
+		c.GetRegionStats().Observe(region, c.GetBasicCluster().GetRegionStores(region))
 	}
 }

pkg/core/region.go

@@ -16,6 +16,7 @@ package core
 
 import (
 	"bytes"
+	"context"
 	"encoding/hex"
 	"fmt"
 	"math"
@@ -35,6 +36,8 @@ import (
 	"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/ctxutil"
 	"github.com/tikv/pd/pkg/utils/logutil"
 	"github.com/tikv/pd/pkg/utils/syncutil"
 	"github.com/tikv/pd/pkg/utils/typeutil"
@@ -711,20 +714,51 @@ func (r *RegionInfo) isRegionRecreated() bool {
 
 // 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 context.Context, 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(msg string, fields ...zap.Field) {}
-	debug, info := noLog, noLog
+	d, i := noLog, noLog
+	debug, info := d, i
 	if enableLog {
-		debug = log.Debug
-		info = log.Info
+		d = log.Debug
+		i = log.Info
+		debug, info = d, i
 	}
 	// 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 context.Context, region, origin *RegionInfo) (saveKV, saveCache, needSync bool) {
+		taskRunner, ok := ctx.Value("taskRunner").(ratelimit.Runner)
+		limiter, _ := ctx.Value("limiter").(*ratelimit.ConcurrencyLimiter)
+		// print log asynchronously
+		if ok {
+			debug = func(msg string, fields ...zap.Field) {
+				taskRunner.RunTask(
+					ctx,
+					ratelimit.TaskOpts{
+						TaskName: "Log",
+						Limit:    limiter,
+					},
+					func(ctx context.Context) {
+						d(msg, fields...)
+					},
+				)
+			}
+			info = func(msg string, fields ...zap.Field) {
+				taskRunner.RunTask(
+					ctx,
+					ratelimit.TaskOpts{
+						TaskName: "Log",
+						Limit:    limiter,
+					},
+					func(ctx context.Context) {
+						i(msg, fields...)
+					},
+				)
+			}
+		}
 		if origin == nil {
 			if log.GetLevel() <= zap.DebugLevel {
 				debug("insert new region",
@@ -789,7 +823,7 @@ func GenerateRegionGuideFunc(enableLog bool) RegionGuideFunc {
 			}
 			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 +946,7 @@ func (r *RegionsInfo) CheckAndPutRegion(region *RegionInfo) []*RegionInfo {
 	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 +967,116 @@ func (r *RegionsInfo) PutRegion(region *RegionInfo) []*RegionInfo {
 }
 
 // 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 context.Context, region *RegionInfo) ([]*RegionInfo, error) {
+	tracer, ok := ctx.Value("tracer").(RegionHeartbeatProcessTracer)
+	if !ok {
+		tracer = NewNoopHeartbeatProcessTracer()
+	}
 	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
+}
+
+// CheckAndPutRootTree checks if the region is valid to put to the root, if valid then return error.
+// Usually used with AtomicCheckAndPutSubTree together.
+func (r *RegionsInfo) CheckAndPutRootTree(ctx context.Context, region *RegionInfo) ([]*RegionInfo, error) {
+	tracer, ok := ctx.Value(ctxutil.HeartbeatTracerKey).(RegionHeartbeatProcessTracer)
+	if !ok {
+		tracer = NewNoopHeartbeatProcessTracer()
+	}
+	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})
+	}
+	tracer.OnCheckOverlapsFinished()
+	err := check(region, origin, convertItemsToRegions(ols))
+	if err != nil {
+		r.t.Unlock()
+		tracer.OnValidateRegionFinished()
+		return nil, err
+	}
+	tracer.OnValidateRegionFinished()
+	_, overlaps, _ := r.setRegionLocked(region, true, ols...)
+	r.t.Unlock()
+	tracer.OnSetRegionFinished()
 	return overlaps, nil
 }
 
+// CheckAndPutSubTree checks if the region is valid to put to the sub tree, if valid then return error.
+// Usually used with AtomicCheckAndPutRootTree together.
+func (r *RegionsInfo) CheckAndPutSubTree(ctx context.Context, region *RegionInfo) error {
+	var origin *RegionInfo
+	r.st.RLock()
+	// get origin from sub tree.
+	originItem, ok := r.subRegions[region.GetID()]
+	if ok {
+		origin = originItem.RegionInfo
+	}
+	r.st.RUnlock()
+	// new region get from root tree again
+	var (
+		overlaps  []*RegionInfo
+		newRegion *RegionInfo
+	)
+	rangeChanged := true
+	if origin != nil {
+		newRegion, overlaps = r.GetRelevantRegions(origin)
+		rangeChanged = !origin.rangeEqualsTo(region)
+	} else {
+		newRegion = r.GetRegion(region.GetID())
+	}
+	r.UpdateSubTree(newRegion, origin, overlaps, rangeChanged)
+	return nil
+}
+
 // GetRelevantRegions returns the relevant regions for a given region.
-func (r *RegionsInfo) GetRelevantRegions(region *RegionInfo, trace 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() {
@@ -1065,6 +1167,9 @@ func (r *RegionsInfo) UpdateSubTree(region, origin *RegionInfo, overlaps []*Regi
 	r.st.Lock()
 	defer r.st.Unlock()
 	if origin != nil {
+		// TODO:  Check if the pointer address is consistent？
+		// eg &origin == &subtree[origin.GetID].
+		// compare the pointer of subtree with the pointer of origin to ensure no changes have occurred before update subtree.
 		if rangeChanged || !origin.peersEqualTo(region) {
 			// If the range or peers have changed, the sub regionTree needs to be cleaned up.
 			// TODO: Improve performance by deleting only the different peers.

pkg/core/region_test.go

@@ -15,6 +15,7 @@
 package core
 
 import (
+	"context"
 	"crypto/rand"
 	"fmt"
 	"math"
@@ -363,7 +364,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(context.TODO(), regionA, regionB)
 		re.Equal(testCase.needSync, needSync)
 	}
 }
@@ -459,9 +460,9 @@ func TestSetRegionConcurrence(t *testing.T) {
 	regions := NewRegionsInfo()
 	region := NewTestRegionInfo(1, 1, []byte("a"), []byte("b"))
 	go func() {
-		regions.AtomicCheckAndPutRegion(region, NewNoopHeartbeatProcessTracer())
+		regions.AtomicCheckAndPutRegion(context.TODO(), region)
 	}()
-	regions.AtomicCheckAndPutRegion(region, NewNoopHeartbeatProcessTracer())
+	regions.AtomicCheckAndPutRegion(context.TODO(), region)
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/core/UpdateSubTree"))
 }
 

pkg/core/region_tree.go

@@ -35,6 +35,10 @@ func (r *regionItem) GetStartKey() []byte {
 	return r.meta.StartKey
 }
 
+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