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

Closes grafana#1352.

api/v1/status.go

@@ -32,6 +32,7 @@ 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"`
+	Stopped bool      `json:"stopped" yaml:"stopped"`
 	Running bool      `json:"running" yaml:"running"`
 	Tainted bool      `json:"tainted" yaml:"tainted"`
 }
@@ -42,6 +43,7 @@ func NewStatus(engine *core.Engine) Status {
 		Status:  executionState.GetCurrentExecutionStatus(),
 		Running: executionState.HasStarted() && !executionState.HasEnded(),
 		Paused:  null.BoolFrom(executionState.IsPaused()),
+		Stopped: engine.IsStopped(),
 		VUs:     null.IntFrom(executionState.GetCurrentlyActiveVUsCount()),
 		VUsMax:  null.IntFrom(executionState.GetInitializedVUsCount()),
 		Tainted: engine.IsTainted(),

api/v1/status_routes.go

@@ -71,34 +71,39 @@ func HandlePatchStatus(rw http.ResponseWriter, r *http.Request, p httprouter.Par
 		return
 	}
 
-	if status.Paused.Valid {
-		if err = engine.ExecutionScheduler.SetPaused(status.Paused.Bool); err != nil {
-			apiError(rw, "Pause error", err.Error(), http.StatusInternalServerError)
-			return
+	if status.Stopped {
+		engine.Stop()
+	} else {
+		if status.Paused.Valid {
+			if err = engine.ExecutionScheduler.SetPaused(status.Paused.Bool); err != nil {
+				apiError(rw, "Pause error", err.Error(), http.StatusInternalServerError)
+				return
+			}
 		}
-	}
 
-	if status.VUsMax.Valid || status.VUs.Valid {
-		//TODO: add ability to specify the actual executor id? Though this should
-		//likely be in the v2 REST API, where we could implement it in a way that
-		//may allow us to eventually support other executor types.
-		executor, updateErr := getFirstExternallyControlledExecutor(engine.ExecutionScheduler)
-		if updateErr != nil {
-			apiError(rw, "Execution config error", updateErr.Error(), http.StatusInternalServerError)
-			return
-		}
-		newConfig := executor.GetCurrentConfig().ExternallyControlledConfigParams
-		if status.VUsMax.Valid {
-			newConfig.MaxVUs = status.VUsMax
-		}
-		if status.VUs.Valid {
-			newConfig.VUs = status.VUs
-		}
-		if updateErr := executor.UpdateConfig(r.Context(), newConfig); err != nil {
-			apiError(rw, "Config update error", updateErr.Error(), http.StatusInternalServerError)
-			return
+		if status.VUsMax.Valid || status.VUs.Valid {
+			//TODO: add ability to specify the actual executor id? Though this should
+			//likely be in the v2 REST API, where we could implement it in a way that
+			//may allow us to eventually support other executor types.
+			executor, updateErr := getFirstExternallyControlledExecutor(engine.ExecutionScheduler)
+			if updateErr != nil {
+				apiError(rw, "Execution config error", updateErr.Error(), http.StatusInternalServerError)
+				return
+			}
+			newConfig := executor.GetCurrentConfig().ExternallyControlledConfigParams
+			if status.VUsMax.Valid {
+				newConfig.MaxVUs = status.VUsMax
+			}
+			if status.VUs.Valid {
+				newConfig.VUs = status.VUs
+			}
+			if updateErr := executor.UpdateConfig(r.Context(), newConfig); err != nil {
+				apiError(rw, "Config update error", updateErr.Error(), http.StatusInternalServerError)
+				return
+			}
 		}
 	}
+
 	data, err := jsonapi.Marshal(NewStatus(engine))
 	if err != nil {
 		apiError(rw, "Encoding error", err.Error(), http.StatusInternalServerError)

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)
 	})
 }

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 {

core/engine_test.go

@@ -156,6 +156,20 @@ func TestEngineAtTime(t *testing.T) {
 	assert.NoError(t, e.Run(ctx))
 }
 
+func TestEngineStopped(t *testing.T) {
+	e := newTestEngine(t, nil, nil, lib.Options{
+		VUs:      null.IntFrom(1),
+		Duration: types.NullDurationFrom(20 * time.Second),
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+	assert.NoError(t, e.Run(ctx))
+	assert.Equal(t, false, e.IsStopped(), "engine should be running")
+	e.Stop()
+	assert.Equal(t, true, e.IsStopped(), "engine should be stopped")
+}
+
 func TestEngineCollector(t *testing.T) {
 	testMetric := stats.New("test_metric", stats.Trend)