Merge pull request #194 from zezhehh/zh/logger

feat: support customized logger for Boomer, Output and runner

boomer.go

@@ -10,7 +10,7 @@ import (
 	"time"
 )
 
-var defaultBoomer = &Boomer{}
+var defaultBoomer = &Boomer{logger: log.Default()}
 
 // Mode is the running mode of boomer, both standalone and distributed are supported.
 type Mode int
@@ -24,6 +24,7 @@ const (
 
 // A Boomer is used to run tasks.
 // This type is exposed, so users can create and control a Boomer instance programmatically.
+// A non-nil logger is supposed to be set.
 type Boomer struct {
 	masterHost  string
 	masterPort  int
@@ -42,6 +43,8 @@ type Boomer struct {
 	memoryProfileDuration time.Duration
 
 	outputs []Output
+
+	logger *log.Logger
 }
 
 // NewBoomer returns a new Boomer.
@@ -50,6 +53,7 @@ func NewBoomer(masterHost string, masterPort int) *Boomer {
 		masterHost: masterHost,
 		masterPort: masterPort,
 		mode:       DistributedMode,
+		logger:     log.Default(),
 	}
 }
 
@@ -59,7 +63,24 @@ func NewStandaloneBoomer(spawnCount int, spawnRate float64) *Boomer {
 		spawnCount: spawnCount,
 		spawnRate:  spawnRate,
 		mode:       StandaloneMode,
+		logger:     log.Default(),
+	}
+}
+
+// WithLogger allows user to use their own logger.
+// If the logger is nil, it will not take effect.
+func (b *Boomer) WithLogger(logger *log.Logger) *Boomer {
+	if logger == nil {
+		return b
+	}
+	b.logger = logger
+	if b.slaveRunner != nil {
+		b.slaveRunner.setLogger(logger)
+	}
+	if b.localRunner != nil {
+		b.localRunner.setLogger(logger)
 	}
+	return b
 }
 
 // SetRateLimiter allows user to use their own rate limiter.
@@ -76,7 +97,7 @@ func (b *Boomer) SetMode(mode Mode) {
 	case StandaloneMode:
 		b.mode = StandaloneMode
 	default:
-		log.Println("Invalid mode, ignored!")
+		b.logger.Println("Invalid mode, ignored!")
 	}
 }
 
@@ -102,33 +123,35 @@ func (b *Boomer) Run(tasks ...*Task) {
 	if b.cpuProfileFile != "" {
 		err := StartCPUProfile(b.cpuProfileFile, b.cpuProfileDuration)
 		if err != nil {
-			log.Printf("Error starting cpu profiling, %v", err)
+			b.logger.Printf("Error starting cpu profiling, %v", err)
 		}
 	}
 	if b.memoryProfileFile != "" {
 		err := StartMemoryProfile(b.memoryProfileFile, b.memoryProfileDuration)
 		if err != nil {
-			log.Printf("Error starting memory profiling, %v", err)
+			b.logger.Printf("Error starting memory profiling, %v", err)
 		}
 	}
 
 	switch b.mode {
 	case DistributedMode:
 		b.slaveRunner = newSlaveRunner(b.masterHost, b.masterPort, tasks, b.rateLimiter)
-		println("new slave runner")
+		b.slaveRunner.setLogger(b.logger)
+		b.logger.Println("new slave runner")
 		for _, o := range b.outputs {
 			b.slaveRunner.addOutput(o)
 		}
 		b.slaveRunner.run()
 	case StandaloneMode:
 		b.localRunner = newLocalRunner(tasks, b.rateLimiter, b.spawnCount, b.spawnRate)
-		println("new local runner")
+		b.localRunner.setLogger(b.logger)
+		b.logger.Println("new local runner")
 		for _, o := range b.outputs {
 			b.localRunner.addOutput(o)
 		}
 		b.localRunner.run()
 	default:
-		log.Println("Invalid mode, expected boomer.DistributedMode or boomer.StandaloneMode")
+		b.logger.Println("Invalid mode, expected boomer.DistributedMode or boomer.StandaloneMode")
 	}
 }
 
@@ -202,7 +225,7 @@ func (b *Boomer) Quit() {
 		case <-b.slaveRunner.client.disconnectedChannel():
 			break
 		case <-ticker.C:
-			log.Println("Timeout waiting for sending quit message to master, boomer will quit any way.")
+			b.logger.Println("Timeout waiting for sending quit message to master, boomer will quit any way.")
 			break
 		}
 		b.slaveRunner.shutdown()

boomer_test.go

