resource limiting: complete rework to handle context cancelation and native rate-limiting on a template (#196)

This introduces a complete rework of resource usage by implementing
contexts, so the lock can be cancelled, in case the lock takes too long
to be obtained, or if the instance is shutting-down.

It also introduces a native resource on all resolution based on the
template name: it allows to reduce the concurrent execution of a given
template, or to completely stop the execution of a template if it behave
incorrectly.

Signed-off-by: Romain Beuque <[email protected]>

Author: rbeuque74

README.md

@@ -356,6 +356,7 @@ Note that the operators `IN` and `NOTIN` expect a list of acceptable values in t
 - `dependencies`: a list of step names on which this step waits before running
 - `custom_states`: a list of personnalised allowed state for this step (can be assigned to the state's step using `conditions`)
 - `retry_pattern`: (`seconds`, `minutes`, `hours`) define on what temporal order of magnitude the re-runs of this step should be spread (default = `seconds`)
+- `resources`: a list of resources that will be used during the step execution, to control and limit the concurrent execution of the step (more information in [the resources section](#resources)).
 
 <p align="center">
 <img src="./assets/img/utask_backoff.png" width="70%">
@@ -480,14 +481,16 @@ The `pre_hook` field of a step can be set to define an action that is executed b
 doSomeAuthPost:
   pre_hook:
     type: http
-    method: "GET"
-    url: "https://myAwesomeApi/otp"
+    configuration:
+      method: "GET"
+      url: "https://example.org/otp"
   action:
     type: http
-    method: "POST"
-    url: "https://myAwesomeApi/doSomePost"
-    headers:
-      X-Otp: "{{ .pre_hook.output }}"
+    configuration:
+      method: "POST"
+      url: "https://example.org/doSomePost"
+      headers:
+        X-Otp: "{{ .pre_hook.output }}"
 ```
 
 #### Functions <a name="functions"></a>
@@ -587,6 +590,39 @@ This output can be then passed to another step in json format:
 foreach: '{{.step.prefixStrings.children | toJson}}'
 ```
 
+#### Resources <a name="resources"></a>
+
+Resources are a way to restrict the concurrency factor of certain operations, to control the throughput and avoid dangerous behavior e.g. flooding the targets.
+
+High level view:
+
+- For each action to execute, a list of target `resources` is determined. (see later)
+- In the µTask configuration, numerical limits can be set to each _resource_ label. This acts as a semaphore, allowing a certain number of concurrent slots for the given _resource_ label. If no limit is set for a resource label, the previously mentionned target resources have no effect. Limits are declared in the `resource_limits` property.
+
+The target _resources_ for a step can be defined in its YAML definition, using the `resources` property.
+
+```yaml
+steps:
+  foobar:
+    description: A dummy step, that should not execute in parallel
+    resources: ["myLimitedResource"]
+    action:
+      type: echo
+      configuration:
+        output:
+          foobar: fuzz
+```
+
+Alternatively, some target resources are determined automatically by µTask Engine:
+
+- When a task is run, the resource `template:my-template-name` is used automatically.
+- When a step is run, the plugin in charge of the execution automatically generates a list of resources. This includes generic resources such as `socket`, `url:www.example.org`, `fork`...
+ allowing the µTask administrator to set-up generic limits such as `"socket": 48` or `"url:www.example.org": 1`.
+
+Each builtin plugins declares resources which can be discovered using the _README_ of the plugin (example for [_http_ plugin](./pkg/plugins/builtin/script/README.md#Resources)).
+
+Declared `resource_limits` must be positive integers. When a step is executed, if the number of concurrent executions is reached, the µTask Engine will wait for a slot to be released. If the resource is limited to the `0` value, then the step will not be executed and is set to `TO_RETRY` state, it will be run once the instance allows the execution of its resources. The default time that µTask Engine will wait for a resource to become available is `1 minute`, but it can be configured using the `resource_acquire_timeout` property.
+
 ### Task templates validation
 
 A JSON-schema file is available to validate the syntax of task templates, it's available in `hack/template-schema.json`.

api/handler/resolution.go

@@ -281,7 +281,24 @@ func RunResolution(c *gin.Context, in *runResolutionIn) error {
 
 	logrus.WithFields(logrus.Fields{"resolution_id": r.PublicID}).Debugf("Handler RunResolution: manual resolve %s", r.PublicID)
 
-	return engine.GetEngine().Resolve(in.PublicID, nil)
+	ch := make(chan struct{})
+	go func() {
+		err = engine.GetEngine().Resolve(in.PublicID, nil)
+		close(ch)
+	}()
+
+	timeout := time.NewTicker(5 * time.Second)
+	defer timeout.Stop()
+
+	// manual resolution can be blocked by a lock acquisition on the Execution pool
+	// waiting for 5 seconds to get a response, otherwise let's consider the task will
+	// start correctly when the Execution pool gets available, and prevent API thread to be blocked
+	select {
+	case <-ch:
+		return err
+	case <-timeout.C:
+		return nil
+	}
 }
 
 type extendResolutionIn struct {

api/server.go

@@ -106,7 +106,7 @@ func (s *Server) ListenAndServe() error {
 		logrus.Info("Shutting down...")
 		cancel()
 
-		if err := srv.Shutdown(ctx); err != nil {
+		if err := srv.Shutdown(context.Background()); err != nil {
 			logrus.Fatal(err)
 		}
 	}()

engine/engine.go

@@ -36,8 +36,8 @@ var (
 	eng Engine
 
 	// Used for stopping the current Engine
-	stopRunningSteps chan struct{}
-	gracePeriodEnd   chan struct{}
+	shutdownCtx    context.Context
+	gracePeriodEnd chan struct{}
 )
 
 // Engine is the heart of utask: it is the active process
@@ -97,13 +97,11 @@ func Init(ctx context.Context, wg *sync.WaitGroup, store *configstore.Store) err
 	}
 
 	// channels for handling graceful shutdown
-	stopRunningSteps = make(chan struct{})
+	shutdownCtx = ctx
 	gracePeriodEnd = make(chan struct{})
 	eng.wg = wg
 	go func() {
-		<-ctx.Done()
-		// Stop running new steps
-		close(stopRunningSteps)
+		<-shutdownCtx.Done()
 
 		// Wait for the grace period to end
 		time.Sleep(3 * time.Second)
@@ -215,14 +213,43 @@ func (e Engine) launchResolution(publicID string, async bool, sm *semaphore.Weig
 
 	res.Values.SetConfig(e.config)
 
-	// all ready, run remaining steps
-
-	utask.AcquireExecutionSlot()
-
+	// check if all resources are available before starting the resolution
+	// first, check if we have a custom semaphore, for example, a semaphore that limits the concurrent execution of tasks recovery from a crashed instance.
+	// This semaphore needs to go first, because it will always be smaller than the global execution pool.
 	if sm != nil {
-		sm.Acquire(context.Background(), 1)
+		if err := sm.Acquire(shutdownCtx, 1); err != nil {
+			debugLogger.Debugf("Engine: launchResolution() %s acquire resource: instance is shutting down", res.PublicID)
+			return nil, errors.New("instance is shutting down")
+		}
+	}
+	// second, check if we have a resource limit on the current template
+	// template could be completely deactivated as a "dead resource". If it's the case, we need to exit because
+	// the limit won't change until next instance's reboot.
+	if acquiredErr := utask.AcquireResource(shutdownCtx, "template:"+t.TemplateName); acquiredErr != nil {
+		if shutdownCtx.Err() != nil {
+			debugLogger.Debugf("Engine: launchResolution() %s acquire resource: instance is shutting down", res.PublicID)
+			return nil, errors.New("instance is shutting down")
+		}
+		debugLogger.Debugf("Engine: launchResolution() %s acquire resource %q: failed to acquire resource: %s", res.PublicID, "template:"+t.TemplateName, err)
+		// otherwise, we either reached timeout on the lock for template, or the template is a "dead resource"
+		t.SetState(task.StateBlocked)
+		res.SetNextRetry(time.Now().Add(10 * time.Minute))
+		res.SetState(resolution.StateToAutorunDelayed)
+		if err := commit(dbp, res, t); err != nil {
+			debugLogger.Debugf("Engine: launchResolution() %s acquire resource, FAILED TO COMMIT RESOLUTION: %s", res.PublicID, err)
+		}
+		if sm != nil {
+			sm.Release(1)
+		}
+		return nil, fmt.Errorf("can't acquire lock for template %q: %s", t.TemplateName, err)
+	}
+	// finally, acquire the execution slot
+	if acquiredErr := utask.AcquireExecutionSlot(shutdownCtx); acquiredErr != nil {
+		debugLogger.Debugf("Engine: launchResolution() %s acquire resource: instance is shutting down", res.PublicID)
+		return nil, errors.New("instance is shutting down")
 	}
 
+	// all ready, run remaining steps
 	recap := make([]string, 0)
 	for name, s := range res.Steps {
 		recap = append(recap, fmt.Sprintf("step %s = %s", name, s.State))
@@ -427,7 +454,7 @@ func resolve(dbp zesty.DBProvider, res *resolution.Resolution, t *task.Task, sm
 
 	inShutdown := false
 	select {
-	case <-stopRunningSteps:
+	case <-shutdownCtx.Done():
 		inShutdown = true
 	default:
 	}
@@ -537,6 +564,7 @@ func resolve(dbp zesty.DBProvider, res *resolution.Resolution, t *task.Task, sm
 		sm.Release(1)
 	}
 
+	utask.ReleaseResource("template:" + t.TemplateName)
 	utask.ReleaseExecutionSlot()
 	if err := resumeParentTask(dbp, t, sm, debugLogger); err != nil {
 		debugLogger.WithError(err).Debugf("Engine: resolver(): failed to resume parent task: %s", err)
@@ -602,7 +630,7 @@ func runAvailableSteps(dbp zesty.DBProvider, modifiedSteps map[string]bool, res
 	expanded := 0
 
 	select {
-	case <-stopRunningSteps:
+	case <-shutdownCtx.Done():
 		return 0
 	default:
 		for name, s := range av {
@@ -660,7 +688,7 @@ func runAvailableSteps(dbp zesty.DBProvider, modifiedSteps map[string]bool, res
 
 				// run
 				stepCopy := *s
-				step.Run(&stepCopy, res.BaseConfigurations, res.Values, stepChan, wg, stopRunningSteps)
+				step.Run(&stepCopy, res.BaseConfigurations, res.Values, stepChan, wg, shutdownCtx)
 			}
 		}
 	}

engine/step/step.go

@@ -2,6 +2,7 @@ package step
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"sort"
@@ -144,12 +145,12 @@ func uniqueSortedList(s []string) []string {
 }
 
 type execution struct {
-	baseCfgRaw       json.RawMessage
-	outputs          []*executor.Output
-	config           json.RawMessage
-	runner           Runner
-	ctx              interface{}
-	stopRunningSteps <-chan struct{}
+	baseCfgRaw  json.RawMessage
+	outputs     []*executor.Output
+	config      json.RawMessage
+	runner      Runner
+	ctx         interface{}
+	shutdownCtx context.Context
 }
 
 func (e *execution) generateOutput(st *Step, v *values.Values) error {
@@ -190,10 +191,10 @@ func (e *execution) generateOutput(st *Step, v *values.Values) error {
 	return nil
 }
 
-func (st *Step) generateExecution(action executor.Executor, baseConfig map[string]json.RawMessage, values *values.Values, stopRunningSteps <-chan struct{}) (*execution, error) {
+func (st *Step) generateExecution(action executor.Executor, baseConfig map[string]json.RawMessage, values *values.Values, shutdownCtx context.Context) (*execution, error) {
 	var ret = execution{
-		config:           action.Configuration,
-		stopRunningSteps: stopRunningSteps,
+		config:      action.Configuration,
+		shutdownCtx: shutdownCtx,
 	}
 	var err error
 
@@ -294,7 +295,7 @@ func (st *Step) generateExecution(action executor.Executor, baseConfig map[strin
 func (st *Step) execute(execution *execution, callback func(interface{}, interface{}, map[string]string, error)) {
 
 	select {
-	case <-execution.stopRunningSteps:
+	case <-execution.shutdownCtx.Done():
 		st.State = StateToRetry
 		return
 	default:
@@ -303,7 +304,12 @@ func (st *Step) execute(execution *execution, callback func(interface{}, interfa
 
 	resources := append(execution.runner.Resources(execution.baseCfgRaw, execution.config), st.Resources...)
 	limits := uniqueSortedList(resources)
-	utask.AcquireResources(limits)
+	if acquiredErr := utask.AcquireResources(execution.shutdownCtx, limits); acquiredErr != nil {
+		// if resource acquisition takes too long (timeout or shutdown), let's put the step in ToRetry state
+		// to release the Execution pool, or let the instance shutdowns correctly, as the step execution didn't started yet
+		callback(`{}`, "", map[string]string{}, errors.NotProvisionedf("failed to acquire resources"))
+		return
+	}
 	defer utask.ReleaseResources(limits)
 
 	output, metadata, tags, err := execution.runner.Exec(st.Name, execution.baseCfgRaw, execution.config, execution.ctx)
@@ -312,9 +318,9 @@ func (st *Step) execute(execution *execution, callback func(interface{}, interfa
 
 // Run carries out the action defined by a Step, by providing values to its configuration
 // - a stepChan channel is provided for committing the result back
-// - a stopRunningSteps channel is provided to interrupt execution in flight
+// - a shutdownCtx context is provided to interrupt execution in flight
 // values IS NOT CONCURRENT SAFE, DO NOT SHARE WITH OTHER GOROUTINES
-func Run(st *Step, baseConfig map[string]json.RawMessage, stepValues *values.Values, stepChan chan<- *Step, wg *sync.WaitGroup, stopRunningSteps <-chan struct{}) {
+func Run(st *Step, baseConfig map[string]json.RawMessage, stepValues *values.Values, stepChan chan<- *Step, wg *sync.WaitGroup, shutdownCtx context.Context) {
 
 	// Step already ran, directly going to afterrun process
 	if st.State == StateAfterrunError {
@@ -351,7 +357,7 @@ func Run(st *Step, baseConfig map[string]json.RawMessage, stepValues *values.Val
 	preHookValues := stepValues.Clone()
 	var preHookWg sync.WaitGroup
 	if prehook != nil {
-		preHookExecution, err := st.generateExecution(*prehook, baseConfig, stepValues, stopRunningSteps)
+		preHookExecution, err := st.generateExecution(*prehook, baseConfig, stepValues, shutdownCtx)
 		if err != nil {
 			st.State = StateFatalError
 			st.Error = fmt.Sprintf("prehook: %s", err)
@@ -384,8 +390,18 @@ func Run(st *Step, baseConfig map[string]json.RawMessage, stepValues *values.Val
 		// Wait the prehook execution is done
 		preHookWg.Wait()
 
+		// after pre-hook execution, let's verify if we are not shutting down the instance
+		// in that case, we can just put the step in ToRetry instead of starting the main execution of the step
+		select {
+		case <-shutdownCtx.Done():
+			st.State = StateToRetry
+			go noopStep(st, stepChan)
+			return
+		default:
+		}
+
 		// Generate the execution
-		execution, err := st.generateExecution(st.Action, baseConfig, preHookValues, stopRunningSteps)
+		execution, err := st.generateExecution(st.Action, baseConfig, preHookValues, shutdownCtx)
 		if err != nil {
 			st.State = StateFatalError
 			st.Error = err.Error()