Add a method to stop a running Engine via the HTTP API

Closes grafana#1352.
hynd · Mar 11, 2020 · 31afae9 · 31afae9
1 parent d1f14a3
commit 31afae9
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 8 deletions.
diff --git a/api/v1/status.go b/api/v1/status.go
@@ -26,9 +26,10 @@ import (
 )
 
 type Status struct {
-	Paused null.Bool `json:"paused" yaml:"paused"`
-	VUs    null.Int  `json:"vus" yaml:"vus"`
-	VUsMax null.Int  `json:"vus-max" yaml:"vus-max"`
+	Paused  null.Bool `json:"paused" yaml:"paused"`
+	VUs     null.Int  `json:"vus" yaml:"vus"`
+	VUsMax  null.Int  `json:"vus-max" yaml:"vus-max"`
+	Stopped bool      `json:"stopped" yaml:"stopped"`
 
 	// Readonly.
 	Running bool `json:"running" yaml:"running"`
@@ -40,6 +41,7 @@ func NewStatus(engine *core.Engine) Status {
 	return Status{
 		Running: executionState.HasStarted(),
 		Paused:  null.BoolFrom(executionState.IsPaused()),
+		Stopped: engine.IsStopped(),
 		VUs:     null.IntFrom(executionState.GetCurrentlyActiveVUsCount()),
 		VUsMax:  null.IntFrom(executionState.GetInitializedVUsCount()),
 		Tainted: engine.IsTainted(),

diff --git a/api/v1/status_routes.go b/api/v1/status_routes.go
@@ -71,6 +71,10 @@ func HandlePatchStatus(rw http.ResponseWriter, r *http.Request, p httprouter.Par
 		return
 	}
 
+	if status.Stopped {
+		engine.Stop()
+	}
+
 	if status.Paused.Valid {
 		if err = engine.ExecutionScheduler.SetPaused(status.Paused.Bool); err != nil {
 			apiError(rw, "Pause error", err.Error(), http.StatusInternalServerError)

diff --git a/api/v1/status_routes_test.go b/api/v1/status_routes_test.go
@@ -62,6 +62,7 @@ func TestGetStatus(t *testing.T) {
 		assert.True(t, status.Paused.Valid)
 		assert.True(t, status.VUs.Valid)
 		assert.True(t, status.VUsMax.Valid)
+		assert.False(t, status.Stopped)
 		assert.False(t, status.Tainted)
 	})
 }

diff --git a/core/engine.go b/core/engine.go
@@ -59,7 +59,8 @@ type Engine struct {
 	NoSummary     bool
 	SummaryExport bool
 
-	logger *logrus.Logger
+	logger   *logrus.Logger
+	stopChan chan struct{}
 
 	Metrics     map[string]*stats.Metric
 	MetricsLock sync.Mutex
@@ -84,10 +85,11 @@ func NewEngine(ex lib.ExecutionScheduler, o lib.Options, logger *logrus.Logger)
 		ExecutionScheduler: ex,
 		executionState:     ex.GetState(),
 
-		Options: o,
-		Metrics: make(map[string]*stats.Metric),
-		Samples: make(chan stats.SampleContainer, o.MetricSamplesBufferSize.Int64),
-		logger:  logger,
+		Options:  o,
+		Metrics:  make(map[string]*stats.Metric),
+		Samples:  make(chan stats.SampleContainer, o.MetricSamplesBufferSize.Int64),
+		stopChan: make(chan struct{}),
+		logger:   logger,
 	}
 
 	e.thresholds = o.Thresholds
@@ -218,6 +220,10 @@ func (e *Engine) Run(ctx context.Context) error {
 			e.logger.Debug("run: context expired; exiting...")
 			e.setRunStatus(lib.RunStatusAbortedUser)
 			return nil
+		case <-e.stopChan:
+			e.logger.Debug("run: stopped by user; exiting...")
+			e.setRunStatus(lib.RunStatusAbortedUser)
+			return nil
 		}
 	}
 }
@@ -226,6 +232,21 @@ func (e *Engine) IsTainted() bool {
 	return e.thresholdsTainted
 }
 
+// Stop closes a signal channel, forcing a running Engine to return
+func (e *Engine) Stop() {
+	close(e.stopChan)
+}
+
+// IsStopped returns a bool indicating whether the Engine has been stopped
+func (e *Engine) IsStopped() bool {
+	select {
+	case <-e.stopChan:
+		return true
+	default:
+		return false
+	}
+}
+
 func (e *Engine) runMetricsEmission(ctx context.Context) {
 	ticker := time.NewTicker(MetricsRate)
 	for {