@@ -2,6 +2,7 @@ package boomer
 
 import (
 	"flag"
+	"log"
 	"math"
 	"os"
 	"runtime"
@@ -215,7 +216,7 @@ var _ = Describe("Test Boomer", func() {
 
 	It("test record success", func() {
 		defer func() {
-			defaultBoomer = &Boomer{}
+			defaultBoomer = &Boomer{logger: log.Default()}
 		}()
 
 		// called before runner instance created
@@ -245,7 +246,7 @@ var _ = Describe("Test Boomer", func() {
 
 	It("test record failure", func() {
 		defer func() {
-			defaultBoomer = &Boomer{}
+			defaultBoomer = &Boomer{logger: log.Default()}
 		}()
 
 		// called before runner instance created
@@ -274,4 +275,20 @@ var _ = Describe("Test Boomer", func() {
 		Expect(requestFailureMsg.responseTime).To(BeEquivalentTo(2))
 		Expect(requestFailureMsg.error).To(Equal("udp error"))
 	})
+
+	It("test loggers", func() {
+		defer func() {
+			defaultBoomer = &Boomer{logger: log.Default()}
+		}()
+
+		logger := log.New(os.Stdout, "[boomer]", log.LstdFlags)
+
+		defaultBoomer = &Boomer{}
+		defaultBoomer.WithLogger(nil)
+		defaultBoomer.WithLogger(logger)
+
+		defaultBoomer.slaveRunner = &slaveRunner{}
+		defaultBoomer.localRunner = &localRunner{}
+		defaultBoomer.WithLogger(logger)
+	})
 })

output.go

@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
-	"os"
 	"sort"
 	"strconv"
 	"time"
@@ -35,11 +34,21 @@ type Output interface {
 
 // ConsoleOutput is the default output for standalone mode.
 type ConsoleOutput struct {
+	logger *log.Logger
 }
 
 // NewConsoleOutput returns a ConsoleOutput.
 func NewConsoleOutput() *ConsoleOutput {
-	return &ConsoleOutput{}
+	return &ConsoleOutput{logger: log.Default()}
+}
+
+// WithLogger allows user to use their own logger.
+// If the logger is nil, it will not take effect.
+func (o *ConsoleOutput) WithLogger(logger *log.Logger) *ConsoleOutput {
+	if logger != nil {
+		o.logger = logger
+	}
+	return o
 }
 
 func getMedianResponseTime(numRequests int64, responseTimes map[int64]int64) int64 {
@@ -119,14 +128,15 @@ func (o *ConsoleOutput) OnStop() {
 func (o *ConsoleOutput) OnEvent(data map[string]interface{}) {
 	output, err := convertData(data)
 	if err != nil {
-		log.Printf("convert data error: %v\n", err)
+		o.logger.Printf("convert data error: %v\n", err)
 		return
 	}
 
 	currentTime := time.Now()
-	println(fmt.Sprintf("Current time: %s, Users: %d, Total RPS: %d, Total Fail Ratio: %.1f%%",
+	o.logger.Println(fmt.Sprintf("Current time: %s, Users: %d, Total RPS: %d, Total Fail Ratio: %.1f%%",
 		currentTime.Format("2006/01/02 15:04:05"), output.UserCount, output.TotalRPS, output.TotalFailRatio*100))
-	table := tablewriter.NewWriter(os.Stdout)
+	noPrefixLogger := log.New(o.logger.Writer(), "", 0)
+	table := tablewriter.NewWriter(noPrefixLogger.Writer())
 	table.SetHeader([]string{"Type", "Name", "# requests", "# fails", "Median", "Average", "Min", "Max", "Content Size", "# reqs/sec", "# fails/sec"})
 
 	for _, stat := range output.Stats {
@@ -145,7 +155,7 @@ func (o *ConsoleOutput) OnEvent(data map[string]interface{}) {
 		table.Append(row)
 	}
 	table.Render()
-	println()
+	o.logger.Println()
 }
 
 type statsEntryOutput struct {
@@ -334,17 +344,28 @@ var (
 func NewPrometheusPusherOutput(gatewayURL, jobName string) *PrometheusPusherOutput {
 	return &PrometheusPusherOutput{
 		pusher: push.New(gatewayURL, jobName),
+		logger: log.Default(),
+	}
+}
+
+// WithLogger allows user to use their own logger.
+// If the logger is nil, it will not take effect.
+func (o *PrometheusPusherOutput) WithLogger(logger *log.Logger) *PrometheusPusherOutput {
+	if logger != nil {
+		o.logger = logger
 	}
+	return o
 }
 
 // PrometheusPusherOutput pushes boomer stats to Prometheus Pushgateway.
 type PrometheusPusherOutput struct {
 	pusher *push.Pusher // Prometheus Pushgateway Pusher
+	logger *log.Logger
 }
 
 // OnStart will register all prometheus metric collectors
 func (o *PrometheusPusherOutput) OnStart() {
-	log.Println("register prometheus metric collectors")
+	o.logger.Println("register prometheus metric collectors")
 	registry := prometheus.NewRegistry()
 	registry.MustRegister(
 		// gauge vectors for requests
@@ -374,7 +395,7 @@ func (o *PrometheusPusherOutput) OnStop() {
 func (o *PrometheusPusherOutput) OnEvent(data map[string]interface{}) {
 	output, err := convertData(data)
 	if err != nil {
-		log.Printf("convert data error: %v\n", err)
+		o.logger.Printf("convert data error: %v\n", err)
 		return
 	}
 
@@ -402,6 +423,6 @@ func (o *PrometheusPusherOutput) OnEvent(data map[string]interface{}) {
 	}
 
 	if err := o.pusher.Push(); err != nil {
-		log.Printf("Could not push to Pushgateway: error: %v\n", err)
+		o.logger.Printf("Could not push to Pushgateway: error: %v\n", err)
 	}
 }

output_test.go

@@ -1,6 +1,9 @@
 package boomer
 
 import (
+	"log"
+	"os"
+
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 )
@@ -94,4 +97,17 @@ var _ = Describe("test output", func() {
 
 		o.OnStop()
 	})
+
+	It("test loggers", func() {
+		o := NewConsoleOutput()
+
+		logger := log.New(os.Stdout, "[boomer]", log.LstdFlags)
+
+		o.WithLogger(nil)
+		o.WithLogger(logger)
+
+		o2 := NewPrometheusPusherOutput("", "")
+		o2.WithLogger(nil)
+		o2.WithLogger(logger)
+	})
 })