WIP Add Cleanup() hook for collectors

cmd/run.go

@@ -278,6 +278,11 @@ a commandline interface for interacting with it.`,
 			runCancel() // stop the test run, metric processing is cancelled below
 			lingerCancel()
 
+			for _, c := range engine.Collectors {
+				c.SetRunStatus(lib.RunStatusAbortedUser)
+				c.Cleanup()
+			}
+
 			// If we get a second signal, we immediately exit, so something like
 			// https://github.com/loadimpact/k6/issues/971 never happens again
 			sig = <-sigC

lib/collector.go

@@ -53,6 +53,10 @@ type Collector interface {
 	// You should do any lengthy setup here rather than in New.
 	Init() error
 
+	// Cleanup is run on premature engine interruption (e.g. Ctrl+C).
+	// It can be used to perform optional maintenance.
+	Cleanup()
+
 	// Run is called in a goroutine and starts the collector. Should commit samples to the backend
 	// at regular intervals and when the context is terminated.
 	Run(ctx context.Context)

stats/cloud/collector.go

@@ -200,6 +200,11 @@ func (c *Collector) Link() string {
 	return URLForResults(c.referenceID, c.config)
 }
 
+// Cleanup updates the test status in the backend
+func (c *Collector) Cleanup() {
+	c.testFinished()
+}
+
 // Run is called in a goroutine and starts the collector. Should commit samples to the backend
 // at regular intervals and when the context is terminated.
 func (c *Collector) Run(ctx context.Context) {

stats/csv/collector.go

@@ -31,10 +31,11 @@ import (
 	"sync"
 	"time"
 
-	"github.com/loadimpact/k6/lib"
-	"github.com/loadimpact/k6/stats"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/afero"
+
+	"github.com/loadimpact/k6/lib"
+	"github.com/loadimpact/k6/stats"
 )
 
 // Collector saving output to csv implements the lib.Collector interface
@@ -118,6 +119,9 @@ func (c *Collector) Init() error {
 	return nil
 }
 
+// Cleanup does nothing, it's only included to satisfy the lib.Collector interface
+func (c *Collector) Cleanup() {}
+
 // SetRunStatus does nothing
 func (c *Collector) SetRunStatus(status lib.RunStatus) {}
 

stats/dummy/collector.go

@@ -43,6 +43,9 @@ var _ lib.Collector = &Collector{}
 // Init does nothing, it's only included to satisfy the lib.Collector interface
 func (c *Collector) Init() error { return nil }
 
+// Cleanup does nothing, it's only included to satisfy the lib.Collector interface
+func (c *Collector) Cleanup() {}
+
 // MakeConfig does nothing, it's only included to satisfy the lib.Collector interface
 func (c *Collector) MakeConfig() interface{} { return nil }
 

stats/influxdb/collector.go

@@ -27,9 +27,10 @@ import (
 	"time"
 
 	client "github.com/influxdata/influxdb1-client/v2"
+	"github.com/sirupsen/logrus"
+
 	"github.com/loadimpact/k6/lib"
 	"github.com/loadimpact/k6/stats"
-	"github.com/sirupsen/logrus"
 )
 
 // Verify that Collector implements lib.Collector
@@ -74,6 +75,9 @@ func (c *Collector) Init() error {
 	return nil
 }
 
+// Cleanup does nothing, it's only included to satisfy the lib.Collector interface
+func (c *Collector) Cleanup() {}
+
 func (c *Collector) Run(ctx context.Context) {
 	logrus.Debug("InfluxDB: Running!")
 	ticker := time.NewTicker(time.Duration(c.Config.PushInterval.Duration))

stats/json/collector.go

@@ -91,9 +91,11 @@ func New(fs afero.Fs, fname string) (*Collector, error) {
 	return c, nil
 }
 
-func (c *Collector) Init() error {
-	return nil
-}
+// Init does nothing, it's only included to satisfy the lib.Collector interface
+func (c *Collector) Init() error { return nil }
+
+// Cleanup does nothing, it's only included to satisfy the lib.Collector interface
+func (c *Collector) Cleanup() {}
 
 func (c *Collector) SetRunStatus(status lib.RunStatus) {}
 

stats/kafka/collector.go

@@ -60,6 +60,9 @@ func New(conf Config) (*Collector, error) {
 // Init does nothing, it's only included to satisfy the lib.Collector interface
 func (c *Collector) Init() error { return nil }
 
+// Cleanup does nothing, it's only included to satisfy the lib.Collector interface
+func (c *Collector) Cleanup() {}
+
 // Run just blocks until the context is done
 func (c *Collector) Run(ctx context.Context) {
 	logrus.Debug("Kafka: Running!")

stats/statsd/common/collector.go

@@ -82,6 +82,9 @@ func (c *Collector) Link() string {
 	return c.Config.Addr.String
 }
 
+// Cleanup does nothing, it's only included to satisfy the lib.Collector interface
+func (c *Collector) Cleanup() {}
+
 // Run the collector
 func (c *Collector) Run(ctx context.Context) {
 	c.logger.Debugf("%s: Running!", c.Type)