refactor(traffic): replace package-vars with Scheduler interface and struct

arbiter.go

@@ -273,8 +273,13 @@ func (a *abtr) run(metadata module.Metadata) error {
 	defer reporterCancel()
 	reporter.Start(reporterCtx)
 
+	sched := traffic.New(&traffic.Opts{
+		Logger:      logger,
+		WorkerLimit: workerLimit.Value(),
+	})
+
 	// Run traffic.
-	if err := traffic.Run(timeoutCtx, metadata, reporter, workerLimit.Value()); err != nil {
+	if err := sched.Run(timeoutCtx, metadata, reporter); err != nil {
 		reporter.ReportError(err) // Report is done in case of early traffic failure, to highlight issues in the TUI.
 		logger.Error(err, "Failed to start traffic")
 		return err
@@ -296,7 +301,7 @@ func (a *abtr) run(metadata module.Metadata) error {
 	// stopErr accumulates any errors from stopping traffic and modules, and finalising the report,
 	// to be returned at the end of the function.
 	var stopErr error
-	if stopErr = traffic.Stop(); stopErr != nil {
+	if stopErr = sched.Stop(); stopErr != nil {
 		logger.Error(stopErr, "Error when stopping traffic")
 		stopErr = fmt.Errorf("%w: traffic stop: %w", ErrStopping, stopErr)
 	}

pkg/traffic/scheduler.go

@@ -8,28 +8,15 @@ import (
 	"time"
 
 	"github.com/go-logr/logr"
-	abtrlog "github.com/maansaake/arbiter/internal/log"
 	"github.com/maansaake/arbiter/pkg/module"
 	"github.com/maansaake/arbiter/pkg/report"
 )
 
 var (
-	reporter report.Reporter //nolint:gochecknoglobals // package-level state for traffic scheduler
-
-	workloads []*workload //nolint:gochecknoglobals // package-level state for traffic scheduler
-
-	// Stop stuff.
-	stopChan chan *workload //nolint:gochecknoglobals // package-level state for traffic scheduler
-
-	// logger is the package logger for the traffic package.
-	logger logr.Logger //nolint:gochecknoglobals // package-level state for traffic scheduler
-
 	ErrNoOpsToSchedule = errors.New("there were no operations to schedule")
 	ErrZeroRate        = errors.New("operation has a zero rate")
 	ErrCleanupTimeout  = errors.New("cleanup timed out")
 	ErrRateIssue       = errors.New("rate issue")
-
-	SampleTolerancePerc = 0.05 //nolint:gochecknoglobals // exported config var for tests
 )
 
 const (
@@ -38,75 +25,125 @@ const (
 	defaultSampleIntervalSeconds = 10
 	cleanupTimeout               = 5 * time.Second
 	minRateForDefaultSample      = 30
+	defaultSampleTolerancePerc   = 0.05
 )
 
+// Opts configures a Scheduler.
+type Opts struct {
+	// Logger is used for traffic scheduler logs. Defaults to a discard logger if not set.
+	Logger logr.Logger
+	// WorkerLimit is the maximum number of concurrent workers per workload. Defaults to DefaultWorkerLimit.
+	WorkerLimit int
+	// SampleTolerancePerc is the tolerance percentage used when comparing sampled rates in tests.
+	// Defaults to 0.05 (5%).
+	SampleTolerancePerc float64
+}
+
+// Scheduler runs traffic against registered modules.
+type Scheduler interface {
+	// Run starts traffic generation for the given modules, reporting results to reporter.
+	// It is asynchronous: it returns once the goroutines are launched and monitors ctx
+	// to stop gracefully when it is cancelled.
+	Run(ctx context.Context, metadata module.Metadata, reporter report.Reporter) error
+	// Stop waits for all workloads to finish after the context passed to Run is cancelled.
+	Stop() error
+}
+
+type scheduler struct {
+	logger              logr.Logger
+	workerLimit         int
+	sampleTolerancePerc float64
+
+	workloads []*workload
+	stopChan  chan *workload
+}
+
+// New creates a Scheduler with the given options. A nil opts uses all defaults.
+func New(opts *Opts) Scheduler {
+	if opts == nil {
+		opts = &Opts{}
+	}
+	if opts.WorkerLimit == 0 {
+		opts.WorkerLimit = DefaultWorkerLimit
+	}
+	if opts.SampleTolerancePerc == 0 {
+		opts.SampleTolerancePerc = defaultSampleTolerancePerc
+	}
+	return &scheduler{
+		logger:              opts.Logger,
+		workerLimit:         opts.WorkerLimit,
+		sampleTolerancePerc: opts.SampleTolerancePerc,
+	}
+}
+
 // Run traffic for the input modules using their exposed operations. Traffic
 // generation will make operation calls at the specified rates and report
 // problems to the reporter. Run() is asynchronous and returns once the main
 // go-routine has been started. Run() will monitor the context's done channel
 // and stop gracefully once it's closed.
-func Run(
+func (s *scheduler) Run(
 	ctx context.Context,
 	metadata module.Metadata,
-	r report.Reporter,
-	workerLimit int,
+	reporter report.Reporter,
 ) error {
-	logger = abtrlog.GetLogger()
-
-	logger.Info("Running traffic generator")
-	// Run initialisation of traffic synchronously
-	reporter = r
+	s.logger.Info("Running traffic generator")
 
-	workloads = make([]*workload, 0, len(metadata))
+	s.workloads = make([]*workload, 0, len(metadata))
 	for _, meta := range metadata {
 		for _, op := range meta.Ops() {
 			if op.Disabled {
-				logger.Info("Skipping disabled operation", "mod", meta.Name(), "op", op.Name)
+				s.logger.Info("Skipping disabled operation", "mod", meta.Name(), "op", op.Name)
 				continue
 			}
 
 			if op.Rate == 0 {
 				return fmt.Errorf("%w: %s", ErrZeroRate, op.Name)
 			}
 
-			workloads = append(workloads, &workload{
-				workerLimit: workerLimit,
+			s.workloads = append(s.workloads, &workload{
+				workerLimit: s.workerLimit,
 				statLock:    &sync.Mutex{},
 				mod:         meta.Name(),
 				op:          op,
+				reporter:    reporter,
+				logger:      s.logger,
 			})
 		}
 	}
 
-	if len(workloads) == 0 {
+	if len(s.workloads) == 0 {
 		return ErrNoOpsToSchedule
 	}
 
 	// Create stop channel that workloads will report to when stopping.
-	stopChan = make(chan *workload, len(workloads))
+	s.stopChan = make(chan *workload, len(s.workloads))
+	for _, wl := range s.workloads {
+		wl.stopChan = s.stopChan
+	}
 
-	// Run the workload in a separate go-routine, runs until context is done
-	for _, workload := range workloads {
-		go workload.run(ctx)
+	// Run the workloads in separate go-routines, each runs until context is done.
+	for _, wl := range s.workloads {
+		go wl.run(ctx)
 	}
 
 	return nil
 }
 
-func Stop() error {
-	logger.Info("Stopping traffic generator", "workload_count", len(workloads))
+// Stop waits for all workloads to finish and returns any error encountered.
+func (s *scheduler) Stop() error {
+	s.logger.Info("Stopping traffic generator", "workload_count", len(s.workloads))
 
 	stopCount := 0
 	for {
 		select {
 		case <-time.After(cleanupTimeout):
-			logger.Error(ErrCleanupTimeout, "Cleanup timed out after "+cleanupTimeout.String())
+			s.logger.Error(ErrCleanupTimeout, "Cleanup timed out after "+cleanupTimeout.String())
 			return ErrCleanupTimeout
-		case workload := <-stopChan:
-			logger.Info("Workload stopped", "mod", workload.mod, "op", workload.op.Name)
+		case wl := <-s.stopChan:
+			s.logger.Info("Workload stopped", "mod", wl.mod, "op", wl.op.Name)
 			stopCount++
-			if stopCount == len(workloads) {
-				logger.Info("All workloads have stopped")
+			if stopCount == len(s.workloads) {
+				s.logger.Info("All workloads have stopped")
 				return nil
 			}
 		}
@@ -117,7 +154,6 @@ func getSampleInterval(op *module.Op) time.Duration {
 	if op.Rate < minRateForDefaultSample {
 		// Minimum 5 samples, this should be a super corner case. Add some time
 		// to allow the 5th invocation to fire.
-
 		return time.Minute/time.Duration(op.Rate)*5 + 250*time.Millisecond
 	}
 

pkg/traffic/scheduler_test.go

@@ -7,12 +7,21 @@ import (
 	"testing"
 	"time"
 
+	"github.com/go-logr/logr"
 	"github.com/maansaake/arbiter/pkg/module"
 	modulemock "github.com/maansaake/arbiter/pkg/module/mock"
 	reportmock "github.com/maansaake/arbiter/pkg/report/mock"
 	log "github.com/trebent/zerologr"
 )
 
+// newTestScheduler creates a Scheduler pre-configured for tests.
+func newTestScheduler() Scheduler {
+	return New(&Opts{
+		Logger:      logr.Discard(),
+		WorkerLimit: DefaultWorkerLimit,
+	})
+}
+
 func TestRunAndAwaitStop(t *testing.T) {
 	opWg := sync.WaitGroup{}
 	opWg.Add(2)
@@ -30,7 +39,8 @@ func TestRunAndAwaitStop(t *testing.T) {
 		},
 	}
 	ctx, cancel := context.WithCancel(context.Background())
-	err := Run(ctx, []*module.Meta{{Module: mod}}, reportmock.NewMock(), DefaultWorkerLimit)
+	sched := newTestScheduler()
+	err := sched.Run(ctx, []*module.Meta{{Module: mod}}, reportmock.NewMock())
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -39,14 +49,15 @@ func TestRunAndAwaitStop(t *testing.T) {
 	opWg.Wait()
 
 	cancel()
-	err = Stop()
+	err = sched.Stop()
 	if err != nil {
 		t.Fatal(err)
 	}
 }
 
 func TestRunNoOps(t *testing.T) {
-	err := Run(context.TODO(), []*module.Meta{{Module: modulemock.NewMock()}}, nil, DefaultWorkerLimit)
+	sched := newTestScheduler()
+	err := sched.Run(context.TODO(), []*module.Meta{{Module: modulemock.NewMock()}}, nil)
 	if err != nil && !errors.Is(err, ErrNoOpsToSchedule) {
 		t.Fatal("unexpected error type")
 	}
@@ -63,7 +74,8 @@ func TestRunZeroRate(t *testing.T) {
 			Rate: 0,
 		},
 	}
-	err := Run(context.TODO(), []*module.Meta{{Module: mod}}, nil, DefaultWorkerLimit)
+	sched := newTestScheduler()
+	err := sched.Run(context.TODO(), []*module.Meta{{Module: mod}}, nil)
 	if err != nil && !errors.Is(err, ErrZeroRate) {
 		t.Fatal("unexpected error type")
 	}
@@ -90,15 +102,16 @@ func TestReportOpToReporter(t *testing.T) {
 
 	reporter := reportmock.NewMock()
 	ctx, cancel := context.WithCancel(context.Background())
-	if err := Run(ctx, []*module.Meta{{Module: mod}}, reporter, DefaultWorkerLimit); err != nil {
+	sched := newTestScheduler()
+	if err := sched.Run(ctx, []*module.Meta{{Module: mod}}, reporter); err != nil {
 		t.Fatal(err)
 	}
 
 	wg.Wait()
 	cancel()
 
 	// This should ensure the reporter mock has received the Op report.
-	Stop()
+	sched.Stop()
 
 	if reporter.OpResults[0].Duration == 0 {
 		t.Fatal("duration was not reported to reporter")
@@ -123,15 +136,16 @@ func TestReportOpDurationOverrideToReporter(t *testing.T) {
 
 	reporter := reportmock.NewMock()
 	ctx, cancel := context.WithCancel(context.Background())
-	if err := Run(ctx, []*module.Meta{{Module: mod}}, reporter, DefaultWorkerLimit); err != nil {
+	sched := newTestScheduler()
+	if err := sched.Run(ctx, []*module.Meta{{Module: mod}}, reporter); err != nil {
 		t.Fatal(err)
 	}
 
 	wg.Wait()
 	cancel()
 
 	// This should ensure the reporter mock has received the Op report.
-	Stop()
+	sched.Stop()
 
 	if reporter.OpResults[0].Duration != 12*time.Millisecond {
 		t.Fatal("duration override was not used")
@@ -156,15 +170,16 @@ func TestReportOpErr(t *testing.T) {
 
 	reporter := reportmock.NewMock()
 	ctx, cancel := context.WithCancel(context.Background())
-	if err := Run(ctx, []*module.Meta{{Module: mod}}, reporter, DefaultWorkerLimit); err != nil {
+	sched := newTestScheduler()
+	if err := sched.Run(ctx, []*module.Meta{{Module: mod}}, reporter); err != nil {
 		t.Fatal(err)
 	}
 
 	wg.Wait()
 	cancel()
 
 	// This should ensure the reporter mock has received the Op report.
-	Stop()
+	sched.Stop()
 
 	if len(reporter.OpResults) != 0 {
 		t.Fatal("unexpected op results found")

pkg/traffic/worker.go

@@ -14,27 +14,33 @@ type worker struct {
 }
 
 func (worker *worker) run(ctx context.Context) {
-	logger.Info("Starting worker", "mod", worker.parent.mod, "op", worker.parent.op.Name)
+	worker.parent.logger.Info("Starting worker", "mod", worker.parent.mod, "op", worker.parent.op.Name)
 	worker.done = make(chan bool)
 
 	for {
 		select {
 		case <-ctx.Done():
-			logger.Info("Context closed, stopping worker", "mod", worker.parent.mod, "op", worker.parent.op.Name)
+			worker.parent.logger.Info(
+				"Context closed, stopping worker",
+				"mod",
+				worker.parent.mod,
+				"op",
+				worker.parent.op.Name,
+			)
 			worker.ticker.Stop()
 
 			close(worker.done)
 			return
 		case t := <-worker.ticker.C:
-			logger.V(workerVerboseLogLevel).
+			worker.parent.logger.V(workerVerboseLogLevel).
 				Info("Worker tick", "time", t, "mod", worker.parent.mod, "op", worker.parent.op.Name)
 			worker.parent.doOp()
 		}
 	}
 }
 
 func (worker *worker) reset(tickerInterval time.Duration) {
-	logger.Info(
+	worker.parent.logger.Info(
 		"Resetting worker ticker",
 		"mod",
 		worker.parent.mod,