Enable cancellation of pipelines

Author: ccremer

examples/context_test.go

@@ -19,12 +19,11 @@ type Data struct {
 func TestExample_Context(t *testing.T) {
 	// Create pipeline with defaults
 	p := pipeline.NewPipeline()
-	p.WithContext(context.WithValue(context.Background(), "data", &Data{}))
 	p.WithSteps(
 		pipeline.NewStep("define random number", defineNumber),
 		pipeline.NewStepFromFunc("print number", printNumber),
 	)
-	result := p.Run()
+	result := p.RunWithContext(context.WithValue(context.Background(), "data", &Data{}))
 	if !result.IsSuccessful() {
 		t.Fatal(result.Err)
 	}

parallel/fanout.go

@@ -12,8 +12,6 @@ NewFanOutStep creates a pipeline step that runs nested pipelines in their own Go
 The function provided as PipelineSupplier is expected to close the given channel when no more pipelines should be executed, otherwise this step blocks forever.
 The step waits until all pipelines are finished.
 If the given ResultHandler is non-nil it will be called after all pipelines were run, otherwise the step is considered successful.
-The given pipelines have to define their own context.Context, it's not passed "down" from parent pipeline.
-However, The context.Context for the ResultHandler will be the one from parent pipeline.
 */
 func NewFanOutStep(name string, pipelineSupplier PipelineSupplier, handler ResultHandler) pipeline.Step {
 	step := pipeline.Step{Name: name}

parallel/pool.go

@@ -16,8 +16,6 @@ The step waits until all pipelines are finished.
  * The pipelines are executed in a pool of a number of Go routines indicated by size.
  * If size is 1, the pipelines are effectively run in sequence.
  * If size is 0 or less, the function panics.
-The given pipelines have to define their own context.Context, it's not passed "down" from parent pipeline.
-However, The context.Context for the ResultHandler will be the one from parent pipeline.
 */
 func NewWorkerPoolStep(name string, size int, pipelineSupplier PipelineSupplier, handler ResultHandler) pipeline.Step {
 	if size < 1 {
@@ -33,7 +31,7 @@ func NewWorkerPoolStep(name string, size int, pipelineSupplier PipelineSupplier,
 		go pipelineSupplier(pipelineChan)
 		for i := 0; i < size; i++ {
 			wg.Add(1)
-			go poolWork(pipelineChan, &wg, &count, &m)
+			go poolWork(ctx, pipelineChan, &wg, &count, &m)
 		}
 
 		wg.Wait()
@@ -42,10 +40,10 @@ func NewWorkerPoolStep(name string, size int, pipelineSupplier PipelineSupplier,
 	return step
 }
 
-func poolWork(pipelineChan chan *pipeline.Pipeline, wg *sync.WaitGroup, i *uint64, m *sync.Map) {
+func poolWork(ctx context.Context, pipelineChan chan *pipeline.Pipeline, wg *sync.WaitGroup, i *uint64, m *sync.Map) {
 	defer wg.Done()
 	for pipe := range pipelineChan {
 		n := atomic.AddUint64(i, 1) - 1
-		m.Store(n, pipe.Run())
+		m.Store(n, pipe.RunWithContext(ctx))
 	}
 }

parallel/pool_test.go

@@ -46,7 +46,7 @@ func TestNewWorkerPoolStep(t *testing.T) {
 				assert.Error(t, results[0].Err)
 				return pipeline.Result{Err: results[0].Err}
 			})
-			result := step.F(nil)
+			result := step.F(context.Background())
 			assert.Error(t, result.Err)
 		})
 	}

pipeline.go

@@ -14,18 +14,6 @@ type Pipeline struct {
 	options     options
 }
 
-// Result is the object that is returned after each step and after running a pipeline.
-type Result struct {
-	// Err contains the step's returned error, nil otherwise.
-	// In an aborted pipeline with ErrAbort it will still be nil.
-	Err error
-	// Name is an optional identifier for a result.
-	// ActionFunc may set this property before returning to help a ResultHandler with further processing.
-	Name string
-
-	aborted bool
-}
-
 // Step is an intermediary action and part of a Pipeline.
 type Step struct {
 	// Name describes the step's human-readable name.
@@ -84,24 +72,24 @@ func (p *Pipeline) WithSteps(steps ...Step) *Pipeline {
 
 // WithNestedSteps is similar to AsNestedStep, but it accepts the steps given directly as parameters.
 func (p *Pipeline) WithNestedSteps(name string, steps ...Step) Step {
-	return NewStep(name, func(_ context.Context) Result {
+	return NewStep(name, func(ctx context.Context) Result {
 		nested := &Pipeline{beforeHooks: p.beforeHooks, steps: steps, options: p.options}
-		return nested.Run()
+		return nested.RunWithContext(ctx)
 	})
 }
 
 // AsNestedStep converts the Pipeline instance into a Step that can be used in other pipelines.
 // The properties are passed to the nested pipeline.
 func (p *Pipeline) AsNestedStep(name string) Step {
-	return NewStep(name, func(_ context.Context) Result {
+	return NewStep(name, func(ctx context.Context) Result {
 		nested := &Pipeline{beforeHooks: p.beforeHooks, steps: p.steps, options: p.options}
-		return nested.Run()
+		return nested.RunWithContext(ctx)
 	})
 }
 
 // WithFinalizer returns itself while setting the finalizer for the pipeline.
 // The finalizer is a handler that gets called after the last step is in the pipeline is completed.
-// If a pipeline aborts early then it is also called.
+// If a pipeline aborts early or gets canceled then it is also called.
 func (p *Pipeline) WithFinalizer(handler ResultHandler) *Pipeline {
 	p.finalizer = handler
 	return p
@@ -116,6 +104,9 @@ func (p *Pipeline) Run() Result {
 }
 
 // RunWithContext is like Run but with a given context.Context.
+// Upon cancellation of the context, the pipeline does not terminate a currently running step, instead it skips the remaining steps in the execution order.
+// The context is passed to each Step.F and each Step may need to listen to the context cancellation event to truly cancel a long-running step.
+// If the pipeline gets canceled, Result.IsCanceled returns true and Result.Err contains the context's error.
 func (p *Pipeline) RunWithContext(ctx context.Context) Result {
 	result := p.doRun(ctx)
 	if p.finalizer != nil {
@@ -125,28 +116,45 @@ func (p *Pipeline) RunWithContext(ctx context.Context) Result {
 }
 
 func (p *Pipeline) doRun(ctx context.Context) Result {
+	name := ""
 	for _, step := range p.steps {
-		for _, hooks := range p.beforeHooks {
-			hooks(step)
-		}
+		name = step.Name
+		select {
+		case <-ctx.Done():
+			result := p.fail(ctx.Err(), step)
+			return result
+		default:
+			for _, hooks := range p.beforeHooks {
+				hooks(step)
+			}
 
-		result := step.F(ctx)
-		var err error
-		if step.H != nil {
-			err = step.H(ctx, result)
-		} else {
-			err = result.Err
-		}
-		if err != nil {
-			if errors.Is(err, ErrAbort) {
-				// Abort pipeline without error
-				return Result{aborted: true}
+			result := step.F(ctx)
+			var err error
+			if step.H != nil {
+				result.name = step.Name
+				err = step.H(ctx, result)
+			} else {
+				err = result.Err
 			}
-			if p.options.disableErrorWrapping {
-				return Result{Err: err}
+			if err != nil {
+				if errors.Is(err, ErrAbort) {
+					// Abort pipeline without error
+					return Result{aborted: true, name: step.Name}
+				}
+				return p.fail(err, step)
 			}
-			return Result{Err: fmt.Errorf("step '%s' failed: %w", step.Name, err)}
 		}
 	}
-	return Result{}
+	return Result{name: name}
+}
+
+func (p *Pipeline) fail(err error, step Step) Result {
+	result := Result{name: step.Name}
+	if p.options.disableErrorWrapping {
+		result.Err = err
+	} else {
+		result.Err = fmt.Errorf("step %q failed: %w", step.Name, err)
+	}
+	result.canceled = errors.Is(err, ErrCanceled) || errors.Is(err, context.DeadlineExceeded)
+	return result
 }

pipeline_test.go

@@ -3,7 +3,9 @@ package pipeline
 import (
 	"context"
 	"errors"
+	"fmt"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -165,6 +167,55 @@ func TestPipeline_Run(t *testing.T) {
 	}
 }
 
+func TestPipeline_RunWithContext_CancelLongRunningStep(t *testing.T) {
+	p := NewPipeline().WithSteps(
+		NewStepFromFunc("long running", func(ctx context.Context) error {
+			for {
+				select {
+				case <-ctx.Done():
+					return ErrCanceled
+				default:
+					// doing nothing
+				}
+			}
+			return nil
+		}),
+	)
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	result := p.RunWithContext(ctx)
+	cancel()
+	assert.True(t, result.IsCanceled(), "IsCanceled()")
+	assert.Equal(t, "long running", result.Name())
+	assert.EqualError(t, result.Err, "step \"long running\" failed: canceled")
+}
+
+func ExamplePipeline_RunWithContext() {
+	// prepare pipeline
+	p := NewPipeline().WithSteps(
+		NewStepFromFunc("short step", func(ctx context.Context) error {
+			fmt.Println("short step")
+			return nil
+		}),
+		NewStepFromFunc("long running step", func(ctx context.Context) error {
+			time.Sleep(100 * time.Millisecond)
+			return nil
+		}),
+		NewStepFromFunc("canceled step", func(ctx context.Context) error {
+			return errors.New("shouldn't execute")
+		}),
+	)
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	result := p.RunWithContext(ctx)
+	// cancel the pipeline
+	cancel()
+	// inspect the result
+	fmt.Println(result.IsCanceled())
+	fmt.Println(result.Err)
+	// Output: short step
+	// true
+	// step "canceled step" failed: context deadline exceeded
+}
+
 func TestNewStepFromFunc(t *testing.T) {
 	called := false
 	step := NewStepFromFunc("name", func(ctx context.Context) error {

result.go

@@ -2,9 +2,28 @@ package pipeline
 
 import "errors"
 
-// ErrAbort indicates that the pipeline should be terminated immediately without returning an error.
+// ErrAbort indicates that the pipeline should be terminated immediately without being marked as failed (returning an error).
 var ErrAbort = errors.New("abort")
 
+// ErrCanceled indicates that the pipeline has been canceled.
+var ErrCanceled = errors.New("canceled")
+
+// Result is the object that is returned after each step and after running a pipeline.
+type Result struct {
+	// Err contains the step's returned error, nil otherwise.
+	// In an aborted pipeline with ErrAbort it will still be nil.
+	Err error
+
+	name     string
+	aborted  bool
+	canceled bool
+}
+
+// Name retrieves the name of the (last) step that has been executed.
+func (r Result) Name() string {
+	return r.name
+}
+
 // IsSuccessful returns true if the contained error is nil.
 // Aborted pipelines (with ErrAbort) are still reported as success.
 // To query if a pipeline is aborted early, use IsAborted.
@@ -21,3 +40,8 @@ func (r Result) IsFailed() bool {
 func (r Result) IsAborted() bool {
 	return r.aborted
 }
+
+// IsCanceled returns true if the pipeline's context has been canceled.
+func (r Result) IsCanceled() bool {
+	return r.canceled
+}

step.go

@@ -14,7 +14,7 @@ func NewStep(name string, action ActionFunc) Step {
 func NewStepFromFunc(name string, fn func(ctx context.Context) error) Step {
 	return NewStep(name, func(ctx context.Context) Result {
 		err := fn(ctx)
-		return Result{Err: err, Name: name}
+		return Result{Err: err, name: name}
 	})
 }