diff --git a/errext/cancellable_context.go b/errext/cancellable_context.go new file mode 100644 index 000000000000..1453013615dd --- /dev/null +++ b/errext/cancellable_context.go @@ -0,0 +1,52 @@ +package errext + +import ( + "context" +) + +// cancelReasonKey is the key used to store both the cancel function for the +// context, and the reason it of an executor. This is a work around to avoid +// excessive changes for the ability of nested functions to cancel the passed +// context and provide the reason they did so. +type cancelAndReasonKey struct{} + +type cancelAndReasonCtxVal struct { + cancel context.CancelFunc + reason error +} + +// CancellableContext returns context.Context that can be cancelled by calling +// CancelContextWithError. It is used to initialize contexts that will be passed +// to executors. +// +// This allows executors to globally halt any executions that uses this context. +// Example use case is when a script calls test.abort(). +// +// TODO: maybe make this into a custom type on top of context.Context? A +// k6-specific interface that has an extra method to cancel it with a reason? +func CancellableContext(ctx context.Context) (context.Context, func() error) { + ctx, cancel := context.WithCancel(ctx) + val := &cancelAndReasonCtxVal{cancel: cancel} + reasonGetter := func() error { + return val.reason + } + return context.WithValue(ctx, cancelAndReasonKey{}, val), reasonGetter +} + +// CancelContextWithError cancels the executor context found in ctx and saves +// the given error inside. +// +// ctx can only be a context created by the execution.Scheduler and passed to an +// executor's Run() method, this function will panic for any other context +func CancelContextWithError(ctx context.Context, err error) { + x := ctx.Value(cancelAndReasonKey{}) + if x == nil { + panic("invalid context value, not cancellable") + } + v, ok := x.(*cancelAndReasonCtxVal) + if !ok { + panic("invalid context value, missing cancel reason and function") + } + v.reason = err + v.cancel() +} diff --git a/execution/scheduler.go b/execution/scheduler.go index a2e7214330ef..a2d1386f283a 100644 --- a/execution/scheduler.go +++ b/execution/scheduler.go @@ -11,7 +11,6 @@ import ( "go.k6.io/k6/errext" "go.k6.io/k6/lib" - "go.k6.io/k6/lib/executor" "go.k6.io/k6/metrics" "go.k6.io/k6/ui/pb" ) @@ -451,7 +450,7 @@ func (e *Scheduler) Run(globalCtx, runCtx context.Context, engineOut chan<- metr // this context effectively stopping all executions. // // This is for addressing test.abort(). - execCtx := executor.Context(runSubCtx) + execCtx, getAbortReason := errext.CancellableContext(runSubCtx) for _, exec := range e.executors { go e.runExecutor(execCtx, runResults, engineOut, exec) } @@ -479,7 +478,7 @@ func (e *Scheduler) Run(globalCtx, runCtx context.Context, engineOut chan<- metr return err } } - if err := executor.CancelReason(execCtx); err != nil && errext.IsInterruptError(err) { + if err := getAbortReason(); err != nil && errext.IsInterruptError(err) { interrupted = true return err } diff --git a/lib/executor/helpers.go b/lib/executor/helpers.go index 13b5df97cdbb..c205f9cf0fa2 100644 --- a/lib/executor/helpers.go +++ b/lib/executor/helpers.go @@ -56,56 +56,12 @@ func validateStages(stages []Stage) []error { return errors } -// cancelKey is the key used to store the cancel function for the context of an -// executor. This is a work around to avoid excessive changes for the ability of -// nested functions to cancel the passed context. -type cancelKey struct{} - -type cancelExec struct { - cancel context.CancelFunc - reason error -} - -// Context returns context.Context that can be cancelled by calling -// CancelExecutorContext. Use this to initialize context that will be passed to -// executors. -// -// This allows executors to globally halt any executions that uses this context. -// Example use case is when a script calls test.abort(). -func Context(ctx context.Context) context.Context { - ctx, cancel := context.WithCancel(ctx) - return context.WithValue(ctx, cancelKey{}, &cancelExec{cancel: cancel}) -} - -// cancelExecutorContext cancels executor context found in ctx, ctx can be a -// child of a context that was created with Context function. -func cancelExecutorContext(ctx context.Context, err error) { - if x := ctx.Value(cancelKey{}); x != nil { - if v, ok := x.(*cancelExec); ok { - v.reason = err - v.cancel() - } - } -} - -// CancelReason returns a reason the executor context was cancelled. This will -// return nil if ctx is not an executor context(ctx or any of its parents was -// never created by Context function). -func CancelReason(ctx context.Context) error { - if x := ctx.Value(cancelKey{}); x != nil { - if v, ok := x.(*cancelExec); ok { - return v.reason - } - } - return nil -} - // handleInterrupt returns true if err is InterruptError and if so it // cancels the executor context passed with ctx. func handleInterrupt(ctx context.Context, err error) bool { if err != nil { if errext.IsInterruptError(err) { - cancelExecutorContext(ctx, err) + errext.CancelContextWithError(ctx, err) return true } }