diff --git a/hrp/cmd/boom.go b/hrp/cmd/boom.go index 66d2dec9..22d8782d 100644 --- a/hrp/cmd/boom.go +++ b/hrp/cmd/boom.go @@ -37,7 +37,7 @@ var boomCmd = &cobra.Command{ hrpBoomer.AddOutput(boomer.NewConsoleOutput()) } if prometheusPushgatewayURL != "" { - hrpBoomer.AddOutput(boomer.NewPrometheusPusherOutput(prometheusPushgatewayURL, "hrp")) + hrpBoomer.AddOutput(boomer.NewPrometheusPusherOutput(prometheusPushgatewayURL, "hrp", hrpBoomer.GetMode())) } hrpBoomer.SetDisableKeepAlive(disableKeepalive) hrpBoomer.SetDisableCompression(disableCompression) diff --git a/hrp/internal/boomer/boomer.go b/hrp/internal/boomer/boomer.go index 6c67007d..cc424b12 100644 --- a/hrp/internal/boomer/boomer.go +++ b/hrp/internal/boomer/boomer.go @@ -7,11 +7,26 @@ import ( "syscall" "time" + "github.com/pkg/errors" "github.com/rs/zerolog/log" ) +// Mode is the running mode of boomer, both standalone and distributed are supported. +type Mode int + +const ( + // DistributedMasterMode requires being connected by each worker. + DistributedMasterMode Mode = iota + // DistributedWorkerMode requires connecting to a master. + DistributedWorkerMode + // StandaloneMode will run without a master. + StandaloneMode +) + // A Boomer is used to run tasks. type Boomer struct { + mode Mode + localRunner *localRunner cpuProfile string @@ -24,9 +39,39 @@ type Boomer struct { disableCompression bool } +// SetMode only accepts boomer.DistributedMasterMode态boomer.DistributedWorkerMode and boomer.StandaloneMode. +func (b *Boomer) SetMode(mode Mode) { + switch mode { + case DistributedMasterMode: + b.mode = DistributedMasterMode + case DistributedWorkerMode: + b.mode = DistributedWorkerMode + case StandaloneMode: + b.mode = StandaloneMode + default: + log.Error().Err(errors.New("Invalid mode, ignored!")) + } +} + +// GetMode returns boomer operating mode +func (b *Boomer) GetMode() string { + switch b.mode { + case DistributedMasterMode: + return "master" + case DistributedWorkerMode: + return "worker" + case StandaloneMode: + return "standalone" + default: + log.Error().Err(errors.New("Invalid mode, ignored!")) + return "" + } +} + // NewStandaloneBoomer returns a new Boomer, which can run without master. func NewStandaloneBoomer(spawnCount int, spawnRate float64) *Boomer { return &Boomer{ + mode: StandaloneMode, localRunner: newLocalRunner(spawnCount, spawnRate), } } diff --git a/hrp/internal/boomer/output.go b/hrp/internal/boomer/output.go index 3bcca9d9..db77b053 100644 --- a/hrp/internal/boomer/output.go +++ b/hrp/internal/boomer/output.go @@ -471,10 +471,12 @@ var ( ) // NewPrometheusPusherOutput returns a PrometheusPusherOutput. -func NewPrometheusPusherOutput(gatewayURL, jobName string) *PrometheusPusherOutput { +func NewPrometheusPusherOutput(gatewayURL, jobName string, mode string) *PrometheusPusherOutput { nodeUUID, _ := uuid.NewUUID() return &PrometheusPusherOutput{ - pusher: push.New(gatewayURL, jobName).Grouping("instance", nodeUUID.String()), + pusher: push.New(gatewayURL, jobName). + Grouping("instance", nodeUUID.String()). + Grouping("mode", mode), } }