fix: Clean up process context if available (#161)

* fix: Remove confusing log statement

This is only relevant for runs where the client is waiting for the
results. Once the request is complete (prematurely or after the run has
completed successfully) the log line is printed, which might confuse
admins.

* fix: Clean up process context if available

* chore: Temporarily increase cyclomatic complexity

.golangci.yaml

@@ -25,4 +25,6 @@ linters:
     - revive
 linters-settings:
   gocyclo:
-    min-complexity: 15
+    # Original value: 15. Temporarily increased before a refactoring of
+    # launcher.ServeHTTP.
+    min-complexity: 18

pkg/handlers/launch.go

@@ -263,6 +263,9 @@ func (h *launchHandler) waitForProcess(cmd k6.TestRun) {
 	h.trackExecutionDuration(cmd)
 	log.WithField("pid", pid).Debugf("testrun exited")
 
+	// Also clean up the context attached to this process if present:
+	cmd.CleanupContext()
+
 	h.availableTestRuns <- struct{}{}
 }
 
@@ -407,6 +410,9 @@ func (h *launchHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 	cmd, err := h.client.Start(ctx, payload.Metadata.Script, payload.Metadata.UploadToCloud, envVars, &buf)
 	if err != nil {
 		fail(fmt.Sprintf("error while launching the test: %v", err))
+		if !payload.Metadata.WaitForResults {
+			cancelCtx()
+		}
 		return
 	}
 
@@ -418,6 +424,9 @@ func (h *launchHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 		logIfError(err)
 		logIfError(h.slackClient.AddFileToThreads(slackMessages, "k6-results.txt", buf.String()))
 		fail(fmt.Sprintf("error while waiting for test to start: %v", waitErr))
+		if !payload.Metadata.WaitForResults {
+			cancelCtx()
+		}
 		h.registerProcessCleanup(cmd)
 		return
 	}
@@ -426,6 +435,9 @@ func (h *launchHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 		url, err := getCloudURL(buf.String())
 		if err != nil {
 			fail(err.Error())
+			if !payload.Metadata.WaitForResults {
+				cancelCtx()
+			}
 			h.registerProcessCleanup(cmd)
 			return
 		}
@@ -439,6 +451,10 @@ func (h *launchHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 
 	if !payload.Metadata.WaitForResults {
 		cmdLog.Infof("the load test for %s.%s was launched successfully!", payload.Name, payload.Namespace)
+		// We also need to register the cancelCtx func for asynchronous
+		// cleanup. In the synchronous cases we can cancel that context right
+		// away.
+		cmd.SetCancelFunc(cancelCtx)
 		h.registerProcessCleanup(cmd)
 		return
 	}
@@ -486,7 +502,6 @@ func (h *launchHandler) propagateCancel(req *http.Request, payload *launchPayloa
 		<-h.ctx.Done()
 		cancelCtx()
 	}
-	log.Info("canceling process")
 }
 
 func (h *launchHandler) waitForOutputPath(cmdLog *log.Entry, buf *bytes.Buffer) error {

pkg/handlers/launch_test.go

@@ -671,6 +671,8 @@ func TestProcessHandler(t *testing.T) {
 			tr.EXPECT().Wait().Return(nil).Times(1)
 			tr.EXPECT().Exited().Return(true).AnyTimes()
 			tr.EXPECT().ExitCode().Return(0).AnyTimes()
+			tr.EXPECT().SetCancelFunc(gomock.Any()).Return().AnyTimes()
+			tr.EXPECT().CleanupContext().Return().AnyTimes()
 			tr.EXPECT().ExecutionDuration().Return(time.Minute).AnyTimes()
 			handler.registerProcessCleanup(tr)
 		}
@@ -773,6 +775,8 @@ func setupHandlerWithKubernetesObjects(t *testing.T, maxConcurrentTests int, exp
 	// value here:
 	testRun.EXPECT().ExecutionDuration().Return(time.Minute).AnyTimes()
 	testRun.EXPECT().ExitCode().Return(0).AnyTimes()
+	testRun.EXPECT().SetCancelFunc(gomock.Any()).Return().AnyTimes()
+	testRun.EXPECT().CleanupContext().Return().AnyTimes()
 
 	ctx, cancel := context.WithCancel(context.Background())
 	handler, err := NewLaunchHandler(ctx, k6Client, kubeClient, slackClient, maxConcurrentTests)

pkg/k6/cmd.go

@@ -23,8 +23,9 @@ func NewLocalRunnerClient(token string) (Client, error) {
 
 type DefaultTestRun struct {
 	*exec.Cmd
-	startedAt time.Time
-	exitedAt  time.Time
+	startedAt     time.Time
+	exitedAt      time.Time
+	cancelContext context.CancelFunc
 }
 
 func (tr *DefaultTestRun) Start() error {
@@ -46,6 +47,12 @@ func (tr *DefaultTestRun) ExitCode() int {
 	return -1
 }
 
+func (tr *DefaultTestRun) CleanupContext() {
+	if tr.cancelContext != nil {
+		tr.cancelContext()
+	}
+}
+
 func (tr *DefaultTestRun) ExecutionDuration() time.Duration {
 	if tr.startedAt.IsZero() || tr.exitedAt.IsZero() {
 		return time.Duration(0)
@@ -74,6 +81,10 @@ func (tr *DefaultTestRun) Exited() bool {
 	return false
 }
 
+func (tr *DefaultTestRun) SetCancelFunc(fn context.CancelFunc) {
+	tr.cancelContext = fn
+}
+
 func (c *LocalRunnerClient) Start(ctx context.Context, scriptContent string, upload bool, envVars map[string]string, outputWriter io.Writer) (TestRun, error) {
 	tempFile, err := os.CreateTemp("", "k6-script")
 	if err != nil {

pkg/k6/interface.go

@@ -19,4 +19,6 @@ type TestRun interface {
 	Exited() bool
 	ExitCode() int
 	ExecutionDuration() time.Duration
+	CleanupContext()
+	SetCancelFunc(context.CancelFunc)
 }