Add abortTest() helper function

Co-authored-by: Ivan Mirić <[email protected]>

Closes #1001
This adds abortTest() helper function to the k6 module. This function when
called inside a script it will

- stop the whole test run and k6 will exit with 107 status code
- stops immediately the VU that called it  and no more iterations are started
- make sure the teardown() is called
- the engine run status is 7 (RunStatusAbortedScriptError)

`(*goja.Runtime).Interrupt` is used for halting script execution and capturing
stack traces for better error message of what is happening with the script.

We introduce InterruptError which is used with `(*goja.Runtime).Interrupt` to
identify interrupts emitted by abortTest(). This way we use special handling of
this type.

Example script is

```js
import { abortTest, sleep } from 'k6';

export default function () {
    // We abort the test on second iteration
    if (__ITER == 1) {
        abortTest();
    }
    sleep(1);
}

export function teardown() {
    console.log('This function will be called even when we abort script');
}
```

abortTest() can be called in both default and setup functions, however you can't use it
in the init context

The following script will fail with the error
```
ERRO[0000] Using abortTest() in the init context is not supported at (...path to the script )init.js:13:43(34)
```

```js
import {
    abortTest
} from 'k6';

abortTest();

export function setup() {
}

export default function () {
    // ... some test logic ...
    console.log('mayday, mayday');
}
```

You can customize the reason for abortTest() by passing values to the function

```js
abortTest("Exceeded expectations");
```

Will emit `"Exceeded expectations"` on the logs

cmd/run.go

