Release executor slot if not executing

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)
 

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