From b970b7b0fc834e9d671846b05e6375b5b0be56d5 Mon Sep 17 00:00:00 2001
From: Theerapat Chawannakul <theerapat.cha@gmail.com>
Date: Mon, 2 Nov 2020 18:03:17 +0700
Subject: [PATCH] Immediately shutdown `k6 cloud` on second Ctrl+C signal
 (#1647)

---
 cmd/cloud.go | 36 ++++++++++++++++++++++--------------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/cmd/cloud.go b/cmd/cloud.go
index b4d136efd89..9d0a709e88c 100644
--- a/cmd/cloud.go
+++ b/cmd/cloud.go
@@ -199,6 +199,9 @@ This will execute the test on the k6 cloud service. Use "k6 login cloud" to auth
 			name = filepath.Base(filename)
 		}
 
+		globalCtx, globalCancel := context.WithCancel(context.Background())
+		defer globalCancel()
+
 		// Start cloud test run
 		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Validating script options"))
 		client := cloud.NewClient(logger, cloudConfig.Token.String, cloudConfig.Host.String, consts.Version)
@@ -226,7 +229,7 @@ This will execute the test on the k6 cloud service. Use "k6 login cloud" to auth
 			pb.WithConstProgress(0, "Initializing the cloud test"),
 		)
 
-		progressCtx, progressCancel := context.WithCancel(context.Background())
+		progressCtx, progressCancel := context.WithCancel(globalCtx)
 		progressBarWG := &sync.WaitGroup{}
 		progressBarWG.Add(1)
 		defer progressBarWG.Wait()
@@ -240,6 +243,19 @@ This will execute the test on the k6 cloud service. Use "k6 login cloud" to auth
 		sigC := make(chan os.Signal, 1)
 		signal.Notify(sigC, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 		defer signal.Stop(sigC)
+		go func() {
+			sig := <-sigC
+			logger.WithField("sig", sig).Print("Stopping k6 in response to signal...")
+			err := client.StopCloudTestRun(refID)
+			if err != nil {
+				logger.WithError(err).Error("Stop cloud test error")
+			}
+			globalCancel()
+
+			sig = <-sigC
+			logger.WithField("sig", sig).Error("Aborting k6 in response to signal")
+			os.Exit(externalAbortErrorCode)
+		}()
 
 		var (
 			startTime   time.Time
@@ -279,12 +295,11 @@ This will execute the test on the k6 cloud service. Use "k6 login cloud" to auth
 		)
 
 		ticker := time.NewTicker(time.Millisecond * 2000)
-		shouldExitLoop := false
 		if showCloudLogs {
 			go func() {
 				logger.Debug("Connecting to cloud logs server...")
 				// TODO replace with another context
-				if err := cloudConfig.StreamLogsToLogger(context.Background(), logger, refID, 0); err != nil {
+				if err := cloudConfig.StreamLogsToLogger(globalCtx, logger, refID, 0); err != nil {
 					logger.WithError(err).Error("error while tailing cloud logs")
 				}
 			}()
@@ -298,7 +313,8 @@ This will execute the test on the k6 cloud service. Use "k6 login cloud" to auth
 				if progressErr == nil {
 					if (newTestProgress.RunStatus > lib.RunStatusRunning) ||
 						(exitOnRunning && newTestProgress.RunStatus == lib.RunStatusRunning) {
-						shouldExitLoop = true
+						globalCancel()
+						break runningLoop
 					}
 					testProgressLock.Lock()
 					testProgress = newTestProgress
@@ -306,16 +322,8 @@ This will execute the test on the k6 cloud service. Use "k6 login cloud" to auth
 				} else {
 					logger.WithError(progressErr).Error("Test progress error")
 				}
-				if shouldExitLoop {
-					break runningLoop
-				}
-			case sig := <-sigC:
-				logger.WithField("sig", sig).Print("Exiting in response to signal...")
-				err := client.StopCloudTestRun(refID)
-				if err != nil {
-					logger.WithError(err).Error("Stop cloud test error")
-				}
-				shouldExitLoop = true // Exit after the next GetTestProgress call
+			case <-globalCtx.Done():
+				break runningLoop
 			}
 		}