From 60e3dd38f9166229169436fa2ad14f82ea219936 Mon Sep 17 00:00:00 2001
From: Mihail Stoykov <M.Stoikov@gmail.com>
Date: Wed, 14 Apr 2021 19:19:00 +0300
Subject: [PATCH] Add cloud output specific option to stop engine on error

---
 cloudapi/config.go          |  4 ++++
 cloudapi/config_test.go     |  1 +
 output/cloud/output.go      |  4 +++-
 output/cloud/output_test.go | 23 ++++++++++++++++++++---
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/cloudapi/config.go b/cloudapi/config.go
index e5590babb393..cec1d3df4228 100644
--- a/cloudapi/config.go
+++ b/cloudapi/config.go
@@ -44,6 +44,7 @@ type Config struct {
 	PushRefID   null.String `json:"pushRefID" envconfig:"K6_CLOUD_PUSH_REF_ID"`
 	WebAppURL   null.String `json:"webAppURL" envconfig:"K6_CLOUD_WEB_APP_URL"`
 	NoCompress  null.Bool   `json:"noCompress" envconfig:"K6_CLOUD_NO_COMPRESS"`
+	StopOnError null.Bool   `json:"stopOnError" envconfig:"K6_CLOUD_STOP_ON_ERROR"`
 
 	MaxMetricSamplesPerPackage null.Int `json:"maxMetricSamplesPerPackage" envconfig:"K6_CLOUD_MAX_METRIC_SAMPLES_PER_PACKAGE"`
 
@@ -209,6 +210,9 @@ func (c Config) Apply(cfg Config) Config {
 	if cfg.NoCompress.Valid {
 		c.NoCompress = cfg.NoCompress
 	}
+	if cfg.StopOnError.Valid {
+		c.StopOnError = cfg.StopOnError
+	}
 	if cfg.MaxMetricSamplesPerPackage.Valid {
 		c.MaxMetricSamplesPerPackage = cfg.MaxMetricSamplesPerPackage
 	}
diff --git a/cloudapi/config_test.go b/cloudapi/config_test.go
index 577d27ac992a..33bf2f9b377f 100644
--- a/cloudapi/config_test.go
+++ b/cloudapi/config_test.go
@@ -48,6 +48,7 @@ func TestConfigApply(t *testing.T) {
 		PushRefID:                       null.NewString("PushRefID", true),
 		WebAppURL:                       null.NewString("foo", true),
 		NoCompress:                      null.NewBool(true, true),
+		StopOnError:                     null.NewBool(true, true),
 		MaxMetricSamplesPerPackage:      null.NewInt(2, true),
 		MetricPushInterval:              types.NewNullDuration(1*time.Second, true),
 		MetricPushConcurrency:           null.NewInt(3, true),
diff --git a/output/cloud/output.go b/output/cloud/output.go
index 4e51cd6f14b9..84eae47591c0 100644
--- a/output/cloud/output.go
+++ b/output/cloud/output.go
@@ -667,7 +667,9 @@ func (out *Output) pushMetrics() {
 		if err != nil {
 			if out.shouldStopSendingMetrics(err) {
 				out.logger.WithError(err).Warn("Stopped sending metrics to cloud due to an error")
-				out.engineStopFunc(err)
+				if out.config.StopOnError.Bool {
+					out.engineStopFunc(err)
+				}
 				close(out.stopSendingMetrics)
 				break
 			}
diff --git a/output/cloud/output_test.go b/output/cloud/output_test.go
index 719587abd5fd..2958fc106641 100644
--- a/output/cloud/output_test.go
+++ b/output/cloud/output_test.go
@@ -388,6 +388,18 @@ func TestCloudOutputMaxPerPacket(t *testing.T) {
 
 func TestCloudOutputStopSendingMetric(t *testing.T) {
 	t.Parallel()
+	t.Run("stop engine on error", func(t *testing.T) {
+		t.Parallel()
+		testCloudOutputStopSendingMetric(t, true)
+	})
+
+	t.Run("don't stop engine on error", func(t *testing.T) {
+		t.Parallel()
+		testCloudOutputStopSendingMetric(t, false)
+	})
+}
+
+func testCloudOutputStopSendingMetric(t *testing.T, stopOnError bool) {
 	tb := httpmultibin.NewHTTPMultiBin(t)
 	tb.Mux.HandleFunc("/v1/tests", http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
 		body, err := ioutil.ReadAll(req.Body)
@@ -417,8 +429,9 @@ func TestCloudOutputStopSendingMetric(t *testing.T) {
 		JSONConfig: json.RawMessage(fmt.Sprintf(`{
 			"host": "%s", "noCompress": true,
 			"maxMetricSamplesPerPackage": 50,
-			"name": "something-that-should-be-overwritten"
-		}`, tb.ServerHTTP.URL)),
+			"name": "something-that-should-be-overwritten",
+			"stopOnError": %t
+		}`, tb.ServerHTTP.URL, stopOnError)),
 		ScriptOptions: lib.Options{
 			Duration:   types.NullDurationFrom(1 * time.Second),
 			SystemTags: &stats.DefaultSystemTagSet,
@@ -428,6 +441,10 @@ func TestCloudOutputStopSendingMetric(t *testing.T) {
 		},
 		ScriptPath: &url.URL{Path: "/script.js"},
 	})
+	var expectedEngineStopFuncCalled int64
+	if stopOnError {
+		expectedEngineStopFuncCalled = 1
+	}
 	var engineStopFuncCalled int64
 	out.engineStopFunc = func(error) {
 		atomic.AddInt64(&engineStopFuncCalled, 1)
@@ -500,7 +517,7 @@ func TestCloudOutputStopSendingMetric(t *testing.T) {
 		t.Fatal("sending metrics wasn't stopped")
 	}
 	require.Equal(t, max, count)
-	require.EqualValues(t, engineStopFuncCalled, 1)
+	require.Equal(t, expectedEngineStopFuncCalled, engineStopFuncCalled)
 
 	nBufferSamples := len(out.bufferSamples)
 	nBufferHTTPTrails := len(out.bufferHTTPTrails)