From e31ef1f3162c9fe920ad0f8d526caffed3c8471c Mon Sep 17 00:00:00 2001
From: xucong053 <xucong053@outlook.com>
Date: Thu, 24 Mar 2022 23:14:36 +0800
Subject: [PATCH] fix: modify loop count

---
 hrp/internal/boomer/boomer.go      |  3 ++-
 hrp/internal/boomer/runner.go      | 15 ++++++++++++---
 hrp/internal/boomer/runner_test.go |  6 ++----
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/hrp/internal/boomer/boomer.go b/hrp/internal/boomer/boomer.go
index 6d22a10c..83459a4f 100644
--- a/hrp/internal/boomer/boomer.go
+++ b/hrp/internal/boomer/boomer.go
@@ -78,7 +78,8 @@ func (b *Boomer) GetDisableCompression() bool {
 
 // SetLoopCount set loop count for test.
 func (b *Boomer) SetLoopCount(loopCount int64) {
-	b.localRunner.loop = &Loop{loopCount: loopCount}
+	// total loop count for testcase, it will be evenly distributed to each worker
+	b.localRunner.loop = &Loop{loopCount: loopCount * int64(b.localRunner.spawnCount)}
 }
 
 // AddOutput accepts outputs which implements the boomer.Output interface.
diff --git a/hrp/internal/boomer/runner.go b/hrp/internal/boomer/runner.go
index c4da7cfd..25ad1bbf 100644
--- a/hrp/internal/boomer/runner.go
+++ b/hrp/internal/boomer/runner.go
@@ -65,7 +65,7 @@ type runner struct {
 	currentClientsNum int32 // current clients count
 	spawnCount        int   // target clients to spawn
 	spawnRate         float64
-	loop              *Loop // specify running cycles
+	loop              *Loop // specify loop count for testcase, count = loopCount * spawnCount
 	spawnDone         chan struct{}
 
 	outputs []Output
@@ -189,6 +189,12 @@ func (r *localRunner) spawnWorkers(spawnCount int, spawnRate float64, quit chan
 		sleepTime := time.Duration(1000000/r.spawnRate) * time.Microsecond
 		time.Sleep(sleepTime)
 
+		// loop count per worker
+		var workerLoop *Loop
+		if r.loop != nil {
+			workerLoop = &Loop{loopCount: atomic.LoadInt64(&r.loop.loopCount) / int64(r.spawnCount)}
+		}
+
 		select {
 		case <-quit:
 			// quit spawning goroutine
@@ -202,7 +208,7 @@ func (r *localRunner) spawnWorkers(spawnCount int, spawnRate float64, quit chan
 					case <-quit:
 						return
 					default:
-						if r.loop != nil && !r.loop.acquire() {
+						if workerLoop != nil && !workerLoop.acquire() {
 							return
 						}
 						if r.rateLimitEnabled {
@@ -215,8 +221,11 @@ func (r *localRunner) spawnWorkers(spawnCount int, spawnRate float64, quit chan
 							task := r.getTask()
 							r.safeRun(task.Fn)
 						}
-						if r.loop != nil {
+						if workerLoop != nil {
+							// finished count of total
 							r.loop.increaseFinishedCount()
+							// finished count of single worker
+							workerLoop.increaseFinishedCount()
 							if r.loop.isFinished() {
 								r.stop()
 							}
diff --git a/hrp/internal/boomer/runner_test.go b/hrp/internal/boomer/runner_test.go
index e6e15992..a4f674b4 100644
--- a/hrp/internal/boomer/runner_test.go
+++ b/hrp/internal/boomer/runner_test.go
@@ -98,7 +98,7 @@ func TestLoopCount(t *testing.T) {
 	taskA := &Task{
 		Weight: 10,
 		Fn: func() {
-			time.Sleep(time.Second)
+			time.Sleep(time.Millisecond)
 		},
 		Name: "TaskA",
 	}
@@ -107,9 +107,7 @@ func TestLoopCount(t *testing.T) {
 	runner.loop = &Loop{loopCount: 4}
 	runner.setTasks(tasks)
 	go runner.start()
-	ticker := time.NewTicker(4 * time.Second)
-	defer ticker.Stop()
-	<-ticker.C
+	<-runner.stopChan
 	if !assert.Equal(t, runner.loop.loopCount, atomic.LoadInt64(&runner.loop.finishedCount)) {
 		t.Fail()
 	}