Merge pull request #188 from cschleiden/fix-cmd-cancellation

Support invalid state transition for finished commands during workflow cancellation

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)

internal/sync/channel.go

@@ -12,14 +12,20 @@ type Channel[T any] interface {
 	Close()
 }
 
+type Receiver[T any] struct {
+	Receive func(v T, ok bool)
+}
+
 type ChannelInternal[T any] interface {
 	Closed() bool
 
 	ReceiveNonBlocking() (v T, ok bool)
 
 	// AddReceiveCallback adds a callback that is called once when a value is sent to the channel. This is similar
 	// to the blocking `Receive` method, but is not blocking a coroutine.
-	AddReceiveCallback(cb func(v T, ok bool))
+	AddReceiveCallback(rcb *Receiver[T])
+
+	RemoveReceiveCallback(rcb *Receiver[T])
 }
 
 // Ensure channel implementation support internal interface
@@ -40,7 +46,7 @@ func NewBufferedChannel[T any](size int) Channel[T] {
 
 type channel[T any] struct {
 	c         []T
-	receivers []func(value T, ok bool)
+	receivers []*Receiver[T]
 	senders   []func() T
 	closed    bool
 	size      int
@@ -62,7 +68,7 @@ func (c *channel[T]) Close() {
 
 		// Send zero value to pending receiver
 		var v T
-		r(v, false)
+		r.Receive(v, false)
 	}
 }
 
@@ -119,10 +125,12 @@ func (c *channel[T]) Receive(ctx Context) (v T, ok bool) {
 
 		// Register handler to receive value once
 		if !addedListener {
-			cb := func(rv T, rok bool) {
-				receivedValue = true
-				v = rv
-				ok = rok
+			cb := &Receiver[T]{
+				Receive: func(rv T, rok bool) {
+					receivedValue = true
+					v = rv
+					ok = rok
+				},
 			}
 
 			c.receivers = append(c.receivers, cb)
@@ -176,7 +184,7 @@ func (c *channel[T]) trySend(v T) bool {
 		c.receivers[0] = nil
 		c.receivers = c.receivers[1:]
 
-		r(v, true)
+		r.Receive(v, true)
 
 		return true
 	}
@@ -223,10 +231,20 @@ func (c *channel[T]) hasCapacity() bool {
 	return len(c.c) < c.size
 }
 
-func (c *channel[T]) AddReceiveCallback(cb func(v T, ok bool)) {
+func (c *channel[T]) AddReceiveCallback(cb *Receiver[T]) {
 	c.receivers = append(c.receivers, cb)
 }
 
+func (c *channel[T]) RemoveReceiveCallback(cb *Receiver[T]) {
+	for i, r := range c.receivers {
+		if r == cb {
+			c.receivers[i] = nil
+			c.receivers = append(c.receivers[:i], c.receivers[i+1:]...)
+			return
+		}
+	}
+}
+
 func (c *channel[T]) Closed() bool {
 	return c.closed
 }

internal/sync/channel_test.go

@@ -413,3 +413,60 @@ func Test_Channel_Buffered(t *testing.T) {
 		})
 	}
 }
+
+func Test_CancellationHandler_Add(t *testing.T) {
+	ctx, cancel := WithCancel(Background())
+
+	f := 1
+
+	c := NewChannel[int]()
+	ic := c.(*channel[int])
+
+	cr := NewCoroutine(ctx, func(ctx Context) error {
+		r := &Receiver[int]{
+			Receive: func(_ int, ok bool) {
+				f++
+			},
+		}
+
+		ic.AddReceiveCallback(r)
+
+		c.Send(ctx, 42)
+		return nil
+	})
+
+	cr.Execute()
+
+	cancel()
+
+	require.Equal(t, 2, f)
+}
+
+func Test_CancellationHandler_Remove(t *testing.T) {
+	ctx, cancel := WithCancel(Background())
+
+	f := 1
+
+	c := NewChannel[int]()
+	ic := c.(*channel[int])
+
+	cr := NewCoroutine(ctx, func(ctx Context) error {
+		r := &Receiver[int]{
+			Receive: func(_ int, ok bool) {
+				f++
+			},
+		}
+
+		ic.AddReceiveCallback(r)
+		ic.RemoveReceiveCallback(r)
+
+		c.Send(ctx, 42)
+		return nil
+	})
+
+	cr.Execute()
+
+	cancel()
+
+	require.Equal(t, 1, f)
+}

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