tools: support more concurrency for API bench (#7867)

ref #7703 Signed-off-by: Cabinfever_B <[email protected]>
tikv · Mar 4, 2024 · 2851db2 · 2851db2
1 parent 1833ce2
commit 2851db2
Showing 1 changed file with 51 additions and 39 deletions.
diff --git 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 error 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 error 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 error 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)
 	}