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

@@ -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
+}

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
 	}

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
 		}
 	}