fix(runtime): support scoped cancellation acknowledgments

.github/workflows/rpc-cancellation.yml

@@ -0,0 +1,50 @@
+name: Scoped RPC cancellation
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+    paths:
+      - .github/workflows/rpc-cancellation.yml
+      - src/Orleans.Core/Runtime/CancellationAcknowledgementScope.cs
+      - src/Orleans.Core/Runtime/CallbackData.cs
+      - src/Orleans.Core/Runtime/GrainReferenceRuntime.cs
+      - src/Orleans.Core/Runtime/IRuntimeClient.cs
+      - src/Orleans.Core/Runtime/OutgoingCallInvoker.cs
+      - src/Orleans.Core/Runtime/OutsideRuntimeClient.cs
+      - src/Orleans.Runtime/Core/InsideRuntimeClient.cs
+      - test/Orleans.Runtime.Tests/CallbackDataTests.cs
+      - test/Orleans.Runtime.Tests/CancellationTests/**
+permissions:
+  contents: read
+jobs:
+  scoped-cancellation:
+    name: Scoped cancellation (${{ matrix.framework }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    strategy:
+      fail-fast: false
+      matrix:
+        framework: [net8.0, net10.0]
+    steps:
+      - uses: actions/checkout@93cb6efe18208431cddfb8368fd83d5badbf9bfd # v5
+      - uses: actions/setup-dotnet@c2fa09f4bde5ebb9d1777cf28262a3eb3db3ced7 # v5
+        with:
+          global-json-file: global.json
+          dotnet-version: 8.0.x
+      - name: Run focused callback and generated-proxy regressions
+        shell: pwsh
+        run: >-
+          dotnet test --project test/Orleans.Runtime.Tests/Orleans.Runtime.Tests.csproj
+          --framework "${{ matrix.framework }}"
+          --filter-query "/*/Tester/(CallbackDataTests)|(CancellationAcknowledgementScopeTests)|(ScopedCancellationRpcTests)"
+          --minimum-expected-tests 29
+          --report-trx
+          --report-trx-filename "scoped-cancellation-${{ matrix.framework }}.trx"
+      - uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
+        if: always()
+        with:
+          name: scoped-cancellation-${{ matrix.framework }}
+          retention-days: 7
+          path: |
+            **/TestResults/*.trx

src/Orleans.Core/Runtime/CallbackData.cs

@@ -16,6 +16,7 @@ internal sealed partial class CallbackData
         private readonly IResponseCompletionSource context;
         private readonly ApplicationRequestInstruments _applicationRequestInstruments;
         private readonly long _startTimestamp;
+        private readonly bool _waitForCancellationAcknowledgement;
         private int _state;
         private StatusResponse? lastKnownStatus;
         private CancellationTokenRegistration _cancellationTokenRegistration;
@@ -24,13 +25,15 @@ public CallbackData(
             SharedCallbackData shared,
             IResponseCompletionSource ctx,
             Message msg,
-            ApplicationRequestInstruments applicationRequestInstruments)
+            ApplicationRequestInstruments applicationRequestInstruments,
+            bool waitForCancellationAcknowledgement = false)
         {
             this.shared = shared;
             this.context = ctx;
             this.Message = msg;
             _applicationRequestInstruments = applicationRequestInstruments;
             _startTimestamp = shared.TimeProvider.GetTimestamp();
+            _waitForCancellationAcknowledgement = shared.WaitForCancellationAcknowledgement || waitForCancellationAcknowledgement;
         }
 
         public Message Message { get; } // might hold metadata used by response pipeline
@@ -113,7 +116,7 @@ private void OnCancellation(CancellationToken cancellationToken)
         {
             // If waiting for acknowledgement is enabled, simply signal to the remote grain that cancellation
             // is requested and return.
-            if (shared.WaitForCancellationAcknowledgement)
+            if (_waitForCancellationAcknowledgement)
             {
                 SignalCancellation();
                 return;

src/Orleans.Core/Runtime/CancellationAcknowledgementScope.cs

@@ -0,0 +1,47 @@
+using System;
+using System.Threading;
+
+namespace Orleans.Runtime;
+
+/// <summary>
+/// Keeps calls initiated in this scope pending after caller cancellation until a response or another
+/// terminal runtime outcome, including the configured response timeout or host shutdown.
+/// </summary>
+/// <remarks>
+/// Enclose creation of the proxy task in <see cref="Enter"/> and dispose the scope before awaiting an
+/// independently bounded local wait. Cancellation still signals the callee. Pending requests describe
+/// the caller's observable lifetime; remote execution can continue after a runtime timeout.
+/// The scope flows with the execution context and restores the previous value when disposed in lexical
+/// order. The invocation boundary captures and suppresses it before running outgoing filters, so each
+/// filter's additional calls select their own policy.
+/// </remarks>
+internal readonly struct CancellationAcknowledgementScope : IDisposable
+{
+    private static readonly AsyncLocal<bool> Current = new();
+    private readonly bool _previous;
+
+    private CancellationAcknowledgementScope(bool enabled)
+    {
+        _previous = Current.Value;
+        if (_previous != enabled)
+        {
+            Current.Value = enabled;
+        }
+    }
+
+    /// <summary>Enables cancellation acknowledgement waiting for calls initiated in this scope.</summary>
+    public static CancellationAcknowledgementScope Enter() => new(enabled: true);
+
+    /// <summary>Captures the caller's policy and clears it for invocation infrastructure.</summary>
+    internal static CancellationAcknowledgementScope Capture() => new(enabled: false);
+
+    internal bool WaitForAcknowledgement => _previous;
+
+    public void Dispose()
+    {
+        if (Current.Value != _previous)
+        {
+            Current.Value = _previous;
+        }
+    }
+}

src/Orleans.Core/Runtime/GrainReferenceRuntime.cs

@@ -16,7 +16,7 @@ internal class GrainReferenceRuntime : IGrainReferenceRuntime
         private readonly GrainInterfaceTypeResolver interfaceTypeResolver;
         private readonly IGrainCancellationTokenRuntime cancellationTokenRuntime;
         private readonly IOutgoingGrainCallFilter[] filters;
-        private readonly Action<GrainReference, IResponseCompletionSource, IInvokable, InvokeMethodOptions> sendRequest;
+        private readonly Action<GrainReference, IResponseCompletionSource, IInvokable, InvokeMethodOptions, bool> sendRequest;
 
         public GrainReferenceRuntime(
             IRuntimeClient runtimeClient,
@@ -30,71 +30,75 @@ public GrainReferenceRuntime(
             this.referenceActivator = referenceActivator;
             this.interfaceTypeResolver = interfaceTypeResolver;
             this.filters = outgoingCallFilters.ToArray();
-            this.sendRequest = (GrainReference reference, IResponseCompletionSource callback, IInvokable body, InvokeMethodOptions options) => RuntimeClient.SendRequest(reference, body, callback, options);
+            this.sendRequest = (reference, callback, body, options, waitForCancellationAcknowledgement) =>
+                RuntimeClient.SendRequest(reference, body, callback, options, waitForCancellationAcknowledgement);
         }
 
         public IRuntimeClient RuntimeClient { get; private set; }
 
         public ValueTask<TResult?> InvokeMethodAsync<TResult>(GrainReference reference, IInvokable request, InvokeMethodOptions options)
         {
+            using var cancellationScope = CancellationAcknowledgementScope.Capture();
             // TODO: Remove expensive interface type check
             if (this.filters.Length == 0 && request is not IOutgoingGrainCallFilter)
             {
                 SetGrainCancellationTokensTarget(reference, request);
                 var responseCompletionSource = ResponseCompletionSourcePool.Get<TResult>();
-                this.RuntimeClient.SendRequest(reference, request, responseCompletionSource, options);
+                this.RuntimeClient.SendRequest(reference, request, responseCompletionSource, options, cancellationScope.WaitForAcknowledgement);
                 return responseCompletionSource.AsValueTask();
             }
             else
             {
-                return InvokeMethodWithFiltersAsync<TResult>(reference, request, options);
+                return InvokeMethodWithFiltersAsync<TResult>(reference, request, options, cancellationScope.WaitForAcknowledgement);
             }
         }
 
         public ValueTask InvokeMethodAsync(GrainReference reference, IInvokable request, InvokeMethodOptions options)
         {
+            using var cancellationScope = CancellationAcknowledgementScope.Capture();
             // TODO: Remove expensive interface type check
             if (filters.Length == 0 && request is not IOutgoingGrainCallFilter)
             {
                 SetGrainCancellationTokensTarget(reference, request);
                 var responseCompletionSource = ResponseCompletionSourcePool.Get();
-                this.RuntimeClient.SendRequest(reference, request, responseCompletionSource, options);
+                this.RuntimeClient.SendRequest(reference, request, responseCompletionSource, options, cancellationScope.WaitForAcknowledgement);
                 return responseCompletionSource.AsVoidValueTask();
             }
             else
             {
-                return InvokeMethodWithFiltersAsync(reference, request, options);
+                return InvokeMethodWithFiltersAsync(reference, request, options, cancellationScope.WaitForAcknowledgement);
             }
         }
 
         public void InvokeMethod(GrainReference reference, IInvokable request, InvokeMethodOptions options)
         {
+            using var cancellationScope = CancellationAcknowledgementScope.Capture();
             Debug.Assert((options & InvokeMethodOptions.OneWay) != 0);
 
             // TODO: Remove expensive interface type check
             if (filters.Length == 0 && request is not IOutgoingGrainCallFilter)
             {
                 SetGrainCancellationTokensTarget(reference, request);
-                this.RuntimeClient.SendRequest(reference, request, context: null, options);
+                this.RuntimeClient.SendRequest(reference, request, context: null, options, cancellationScope.WaitForAcknowledgement);
             }
             else
             {
-                InvokeMethodWithFiltersAsync(reference, request, options).AsTask().Ignore();
+                InvokeMethodWithFiltersAsync(reference, request, options, cancellationScope.WaitForAcknowledgement).AsTask().Ignore();
             }
         }
 
-        private async ValueTask<TResult?> InvokeMethodWithFiltersAsync<TResult>(GrainReference reference, IInvokable request, InvokeMethodOptions options)
+        private async ValueTask<TResult?> InvokeMethodWithFiltersAsync<TResult>(GrainReference reference, IInvokable request, InvokeMethodOptions options, bool waitForCancellationAcknowledgement)
         {
             SetGrainCancellationTokensTarget(reference, request);
-            var invoker = new OutgoingCallInvoker<TResult>(reference, request, options, this.sendRequest, this.filters);
+            var invoker = new OutgoingCallInvoker<TResult>(reference, request, options, this.sendRequest, this.filters, waitForCancellationAcknowledgement);
             await invoker.Invoke();
             return invoker.TypedResult;
         }
 
-        private async ValueTask InvokeMethodWithFiltersAsync(GrainReference reference, IInvokable request, InvokeMethodOptions options)
+        private async ValueTask InvokeMethodWithFiltersAsync(GrainReference reference, IInvokable request, InvokeMethodOptions options, bool waitForCancellationAcknowledgement)
         {
             SetGrainCancellationTokensTarget(reference, request);
-            var invoker = new OutgoingCallInvoker<object>(reference, request, options, this.sendRequest, this.filters);
+            var invoker = new OutgoingCallInvoker<object>(reference, request, options, this.sendRequest, this.filters, waitForCancellationAcknowledgement);
             await invoker.Invoke();
         }
 

src/Orleans.Core/Runtime/IRuntimeClient.cs

@@ -44,7 +44,7 @@ internal interface IRuntimeClient
         /// <param name="timeout">New response timeout value</param>
         void SetResponseTimeout(TimeSpan timeout);
 
-        void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource? context, InvokeMethodOptions options);
+        void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource? context, InvokeMethodOptions options, bool waitForCancellationAcknowledgement);
 
         void SendResponse(Message request, Response response);
 

src/Orleans.Core/Runtime/OutgoingCallInvoker.cs

@@ -13,7 +13,8 @@ internal sealed class OutgoingCallInvoker<TResult> : IOutgoingGrainCallContext
     {
         private readonly IInvokable request;
         private readonly InvokeMethodOptions options;
-        private readonly Action<GrainReference, IResponseCompletionSource, IInvokable, InvokeMethodOptions> sendRequest;
+        private readonly Action<GrainReference, IResponseCompletionSource, IInvokable, InvokeMethodOptions, bool> sendRequest;
+        private readonly bool _waitForCancellationAcknowledgement;
         private readonly IOutgoingGrainCallFilter[] filters;
         private readonly int stages;
         private readonly GrainReference grainReference;
@@ -28,18 +29,21 @@ internal sealed class OutgoingCallInvoker<TResult> : IOutgoingGrainCallContext
         /// <param name="options"></param>
         /// <param name="sendRequest"></param>
         /// <param name="filters">The invocation interceptors.</param>
+        /// <param name="waitForCancellationAcknowledgement">The cancellation policy captured at invocation entry.</param>
         public OutgoingCallInvoker(
             GrainReference grain,
             IInvokable request,
             InvokeMethodOptions options,
-            Action<GrainReference, IResponseCompletionSource, IInvokable, InvokeMethodOptions> sendRequest,
-            IOutgoingGrainCallFilter[] filters)
+            Action<GrainReference, IResponseCompletionSource, IInvokable, InvokeMethodOptions, bool> sendRequest,
+            IOutgoingGrainCallFilter[] filters,
+            bool waitForCancellationAcknowledgement)
         {
             this.request = request;
             this.options = options;
             this.sendRequest = sendRequest;
             this.grainReference = grain;
             this.filters = filters;
+            _waitForCancellationAcknowledgement = waitForCancellationAcknowledgement;
             this.stages = filters.Length;
             SourceContext = RuntimeContext.Current;
 
@@ -113,7 +117,7 @@ public async Task Invoke()
                     // Finally call the root-level invoker.
                     stage++;
                     var responseCompletionSource = ResponseCompletionSourcePool.Get();
-                    this.sendRequest(this.grainReference, responseCompletionSource, this.request, this.options);
+                    this.sendRequest(this.grainReference, responseCompletionSource, this.request, this.options, _waitForCancellationAcknowledgement);
                     this.Response = await responseCompletionSource.AsValueTask().ConfigureAwait(false);
 
                     return;

src/Orleans.Core/Runtime/OutsideRuntimeClient.cs

@@ -265,7 +265,7 @@ public void SendResponse(Message request, Response response)
             MessageCenter!.SendMessage(message);
         }
 
-        public void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource? context, InvokeMethodOptions options)
+        public void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource? context, InvokeMethodOptions options, bool waitForCancellationAcknowledgement)
         {
             ThrowIfDisposed();
             var cancellationToken = request.GetCancellationToken();
@@ -295,7 +295,7 @@ public void SendRequest(GrainReference target, IInvokable request, IResponseComp
 
             if (!oneWay)
             {
-                var callbackData = new CallbackData(this.sharedCallbackData, context!, message, _applicationRequestInstruments);
+                var callbackData = new CallbackData(this.sharedCallbackData, context!, message, _applicationRequestInstruments, waitForCancellationAcknowledgement);
                 if (Volatile.Read(ref _isStopping) != 0)
                 {
                     callbackData.OnHostShutdown();

src/Orleans.Runtime/Core/InsideRuntimeClient.cs

@@ -138,7 +138,8 @@ public void SendRequest(
             GrainReference target,
             IInvokable request,
             IResponseCompletionSource? context,
-            InvokeMethodOptions options)
+            InvokeMethodOptions options,
+            bool waitForCancellationAcknowledgement)
         {
             var cancellationToken = request.GetCancellationToken();
             cancellationToken.ThrowIfCancellationRequested();
@@ -190,7 +191,7 @@ public void SendRequest(
                 Debug.Assert(context is not null);
 
                 // Register a callback for the request.
-                callbackData = new CallbackData(sharedData, context, message, _applicationRequestInstruments);
+                callbackData = new CallbackData(sharedData, context, message, _applicationRequestInstruments, waitForCancellationAcknowledgement);
                 if (Volatile.Read(ref _isStopping) != 0)
                 {
                     callbackData.OnHostShutdown();