Support workflow cancellation after commands are done

Author: cschleiden

internal/command/cancelablecommand.go

@@ -38,6 +38,10 @@ func (c *cancelableCommand) Done() {
 	switch c.state {
 	case CommandState_Committed, CommandState_Canceled:
 		c.state = CommandState_Done
+		if c.whenDone != nil {
+			c.whenDone()
+		}
+
 	default:
 		c.invalidStateTransition(CommandState_Done)
 	}

internal/command/command.go

@@ -24,6 +24,8 @@ type Command interface {
 	// Done marks the command as done. This transitions the state to done and indicates that the result
 	// of this command has been applied.
 	Done()
+
+	WhenDone(fn func())
 }
 
 type CommandResult struct {
@@ -40,6 +42,8 @@ type command struct {
 	state CommandState
 
 	id int64
+
+	whenDone func()
 }
 
 func (c *command) ID() int64 {
@@ -67,11 +71,19 @@ func (c *command) Done() {
 	switch c.state {
 	case CommandState_Committed:
 		c.state = CommandState_Done
+
+		if c.whenDone != nil {
+			c.whenDone()
+		}
 	default:
 		c.invalidStateTransition(CommandState_Done)
 	}
 }
 
+func (c *command) WhenDone(fn func()) {
+	c.whenDone = fn
+}
+
 func (c *command) invalidStateTransition(state CommandState) {
 	panic(fmt.Errorf("invalid state transition for command %s: %s -> %s", c.name, c.State().String(), state.String()))
 }

internal/command/sideeffect.go

@@ -65,6 +65,9 @@ func (c *SideEffectCommand) Done() {
 	switch c.state {
 	case CommandState_Pending, CommandState_Committed:
 		c.state = CommandState_Done
+		if c.whenDone != nil {
+			c.whenDone()
+		}
 
 	default:
 		c.invalidStateTransition(CommandState_Done)

workflow/subworkflow.go

@@ -85,22 +85,31 @@ func createSubWorkflowInstance[TResult any](ctx sync.Context, options SubWorkflo
 	span.Marshal(metadata)
 
 	cmd := command.NewScheduleSubWorkflowCommand(scheduleEventID, wfState.Instance(), options.InstanceID, name, inputs, metadata)
+
 	wfState.AddCommand(cmd)
 	wfState.TrackFuture(scheduleEventID, workflowstate.AsDecodingSettable(cv, f))
 
 	// Check if the channel is cancelable
 	if c, cancelable := ctx.Done().(sync.CancelChannel); cancelable {
-		c.AddReceiveCallback(func(v struct{}, ok bool) {
-			cmd.Cancel()
-			if cmd.State() == command.CommandState_Canceled {
-				// Remove the sub-workflow future from the workflow state and mark it as canceled if it hasn't already fired
-				if fi, ok := f.(sync.FutureInternal[TResult]); ok {
-					if !fi.Ready() {
-						wfState.RemoveFuture(scheduleEventID)
-						f.Set(*new(TResult), sync.Canceled)
+		cancelReceiver := &sync.Receiver[struct{}]{
+			Receive: func(v struct{}, ok bool) {
+				cmd.Cancel()
+				if cmd.State() == command.CommandState_Canceled {
+					// Remove the sub-workflow future from the workflow state and mark it as canceled if it hasn't already fired
+					if fi, ok := f.(sync.FutureInternal[TResult]); ok {
+						if !fi.Ready() {
+							wfState.RemoveFuture(scheduleEventID)
+							f.Set(*new(TResult), sync.Canceled)
+						}
 					}
 				}
-			}
+			},
+		}
+
+		c.AddReceiveCallback(cancelReceiver)
+
+		cmd.WhenDone(func() {
+			c.RemoveReceiveCallback(cancelReceiver)
 		})
 	}
 

workflow/timer.go

@@ -25,11 +25,27 @@ func ScheduleTimer(ctx Context, delay time.Duration) Future[struct{}] {
 
 	scheduleEventID := wfState.GetNextScheduleEventID()
 	at := Now(ctx).Add(delay)
+
 	timerCmd := command.NewScheduleTimerCommand(scheduleEventID, at)
 	wfState.AddCommand(timerCmd)
-
 	wfState.TrackFuture(scheduleEventID, workflowstate.AsDecodingSettable(converter.GetConverter(ctx), f))
 
+	cancelReceiver := &sync.Receiver[struct{}]{
+		Receive: func(v struct{}, ok bool) {
+			timerCmd.Cancel()
+
+			// Remove the timer future from the workflow state and mark it as canceled if it hasn't already fired. This is different
+			// from subworkflow behavior, where we want to wait for the subworkflow to complete before proceeding. Here we can
+			// continue right away.
+			if fi, ok := f.(sync.FutureInternal[struct{}]); ok {
+				if !fi.Ready() {
+					wfState.RemoveFuture(scheduleEventID)
+					f.Set(v, sync.Canceled)
+				}
+			}
+		},
+	}
+
 	ctx, span := workflowtracer.Tracer(ctx).Start(ctx, "ScheduleTimer",
 		trace.WithAttributes(
 			attribute.Int64("duration_ms", int64(delay/time.Millisecond)),
@@ -42,26 +58,10 @@ func ScheduleTimer(ctx Context, delay time.Duration) Future[struct{}] {
 	if c, cancelable := ctx.Done().(sync.CancelChannel); cancelable {
 		// Register a callback for when it's canceled. The only operation on the `Done` channel
 		// is that it's closed when the context is canceled.
-		canceled := false
-
-		c.AddReceiveCallback(func(v struct{}, ok bool) {
-			// Ignore any future cancelation events for this timer
-			if canceled {
-				return
-			}
-			canceled = true
+		c.AddReceiveCallback(cancelReceiver)
 
-			timerCmd.Cancel()
-
-			// Remove the timer future from the workflow state and mark it as canceled if it hasn't already fired. This is different
-			// from subworkflow behavior, where we want to wait for the subworkflow to complete before proceeding. Here we can
-			// continue right away.
-			if fi, ok := f.(sync.FutureInternal[struct{}]); ok {
-				if !fi.Ready() {
-					wfState.RemoveFuture(scheduleEventID)
-					f.Set(v, sync.Canceled)
-				}
-			}
+		timerCmd.WhenDone(func() {
+			c.RemoveReceiveCallback(cancelReceiver)
 		})
 	}
 

workflow/timer_test.go

@@ -0,0 +1,50 @@
+package workflow
+
+import (
+	"testing"
+	"time"
+
+	"github.com/benbjohnson/clock"
+	"github.com/cschleiden/go-workflows/internal/converter"
+	"github.com/cschleiden/go-workflows/internal/core"
+	"github.com/cschleiden/go-workflows/internal/logger"
+	"github.com/cschleiden/go-workflows/internal/sync"
+	"github.com/cschleiden/go-workflows/internal/workflowstate"
+	"github.com/cschleiden/go-workflows/internal/workflowtracer"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace"
+)
+
+func Test_Timer_Cancellation(t *testing.T) {
+	state := workflowstate.NewWorkflowState(core.NewWorkflowInstance("a", ""), logger.NewDefaultLogger(), clock.New())
+
+	ctx, cancel := sync.WithCancel(sync.Background())
+	ctx = converter.WithConverter(ctx, converter.DefaultConverter)
+	ctx = workflowstate.WithWorkflowState(ctx, state)
+	ctx = workflowtracer.WithWorkflowTracer(ctx, workflowtracer.New(trace.NewNoopTracerProvider().Tracer("test")))
+
+	c := sync.NewCoroutine(ctx, func(ctx sync.Context) error {
+		f := ScheduleTimer(ctx, time.Second*1)
+		f.Get(ctx)
+
+		// Block workflow
+		sync.NewFuture[int]().Get(ctx)
+
+		return nil
+	})
+	c.Execute()
+	require.False(t, c.Finished())
+
+	// Fire timer
+	cmd := state.CommandByScheduleEventID(1)
+	cmd.Commit()
+	cmd.Done()
+	fs, ok := state.FutureByScheduleEventID(1)
+	require.True(t, ok)
+	fs(nil, nil)
+
+	c.Execute()
+	require.False(t, c.Finished())
+
+	cancel()
+}