feat: specify running cycles for load testing.

cli/hrp/cmd/boom.go

@@ -29,6 +29,9 @@ var boomCmd = &cobra.Command{
 		}
 		hrpBoomer := hrp.NewBoomer(spawnCount, spawnRate)
 		hrpBoomer.SetRateLimiter(maxRPS, requestIncreaseRate)
+		if loopCount > 0 {
+			hrpBoomer.SetLoopCount(loopCount)
+		}
 		if !disableConsoleOutput {
 			hrpBoomer.AddOutput(boomer.NewConsoleOutput())
 		}
@@ -45,6 +48,7 @@ var (
 	spawnCount               int
 	spawnRate                float64
 	maxRPS                   int64
+	loopCount                int64
 	requestIncreaseRate      string
 	memoryProfile            string
 	memoryProfileDuration    time.Duration

internal/boomer/boomer.go

@@ -52,6 +52,11 @@ func (b *Boomer) SetRateLimiter(maxRPS int64, requestIncreaseRate string) {
 	}
 }
 
+// SetLoopCount set loop count for test.
+func (b *Boomer) SetLoopCount(loopCount int64) {
+	b.localRunner.loop = &Loop{loopCount: loopCount}
+}
+
 // AddOutput accepts outputs which implements the boomer.Output interface.
 func (b *Boomer) AddOutput(o Output) {
 	b.localRunner.addOutput(o)

internal/boomer/runner.go

@@ -24,6 +24,31 @@ const (
 	reportStatsInterval = 3 * time.Second
 )
 
+type Loop struct {
+	loopCount     int64 // more than 0
+	acquiredCount int64 // count acquired of load testing
+	finishedCount int64 // count finished of load testing
+}
+
+func (l *Loop) isFinished() bool {
+	// return true when there are no remaining loop count to test
+	return atomic.LoadInt64(&l.finishedCount) == l.loopCount
+}
+
+func (l *Loop) acquire() bool {
+	// get one ticket when there are still remaining loop count to test
+	// return true when getting ticket successfully
+	if atomic.LoadInt64(&l.acquiredCount) < l.loopCount {
+		atomic.AddInt64(&l.acquiredCount, 1)
+		return true
+	}
+	return false
+}
+
+func (l *Loop) increaseFinishedCount() {
+	atomic.AddInt64(&l.finishedCount, 1)
+}
+
 type runner struct {
 	state int32
 
@@ -37,6 +62,7 @@ type runner struct {
 	currentClientsNum int32 // current clients count
 	spawnCount        int   // target clients to spawn
 	spawnRate         float64
+	loop              *Loop // specify running cycles
 
 	outputs []Output
 }
@@ -78,7 +104,7 @@ func (r *runner) outputOnStart() {
 	wg.Wait()
 }
 
-func (r *runner) outputOnEevent(data map[string]interface{}) {
+func (r *runner) outputOnEvent(data map[string]interface{}) {
 	size := len(r.outputs)
 	if size == 0 {
 		return
@@ -110,7 +136,14 @@ func (r *runner) outputOnStop() {
 	wg.Wait()
 }
 
-func (r *runner) spawnWorkers(spawnCount int, spawnRate float64, quit chan bool, spawnCompleteFunc func()) {
+func (r *runner) reportStats() {
+	data := r.stats.collectReportData()
+	data["user_count"] = atomic.LoadInt32(&r.currentClientsNum)
+	data["state"] = atomic.LoadInt32(&r.state)
+	r.outputOnEvent(data)
+}
+
+func (r *localRunner) spawnWorkers(spawnCount int, spawnRate float64, quit chan bool, spawnCompleteFunc func()) {
 	log.Info().
 		Int("spawnCount", spawnCount).
 		Float64("spawnRate", spawnRate).
@@ -135,6 +168,9 @@ func (r *runner) spawnWorkers(spawnCount int, spawnRate float64, quit chan bool,
 					case <-quit:
 						return
 					default:
+						if r.loop != nil && !r.loop.acquire() {
+							return
+						}
 						if r.rateLimitEnabled {
 							blocked := r.rateLimiter.Acquire()
 							if !blocked {
@@ -145,6 +181,12 @@ func (r *runner) spawnWorkers(spawnCount int, spawnRate float64, quit chan bool,
 							task := r.getTask()
 							r.safeRun(task.Fn)
 						}
+						if r.loop != nil {
+							r.loop.increaseFinishedCount()
+							if r.loop.isFinished() {
+								r.stop()
+							}
+						}
 					}
 				}
 			}()
@@ -250,10 +292,7 @@ func (r *localRunner) start() {
 			r.stats.logError(n.requestType, n.name, n.errMsg)
 		// report stats
 		case <-ticker.C:
-			data := r.stats.collectReportData()
-			data["user_count"] = atomic.LoadInt32(&r.currentClientsNum)
-			data["state"] = atomic.LoadInt32(&r.state)
-			r.outputOnEevent(data)
+			r.reportStats()
 		// stop
 		case <-r.stopChan:
 			atomic.StoreInt32(&r.state, stateQuitting)
@@ -267,6 +306,10 @@ func (r *localRunner) start() {
 				r.rateLimiter.Stop()
 			}
 
+			// report last stats
+			<-ticker.C
+			r.reportStats()
+
 			// output teardown
 			r.outputOnStop()
 

internal/boomer/runner_test.go

@@ -3,6 +3,8 @@ package boomer
 import (
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 type HitOutput struct {
@@ -45,13 +47,13 @@ func TestOutputOnStart(t *testing.T) {
 	}
 }
 
-func TestOutputOnEevent(t *testing.T) {
+func TestOutputOnEvent(t *testing.T) {
 	hitOutput := &HitOutput{}
 	hitOutput2 := &HitOutput{}
 	runner := &runner{}
 	runner.addOutput(hitOutput)
 	runner.addOutput(hitOutput2)
-	runner.outputOnEevent(nil)
+	runner.outputOnEvent(nil)
 	if !hitOutput.onEvent {
 		t.Error("hitOutput's OnEvent has not been called")
 	}
@@ -90,3 +92,24 @@ func TestLocalRunner(t *testing.T) {
 	time.Sleep(4 * time.Second)
 	runner.stop()
 }
+
+func TestLoopCount(t *testing.T) {
+	taskA := &Task{
+		Weight: 10,
+		Fn: func() {
+			time.Sleep(time.Second)
+		},
+		Name: "TaskA",
+	}
+	tasks := []*Task{taskA}
+	runner := newLocalRunner(2, 2)
+	runner.loop = &Loop{loopCount: 4}
+	runner.setTasks(tasks)
+	go runner.start()
+	ticker := time.NewTicker(4 * time.Second)
+	defer ticker.Stop()
+	<-ticker.C
+	if !assert.Equal(t, runner.loop.loopCount, runner.loop.finishedCount) {
+		t.Fail()
+	}
+}