diff --git a/loadtest.go b/loadtest.go index 50914c5..7b90248 100644 --- a/loadtest.go +++ b/loadtest.go @@ -37,6 +37,8 @@ type LoadTestOptions struct { resultHandlers []resultHandler weightOverrides map[string]int reportInterval time.Duration + ctxModifier func(ctx context.Context) context.Context + defaultTimeout time.Duration } type LoadTestOption func(*LoadTestOptions) diff --git a/options.go b/options.go index 0da8186..2946477 100644 --- a/options.go +++ b/options.go @@ -1,6 +1,7 @@ package goload import ( + "context" "github.com/HenriBeck/goload/pacer" "time" ) @@ -53,6 +54,18 @@ func WithAdditionalResultHandler(handler resultHandler) LoadTestOption { } } +func WithContextModifier(fn func(ctx context.Context) context.Context) LoadTestOption { + return func(options *LoadTestOptions) { + options.ctxModifier = fn + } +} + +func WithDefaultTimeout(timeout time.Duration) LoadTestOption { + return func(options *LoadTestOptions) { + options.defaultTimeout = timeout + } +} + func WithWeightOverrides(overrides map[string]int) LoadTestOption { return func(options *LoadTestOptions) { options.weightOverrides = overrides diff --git a/runner.go b/runner.go index 3decad4..ca5bc6e 100644 --- a/runner.go +++ b/runner.go @@ -16,6 +16,9 @@ type Runner struct { maxWorkers int weightOverrides map[string]int + ctxModifier func(ctx context.Context) context.Context + defaultTimeout time.Duration + startedAt *time.Time } @@ -26,6 +29,8 @@ func NewRunner(loadTestOptions LoadTestOptions) *Runner { workers: loadTestOptions.initialWorkers, maxWorkers: loadTestOptions.maxWorkers, weightOverrides: loadTestOptions.weightOverrides, + ctxModifier: loadTestOptions.ctxModifier, + defaultTimeout: loadTestOptions.defaultTimeout, startedAt: nil, } @@ -155,6 +160,14 @@ func (r *Runner) hit(ex Executor, began time.Time) *Result { var cancel context.CancelFunc ctx, cancel = context.WithTimeout(context.Background(), ex.Options().Timeout) defer cancel() + } else if r.defaultTimeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(context.Background(), ex.Options().Timeout) + defer cancel() + } + + if r.ctxModifier != nil { + ctx = r.ctxModifier(ctx) } resp := ex.Execute(ctx)