diff --git a/arbiter.go b/arbiter.go index 3758897..51b30db 100644 --- a/arbiter.go +++ b/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) } diff --git a/pkg/traffic/scheduler.go b/pkg/traffic/scheduler.go index 8038471..51d877f 100644 --- a/pkg/traffic/scheduler.go +++ b/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,30 +25,74 @@ 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 } @@ -69,44 +100,50 @@ func Run( 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 } diff --git a/pkg/traffic/scheduler_test.go b/pkg/traffic/scheduler_test.go index c73b379..1e4b4de 100644 --- a/pkg/traffic/scheduler_test.go +++ b/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,7 +102,8 @@ 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) } @@ -98,7 +111,7 @@ func TestReportOpToReporter(t *testing.T) { 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,7 +136,8 @@ 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) } @@ -131,7 +145,7 @@ func TestReportOpDurationOverrideToReporter(t *testing.T) { 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,7 +170,8 @@ 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) } @@ -164,7 +179,7 @@ func TestReportOpErr(t *testing.T) { 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") diff --git a/pkg/traffic/worker.go b/pkg/traffic/worker.go index ac95c4f..95ec516 100644 --- a/pkg/traffic/worker.go +++ b/pkg/traffic/worker.go @@ -14,19 +14,25 @@ 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() } @@ -34,7 +40,7 @@ func (worker *worker) run(ctx context.Context) { } func (worker *worker) reset(tickerInterval time.Duration) { - logger.Info( + worker.parent.logger.Info( "Resetting worker ticker", "mod", worker.parent.mod, diff --git a/pkg/traffic/workload.go b/pkg/traffic/workload.go index b7c336a..3a0ff80 100644 --- a/pkg/traffic/workload.go +++ b/pkg/traffic/workload.go @@ -6,7 +6,9 @@ import ( "sync" "time" + "github.com/go-logr/logr" "github.com/maansaake/arbiter/pkg/module" + "github.com/maansaake/arbiter/pkg/report" ) type workload struct { @@ -19,6 +21,10 @@ type workload struct { statLock *sync.Mutex calls float64 totalDur time.Duration + + reporter report.Reporter + stopChan chan *workload + logger logr.Logger } const workloadVerboseLogLevel = 100 @@ -26,7 +32,7 @@ const workloadVerboseLogLevel = 100 // run runs the workload, which in turn spawns workers to do the actual invocations. The workload // will monitor call-rates and scale the number of workers as needed. func (w *workload) run(ctx context.Context) { - logger.Info("Starting workload", "mod", w.mod, "op", w.op.Name, "rate", w.op.Rate) + w.logger.Info("Starting workload", "mod", w.mod, "op", w.op.Name, "rate", w.op.Rate) // All workload start with exactly one worker. After the first sampling // period, this may be increased. @@ -36,7 +42,7 @@ func (w *workload) run(ctx context.Context) { samplingInterval := getSampleInterval(w.op) expectedCalls := float64(samplingInterval) / float64(time.Minute) * float64(w.op.Rate) - logger.Info( + w.logger.Info( "Setting sampling interval", "mod", w.mod, @@ -52,20 +58,20 @@ func (w *workload) run(ctx context.Context) { for { select { case <-ctx.Done(): - logger.Info("Context closed, stopping workload", "mod", w.mod, "op", w.op.Name) + w.logger.Info("Context closed, stopping workload", "mod", w.mod, "op", w.op.Name) rateCheckTicker.Stop() for i, worker := range w.workers { - logger.Info("Awaiting worker", "mod", w.mod, "op", w.op.Name, "worker", i) + w.logger.Info("Awaiting worker", "mod", w.mod, "op", w.op.Name, "worker", i) // Await the stop of each worker <-worker.done - logger.Info("Worker stopped", "mod", w.mod, "op", w.op.Name, "worker", i) + w.logger.Info("Worker stopped", "mod", w.mod, "op", w.op.Name, "worker", i) } - stopChan <- w + w.stopChan <- w return case <-rateCheckTicker.C: - logger.Info( + w.logger.Info( "Running rate check", "mod", w.mod, @@ -79,7 +85,7 @@ func (w *workload) run(ctx context.Context) { if w.calls > 0 { avgUs := (w.totalDur / time.Duration(w.calls)).Microseconds() - logger.Info("Average exec time", "mod", w.mod, "op", w.op.Name, "avg_µs", avgUs) + w.logger.Info("Average exec time", "mod", w.mod, "op", w.op.Name, "avg_µs", avgUs) w.scale(ctx) } @@ -105,10 +111,10 @@ func (w *workload) withStatLock(f func()) { // rate per worker when this is called. func (w *workload) scale(ctx context.Context) { requiredWorkers := w.getWorkerCount() - logger.Info("Required worker count", "required_workers", requiredWorkers, "mod", w.mod, "op", w.op.Name) + w.logger.Info("Required worker count", "required_workers", requiredWorkers, "mod", w.mod, "op", w.op.Name) if int(requiredWorkers) > len(w.workers) { - logger.Info("Adding workers", "count", int(requiredWorkers)-len(w.workers), "mod", w.mod, "op", w.op.Name) + w.logger.Info("Adding workers", "count", int(requiredWorkers)-len(w.workers), "mod", w.mod, "op", w.op.Name) for i := int(requiredWorkers) - len(w.workers); i > 0; i-- { w.addWorker(ctx) } @@ -156,7 +162,7 @@ func (w *workload) getAverageDuration() time.Duration { } func (w *workload) addWorker(ctx context.Context) { - logger.Info("Adding worker", "mod", w.mod, "op", w.op.Name) + w.logger.Info("Adding worker", "mod", w.mod, "op", w.op.Name) worker := &worker{ parent: w, @@ -171,11 +177,11 @@ func (w *workload) addWorker(ctx context.Context) { // doOp executes the workload operation and reports the result to the reporter. It also updates // the total duration and call count for the workload, which are used to calculate the average execution time. func (w *workload) doOp() { - logger.V(workloadVerboseLogLevel).Info("Triggering workload op", "mod", w.mod, "op", w.op.Name) + w.logger.V(workloadVerboseLogLevel).Info("Triggering workload op", "mod", w.mod, "op", w.op.Name) start := time.Now() res, err := w.op.Do() - logger.V(workloadVerboseLogLevel).Info("Ran op", "mod", w.mod, "op", w.op.Name) + w.logger.V(workloadVerboseLogLevel).Info("Ran op", "mod", w.mod, "op", w.op.Name) if res.Duration == 0 { res.Duration = time.Since(start) @@ -187,10 +193,10 @@ func (w *workload) doOp() { w.calls++ w.totalDur += res.Duration }) - logger.V(workloadVerboseLogLevel).Info("Reporting", "mod", w.mod, "op", w.op.Name) + w.logger.V(workloadVerboseLogLevel).Info("Reporting", "mod", w.mod, "op", w.op.Name) - reporter.ReportOp(w.mod, w.op.Name, &res, err) + w.reporter.ReportOp(w.mod, w.op.Name, &res, err) - logger.V(workloadVerboseLogLevel). + w.logger.V(workloadVerboseLogLevel). Info("Trigger done", "mod", w.mod, "op", w.op.Name, "duration_µs", time.Since(start).Microseconds()) }