From e5b887962b89d0dc1ce50c01ab2176d369459d7b Mon Sep 17 00:00:00 2001 From: Nedyalko Andreev Date: Thu, 8 Dec 2022 00:17:45 +0200 Subject: [PATCH] Move and expand the test-abort context to the execution package The expanded capabilities will be necessary for changes in upcoming PRs --- execution/abort.go | 84 +++++++++++++++++++++++++++++++++++++++++ execution/scheduler.go | 5 +-- lib/executor/helpers.go | 47 +---------------------- 3 files changed, 88 insertions(+), 48 deletions(-) create mode 100644 execution/abort.go diff --git a/execution/abort.go b/execution/abort.go new file mode 100644 index 00000000000..08e9bfa8ea2 --- /dev/null +++ b/execution/abort.go @@ -0,0 +1,84 @@ +package execution + +import ( + "context" + "sync" + + "github.com/sirupsen/logrus" +) + +// testAbortKey is the key used to store the abort function for the context of +// an executor. This allows any users of that context or its sub-contexts to +// cancel the whole execution tree, while at the same time providing all of the +// details for why they cancelled it via the attached error. +type testAbortKey struct{} + +type testAbortController struct { + cancel context.CancelFunc + + logger logrus.FieldLogger + lock sync.Mutex // only the first reason will be kept, other will be logged + reason error // see errext package, you can wrap errors to attach exit status, run status, etc. +} + +func (tac *testAbortController) abort(err error) { + tac.lock.Lock() + defer tac.lock.Unlock() + if tac.reason != nil { + tac.logger.Debugf( + "test abort with reason '%s' was attempted when the test was already aborted due to '%s'", + err.Error(), tac.reason.Error(), + ) + return + } + tac.reason = err + tac.cancel() +} + +func (tac *testAbortController) getReason() error { + tac.lock.Lock() + defer tac.lock.Unlock() + return tac.reason +} + +// NewTestRunContext returns context.Context that can be aborted by calling the +// returned TestAbortFunc or by calling CancelTestRunContext() on the returned +// context or a sub-context of it. Use this to initialize the context that will +// be passed to the ExecutionScheduler, so `execution.test.abort()` and the REST +// API test stopping both work. +func NewTestRunContext( + ctx context.Context, logger logrus.FieldLogger, +) (newCtx context.Context, abortTest func(reason error)) { + ctx, cancel := context.WithCancel(ctx) + + controller := &testAbortController{ + cancel: cancel, + logger: logger, + } + + return context.WithValue(ctx, testAbortKey{}, controller), controller.abort +} + +// AbortTestRun will cancel the test run context with the given reason if the +// provided context is actually a TestRuncontext or a child of one. +func AbortTestRun(ctx context.Context, err error) bool { + if x := ctx.Value(testAbortKey{}); x != nil { + if v, ok := x.(*testAbortController); ok { + v.abort(err) + return true + } + } + return false +} + +// GetCancelReasonIfTestAborted returns a reason the Context was cancelled, if it was +// aborted with these functions. It will return nil if ctx is not an +// TestRunContext (or its children) or if it was never aborted. +func GetCancelReasonIfTestAborted(ctx context.Context) error { + if x := ctx.Value(testAbortKey{}); x != nil { + if v, ok := x.(*testAbortController); ok { + return v.getReason() + } + } + return nil +} diff --git a/execution/scheduler.go b/execution/scheduler.go index a2e7214330e..0f5d42a729f 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, _ := NewTestRunContext(runSubCtx, logger) 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 := GetCancelReasonIfTestAborted(execCtx); err != nil && errext.IsInterruptError(err) { interrupted = true return err } diff --git a/lib/executor/helpers.go b/lib/executor/helpers.go index 13b5df97cdb..5ab401c1d5d 100644 --- a/lib/executor/helpers.go +++ b/lib/executor/helpers.go @@ -10,6 +10,7 @@ import ( "github.com/sirupsen/logrus" "go.k6.io/k6/errext" + "go.k6.io/k6/execution" "go.k6.io/k6/lib" "go.k6.io/k6/lib/types" "go.k6.io/k6/ui/pb" @@ -56,56 +57,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) + execution.AbortTestRun(ctx, err) return true } }