bench/heartbeat: make heartbeat bench better

tools/pd-heartbeat-bench/README.md

@@ -0,0 +1,60 @@
+# pd-heartbeat-bench
+
+`pd-heartbeat-bench` is used to bench heartbeart.
+
+1. You need to deploy a cluster that only contain pd firstly, like `tiup playground nightly --pd 3 --kv 0 --db 0`.
+2. Then, execute `pd-heartbeart-bench` and set the pd leader as `--pd-endpoints` 
+
+## HTTP Server
+The tool starts an HTTP server based on the StatusAddr field in the configuration file. Ensure that the StatusAddr is correctly configured before starting the server.
+
+## API Endpoints
+1. Get Current Configuration
+
+**Endpoint**: `GET /config`
+
+**Description**: Returns the current configuration of the benchmark tool.
+
+Response:
+
+- Status Code: `200 OK`
+- Content: JSON representation of the configuration.
+
+Example Response:
+
+```json
+{
+  "HotStoreCount": 10,
+  "FlowUpdateRatio": 0.5,
+  "LeaderUpdateRatio": 0.3,
+  "EpochUpdateRatio": 0.2,
+  "SpaceUpdateRatio": 0.1,
+  "ReportRatio": 0.05
+}
+```
+
+2. Update Configuration
+
+**Endpoint**: `PUT /config`
+
+**Description**: Updates the configuration of the benchmark tool.
+
+Request Body: JSON representation of the new configuration.
+
+Response:
+
+- Status Code: `200 OK` if the update is successful.
+- Status Code: `400 Bad Request` if the request body is invalid or fails validation.
+
+Example Request:
+
+```json
+{
+  "HotStoreCount": 15,
+  "FlowUpdateRatio": 0.6,
+  "LeaderUpdateRatio": 0.4,
+  "EpochUpdateRatio": 0.3,
+  "SpaceUpdateRatio": 0.2,
+  "ReportRatio": 0.1
+}
+```

tools/pd-heartbeat-bench/config-template.toml

@@ -1,14 +1,14 @@
 round = 0
 
 store-count = 100
+hot-store-count = 10
 region-count = 2000000
 
 replica = 3
 
 leader-update-ratio = 0.06
-epoch-update-ratio = 0.0
-space-update-ratio = 0.0
-flow-update-ratio = 0.0
-no-update-ratio = 0.0
+epoch-update-ratio = 0.04
+space-update-ratio = 0.15
+flow-update-ratio = 0.35
 
 sample = false

tools/pd-heartbeat-bench/config/config.go

@@ -79,7 +79,7 @@ func NewConfig() *Config {
 	fs := cfg.flagSet
 	fs.ParseErrorsWhitelist.UnknownFlags = true
 	fs.StringVar(&cfg.configFile, "config", "", "config file")
-	fs.StringVar(&cfg.PDAddr, "pd-endpoints", "127.0.0.1:2379", "pd address")
+	fs.StringVar(&cfg.PDAddr, "pd-endpoints", "127.0.0.1:2379", "pd leader address")
 	fs.StringVar(&cfg.Log.File.Filename, "log-file", "", "log file path")
 	fs.StringVar(&cfg.StatusAddr, "status-addr", "127.0.0.1:20180", "status address")
 	fs.StringVar(&cfg.Security.CAPath, "cacert", "", "path of file that contains list of trusted TLS CAs")

tools/pd-heartbeat-bench/main.go

@@ -121,7 +121,7 @@ func bootstrap(ctx context.Context, cli pdpb.PDClient) {
 	store := &metapb.Store{
 		Id:      1,
 		Address: fmt.Sprintf("mock://tikv-1:%d", 2),
-		Version: "6.4.0-alpha",
+		Version: "9.0.0-alpha.1",
 	}
 	region := &metapb.Region{
 		Id:          1,
@@ -150,7 +150,7 @@ func putStores(ctx context.Context, cfg *config.Config, cli pdpb.PDClient, store
 		store := &metapb.Store{
 			Id:      i,
 			Address: fmt.Sprintf("mock://tikv-%d:%d", i, i),
-			Version: "6.4.0-alpha",
+			Version: "9.0.0-alpha.1",
 		}
 		cctx, cancel := context.WithCancel(ctx)
 		resp, err := cli.PutStore(cctx, &pdpb.PutStoreRequest{Header: header(), Store: store})
@@ -187,9 +187,17 @@ func createHeartbeatStream(ctx context.Context, cfg *config.Config) (pdpb.PDClie
 	}
 
 	go func() {
-		// do nothing
+		failedRequest := 0
 		for {
-			stream.Recv()
+			resp, err := stream.Recv()
+			if err != nil {
+				log.Error("receive error", zap.Error(err), zap.Int("failed-request", failedRequest))
+				failedRequest++
+			}
+			if resp.GetHeader().GetError() != nil {
+				log.Error("receive error", zap.String("err", resp.GetHeader().GetError().String()), zap.Int("failed-request", failedRequest))
+				failedRequest++
+			}
 		}
 	}()
 	return cli, stream
@@ -312,7 +320,7 @@ func main() {
 
 	initClusterID(ctx, cli)
 	go runHTTPServer(cfg, options)
-	regions := utils.NewRegions(cfg.RegionCount, cfg.Replica, header())
+	regions := utils.NewRegions(cfg.RegionCount, cfg.Replica, cfg.StoreCount, header())
 	log.Info("finish init regions")
 	stores := newStores(cfg.StoreCount)
 	stores.update(regions)

tools/pd-region-bench/main.go

@@ -208,7 +208,7 @@
 	defer cancel()
 
 	// Generate the regions info first.
-	s.regions = utils.NewRegions(*regionCount, 3, s.header())
+	s.regions = utils.NewRegions(*regionCount, 3, 3, s.header())
 	// Heartbeat the regions with two rounds to ensure all regions are fulfilled.
 	heartbeatStream := s.createHeartbeatStream(cctx, cli)
 	for range 2 {

tools/utils/meta.go

@@ -139,7 +139,7 @@
 }
 
 // NewRegions initializes the regions with the given region count and replica count.
-func NewRegions(regionCount, replicaCount int, header *pdpb.RequestHeader) *Regions {
+func NewRegions(regionCount, replicaCount, storeCount int, header *pdpb.RequestHeader) *Regions {
 	rs := &Regions{
 		regionCount:  regionCount,
 		replicaCount: replicaCount,
@@ -180,7 +180,7 @@
 
 		peers := make([]*metapb.Peer, 0, replicaCount)
 		for j := range replicaCount {
-			peers = append(peers, &metapb.Peer{Id: id, StoreId: uint64((i+j)%replicaCount + 1)})
+			peers = append(peers, &metapb.Peer{Id: id, StoreId: uint64((i+j)%storeCount + 1)})
 			id += 1
 		}