From b37371612422c989cd4d27c8e0805427502807a0 Mon Sep 17 00:00:00 2001 From: yashnevatia Date: Wed, 12 Aug 2026 17:39:33 +0100 Subject: [PATCH 1/6] Release executor slot if not executing --- .../executable/request/server_request.go | 45 +++++++++++++++---- 1 file changed, 36 insertions(+), 9 deletions(-) diff --git a/core/capabilities/remote/executable/request/server_request.go b/core/capabilities/remote/executable/request/server_request.go index a76ffb07f8c..05272339387 100644 --- a/core/capabilities/remote/executable/request/server_request.go +++ b/core/capabilities/remote/executable/request/server_request.go @@ -7,6 +7,7 @@ import ( "slices" "strconv" "sync" + "sync/atomic" "time" "go.opentelemetry.io/otel/attribute" @@ -133,7 +134,14 @@ type ServerRequest struct { // Metadata.WorkflowDonID to match the authenticated calling DON. workflowDONBindingGate limits.GateLimiter - mux sync.Mutex + // stateMux guards requesters, responseSentToRequester, and response. + // It is held only for short map/field operations, never during capability execution. + stateMux sync.Mutex + + // executionClaimed is set to true exactly once (via CAS) by the message that + // wins the right to execute the capability. All other messages skip execution. + executionClaimed atomic.Bool + lggr logger.Logger metrics *srMetrics @@ -178,9 +186,6 @@ func NewServerRequest(capability capabilities.ExecutableCapability, method strin func (e *ServerRequest) OnMessage(ctx context.Context, msg *types.MessageBody) error { e.metrics.countExecutionRequest(ctx) - e.mux.Lock() - defer e.mux.Unlock() - if msg.Sender == nil { return errors.New("sender missing from message") } @@ -190,22 +195,35 @@ func (e *ServerRequest) OnMessage(ctx context.Context, msg *types.MessageBody) e return fmt.Errorf("failed to convert message sender to PeerID: %w", err) } + e.stateMux.Lock() if err := e.addRequester(requester); err != nil { + e.stateMux.Unlock() return fmt.Errorf("failed to add requester to request: %w", err) } - e.lggr.Debugw("OnMessage called for request", "calls", len(e.requesters), - "hasResponse", e.response != nil, "requester", requester.String(), "minRequsters", e.callingDon.F+1) + quorumReached := e.minimumRequiredRequestsReceived() + hasResponse := e.hasResponse() + e.stateMux.Unlock() + + e.lggr.Debugw("OnMessage called for request", "requester", requester.String(), + "quorumReached", quorumReached, "hasResponse", hasResponse, "minRequesters", e.callingDon.F+1) - if e.minimumRequiredRequestsReceived() && !e.hasResponse() { + // Only one message wins the right to execute. All others skip execution + // and either wait for the executor's fan-out or self-send if the response + // is already available. + if quorumReached && !hasResponse && e.executionClaimed.CompareAndSwap(false, true) { switch e.method { case types.MethodExecute: e.executeRequest(ctx, msg, executeCapabilityRequest) default: + e.stateMux.Lock() e.setError(types.Error_INTERNAL_ERROR, "unknown method %s"+e.method) + e.stateMux.Unlock() } } + e.stateMux.Lock() + defer e.stateMux.Unlock() if err := e.sendResponses(ctx); err != nil { return fmt.Errorf("failed to send responses: %w", err) } @@ -223,8 +241,14 @@ func (e *ServerRequest) Evictable(minRetention time.Duration) bool { } func (e *ServerRequest) Cancel(ctx context.Context, err types.Error, msg string) error { - e.mux.Lock() - defer e.mux.Unlock() + // If execution has been claimed, the executor will set the response when it + // finishes. Cancelling at that point would race with the executor's result. + if e.executionClaimed.Load() { + return nil + } + + e.stateMux.Lock() + defer e.stateMux.Unlock() if !e.hasResponse() { e.setError(err, msg) @@ -245,12 +269,15 @@ func (e *ServerRequest) executeRequest(ctx context.Context, msg *types.MessageBo success := false start := time.Now() responsePayload, err := method(ctxWithTimeout, e.lggr, e.capability, msg.Payload, e.callingDon.ID, e.workflowDONBindingGate) + + e.stateMux.Lock() if err != nil { e.setError(types.Error_INTERNAL_ERROR, err.Error()) } else { success = true e.setResult(responsePayload) } + e.stateMux.Unlock() e.metrics.countExecution(ctx, success) e.metrics.recordExecutionDuration(ctx, time.Since(start), success) From 6ff43f3fa9a9f5afb7beca708135508ad25cb6b6 Mon Sep 17 00:00:00 2001 From: yashnevatia Date: Wed, 12 Aug 2026 17:57:26 +0100 Subject: [PATCH 2/6] Add exec done channel --- .../executable/request/server_request.go | 27 ++++++++++++++++--- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/core/capabilities/remote/executable/request/server_request.go b/core/capabilities/remote/executable/request/server_request.go index 05272339387..efccc459c7d 100644 --- a/core/capabilities/remote/executable/request/server_request.go +++ b/core/capabilities/remote/executable/request/server_request.go @@ -142,6 +142,14 @@ type ServerRequest struct { // wins the right to execute the capability. All other messages skip execution. executionClaimed atomic.Bool + // executionDone is closed by the executor after it writes the response under stateMux. + // Cancel waits on this channel to avoid a race where it could acquire stateMux + // between "executor finishes capability call" and "executor writes response", + // see no response, and set a timeout error. The executor would then overwrite + // with success, but all requesters would already be marked as sent — leaving + // them with the timeout error while the success response is lost. + executionDone chan struct{} + lggr logger.Logger metrics *srMetrics @@ -178,6 +186,7 @@ func NewServerRequest(capability capabilities.ExecutableCapability, method strin requestTimeout: requestTimeout, capMethodName: capMethodName, workflowDONBindingGate: workflowDONBindingGate, + executionDone: make(chan struct{}), lggr: lggr, metrics: m, }, nil @@ -202,11 +211,12 @@ func (e *ServerRequest) OnMessage(ctx context.Context, msg *types.MessageBody) e } quorumReached := e.minimumRequiredRequestsReceived() + calls := len(e.requesters) hasResponse := e.hasResponse() e.stateMux.Unlock() e.lggr.Debugw("OnMessage called for request", "requester", requester.String(), - "quorumReached", quorumReached, "hasResponse", hasResponse, "minRequesters", e.callingDon.F+1) + "quorumReached", quorumReached, "hasResponse", hasResponse, "minRequesters", e.callingDon.F+1, "calls", calls) // Only one message wins the right to execute. All others skip execution // and either wait for the executor's fan-out or self-send if the response @@ -219,6 +229,7 @@ func (e *ServerRequest) OnMessage(ctx context.Context, msg *types.MessageBody) e e.stateMux.Lock() e.setError(types.Error_INTERNAL_ERROR, "unknown method %s"+e.method) e.stateMux.Unlock() + close(e.executionDone) } } @@ -241,15 +252,21 @@ func (e *ServerRequest) Evictable(minRetention time.Duration) bool { } func (e *ServerRequest) Cancel(ctx context.Context, err types.Error, msg string) error { - // If execution has been claimed, the executor will set the response when it - // finishes. Cancelling at that point would race with the executor's result. + // If execution has been claimed, wait for it to complete so we can overwrite + // the response with the cancellation error. if e.executionClaimed.Load() { - return nil + select { + case <-e.executionDone: + // Execution finished, safe to overwrite response + case <-ctx.Done(): + return ctx.Err() + } } e.stateMux.Lock() defer e.stateMux.Unlock() + // Only set cancellation error if no response exists (matches original behavior) if !e.hasResponse() { e.setError(err, msg) if err := e.sendResponses(ctx); err != nil { @@ -279,6 +296,8 @@ func (e *ServerRequest) executeRequest(ctx context.Context, msg *types.MessageBo } e.stateMux.Unlock() + close(e.executionDone) + e.metrics.countExecution(ctx, success) e.metrics.recordExecutionDuration(ctx, time.Since(start), success) } From eb8f9de5c2cef32b6721dd8bf1e60942f2ba3217 Mon Sep 17 00:00:00 2001 From: yashnevatia Date: Wed, 12 Aug 2026 18:13:03 +0100 Subject: [PATCH 3/6] fix comment --- core/capabilities/remote/executable/request/server_request.go | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/core/capabilities/remote/executable/request/server_request.go b/core/capabilities/remote/executable/request/server_request.go index efccc459c7d..cd4400d5558 100644 --- a/core/capabilities/remote/executable/request/server_request.go +++ b/core/capabilities/remote/executable/request/server_request.go @@ -145,9 +145,7 @@ type ServerRequest struct { // executionDone is closed by the executor after it writes the response under stateMux. // Cancel waits on this channel to avoid a race where it could acquire stateMux // between "executor finishes capability call" and "executor writes response", - // see no response, and set a timeout error. The executor would then overwrite - // with success, but all requesters would already be marked as sent — leaving - // them with the timeout error while the success response is lost. + // see no response, and set and send timeout error when actually we want to send the real response. executionDone chan struct{} lggr logger.Logger From 4e2cfed4b8c784abb83e0c28faba228e282c147b Mon Sep 17 00:00:00 2001 From: yashnevatia Date: Wed, 19 Aug 2026 14:11:12 +0100 Subject: [PATCH 4/6] move executeReq upwards and cancel it if Cancel() is called --- .../executable/request/server_request.go | 86 +++++++++---------- 1 file changed, 41 insertions(+), 45 deletions(-) diff --git a/core/capabilities/remote/executable/request/server_request.go b/core/capabilities/remote/executable/request/server_request.go index cd4400d5558..918c6a30600 100644 --- a/core/capabilities/remote/executable/request/server_request.go +++ b/core/capabilities/remote/executable/request/server_request.go @@ -134,7 +134,7 @@ type ServerRequest struct { // Metadata.WorkflowDonID to match the authenticated calling DON. workflowDONBindingGate limits.GateLimiter - // stateMux guards requesters, responseSentToRequester, and response. + // stateMux guards requesters, responseSentToRequester, response, and executionCancel. // It is held only for short map/field operations, never during capability execution. stateMux sync.Mutex @@ -142,11 +142,10 @@ type ServerRequest struct { // wins the right to execute the capability. All other messages skip execution. executionClaimed atomic.Bool - // executionDone is closed by the executor after it writes the response under stateMux. - // Cancel waits on this channel to avoid a race where it could acquire stateMux - // between "executor finishes capability call" and "executor writes response", - // see no response, and set and send timeout error when actually we want to send the real response. - executionDone chan struct{} + // executionCancel cancels the in-flight capability execution context. + // Set under stateMux when execution starts; called by Cancel to stop + // execution early (e.g. on request expiry) instead of waiting for completion. + executionCancel context.CancelFunc lggr logger.Logger @@ -184,7 +183,6 @@ func NewServerRequest(capability capabilities.ExecutableCapability, method strin requestTimeout: requestTimeout, capMethodName: capMethodName, workflowDONBindingGate: workflowDONBindingGate, - executionDone: make(chan struct{}), lggr: lggr, metrics: m, }, nil @@ -222,12 +220,36 @@ func (e *ServerRequest) OnMessage(ctx context.Context, msg *types.MessageBody) e if quorumReached && !hasResponse && e.executionClaimed.CompareAndSwap(false, true) { switch e.method { case types.MethodExecute: - e.executeRequest(ctx, msg, executeCapabilityRequest) + ctxWithTimeout, cancel := context.WithTimeout(ctx, e.requestTimeout) + defer cancel() + + // Expose the cancel func so Cancel can stop the in-flight execution early. + e.stateMux.Lock() + e.executionCancel = cancel + e.stateMux.Unlock() + success := false + start := time.Now() + responsePayload, responseErr := executeCapabilityRequest(ctxWithTimeout, e.lggr, e.capability, msg.Payload, e.callingDon.ID, e.workflowDONBindingGate) + + e.stateMux.Lock() + // Cancel may have already set a timeout error response; never overwrite an + // existing response. + if !e.hasResponse() { + if responseErr != nil { + e.setError(types.Error_INTERNAL_ERROR, responseErr.Error()) + } else { + success = true + e.setResult(responsePayload) + } + } + e.stateMux.Unlock() + + e.metrics.countExecution(ctxWithTimeout, success) + e.metrics.recordExecutionDuration(ctxWithTimeout, time.Since(start), success) default: e.stateMux.Lock() e.setError(types.Error_INTERNAL_ERROR, "unknown method %s"+e.method) e.stateMux.Unlock() - close(e.executionDone) } } @@ -249,21 +271,20 @@ func (e *ServerRequest) Evictable(minRetention time.Duration) bool { return age > e.requestTimeout && age > minRetention } +// Cancel stops any in-flight execution by cancelling its context and, if no +// response has been produced yet, records err as the response and fans it out +// to all requesters. func (e *ServerRequest) Cancel(ctx context.Context, err types.Error, msg string) error { - // If execution has been claimed, wait for it to complete so we can overwrite - // the response with the cancellation error. - if e.executionClaimed.Load() { - select { - case <-e.executionDone: - // Execution finished, safe to overwrite response - case <-ctx.Done(): - return ctx.Err() - } - } - e.stateMux.Lock() defer e.stateMux.Unlock() + // Cancel the in-flight execution, if any. The executor goroutine returns + // early on the cancelled context and skips overwriting the response set + // below (guarded by hasResponse). + if e.executionCancel != nil { + e.executionCancel() + } + // Only set cancellation error if no response exists (matches original behavior) if !e.hasResponse() { e.setError(err, msg) @@ -275,31 +296,6 @@ func (e *ServerRequest) Cancel(ctx context.Context, err types.Error, msg string) return nil } -type executeFn func(ctx context.Context, lggr logger.Logger, capability capabilities.ExecutableCapability, payload []byte, callingDonID uint32, workflowDONBindingGate limits.GateLimiter) ([]byte, error) - -func (e *ServerRequest) executeRequest(ctx context.Context, msg *types.MessageBody, method executeFn) { - ctxWithTimeout, cancel := context.WithTimeout(ctx, e.requestTimeout) - defer cancel() - - success := false - start := time.Now() - responsePayload, err := method(ctxWithTimeout, e.lggr, e.capability, msg.Payload, e.callingDon.ID, e.workflowDONBindingGate) - - e.stateMux.Lock() - if err != nil { - e.setError(types.Error_INTERNAL_ERROR, err.Error()) - } else { - success = true - e.setResult(responsePayload) - } - e.stateMux.Unlock() - - close(e.executionDone) - - e.metrics.countExecution(ctx, success) - e.metrics.recordExecutionDuration(ctx, time.Since(start), success) -} - func (e *ServerRequest) addRequester(from p2ptypes.PeerID) error { fromPeerInCallingDon := slices.Contains(e.callingDon.Members, from) From c8f66a936233b9127326c87b9d6b4ef0393d8572 Mon Sep 17 00:00:00 2001 From: yashnevatia Date: Thu, 20 Aug 2026 18:42:33 +0100 Subject: [PATCH 5/6] Add test --- .../executable/request/server_request_test.go | 195 ++++++++++++++++++ 1 file changed, 195 insertions(+) diff --git a/core/capabilities/remote/executable/request/server_request_test.go b/core/capabilities/remote/executable/request/server_request_test.go index 3c8c376d159..44ab1a5d80b 100644 --- a/core/capabilities/remote/executable/request/server_request_test.go +++ b/core/capabilities/remote/executable/request/server_request_test.go @@ -4,6 +4,8 @@ import ( "context" "crypto/rand" "errors" + "sync" + "sync/atomic" "testing" "time" @@ -362,6 +364,159 @@ func Test_ServerRequest_MessageValidation(t *testing.T) { }) } +func Test_ServerRequest_SingleExecutionNonBlocking(t *testing.T) { + t.Parallel() + + lggr := logger.Test(t) + capabilityPeerID := NewP2PPeerID(t) + + numWorkflowPeers := 2 + workflowPeers := make([]p2ptypes.PeerID, numWorkflowPeers) + for i := range numWorkflowPeers { + workflowPeers[i] = NewP2PPeerID(t) + } + + callingDon := commoncap.DON{ + Members: workflowPeers, + ID: 1, + F: 0, // we want wf peer 0 to execute and peer 1 to return immediately + } + + executeInputs, err := values.NewMap( + map[string]any{ + "executeValue1": "aValue1", + }, + ) + require.NoError(t, err) + + capabilityRequest := commoncap.CapabilityRequest{ + Metadata: commoncap.RequestMetadata{ + WorkflowID: "workflowID", + WorkflowExecutionID: "workflowExecutionID", + }, + Inputs: executeInputs, + } + + rawRequest, err := pb.MarshalCapabilityRequest(capabilityRequest) + require.NoError(t, err) + + t.Run("Execute capability", func(t *testing.T) { + t.Parallel() + + started := make(chan struct{}) + release := make(chan struct{}) + capability := BlockingTestCapability{counter: &atomic.Int32{}, started: started, release: release} + // need a concurrent dispatcher here as we call onMessage for wf peer 0 concurrently + // before this, all onMessage calls were sequential + dispatcher := &concurrentTestDispatcher{} + req, err := request.NewServerRequest(capability, types.MethodExecute, "capabilityID", 2, + capabilityPeerID, callingDon, "requestMessageID", dispatcher, 10*time.Minute, "", limits.NewGateLimiter(false), lggr) + require.NoError(t, err) + + errCh := make(chan error, 1) + + // wf peer 0 calls onMessage + // we have to do this in a goroutine because the capability blocks until we close the release channel + go func() { + errCh <- sendValidRequest( + req, + workflowPeers, + capabilityPeerID, + rawRequest, + ) + }() + + // wait on started to prove that wf peer 0 has started executing the capability + <-started + + // wf peer 1 calls onMessage + err = req.OnMessage(context.Background(), &types.MessageBody{ + Version: 0, + Sender: workflowPeers[1][:], + Receiver: capabilityPeerID[:], + MessageId: []byte("workflowID" + "workflowExecutionID"), + CapabilityId: "capabilityID", + CapabilityDonId: 2, + CallerDonId: 1, + Method: types.MethodExecute, + Payload: rawRequest, + }) + // wf peer 1 should immediately return as its not blocked + require.NoError(t, err) + + t.Log("WF Peer 1 has returned. Closing channel now. This will let capability execution finish and WF Peer 0 can finish.") + close(release) + // ensure wf peer 0 returns without error + errA := <-errCh + require.NoError(t, errA) + + // ensure that the capability was only executed once and that both workflow peers received a response + require.Equal(t, int32(1), capability.counter.Load()) + assert.Len(t, dispatcher.msgs, 2) + assert.Equal(t, types.Error_OK, dispatcher.msgs[0].Error) + assert.Equal(t, types.Error_OK, dispatcher.msgs[1].Error) + }) + + t.Run("Cancel Execute capability", func(t *testing.T) { + t.Parallel() + + started := make(chan struct{}) + release := make(chan struct{}) + capability := BlockingTestCapability{counter: &atomic.Int32{}, started: started, release: release} + // need a concurrent dispatcher here as we call onMessage for wf peer 0 concurrently + // before this, all onMessage calls were sequential + dispatcher := &concurrentTestDispatcher{} + req, err := request.NewServerRequest(capability, types.MethodExecute, "capabilityID", 2, + capabilityPeerID, callingDon, "requestMessageID", dispatcher, 10*time.Minute, "", limits.NewGateLimiter(false), lggr) + require.NoError(t, err) + + errCh := make(chan error, 1) + + // wf peer 0 calls onMessage + // we have to do this in a goroutine because the capability blocks until we close the release channel + go func() { + errCh <- sendValidRequest( + req, + workflowPeers, + capabilityPeerID, + rawRequest, + ) + }() + + // wait on started to prove that wf peer 0 has started executing the capability + <-started + + // wf peer 1 calls onMessage + err = req.OnMessage(context.Background(), &types.MessageBody{ + Version: 0, + Sender: workflowPeers[1][:], + Receiver: capabilityPeerID[:], + MessageId: []byte("workflowID" + "workflowExecutionID"), + CapabilityId: "capabilityID", + CapabilityDonId: 2, + CallerDonId: 1, + Method: types.MethodExecute, + Payload: rawRequest, + }) + // wf peer 1 should immediately return as its not blocked + require.NoError(t, err) + + // cancel request before wf peer 0 finishes executing the capability + req.Cancel(t.Context(), types.Error_TIMEOUT, "cancelled by test") + + // wf peer 0 returns without error + errA := <-errCh + require.NoError(t, errA) + + // ensure that the capability was only executed once + // ensure that both workflow peers received a timeout error + require.Equal(t, int32(1), capability.counter.Load()) + assert.Len(t, dispatcher.msgs, 2) + assert.Equal(t, types.Error_TIMEOUT, dispatcher.msgs[0].Error) + assert.Equal(t, types.Error_TIMEOUT, dispatcher.msgs[1].Error) + }) +} + func Test_ServerRequest_Evictable(t *testing.T) { t.Parallel() @@ -466,6 +621,18 @@ func (t *testDispatcher) Send(peerID p2ptypes.PeerID, msgBody *types.MessageBody return nil } +type concurrentTestDispatcher struct { + testDispatcher + mu sync.Mutex +} + +func (t *concurrentTestDispatcher) Send(peerID p2ptypes.PeerID, msgBody *types.MessageBody) error { + t.mu.Lock() + defer t.mu.Unlock() + t.msgs = append(t.msgs, msgBody) + return nil +} + type abstractTestCapability struct { } @@ -498,6 +665,34 @@ func (t TestCapability) Execute(ctx context.Context, request commoncap.Capabilit }, nil } +type BlockingTestCapability struct { + abstractTestCapability + counter *atomic.Int32 + started chan struct{} + release chan struct{} +} + +func (t BlockingTestCapability) Execute(ctx context.Context, request commoncap.CapabilityRequest) (commoncap.CapabilityResponse, error) { + close(t.started) + value := request.Inputs.Underlying["executeValue1"] + + response, err := values.NewMap(map[string]any{"response": value}) + if err != nil { + return commoncap.CapabilityResponse{}, err + } + t.counter.Add(1) + + // Block until we're released. + select { + case <-t.release: + case <-ctx.Done(): + return commoncap.CapabilityResponse{}, ctx.Err() + } + return commoncap.CapabilityResponse{ + Value: response, + }, nil +} + type TestErrorCapability struct { abstractTestCapability err error From b03c1cf5d3418e1d610aef903c1e4c49cd18c2f4 Mon Sep 17 00:00:00 2001 From: yashnevatia Date: Fri, 21 Aug 2026 14:55:23 +0100 Subject: [PATCH 6/6] lint --- .../remote/executable/request/server_request_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/capabilities/remote/executable/request/server_request_test.go b/core/capabilities/remote/executable/request/server_request_test.go index 44ab1a5d80b..2bc3e7d9561 100644 --- a/core/capabilities/remote/executable/request/server_request_test.go +++ b/core/capabilities/remote/executable/request/server_request_test.go @@ -502,7 +502,8 @@ func Test_ServerRequest_SingleExecutionNonBlocking(t *testing.T) { require.NoError(t, err) // cancel request before wf peer 0 finishes executing the capability - req.Cancel(t.Context(), types.Error_TIMEOUT, "cancelled by test") + cancelErr := req.Cancel(t.Context(), types.Error_TIMEOUT, "cancelled by test") + require.NoError(t, cancelErr) // wf peer 0 returns without error errA := <-errCh