statistics: reduce hot cache task memory retention

pkg/cluster/cluster.go

@@ -37,25 +37,7 @@ type Cluster interface {
 
 // HandleStatsAsync handles the flow asynchronously.
 func HandleStatsAsync(c Cluster, region *core.RegionInfo) {
-	checkWritePeerTask := func(cache *statistics.HotPeerCache) {
-		reportInterval := region.GetInterval()
-		interval := reportInterval.GetEndTimestamp() - reportInterval.GetStartTimestamp()
-		stats := cache.CheckPeerFlow(region, region.GetPeers(), region.GetWriteLoads(), interval)
-		for _, stat := range stats {
-			cache.UpdateStat(stat)
-		}
-	}
-
-	checkExpiredTask := func(cache *statistics.HotPeerCache) {
-		expiredStats := cache.CollectExpiredItems(region)
-		for _, stat := range expiredStats {
-			cache.UpdateStat(stat)
-		}
-	}
-
-	c.GetHotStat().CheckWriteAsync(checkExpiredTask)
-	c.GetHotStat().CheckReadAsync(checkExpiredTask)
-	c.GetHotStat().CheckWriteAsync(checkWritePeerTask)
+	c.GetHotStat().CheckRegionFlowAsync(region)
 	c.GetCoordinator().GetSchedulersController().CheckTransferWitnessLeader(region)
 }
 

pkg/cluster/cluster_test.go

@@ -0,0 +1,90 @@
+// Copyright 2026 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 cluster
+
+import (
+	"context"
+	"testing"
+
+	"github.com/pingcap/kvproto/pkg/metapb"
+
+	"github.com/tikv/pd/pkg/core"
+	"github.com/tikv/pd/pkg/mock/mockcluster"
+	"github.com/tikv/pd/pkg/mock/mockconfig"
+	"github.com/tikv/pd/pkg/schedule"
+	"github.com/tikv/pd/pkg/schedule/hbstream"
+	"github.com/tikv/pd/pkg/statistics"
+	"github.com/tikv/pd/pkg/statistics/utils"
+)
+
+type hotCacheBenchmarkCluster struct {
+	*mockcluster.Cluster
+	coordinator *schedule.Coordinator
+}
+
+func newHotCacheBenchmarkCluster(ctx context.Context) *hotCacheBenchmarkCluster {
+	cluster := mockcluster.NewCluster(ctx, mockconfig.NewTestOptions())
+	return &hotCacheBenchmarkCluster{
+		Cluster:     cluster,
+		coordinator: schedule.NewCoordinator(ctx, cluster, hbstream.NewTestHeartbeatStreams(ctx, cluster, true)),
+	}
+}
+
+func (c *hotCacheBenchmarkCluster) GetHotStat() *statistics.HotStat {
+	return c.HotStat
+}
+
+func (*hotCacheBenchmarkCluster) GetRegionStats() *statistics.RegionStatistics {
+	return nil
+}
+
+func (*hotCacheBenchmarkCluster) GetLabelStats() *statistics.LabelStatistics {
+	return nil
+}
+
+func (c *hotCacheBenchmarkCluster) GetCoordinator() *schedule.Coordinator {
+	return c.coordinator
+}
+
+func BenchmarkHandleStatsAsync(b *testing.B) {
+	ctx, cancel := context.WithCancel(context.Background())
+	b.Cleanup(cancel)
+	cluster := newHotCacheBenchmarkCluster(ctx)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := range b.N {
+		HandleStatsAsync(cluster, newHotCacheBenchmarkRegion(uint64(i+1)))
+	}
+	cluster.GetHotStat().GetHotPeerStats(utils.Write, 0)
+	cluster.GetHotStat().GetHotPeerStats(utils.Read, 0)
+}
+
+func newHotCacheBenchmarkRegion(regionID uint64) *core.RegionInfo {
+	peers := []*metapb.Peer{
+		{Id: regionID*10 + 1, StoreId: 1},
+		{Id: regionID*10 + 2, StoreId: 2},
+		{Id: regionID*10 + 3, StoreId: 3},
+	}
+	return core.NewRegionInfo(
+		&metapb.Region{
+			Id:          regionID,
+			RegionEpoch: &metapb.RegionEpoch{ConfVer: 1, Version: 1},
+			Peers:       peers,
+		},
+		peers[0],
+		core.SetReportInterval(0, utils.RegionHeartBeatReportInterval),
+	)
+}

pkg/mcs/scheduling/server/cluster.go

@@ -593,11 +593,11 @@ func (c *Cluster) HandleStoreHeartbeat(heartbeat *schedulingpb.StoreHeartbeatReq
 	reportInterval := stats.GetInterval()
 	interval := reportInterval.GetEndTimestamp() - reportInterval.GetStartTimestamp()
 
-	regions := make(map[uint64]*core.RegionInfo, len(stats.GetPeerStats()))
+	reportedRegions := make(map[uint64]struct{}, len(stats.GetPeerStats()))
 	for _, peerStat := range stats.GetPeerStats() {
 		regionID := peerStat.GetRegionId()
 		region := c.GetRegion(regionID)
-		regions[regionID] = region
+		reportedRegions[regionID] = struct{}{}
 		if region == nil {
 			log.Warn("discard hot peer stat for unknown region",
 				zap.Uint64("region-id", regionID),
@@ -623,23 +623,11 @@ func (c *Cluster) HandleStoreHeartbeat(heartbeat *schedulingpb.StoreHeartbeatReq
 			utils.RegionReadCPU:       regionReadCPU * float64(interval),
 			utils.RegionWriteCPU:      0,
 		}
-		checkReadPeerTask := func(cache *statistics.HotPeerCache) {
-			stats := cache.CheckPeerFlow(region, []*metapb.Peer{peer}, loads, interval)
-			for _, stat := range stats {
-				cache.UpdateStat(stat)
-			}
-		}
-		c.hotStat.CheckReadAsync(checkReadPeerTask)
+		c.hotStat.CheckReadPeerAsync(region, peer, loads, interval)
 	}
 
 	// Here we will compare the reported regions with the previous hot peers to decide if it is still hot.
-	collectUnReportedPeerTask := func(cache *statistics.HotPeerCache) {
-		stats := cache.CheckColdPeer(storeID, regions, interval)
-		for _, stat := range stats {
-			cache.UpdateStat(stat)
-		}
-	}
-	c.hotStat.CheckReadAsync(collectUnReportedPeerTask)
+	c.hotStat.CheckColdPeerAsync(storeID, reportedRegions, interval)
 	return nil
 }
 

pkg/statistics/hot_cache.go

@@ -73,6 +73,77 @@ func (w *HotCache) CheckReadAsync(task func(cache *HotPeerCache)) bool {
 	}
 }
 
+// CheckRegionFlowAsync checks the expired read and write items and the write
+// flow for a region asynchronously. It preserves the original task count and
+// enqueue order while keeping only the fields used by HotPeerCache.
+func (w *HotCache) CheckRegionFlowAsync(region *core.RegionInfo) {
+	checkExpiredTask, checkWritePeerTask := newRegionFlowTasks(region)
+	w.CheckWriteAsync(checkExpiredTask)
+	w.CheckReadAsync(checkExpiredTask)
+	w.CheckWriteAsync(checkWritePeerTask)
+}
+
+func newRegionFlowTasks(region *core.RegionInfo) (checkExpiredTask, checkWritePeerTask func(*HotPeerCache)) {
+	regionInfo := newHotRegionInfo(region)
+	reportInterval := region.GetInterval()
+	interval := reportInterval.GetEndTimestamp() - reportInterval.GetStartTimestamp()
+	writtenBytes := region.GetBytesWritten()
+	writtenKeys := region.GetKeysWritten()
+	writeQueryNum := region.GetWriteQueryNum()
+	return newExpiredRegionTask(regionInfo), newWriteRegionTask(regionInfo, interval, writtenBytes, writtenKeys, writeQueryNum)
+}
+
+func newExpiredRegionTask(regionInfo *hotRegionInfo) func(*HotPeerCache) {
+	return func(cache *HotPeerCache) {
+		expiredStats := cache.collectExpiredItemsForRegion(regionInfo)
+		for _, stat := range expiredStats {
+			cache.UpdateStat(stat)
+		}
+	}
+}
+
+func newWriteRegionTask(regionInfo *hotRegionInfo, interval, writtenBytes, writtenKeys, writeQueryNum uint64) func(*HotPeerCache) {
+	return func(cache *HotPeerCache) {
+		var loads [utils.RegionStatCount]float64
+		loads[utils.RegionWriteBytes] = float64(writtenBytes)
+		loads[utils.RegionWriteKeys] = float64(writtenKeys)
+		loads[utils.RegionWriteQueryNum] = float64(writeQueryNum)
+		stats := cache.checkPeerFlowForRegion(regionInfo, nil, loads[:], interval)
+		for _, stat := range stats {
+			cache.UpdateStat(stat)
+		}
+	}
+}
+
+// CheckReadPeerAsync checks the read flow for one peer asynchronously without
+// retaining the complete RegionInfo or Peer in the pending task.
+func (w *HotCache) CheckReadPeerAsync(region *core.RegionInfo, peer *metapb.Peer, loads []float64, interval uint64) bool {
+	return w.CheckReadAsync(newReadPeerTask(region, peer, loads, interval))
+}
+
+func newReadPeerTask(region *core.RegionInfo, peer *metapb.Peer, loads []float64, interval uint64) func(*HotPeerCache) {
+	regionInfo := newHotRegionInfo(region)
+	storeID := peer.GetStoreId()
+	return func(cache *HotPeerCache) {
+		stats := cache.checkPeerFlowForRegion(regionInfo, []uint64{storeID}, loads, interval)
+		for _, stat := range stats {
+			cache.UpdateStat(stat)
+		}
+	}
+}
+
+// CheckColdPeerAsync checks peers missing from a store heartbeat
+// asynchronously. The pending task only keeps the reported region IDs.
+func (w *HotCache) CheckColdPeerAsync(storeID uint64, reportedRegions map[uint64]struct{}, interval uint64) bool {
+	checkColdPeerTask := func(cache *HotPeerCache) {
+		stats := cache.checkColdPeerByRegionIDs(storeID, reportedRegions, interval)
+		for _, stat := range stats {
+			cache.UpdateStat(stat)
+		}
+	}
+	return w.CheckReadAsync(checkColdPeerTask)
+}
+
 // GetHotPeerStats returns hot peer stats for the specified kind (read/write).
 // It returns a map where the keys are store IDs and the values are slices of HotPeerStat.
 func (w *HotCache) GetHotPeerStats(kind utils.RWType, minHotDegree int) map[uint64][]*HotPeerStat {

pkg/statistics/hot_cache_test.go

@@ -16,11 +16,14 @@ package statistics
 
 import (
 	"context"
+	"runtime"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 
+	"github.com/pingcap/kvproto/pkg/metapb"
+
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/statistics/utils"
 	"github.com/tikv/pd/pkg/utils/testutil"
@@ -50,3 +53,83 @@ func TestIsHot(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkPendingRegionHeartbeatTasks(b *testing.B) {
+	b.Run("retain-region-info", func(b *testing.B) {
+		runPendingRegionHeartbeatTaskBenchmark(b, newRetainedRegionFlowTasks)
+	})
+	b.Run("compact-region-info", func(b *testing.B) {
+		runPendingRegionHeartbeatTaskBenchmark(b, newRegionFlowTasks)
+	})
+}
+
+func runPendingRegionHeartbeatTaskBenchmark(
+	b *testing.B,
+	newTasks func(*core.RegionInfo) (func(*HotPeerCache), func(*HotPeerCache)),
+) {
+	const maxPendingHeartbeats = 16 * 1024
+	pending := make([][3]func(*HotPeerCache), min(b.N, maxPendingHeartbeats))
+	runtime.GC()
+	var before runtime.MemStats
+	runtime.ReadMemStats(&before)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := range b.N {
+		region := newBenchmarkRegion(uint64(i + 1))
+		checkExpiredTask, checkWritePeerTask := newTasks(region)
+		pending[i%len(pending)] = [3]func(*HotPeerCache){
+			checkExpiredTask,
+			checkExpiredTask,
+			checkWritePeerTask,
+		}
+	}
+	b.StopTimer()
+
+	runtime.GC()
+	var after runtime.MemStats
+	runtime.ReadMemStats(&after)
+	retainedBytes := max(int64(after.HeapAlloc)-int64(before.HeapAlloc), 0)
+	b.ReportMetric(float64(retainedBytes)/float64(len(pending)), "retained-B/heartbeat")
+	b.ReportMetric(float64(retainedBytes)/float64(len(pending)*3), "retained-B/task")
+	runtime.KeepAlive(pending)
+}
+
+func newRetainedRegionFlowTasks(region *core.RegionInfo) (checkExpiredTask, checkWritePeerTask func(*HotPeerCache)) {
+	checkExpiredTask = func(cache *HotPeerCache) {
+		expiredStats := cache.CollectExpiredItems(region)
+		for _, stat := range expiredStats {
+			cache.UpdateStat(stat)
+		}
+	}
+	checkWritePeerTask = func(cache *HotPeerCache) {
+		reportInterval := region.GetInterval()
+		interval := reportInterval.GetEndTimestamp() - reportInterval.GetStartTimestamp()
+		stats := cache.CheckPeerFlow(region, region.GetPeers(), region.GetWriteLoads(), interval)
+		for _, stat := range stats {
+			cache.UpdateStat(stat)
+		}
+	}
+	return checkExpiredTask, checkWritePeerTask
+}
+
+func newBenchmarkRegion(regionID uint64) *core.RegionInfo {
+	peers := []*metapb.Peer{
+		{Id: regionID*10 + 1, StoreId: 1},
+		{Id: regionID*10 + 2, StoreId: 2},
+		{Id: regionID*10 + 3, StoreId: 3},
+	}
+	return core.NewRegionInfo(
+		&metapb.Region{
+			Id:          regionID,
+			StartKey:    make([]byte, 32),
+			EndKey:      make([]byte, 32),
+			RegionEpoch: &metapb.RegionEpoch{ConfVer: 1, Version: 1},
+			Peers:       peers,
+		},
+		peers[0],
+		core.SetWrittenBytes(1<<20),
+		core.SetWrittenKeys(1024),
+		core.SetReportInterval(0, utils.RegionHeartBeatReportInterval),
+	)
+}