From 53bd8003ecd372ddfd360cf1e9c5cf76793285b7 Mon Sep 17 00:00:00 2001
From: Cabinfever_B <cabinfeveroier@gmail.com>
Date: Fri, 1 Mar 2024 17:06:26 +0800
Subject: [PATCH] support more concurrency

Signed-off-by: Cabinfever_B <cabinfeveroier@gmail.com>
---
 tools/pd-api-bench/cases/controller.go | 90 +++++++++++++++-----------
 1 file changed, 51 insertions(+), 39 deletions(-)

diff --git a/tools/pd-api-bench/cases/controller.go b/tools/pd-api-bench/cases/controller.go
index db42c469843..0f2d7bab370 100644
--- a/tools/pd-api-bench/cases/controller.go
+++ b/tools/pd-api-bench/cases/controller.go
@@ -220,21 +220,25 @@ func (c *httpController) run() {
 		c.wg.Add(1)
 		go func(hCli pdHttp.Client) {
 			defer c.wg.Done()
-			var ticker = time.NewTicker(tt)
-			defer ticker.Stop()
-			for {
-				select {
-				case <-ticker.C:
-					for i := int64(0); i < burst; i++ {
-						err := c.Do(c.ctx, hCli)
-						if err != nil {
-							log.Error("meet erorr when doing HTTP request", zap.String("case", c.Name()), zap.Error(err))
+			c.wg.Add(int(burst))
+			for i := int64(0); i < burst; i++ {
+				go func() {
+					defer c.wg.Done()
+					var ticker = time.NewTicker(tt)
+					defer ticker.Stop()
+					for {
+						select {
+						case <-ticker.C:
+							err := c.Do(c.ctx, hCli)
+							if err != nil {
+								log.Error("meet erorr when doing HTTP request", zap.String("case", c.Name()), zap.Error(err))
+							}
+						case <-c.ctx.Done():
+							log.Info("Got signal to exit running HTTP case")
+							return
 						}
 					}
-				case <-c.ctx.Done():
-					log.Info("Got signal to exit running HTTP case")
-					return
-				}
+				}()
 			}
 		}(hCli)
 	}
@@ -285,21 +289,25 @@ func (c *gRPCController) run() {
 		c.wg.Add(1)
 		go func(cli pd.Client) {
 			defer c.wg.Done()
-			var ticker = time.NewTicker(tt)
-			defer ticker.Stop()
-			for {
-				select {
-				case <-ticker.C:
-					for i := int64(0); i < burst; i++ {
-						err := c.Unary(c.ctx, cli)
-						if err != nil {
-							log.Error("meet erorr when doing gRPC request", zap.String("case", c.Name()), zap.Error(err))
+			c.wg.Add(int(burst))
+			for i := int64(0); i < burst; i++ {
+				go func() {
+					defer c.wg.Done()
+					var ticker = time.NewTicker(tt)
+					defer ticker.Stop()
+					for {
+						select {
+						case <-ticker.C:
+							err := c.Unary(c.ctx, cli)
+							if err != nil {
+								log.Error("meet erorr when doing gRPC request", zap.String("case", c.Name()), zap.Error(err))
+							}
+						case <-c.ctx.Done():
+							log.Info("Got signal to exit running gRPC case")
+							return
 						}
 					}
-				case <-c.ctx.Done():
-					log.Info("Got signal to exit running gRPC case")
-					return
-				}
+				}()
 			}
 		}(cli)
 	}
@@ -355,21 +363,25 @@ func (c *etcdController) run() {
 		c.wg.Add(1)
 		go func(cli *clientv3.Client) {
 			defer c.wg.Done()
-			var ticker = time.NewTicker(tt)
-			defer ticker.Stop()
-			for {
-				select {
-				case <-ticker.C:
-					for i := int64(0); i < burst; i++ {
-						err := c.Unary(c.ctx, cli)
-						if err != nil {
-							log.Error("meet erorr when doing etcd request", zap.String("case", c.Name()), zap.Error(err))
+			c.wg.Add(int(burst))
+			for i := int64(0); i < burst; i++ {
+				go func() {
+					defer c.wg.Done()
+					var ticker = time.NewTicker(tt)
+					defer ticker.Stop()
+					for {
+						select {
+						case <-ticker.C:
+							err := c.Unary(c.ctx, cli)
+							if err != nil {
+								log.Error("meet erorr when doing etcd request", zap.String("case", c.Name()), zap.Error(err))
+							}
+						case <-c.ctx.Done():
+							log.Info("Got signal to exit running etcd case")
+							return
 						}
 					}
-				case <-c.ctx.Done():
-					log.Info("Got signal to exit running etcd case")
-					return
-				}
+				}()
 			}
 		}(cli)
 	}