Allow async state transition to be canceled

This adds a context and cancelation facility to the type `AsyncError`.
Async state transitions can now be canceled by calling `CancelTransition`
on the AsyncError returned by `fsm.Event`. The context on that error can
also be handed off as described in looplab#77 (comment).

Author: annismckenzie

errors.go

@@ -14,6 +14,10 @@
 
 package fsm
 
+import (
+	"context"
+)
+
 // InvalidEventError is returned by FSM.Event() when the event cannot be called
 // in the current state.
 type InvalidEventError struct {
@@ -82,6 +86,9 @@ func (e CanceledError) Error() string {
 // asynchronous state transition.
 type AsyncError struct {
 	Err error
+
+	Ctx              context.Context
+	CancelTransition func()
 }
 
 func (e AsyncError) Error() string {

fsm.go

@@ -364,6 +364,16 @@ func (f *FSM) Event(ctx context.Context, event string, args ...interface{}) erro
 	if err = f.leaveStateCallbacks(ctx, e); err != nil {
 		if _, ok := err.(CanceledError); ok {
 			f.transition = nil
+		} else if asyncError, ok := err.(AsyncError); ok {
+			// setup a new context in order for async state transitions to work correctly
+			// this "uncancels" the original context which ignores its cancelation
+			// but keeps the values of the original context available to callers
+			ctx, cancel := uncancelContext(ctx)
+			e.cancelFunc = cancel
+			asyncError.Ctx = ctx
+			asyncError.CancelTransition = cancel
+			f.transition = transitionFunc(ctx, true)
+			return asyncError
 		}
 		return err
 	}
@@ -434,15 +444,15 @@ func (f *FSM) leaveStateCallbacks(ctx context.Context, e *Event) error {
 		if e.canceled {
 			return CanceledError{e.Err}
 		} else if e.async {
-			return AsyncError{e.Err}
+			return AsyncError{Err: e.Err}
 		}
 	}
 	if fn, ok := f.callbacks[cKey{"", callbackLeaveState}]; ok {
 		fn(ctx, e)
 		if e.canceled {
 			return CanceledError{e.Err}
 		} else if e.async {
-			return AsyncError{e.Err}
+			return AsyncError{Err: e.Err}
 		}
 	}
 	return nil

fsm_test.go

@@ -473,6 +473,42 @@ func TestAsyncTransitionNotInProgress(t *testing.T) {
 	}
 }
 
+func TestCancelAsyncTransition(t *testing.T) {
+	fsm := NewFSM(
+		"start",
+		Events{
+			{Name: "run", Src: []string{"start"}, Dst: "end"},
+		},
+		Callbacks{
+			"leave_start": func(_ context.Context, e *Event) {
+				e.Async()
+			},
+		},
+	)
+	err := fsm.Event(context.Background(), "run")
+	asyncError, ok := err.(AsyncError)
+	if !ok {
+		t.Errorf("expected error to be 'AsyncError', got %v", err)
+	}
+	var asyncStateTransitionWasCanceled bool
+	go func() {
+		<-asyncError.Ctx.Done()
+		asyncStateTransitionWasCanceled = true
+	}()
+	asyncError.CancelTransition()
+	time.Sleep(20 * time.Millisecond)
+
+	if err = fsm.Transition(); err != nil {
+		t.Errorf("expected no error, got %v", err)
+	}
+	if !asyncStateTransitionWasCanceled {
+		t.Error("expected async state transition cancelation to have propagated")
+	}
+	if fsm.Current() != "start" {
+		t.Error("expected state to be 'start'")
+	}
+}
+
 func TestCallbackNoError(t *testing.T) {
 	fsm := NewFSM(
 		"start",

uncancel_context.go

@@ -0,0 +1,21 @@
+package fsm
+
+import (
+	"context"
+	"time"
+)
+
+type uncancel struct {
+	context.Context
+}
+
+func (*uncancel) Deadline() (deadline time.Time, ok bool) { return }
+func (*uncancel) Done() <-chan struct{}                   { return nil }
+func (*uncancel) Err() error                              { return nil }
+
+// uncancelContext returns a context which ignores the cancellation of the parent and only keeps the values.
+// Also returns a new cancel function.
+// This is useful to keep a background task running while the initial request is finished.
+func uncancelContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	return context.WithCancel(&uncancel{ctx})
+}

uncancel_context_test.go

@@ -0,0 +1,91 @@
+package fsm
+
+import (
+	"context"
+	"testing"
+)
+
+func TestUncancel(t *testing.T) {
+	t.Run("create a new context", func(t *testing.T) {
+		t.Run("and cancel it", func(t *testing.T) {
+			ctx := context.Background()
+			ctx = context.WithValue(ctx, "key1", "value1")
+			ctx, cancelFunc := context.WithCancel(ctx)
+			cancelFunc()
+
+			if ctx.Err() != context.Canceled {
+				t.Errorf("expected context error 'context canceled', got %v", ctx.Err())
+			}
+			select {
+			case <-ctx.Done():
+			default:
+				t.Error("expected context to be done but it wasn't")
+			}
+
+			t.Run("and uncancel it", func(t *testing.T) {
+				ctx, newCancelFunc := uncancelContext(ctx)
+				if ctx.Err() != nil {
+					t.Errorf("expected context error to be nil, got %v", ctx.Err())
+				}
+				select {
+				case <-ctx.Done():
+					t.Fail()
+				default:
+				}
+
+				t.Run("now it should still contain the values", func(t *testing.T) {
+					if ctx.Value("key1") != "value1" {
+						t.Errorf("expected context value of key 'key1' to be 'value1', got %v", ctx.Value("key1"))
+					}
+				})
+				t.Run("and cancel the child", func(t *testing.T) {
+					newCancelFunc()
+					if ctx.Err() != context.Canceled {
+						t.Errorf("expected context error 'context canceled', got %v", ctx.Err())
+					}
+					select {
+					case <-ctx.Done():
+					default:
+						t.Error("expected context to be done but it wasn't")
+					}
+				})
+			})
+		})
+		t.Run("and uncancel it", func(t *testing.T) {
+			ctx := context.Background()
+			parent := ctx
+			ctx, newCancelFunc := uncancelContext(ctx)
+			if ctx.Err() != nil {
+				t.Errorf("expected context error to be nil, got %v", ctx.Err())
+			}
+			select {
+			case <-ctx.Done():
+				t.Fail()
+			default:
+			}
+
+			t.Run("and cancel the child", func(t *testing.T) {
+				newCancelFunc()
+				if ctx.Err() != context.Canceled {
+					t.Errorf("expected context error 'context canceled', got %v", ctx.Err())
+				}
+				select {
+				case <-ctx.Done():
+				default:
+					t.Error("expected context to be done but it wasn't")
+				}
+
+				t.Run("and ensure the parent is not affected", func(t *testing.T) {
+					if parent.Err() != nil {
+						t.Errorf("expected parent context error to be nil, got %v", ctx.Err())
+					}
+					select {
+					case <-parent.Done():
+						t.Fail()
+					default:
+					}
+				})
+			})
+		})
+	})
+}