@@ -47,6 +47,7 @@ import (
 	"go.k6.io/k6/errext"
 	"go.k6.io/k6/errext/exitcodes"
 	"go.k6.io/k6/js"
+	"go.k6.io/k6/js/common"
 	"go.k6.io/k6/lib"
 	"go.k6.io/k6/lib/consts"
 	"go.k6.io/k6/loader"
@@ -247,6 +248,10 @@ a commandline interface for interacting with it.`,
 			initBar.Modify(pb.WithConstProgress(0, "Init VUs..."))
 			engineRun, engineWait, err := engine.Init(globalCtx, runCtx)
 			if err != nil {
+				var intErr *common.InterruptError
+				if errors.As(err, &intErr) {
+					return errext.WithExitCodeIfNone(err, exitcodes.ScriptException)
+				}
 				// Add a generic engine exit code if we don't have a more specific one
 				return errext.WithExitCodeIfNone(err, exitcodes.GenericEngine)
 			}
@@ -270,8 +275,18 @@ a commandline interface for interacting with it.`,
 
 			// Start the test run
 			initBar.Modify(pb.WithConstProgress(0, "Starting test..."))
-			if err := engineRun(); err != nil {
-				return errext.WithExitCodeIfNone(err, exitcodes.GenericEngine)
+			var interrupt error
+			err = engineRun()
+			if err != nil {
+				if common.IsInterruptError(err) {
+					interrupt = err
+				}
+				if !conf.Linger.Bool {
+					if interrupt == nil {
+						return errext.WithExitCodeIfNone(err, exitcodes.GenericEngine)
+					}
+					return errext.WithExitCodeIfNone(interrupt, exitcodes.ScriptException)
+				}
 			}
 			runCancel()
 			logger.Debug("Engine run terminated cleanly")
@@ -320,6 +335,9 @@ a commandline interface for interacting with it.`,
 			logger.Debug("Waiting for engine processes to finish...")
 			engineWait()
 			logger.Debug("Everything has finished, exiting k6!")
+			if interrupt != nil {
+				return errext.WithExitCodeIfNone(interrupt, exitcodes.ScriptException)
+			}
 			if engine.IsTainted() {
 				return errext.WithExitCodeIfNone(errors.New("some thresholds have failed"), exitcodes.ThresholdsHaveFailed)
 			}

cmd/run_test.go

@@ -22,6 +22,7 @@ package cmd
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"io"
 	"io/ioutil"
@@ -30,11 +31,17 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/afero"
+	"github.com/spf13/pflag"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"go.k6.io/k6/errext"
+	"go.k6.io/k6/errext/exitcodes"
+	"go.k6.io/k6/js/common"
 	"go.k6.io/k6/lib/fsext"
+	"go.k6.io/k6/lib/testutils"
 )
 
 type mockWriter struct {
@@ -125,3 +132,48 @@ func TestHandleSummaryResultError(t *testing.T) {
 	assertEqual(t, "file summary 1", files[filePath1])
 	assertEqual(t, "file summary 2", files[filePath2])
 }
+
+func TestAbortTest(t *testing.T) { //nolint: tparallel
+	t.Parallel()
+
+	t.Run("Check status code is 107", func(t *testing.T) { //nolint: paralleltest
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		logger := testutils.NewLogger(t)
+
+		cmd := getRunCmd(ctx, logger)
+		a, err := filepath.Abs("testdata/abort.js")
+		require.NoError(t, err)
+		cmd.SetArgs([]string{a})
+		err = cmd.Execute()
+		var e errext.HasExitCode
+		require.ErrorAs(t, err, &e)
+		require.Equal(t, exitcodes.ScriptException, e.ExitCode(), "Status code must be 107")
+		require.Contains(t, e.Error(), common.AbortTest)
+	})
+
+	t.Run("Check that teardown is called", func(t *testing.T) { //nolint: paralleltest
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		msg := "Calling teardown function after abortTest()"
+		var buf bytes.Buffer
+		logger := logrus.New()
+		logger.SetOutput(&buf)
+
+		cmd := getRunCmd(ctx, logger)
+		// Redefine the flag to avoid a nil pointer panic on lookup.
+		cmd.Flags().AddFlag(&pflag.Flag{
+			Name:   "address",
+			Hidden: true,
+		})
+		a, err := filepath.Abs("testdata/teardown.js")
+		require.NoError(t, err)
+		cmd.SetArgs([]string{a})
+		err = cmd.Execute()
+		var e errext.HasExitCode
+		require.ErrorAs(t, err, &e)
+		assert.Equal(t, exitcodes.ScriptException, e.ExitCode(), "Status code must be 107")
+		assert.Contains(t, e.Error(), common.AbortTest)
+		assert.Contains(t, buf.String(), msg)
+	})
+}

cmd/testdata/abort.js

@@ -0,0 +1,5 @@
+import { abortTest } from 'k6';
+
+export default function () {
+    abortTest();
+}

cmd/testdata/teardown.js

@@ -0,0 +1,9 @@
+import { abortTest } from 'k6';
+
+export default function () {
+    abortTest();
+}
+
+export function teardown() {
+    console.log('Calling teardown function after abortTest()');
+}

core/local/local.go

@@ -31,6 +31,7 @@ import (
 
 	"go.k6.io/k6/errext"
 	"go.k6.io/k6/lib"
+	"go.k6.io/k6/lib/executor"
 	"go.k6.io/k6/stats"
 	"go.k6.io/k6/ui/pb"
 )
@@ -381,8 +382,14 @@ func (e *ExecutionScheduler) Run(globalCtx, runCtx context.Context, engineOut ch
 	// Start all executors at their particular startTime in a separate goroutine...
 	logger.Debug("Start all executors...")
 	e.state.SetExecutionStatus(lib.ExecutionStatusRunning)
+
+	// We are using this context to allow lib.Executor implementations to cancel
+	// this context effectively stopping all executions.
+	//
+	// This is for addressing abortTest()
+	execCtx := executor.Context(runSubCtx)
 	for _, exec := range e.executors {
-		go e.runExecutor(runSubCtx, runResults, engineOut, exec)
+		go e.runExecutor(execCtx, runResults, engineOut, exec)
 	}
 
 	// Wait for all executors to finish
@@ -409,7 +416,10 @@ func (e *ExecutionScheduler) Run(globalCtx, runCtx context.Context, engineOut ch
 			return err
 		}
 	}
-
+	if err := executor.CancelReason(execCtx); err != nil {
+		// The execution was interupted
+		return err
+	}
 	return firstErr
 }
 

js/common/interrupt_error.go

@@ -0,0 +1,44 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2021 Load Impact
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+package common
+
+import "errors"
+
+// InterruptError is an error that halts engine execution
+type InterruptError struct {
+	Reason string
+}
+
+func (i *InterruptError) Error() string {
+	return i.Reason
+}
+
+// AbortTest is a reason emitted when a test script calls abortTest() without arguments
+const AbortTest = "abortTest() was called in a script"
+
+// IsInterruptError returns true if err is *InterruptError.
+func IsInterruptError(err error) bool {
+	if err == nil {
+		return false
+	}
+	var intErr *InterruptError
+	return errors.As(err, &intErr)
+}

js/modules/k6/k6.go

@@ -118,6 +118,17 @@ func (*K6) Group(ctx context.Context, name string, fn goja.Callable) (goja.Value
 	return ret, err
 }
 
+// AbortTest exposes abortTest function in the k6 module. When called it will
+// interrupt the active goja runtime passed with ctx.
+func (*K6) AbortTest(ctx context.Context, msg goja.Value) {
+	rt := common.GetRuntime(ctx)
+	reason := common.AbortTest
+	if !goja.IsUndefined(msg) {
+		reason = msg.String()
+	}
+	rt.Interrupt(&common.InterruptError{Reason: reason})
+}
+
 // Check will emit check metrics for the provided checks.
 //nolint:cyclop
 func (*K6) Check(ctx context.Context, arg0, checks goja.Value, extras ...goja.Value) (bool, error) {

js/modules/k6/k6_test.go

@@ -406,3 +406,35 @@ func TestCheckTags(t *testing.T) {
 		}, sample.Tags.CloneTags())
 	}
 }
+
+func TestAbortTest(t *testing.T) { //nolint: tparallel
+	t.Parallel()
+
+	rt := goja.New()
+	baseCtx := common.WithRuntime(context.Background(), rt)
+
+	ctx := new(context.Context)
+	*ctx = baseCtx
+	err := rt.Set("k6", common.Bind(rt, New(), ctx))
+	require.Nil(t, err)
+	prove := func(t *testing.T, script, reason string) {
+		_, err := rt.RunString(script)
+		require.NotNil(t, err)
+		var x *goja.InterruptedError
+		assert.ErrorAs(t, err, &x)
+		v, ok := x.Value().(*common.InterruptError)
+		require.True(t, ok)
+		require.Equal(t, v.Reason, reason)
+	}
+	t.Run("Without state", func(t *testing.T) { //nolint: paralleltest
+		prove(t, "k6.abortTest()", common.AbortTest)
+	})
+	t.Run("With state and default reason", func(t *testing.T) { //nolint: paralleltest
+		*ctx = lib.WithState(baseCtx, &lib.State{})
+		prove(t, "k6.abortTest()", common.AbortTest)
+	})
+	t.Run("With state and custom reason", func(t *testing.T) { //nolint: paralleltest
+		*ctx = lib.WithState(baseCtx, &lib.State{})
+		prove(t, `k6.abortTest("mayday")`, "mayday")
+	})
+}

js/runner.go

@@ -691,6 +691,15 @@ func (u *ActiveVU) RunOnce() error {
 
 	// Call the exported function.
 	_, isFullIteration, totalTime, err := u.runFn(u.RunContext, true, fn, u.setupData)
+	if err != nil {
+		var x *goja.InterruptedError
+		if errors.As(err, &x) {
+			if v, ok := x.Value().(*common.InterruptError); ok {
+				v.Reason = x.Error()
+				err = v
+			}
+		}
+	}
 
 	// If MinIterationDuration is specified and the iteration wasn't canceled
 	// and was less than it, sleep for the remainder

js/runner_test.go

@@ -1458,6 +1458,13 @@ func TestInitContextForbidden(t *testing.T) {
 			 exports.default = function() { console.log("p"); }`,
 			k6.ErrCheckInInitContext.Error(),
 		},
+		{
+			"abortTest",
+			`var abortTest = require("k6").abortTest;
+			 abortTest();
+			 exports.default = function() { console.log("p"); }`,
+			common.AbortTest,
+		},
 		{
 			"group",
 			`var group = require("k6").group;

lib/executor/helpers.go

@@ -30,6 +30,7 @@ import (
 	"github.com/sirupsen/logrus"
 
 	"go.k6.io/k6/errext"
+	"go.k6.io/k6/js/common"
 	"go.k6.io/k6/lib"
 	"go.k6.io/k6/lib/types"
 	"go.k6.io/k6/ui/pb"
@@ -76,6 +77,62 @@ 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 abortTest()
+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 common.IsInterruptError(err) {
+			cancelExecutorContext(ctx, err)
+			return true
+		}
+	}
+	return false
+}
+
 // getIterationRunner is a helper function that returns an iteration executor
 // closure. It takes care of updating the execution state statistics and
 // warning messages. And returns whether a full iteration was finished or not
@@ -98,6 +155,11 @@ func getIterationRunner(
 			return false
 		default:
 			if err != nil {
+				if handleInterrupt(ctx, err) {
+					executionState.AddInterruptedIterations(1)
+					return false
+				}
+
 				var exception errext.Exception
 				if errors.As(err, &exception) {
 					// TODO don't count this as a full iteration?