feat: hrp boom support set run time

hrp/boomer.go

@@ -66,6 +66,7 @@ func (b *HRPBoomer) InitBoomer() {
 	}
 	b.SetSpawnCount(b.GetProfile().SpawnCount)
 	b.SetSpawnRate(b.GetProfile().SpawnRate)
+	b.SetRunTime(b.GetProfile().RunTime)
 	if b.GetProfile().LoopCount > 0 {
 		b.SetLoopCount(b.GetProfile().LoopCount)
 	}
@@ -232,6 +233,7 @@ func (b *HRPBoomer) rebalanceBoomer(profile *boomer.Profile) {
 	b.SetProfile(profile)
 	b.SetSpawnCount(b.GetProfile().SpawnCount)
 	b.SetSpawnRate(b.GetProfile().SpawnRate)
+	b.SetRunTime(b.GetProfile().RunTime)
 	b.GetRebalanceChan() <- true
 	log.Info().Interface("profile", profile).Msg("rebalance tasks successful")
 }

hrp/cmd/boom.go

@@ -67,6 +67,7 @@ var boomCmd = &cobra.Command{
 				hrpBoomer.SetExpectWorkers(boomArgs.expectWorkers, boomArgs.expectWorkersMaxWait)
 				hrpBoomer.SetSpawnCount(boomArgs.SpawnCount)
 				hrpBoomer.SetSpawnRate(boomArgs.SpawnRate)
+				hrpBoomer.SetRunTime(boomArgs.RunTime)
 			}
 			if boomArgs.autoStart {
 				hrpBoomer.InitBoomer()
@@ -116,6 +117,7 @@ func init() {
 	boomCmd.Flags().StringVar(&boomArgs.RequestIncreaseRate, "request-increase-rate", "-1", "Request increase rate, disabled by default.")
 	boomCmd.Flags().Int64Var(&boomArgs.SpawnCount, "spawn-count", 1, "The number of users to spawn for load testing")
 	boomCmd.Flags().Float64Var(&boomArgs.SpawnRate, "spawn-rate", 1, "The rate for spawning users")
+	boomCmd.Flags().Int64Var(&boomArgs.RunTime, "run-time", 0, "Stop after the specified amount of time(s), Only used  --autostart. Defaults to run forever.")
 	boomCmd.Flags().Int64Var(&boomArgs.LoopCount, "loop-count", -1, "The specify running cycles for load testing")
 	boomCmd.Flags().StringVar(&boomArgs.MemoryProfile, "mem-profile", "", "Enable memory profiling.")
 	boomCmd.Flags().DurationVar(&boomArgs.MemoryProfileDuration, "mem-profile-duration", 30*time.Second, "Memory profile duration.")

hrp/internal/boomer/boomer.go

@@ -50,6 +50,7 @@ type Boomer struct {
 type Profile struct {
 	SpawnCount               int64         `json:"spawn-count,omitempty" yaml:"spawn-count,omitempty" mapstructure:"spawn-count,omitempty"`
 	SpawnRate                float64       `json:"spawn-rate,omitempty" yaml:"spawn-rate,omitempty" mapstructure:"spawn-rate,omitempty"`
+	RunTime                  int64         `json:"run-time,omitempty" yaml:"run-time,omitempty" mapstructure:"run-time,omitempty"`
 	MaxRPS                   int64         `json:"max-rps,omitempty" yaml:"max-rps,omitempty" mapstructure:"max-rps,omitempty"`
 	LoopCount                int64         `json:"loop-count,omitempty" yaml:"loop-count,omitempty" mapstructure:"loop-count,omitempty"`
 	RequestIncreaseRate      string        `json:"request-increase-rate,omitempty" yaml:"request-increase-rate,omitempty" mapstructure:"request-increase-rate,omitempty"`
@@ -251,6 +252,18 @@ func (b *Boomer) SetSpawnRate(spawnRate float64) {
 	}
 }
 
+// SetRunTime sets run time
+func (b *Boomer) SetRunTime(spawnRate int64) {
+	switch b.mode {
+	case DistributedMasterMode:
+		b.masterRunner.setRunTime(spawnRate)
+	case DistributedWorkerMode:
+		b.workerRunner.setRunTime(spawnRate)
+	default:
+		b.localRunner.setRunTime(spawnRate)
+	}
+}
+
 // SetExpectWorkers sets expect workers while load testing
 func (b *Boomer) SetExpectWorkers(expectWorkers int, expectWorkersMaxWait int) {
 	b.masterRunner.setExpectWorkers(expectWorkers, expectWorkersMaxWait)
@@ -471,6 +484,7 @@ func (b *Boomer) Start(Args *Profile) error {
 	}
 	b.SetSpawnCount(Args.SpawnCount)
 	b.SetSpawnRate(Args.SpawnRate)
+	b.SetRunTime(Args.RunTime)
 	b.SetProfile(Args)
 	err := b.masterRunner.start()
 	return err
@@ -483,6 +497,7 @@ func (b *Boomer) ReBalance(Args *Profile) error {
 	}
 	b.SetSpawnCount(Args.SpawnCount)
 	b.SetSpawnRate(Args.SpawnRate)
+	b.SetRunTime(Args.RunTime)
 	b.SetProfile(Args)
 	err := b.masterRunner.rebalance()
 	if err != nil {

hrp/internal/boomer/runner.go

@@ -195,6 +195,7 @@ type runner struct {
 
 	spawnCount int64 // target clients to spawn
 	spawnRate  float64
+	runTime    int64
 
 	controller *Controller
 	loop       *Loop // specify loop count for testcase, count = loopCount * spawnCount
@@ -238,6 +239,20 @@ func (r *runner) getSpawnRate() float64 {
 	return r.spawnRate
 }
 
+func (r *runner) setRunTime(runTime int64) {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+	if runTime > 0 {
+		r.runTime = time.Now().Unix() + runTime
+	}
+}
+
+func (r *runner) getRunTime() int64 {
+	r.mutex.RLock()
+	defer r.mutex.RUnlock()
+	return r.runTime
+}
+
 func (r *runner) getSpawnCount() int64 {
 	return atomic.LoadInt64(&r.spawnCount)
 }
@@ -364,6 +379,15 @@ func (r *runner) reset() {
 	r.reportedChan = make(chan bool)
 }
 
+func (r *runner) runTimeCheck(runTime int64) {
+	for range time.Tick(time.Second * 3) {
+		nowTime := time.Now().Unix()
+		if nowTime > runTime {
+			r.stop()
+		}
+	}
+}
+
 func (r *runner) spawnWorkers(spawnCount int64, spawnRate float64, quit chan bool, spawnCompleteFunc func()) {
 	r.updateState(StateSpawning)
 	log.Info().
@@ -622,6 +646,11 @@ func (r *localRunner) start() {
 	// output setup
 	r.outputOnStart()
 
+	runTime := r.getRunTime()
+	if runTime != 0 {
+		go r.runTimeCheck(runTime)
+	}
+
 	go r.spawnWorkers(r.getSpawnCount(), r.getSpawnRate(), r.stoppingChan, nil)
 
 	defer func() {
@@ -912,6 +941,11 @@ func (r *workerRunner) start() {
 
 	r.outputOnStart()
 
+	runTime := r.getRunTime()
+	if runTime != 0 {
+		go r.runTimeCheck(runTime)
+	}
+
 	go r.spawnWorkers(r.getSpawnCount(), r.getSpawnRate(), r.stoppingChan, r.spawnComplete)
 
 	defer func() {