Merge pull request #609 from wippyai/fix/contract-future-cancel

fix(contract): cancel async calls at the callee

Author: wolfy-j

runtime/lua/modules/contract/module.go

@@ -797,6 +797,7 @@ func futureCancelImpl(l *lua.LState) int {
 		return 0
 	}
 
+	f.MarkCanceled()
 	yield := AcquireAsyncCancelYield()
 	yield.Topic = f.Topic
 	l.Push(yield)

runtime/lua/modules/contract/module_test.go

@@ -16,6 +16,8 @@ import (
 	"github.com/wippyai/runtime/api/registry"
 	"github.com/wippyai/runtime/api/runtime"
 	secapi "github.com/wippyai/runtime/api/security"
+	"github.com/wippyai/runtime/runtime/lua/engine"
+	"github.com/wippyai/runtime/runtime/lua/modules/future"
 )
 
 type mockInstanceForTest struct {
@@ -430,6 +432,23 @@ func TestAsyncCancelYield_HandleResult(t *testing.T) {
 	assert.Equal(t, lua.LNil, results[1])
 }
 
+func TestFutureCancelImplMarksFutureCanceled(t *testing.T) {
+	l := lua.NewState()
+	defer l.Close()
+
+	f := future.New("@future:test", engine.NewChannel(1))
+	ud := l.NewUserData()
+	ud.Value = f
+	l.Push(ud)
+
+	require.Equal(t, -1, futureCancelImpl(l))
+	require.True(t, f.IsCanceled())
+	yield, ok := l.Get(-1).(*AsyncCancelYield)
+	require.True(t, ok)
+	require.Equal(t, "@future:test", yield.Topic)
+	ReleaseAsyncCancelYield(yield)
+}
+
 func TestPoolConcurrency(_ *testing.T) {
 	const goroutines = 50
 	const iterations = 100

system/contract/dispatcher.go

@@ -4,6 +4,7 @@ package contract
 
 import (
 	"context"
+	"sync"
 
 	ctxapi "github.com/wippyai/runtime/api/context"
 	"github.com/wippyai/runtime/api/contract"
@@ -17,8 +18,18 @@ import (
 
 // Dispatcher handles contract commands.
 type Dispatcher struct {
-	node   relay.Node
-	logger *zap.Logger
+	node       relay.Node
+	logger     *zap.Logger
+	asyncCalls sync.Map // map[asyncCallKey]*asyncCall
+}
+
+type asyncCallKey struct {
+	target string
+	topic  string
+}
+
+type asyncCall struct {
+	cancel context.CancelFunc
 }
 
 // NewDispatcher creates a new contract dispatcher with relay node for async routing.
@@ -34,8 +45,14 @@ func (d *Dispatcher) Start(_ context.Context) error {
 	return nil
 }
 
-// Stop is a no-op for contract dispatcher.
+// Stop cancels calls owned by this dispatcher.
 func (d *Dispatcher) Stop(_ context.Context) error {
+	d.asyncCalls.Range(func(key, value any) bool {
+		if d.asyncCalls.CompareAndDelete(key, value) {
+			value.(*asyncCall).cancel()
+		}
+		return true
+	})
 	return nil
 }
 
@@ -145,10 +162,24 @@ func (d *Dispatcher) handleAsyncCall(ctx context.Context, cmd dispatcher.Command
 	logger := d.logger
 
 	callCtx, fc := ctxapi.ForkFrameContext(ctx)
+	callCtx, cancel := context.WithCancel(callCtx)
+	key := asyncCallKey{target: framePID.String(), topic: topic}
+	ownedCall := &asyncCall{cancel: cancel}
+	if previous, replaced := d.asyncCalls.Swap(key, ownedCall); replaced {
+		previous.(*asyncCall).cancel()
+	}
 
 	go func(callCtx context.Context, callFC ctxapi.FrameContext) {
 		defer ctxapi.ReleaseFrameContext(callFC)
+		defer cancel()
 		result, err := instance.Call(callCtx, method, args, options)
+		// Cancellation and topic replacement both remove ownership before
+		// interrupting the call. Only the still-owned invocation may publish a
+		// result; otherwise its terminal frame could close a newer future that
+		// reused the same caller-owned topic.
+		if !d.asyncCalls.CompareAndDelete(key, ownedCall) {
+			return
+		}
 
 		resultPayload := resultToPayload(result, err)
 		if err := sendAsyncResult(node, framePID, topic, resultPayload); err != nil {
@@ -191,9 +222,15 @@ func (d *Dispatcher) handleAsyncCancel(ctx context.Context, cmd dispatcher.Comma
 		return nil
 	}
 
-	if err := sendAsyncCancel(d.node, framePID, cancelCmd.Topic); err != nil {
+	topic := cancelCmd.Topic
+	key := asyncCallKey{target: framePID.String(), topic: topic}
+	if call, ok := d.asyncCalls.LoadAndDelete(key); ok {
+		call.(*asyncCall).cancel()
+	}
+
+	if err := sendAsyncCancel(d.node, framePID, topic); err != nil {
 		d.logger.Warn("failed to send async cancel",
-			zap.String("topic", cancelCmd.Topic),
+			zap.String("topic", topic),
 			zap.String("target", framePID.String()),
 			zap.Error(err))
 	}

system/contract/dispatcher_test.go

@@ -344,6 +344,173 @@ func TestAsyncCancelHandler(t *testing.T) {
 	}
 }
 
+func startContractAsyncForTest(
+	ctx context.Context,
+	t *testing.T,
+	d *Dispatcher,
+	instance contract.Instance,
+	method string,
+	topic string,
+) {
+	t.Helper()
+	cmd := contract.AcquireAsyncCallCmd()
+	defer cmd.Release()
+	cmd.Instance = instance
+	cmd.Method = method
+	cmd.Topic = topic
+	done := make(chan contract.AsyncCallResult, 1)
+	require.NoError(t, d.handleAsyncCall(ctx, cmd, 0, &testReceiver{cb: func(data any, _ error) {
+		done <- data.(contract.AsyncCallResult)
+	}}))
+	select {
+	case result := <-done:
+		require.NoError(t, result.Error)
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for async contract call")
+	}
+}
+
+func cancelContractAsyncForTest(ctx context.Context, t *testing.T, d *Dispatcher, topic string) {
+	t.Helper()
+	cmd := contract.AcquireAsyncCancelCmd()
+	defer cmd.Release()
+	cmd.Topic = topic
+	done := make(chan struct{}, 1)
+	require.NoError(t, d.handleAsyncCancel(ctx, cmd, 0, &testReceiver{cb: func(_ any, _ error) {
+		done <- struct{}{}
+	}}))
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for async contract cancellation")
+	}
+}
+
+func TestAsyncCancelHandler_CancelsRunningCallContext(t *testing.T) {
+	node := &mockRelayNode{packages: make(chan *relay.Package, 10)}
+	started := make(chan struct{})
+	canceled := make(chan error, 1)
+	instance := &mockInstance{callFn: func(ctx context.Context, _ string, _ payload.Payloads, _ runtime.Options) (*runtime.Result, error) {
+		close(started)
+		<-ctx.Done()
+		canceled <- ctx.Err()
+		return nil, ctx.Err()
+	}}
+	d := NewDispatcher(node, nil)
+	ctx := setupAsyncTestContext()
+	startContractAsyncForTest(ctx, t, d, instance, "run", "@future:cancel")
+
+	select {
+	case <-started:
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for contract call to start")
+	}
+	cancelContractAsyncForTest(ctx, t, d, "@future:cancel")
+	select {
+	case err := <-canceled:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for contract call cancellation")
+	}
+	require.Len(t, node.packages, 1, "canceled call must not publish a second terminal result")
+}
+
+func TestAsyncCallHandler_CleansCancelHandleAfterCompletion(t *testing.T) {
+	node := &mockRelayNode{packages: make(chan *relay.Package, 10)}
+	d := NewDispatcher(node, nil)
+	instance := &mockInstance{callFn: func(_ context.Context, _ string, _ payload.Payloads, _ runtime.Options) (*runtime.Result, error) {
+		return &runtime.Result{Value: payload.New("done")}, nil
+	}}
+	ctx := setupAsyncTestContext()
+	framePID, ok := runtime.GetFramePID(ctx)
+	require.True(t, ok)
+	startContractAsyncForTest(ctx, t, d, instance, "run", "@future:complete")
+
+	select {
+	case <-node.packages:
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for async contract result")
+	}
+	key := asyncCallKey{target: framePID.String(), topic: "@future:complete"}
+	require.Eventually(t, func() bool {
+		_, exists := d.asyncCalls.Load(key)
+		return !exists
+	}, time.Second, 10*time.Millisecond)
+}
+
+func TestAsyncCancelHandler_DoesNotCancelSameTopicForDifferentCallerPID(t *testing.T) {
+	node := &mockRelayNode{packages: make(chan *relay.Package, 10)}
+	started := make(chan context.Context, 1)
+	canceled := make(chan error, 1)
+	instance := &mockInstance{callFn: func(ctx context.Context, _ string, _ payload.Payloads, _ runtime.Options) (*runtime.Result, error) {
+		started <- ctx
+		<-ctx.Done()
+		canceled <- ctx.Err()
+		return nil, ctx.Err()
+	}}
+	d := NewDispatcher(node, nil)
+
+	root := ctxapi.NewRootContext()
+	ownerCtx, _ := ctxapi.OpenFrameContext(root)
+	require.NoError(t, runtime.SetFramePID(ownerCtx, pid.PID{Host: "test", UniqID: "owner"}))
+	startContractAsyncForTest(ownerCtx, t, d, instance, "run", "@future:shared")
+	var callCtx context.Context
+	select {
+	case callCtx = <-started:
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for contract call to start")
+	}
+
+	otherCtx, _ := ctxapi.OpenFrameContext(root)
+	require.NoError(t, runtime.SetFramePID(otherCtx, pid.PID{Host: "test", UniqID: "other"}))
+	cancelContractAsyncForTest(otherCtx, t, d, "@future:shared")
+	require.NoError(t, callCtx.Err())
+
+	cancelContractAsyncForTest(ownerCtx, t, d, "@future:shared")
+	select {
+	case err := <-canceled:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for owning caller cancellation")
+	}
+}
+
+func TestAsyncCallHandler_ReusedOwnedTopicCancelsOnlyTheReplacedCall(t *testing.T) {
+	node := &mockRelayNode{packages: make(chan *relay.Package, 10)}
+	started := make(chan context.Context, 2)
+	canceled := make(chan error, 2)
+	instance := &mockInstance{callFn: func(ctx context.Context, _ string, _ payload.Payloads, _ runtime.Options) (*runtime.Result, error) {
+		started <- ctx
+		<-ctx.Done()
+		canceled <- ctx.Err()
+		return nil, ctx.Err()
+	}}
+	d := NewDispatcher(node, nil)
+	ctx := setupAsyncTestContext()
+
+	startContractAsyncForTest(ctx, t, d, instance, "first", "@future:reused")
+	firstCtx := <-started
+	startContractAsyncForTest(ctx, t, d, instance, "second", "@future:reused")
+	secondCtx := <-started
+	select {
+	case err := <-canceled:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for replaced call cancellation")
+	}
+	require.ErrorIs(t, firstCtx.Err(), context.Canceled)
+	require.NoError(t, secondCtx.Err())
+
+	cancelContractAsyncForTest(ctx, t, d, "@future:reused")
+	select {
+	case err := <-canceled:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for current call cancellation")
+	}
+	require.Len(t, node.packages, 1, "replaced call must not terminate the current future")
+}
+
 func TestDispatcher_RegisterAll(t *testing.T) {
 	d := NewDispatcher(nil, nil)
 
@@ -508,6 +675,32 @@ func TestDispatcher_StartStop(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+func TestDispatcher_StopCancelsOwnedAsyncCalls(t *testing.T) {
+	node := &mockRelayNode{packages: make(chan *relay.Package, 10)}
+	started := make(chan struct{})
+	canceled := make(chan error, 1)
+	instance := &mockInstance{callFn: func(ctx context.Context, _ string, _ payload.Payloads, _ runtime.Options) (*runtime.Result, error) {
+		close(started)
+		<-ctx.Done()
+		canceled <- ctx.Err()
+		return nil, ctx.Err()
+	}}
+	d := NewDispatcher(node, nil)
+	startContractAsyncForTest(setupAsyncTestContext(), t, d, instance, "run", "@future:stop")
+	select {
+	case <-started:
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for contract call to start")
+	}
+	require.NoError(t, d.Stop(context.Background()))
+	select {
+	case err := <-canceled:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("timeout waiting for dispatcher shutdown cancellation")
+	}
+}
+
 func TestOpenHandler_ContextCanceled(t *testing.T) {
 	d := NewDispatcher(nil, nil)
 	mockInst := &mockInstantiator{