From b4e2c553501877afe94c0447c24dd40b5d2bdf97 Mon Sep 17 00:00:00 2001
From: xucong053 <xucong053@outlook.com>
Date: Thu, 12 May 2022 11:26:50 +0800
Subject: [PATCH] fix: add boomer mode(standalone/master/worker)

---
 hrp/cmd/boom.go               |  2 +-
 hrp/internal/boomer/boomer.go | 45 +++++++++++++++++++++++++++++++++++
 hrp/internal/boomer/output.go |  6 +++--
 3 files changed, 50 insertions(+), 3 deletions(-)

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),
 	}
 }