diff --git a/cli/hrp/cmd/boom.go b/cli/hrp/cmd/boom.go index 35a7d6a9..c2765c5a 100644 --- a/cli/hrp/cmd/boom.go +++ b/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 diff --git a/internal/boomer/boomer.go b/internal/boomer/boomer.go index 55a01e7f..0b22f6cb 100644 --- a/internal/boomer/boomer.go +++ b/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) diff --git a/internal/boomer/runner.go b/internal/boomer/runner.go index 2ddd6f30..a205b652 100644 --- a/internal/boomer/runner.go +++ b/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() diff --git a/internal/boomer/runner_test.go b/internal/boomer/runner_test.go index 22f6ad8b..94c2c19b 100644 --- a/internal/boomer/runner_test.go +++ b/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() + } +}