diff --git a/core/capabilities/remote/executable/request/server_request.go b/core/capabilities/remote/executable/request/server_request.go index a76ffb07f8c..918c6a30600 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,19 @@ type ServerRequest struct { // Metadata.WorkflowDonID to match the authenticated calling DON. workflowDONBindingGate limits.GateLimiter - mux sync.Mutex + // stateMux guards requesters, responseSentToRequester, response, and executionCancel. + // 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 + + // 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 metrics *srMetrics @@ -178,9 +191,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 +200,61 @@ 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() + 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, "calls", calls) - 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) + 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() } } + e.stateMux.Lock() + defer e.stateMux.Unlock() if err := e.sendResponses(ctx); err != nil { return fmt.Errorf("failed to send responses: %w", err) } @@ -222,10 +271,21 @@ 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 { - e.mux.Lock() - defer e.mux.Unlock() + 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) if err := e.sendResponses(ctx); err != nil { @@ -236,26 +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) - if err != nil { - e.setError(types.Error_INTERNAL_ERROR, err.Error()) - } else { - success = true - e.setResult(responsePayload) - } - - 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) diff --git a/core/capabilities/remote/executable/request/server_request_test.go b/core/capabilities/remote/executable/request/server_request_test.go index 3c8c376d159..2bc3e7d9561 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,160 @@ 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 + cancelErr := req.Cancel(t.Context(), types.Error_TIMEOUT, "cancelled by test") + require.NoError(t, cancelErr) + + // 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 +622,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 +666,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