fix(runtime): guard pooled observer cancellation

Author: ReubenBond

src/Orleans.CodeGenerator/InvokableGenerator.cs

@@ -911,8 +911,7 @@ static bool IsPoolableBaseType(INamedTypeSymbol type)
                 && type.MetadataName is "Request"
                     or "Request`1"
                     or "TaskRequest"
-                    or "TaskRequest`1"
-                    or "VoidRequest";
+                    or "TaskRequest`1";
     }
 
     internal abstract class InvokerFieldDescription(ITypeSymbol fieldType, string fieldName)

src/Orleans.Core/Runtime/InvokableObjectManager.cs

@@ -90,6 +90,8 @@ public sealed partial class LocalObjectData : IGrainContext, IGrainCallCancellat
             private static readonly Func<object?, Task> HandleFunc = self => ((LocalObjectData)self!).LocalObjectMessagePumpAsync();
             private readonly InvokableObjectManager _manager;
             private readonly Dictionary<Message, Task?> _runningRequests = [];
+            private readonly Dictionary<Message, int> _cancellationLeaseCounts = [];
+            private readonly HashSet<Message> _completedRequests = [];
             private Task? _messagePumpTask;
 
             internal LocalObjectData(IAddressable obj, ObserverGrainId observerId, InvokableObjectManager manager)
@@ -324,12 +326,23 @@ private async Task ProcessMessageAsync(Message message)
                 }
                 finally
                 {
-                    message.DisposeOwnedBody();
-
-                    // Clear the running request when done.
+                    var disposeMessage = false;
                     lock (Messages)
                     {
                         _runningRequests.Remove(message);
+                        if (_cancellationLeaseCounts.ContainsKey(message))
+                        {
+                            _completedRequests.Add(message);
+                        }
+                        else
+                        {
+                            disposeMessage = true;
+                        }
+                    }
+
+                    if (disposeMessage)
+                    {
+                        message.DisposeOwnedBody();
                     }
                 }
             }
@@ -437,6 +450,7 @@ bool TryCancelRequest()
                 {
                     Message? message = null;
                     var wasWaiting = false;
+                    var hasCancellationLease = false;
                     lock (Messages)
                     {
                         // Check the running requests.
@@ -445,6 +459,9 @@ bool TryCancelRequest()
                             if (runningRequest.Id == messageId && runningRequest.SendingGrain == senderGrainId)
                             {
                                 message = runningRequest;
+                                _cancellationLeaseCounts.TryGetValue(message, out var leaseCount);
+                                _cancellationLeaseCounts[message] = leaseCount + 1;
+                                hasCancellationLease = true;
                                 break;
                             }
                         }
@@ -478,34 +495,67 @@ bool TryCancelRequest()
                     }
 
                     var didCancel = false;
-                    if (message is not null)
+                    try
                     {
-                        // The message never began executing, so send a canceled response immediately.
-                        // If the message did begin executing, wait for it to observe the cancellation token and respond itself.
-                        if (wasWaiting)
+                        if (message is not null)
                         {
-                            try
+                            // The message never began executing, so send a canceled response immediately.
+                            // If the message did begin executing, wait for it to observe the cancellation token and respond itself.
+                            if (wasWaiting)
                             {
-                                _manager.runtimeClient.SendResponse(message, Response.FromException(new OperationCanceledException()));
-                                didCancel = true;
+                                try
+                                {
+                                    _manager.runtimeClient.SendResponse(message, Response.FromException(new OperationCanceledException()));
+                                    didCancel = true;
+                                }
+                                finally
+                                {
+                                    message.DisposeOwnedBody();
+                                }
+                            }
+                            else if (message.BodyObject is IInvokable invokableRequest)
+                            {
+                                didCancel = TryCancelInvokable(invokableRequest) || !invokableRequest.IsCancellable;
                             }
-                            finally
+                            else
                             {
-                                message.DisposeOwnedBody();
+                                // Assume the request is not cancellable.
+                                didCancel = true;
                             }
                         }
-                        else if (message.BodyObject is IInvokable invokableRequest)
+                    }
+                    finally
+                    {
+                        if (hasCancellationLease)
                         {
-                            didCancel = TryCancelInvokable(invokableRequest) || !invokableRequest.IsCancellable;
+                            ReleaseCancellationLease(message!);
+                        }
+                    }
+
+                    return didCancel;
+                }
+
+                void ReleaseCancellationLease(Message message)
+                {
+                    var disposeMessage = false;
+                    lock (Messages)
+                    {
+                        var leaseCount = _cancellationLeaseCounts[message];
+                        if (leaseCount == 1)
+                        {
+                            _cancellationLeaseCounts.Remove(message);
+                            disposeMessage = _completedRequests.Remove(message);
                         }
                         else
                         {
-                            // Assume the request is not cancellable.
-                            didCancel = true;
+                            _cancellationLeaseCounts[message] = leaseCount - 1;
                         }
                     }
 
-                    return didCancel;
+                    if (disposeMessage)
+                    {
+                        message.DisposeOwnedBody();
+                    }
                 }
 
                 bool TryCancelInvokable(IInvokable request)

src/Orleans.Serialization/Invocation/Pools/InvokablePool.cs

@@ -10,6 +10,10 @@ namespace Orleans.Serialization.Invocation;
 /// Provides bounded thread-local reuse for <see cref="IInvokable"/> implementations.
 /// </summary>
 /// <typeparam name="T">The invokable type.</typeparam>
+/// <remarks>
+/// A returned instance becomes available to one subsequent rental on the current thread.
+/// Callers reset mutable state before returning an instance and transfer exclusive ownership to the pool.
+/// </remarks>
 public sealed class InvokablePool<T> : IDisposable where T : class, IInvokable
 {
     private const int MaxPoolSizePerThread = 128;
@@ -49,7 +53,7 @@ public bool TryGet([NotNullWhen(true)] out T? item)
     /// <summary>
     /// Makes an instance available for reuse by the current thread.
     /// </summary>
-    /// <param name="item">The instance to return.</param>
+    /// <param name="item">The reset instance whose ownership is transferred to the pool.</param>
     public void Return(T item)
     {
         if (TryGetStack(out var stack) && stack.Count < MaxPoolSizePerThread)

test/Orleans.CodeGenerator.Tests/OrleansSourceGeneratorTests.cs

@@ -739,6 +739,28 @@ public Task<string> SayHello(string name)
     }
 }");
 
+    [Fact]
+    public Task TestGrainRequestReturnShapes() => AssertSuccessfulSourceGeneration(
+@"using Orleans;
+using Orleans.Concurrency;
+using System.Threading.Tasks;
+
+namespace TestProject;
+
+public interface IRequestShapeGrain : IGrainWithIntegerKey
+{
+    ValueTask ValueTaskMethod(int value);
+
+    ValueTask<int> ValueTaskOfTMethod(string value);
+
+    Task TaskMethod(object value);
+
+    Task<int> TaskOfTMethod(byte[] value);
+
+    [OneWay]
+    void OneWayMethod(long value);
+}");
+
     /// <summary>
     /// Tests proxy generation for grains with different key types.
     /// Orleans supports multiple grain key types: