tools: fix store statistics (#7771)

ref #7703 Signed-off-by: Ryan Leung <rleungx@gmail.com>
tikv · Jan 29, 2024 · 7142cf4 · 7142cf4
1 parent 5b939c6
commit 7142cf4
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/tools/pd-heartbeat-bench/main.go b/tools/pd-heartbeat-bench/main.go
@@ -401,7 +401,14 @@ func (s *Stores) update(rs *Regions) {
 			},
 		}
 	}
-	for _, region := range rs.regions {
+	var toUpdate []*pdpb.RegionHeartbeatRequest
+	updatedRegions := rs.awakenRegions.Load()
+	if updatedRegions == nil {
+		toUpdate = rs.regions
+	} else {
+		toUpdate = updatedRegions.([]*pdpb.RegionHeartbeatRequest)
+	}
+	for _, region := range toUpdate {
 		for _, peer := range region.Region.Peers {
 			store := stats[peer.StoreId]
 			store.UsedSize += region.ApproximateSize