make the k6 runner timeout configurable

internal/k6runner/k6runner.go

@@ -270,7 +270,8 @@ func (r requestError) Error() string {
 }
 
 type RunRequest struct {
-	Script []byte `json:"script"`
+	Script  []byte `json:"script"`
+	Timeout int64  `json:"timeout"`
 }
 
 type RunResponse struct {
@@ -287,7 +288,8 @@ func (r HttpRunner) WithLogger(logger *zerolog.Logger) Runner {
 
 func (r HttpRunner) Run(ctx context.Context, script []byte) (*RunResponse, error) {
 	req, err := json.Marshal(&RunRequest{
-		Script: script,
+		Script:  script,
+		Timeout: getTimeout(ctx).Milliseconds(),
 	})
 	if err != nil {
 		return nil, fmt.Errorf("running script: %w", err)
@@ -381,9 +383,7 @@ func (r LocalRunner) Run(ctx context.Context, script []byte) (*RunResponse, erro
 		return nil, fmt.Errorf("cannot find k6 executable: %w", err)
 	}
 
-	timeout := 10 * time.Second // TODO(mem): make this configurable
-	ctx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
+	timeout := getTimeout(ctx)
 
 	// #nosec G204 -- the variables are not user-controlled
 	cmd := exec.CommandContext(
@@ -463,3 +463,12 @@ func mktemp(fs afero.Fs, dir, pattern string) (string, error) {
 	}
 	return f.Name(), nil
 }
+
+func getTimeout(ctx context.Context) time.Duration {
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		return 10 * time.Second
+	}
+
+	return time.Until(deadline)
+}

internal/k6runner/k6runner_test.go

@@ -11,6 +11,7 @@ import (
 	"sort"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/grafana/synthetic-monitoring-agent/internal/prober/logger"
 	"github.com/grafana/synthetic-monitoring-agent/internal/testhelper"
@@ -72,6 +73,7 @@ func TestScriptRun(t *testing.T) {
 
 func TestHttpRunnerRun(t *testing.T) {
 	scriptSrc := testhelper.MustReadFile(t, "testdata/test.js")
+	timeout := 1 * time.Second
 
 	mux := http.NewServeMux()
 	mux.HandleFunc("/run", func(w http.ResponseWriter, r *http.Request) {
@@ -82,6 +84,15 @@ func TestHttpRunnerRun(t *testing.T) {
 		err := json.NewDecoder(r.Body).Decode(&req)
 		require.NoError(t, err)
 		require.Equal(t, scriptSrc, req.Script)
+		// The timeout in the request is not going to be exactly the
+		// original timeout because computers need some time to process
+		// data, and the timeout is set based on the remaining time
+		// until the deadline and the clock starts ticking as soon as
+		// the context is created. Check that the actual timeout is not
+		// greater than the expected value and that it's within 1% of
+		// the expected value.
+		require.LessOrEqual(t, req.Timeout, timeout.Milliseconds())
+		require.InEpsilon(t, timeout.Milliseconds(), req.Timeout, 0.01)
 
 		w.Header().Set("Content-Type", "application/json")
 		w.WriteHeader(http.StatusOK)
@@ -102,7 +113,14 @@ func TestHttpRunnerRun(t *testing.T) {
 	runner := New(srv.URL + "/run")
 	require.IsType(t, &HttpRunner{}, runner)
 
-	ctx, cancel := testhelper.Context(context.Background(), t)
+	ctx := context.Background()
+	ctx, cancel := testhelper.Context(ctx, t)
+	t.Cleanup(cancel)
+
+	// By adding a timeout to the context passed to Run, the expectation is
+	// that the runner extracts the timeout from it and sets the
+	// corresponding field accordingly.
+	ctx, cancel = context.WithTimeout(ctx, timeout)
 	t.Cleanup(cancel)
 
 	_, err := runner.Run(ctx, scriptSrc)

internal/prober/k6/k6.go

@@ -36,7 +36,8 @@ func NewProber(ctx context.Context, check sm.Check, logger zerolog.Logger, runne
 	}
 
 	p.config = settingsToModule(check.Settings.K6)
-	p.config.Timeout = time.Duration(check.Timeout) * time.Millisecond
+	timeout := time.Duration(check.Timeout) * time.Millisecond
+	p.config.Timeout = timeout
 
 	script, err := k6runner.NewScript(check.Settings.K6.Script, runner)
 	if err != nil {

internal/prober/multihttp/multihttp.go

@@ -39,7 +39,8 @@ func NewProber(ctx context.Context, check sm.Check, logger zerolog.Logger, runne
 	}
 
 	p.config = settingsToModule(check.Settings.Multihttp)
-	p.config.Timeout = time.Duration(check.Timeout) * time.Millisecond
+	timeout := time.Duration(check.Timeout) * time.Millisecond
+	p.config.Timeout = timeout
 
 	script, err := settingsToScript(check.Settings.Multihttp)
 	if err != nil {

internal/prober/multihttp/script_test.go

@@ -796,8 +796,9 @@ func TestSettingsToScript(t *testing.T) {
 	require.NotEmpty(t, actual)
 
 	check := sm.Check{
-		Target: settings.Entries[0].Request.Url,
-		Job:    "test",
+		Target:  settings.Entries[0].Request.Url,
+		Job:     "test",
+		Timeout: 10000,
 		Settings: sm.CheckSettings{
 			Multihttp: settings,
 		},