fix(runtime): return completion sources on send failure

Author: ReubenBond

src/Orleans.Core/Runtime/GrainReferenceRuntime.cs

@@ -120,7 +120,7 @@ private async ValueTask InvokeMethodWithFiltersAsync(GrainReference reference, I
 
         private ValueTask<TResult?> InvokeMethodAsyncCore<TResult>(GrainReference reference, IInvokable request, InvokeMethodOptions options)
         {
-            ResponseCompletionSource<TResult> responseCompletionSource;
+            ResponseCompletionSource<TResult>? responseCompletionSource = null;
             try
             {
                 SetGrainCancellationTokensTarget(reference, request);
@@ -130,6 +130,7 @@ private async ValueTask InvokeMethodWithFiltersAsync(GrainReference reference, I
             }
             catch
             {
+                responseCompletionSource?.Reset();
                 DisposeRequest(request);
                 throw;
             }
@@ -149,7 +150,7 @@ private async ValueTask InvokeMethodWithFiltersAsync(GrainReference reference, I
 
         private ValueTask InvokeMethodAsyncCore(GrainReference reference, IInvokable request, InvokeMethodOptions options)
         {
-            ResponseCompletionSource responseCompletionSource;
+            ResponseCompletionSource? responseCompletionSource = null;
             try
             {
                 SetGrainCancellationTokensTarget(reference, request);
@@ -159,6 +160,7 @@ private ValueTask InvokeMethodAsyncCore(GrainReference reference, IInvokable req
             }
             catch
             {
+                responseCompletionSource?.Reset();
                 DisposeRequest(request);
                 throw;
             }

test/Orleans.Core.Tests/Runtime/GrainReferenceRuntimeTests.cs

@@ -0,0 +1,90 @@
+using Microsoft.Extensions.DependencyInjection;
+using Orleans.CodeGeneration;
+using Orleans.Runtime;
+using Orleans.Serialization.Invocation;
+using TestExtensions;
+using UnitTests.Messaging;
+using Xunit;
+
+namespace UnitTests.Runtime;
+
+[TestSuite("BVT")]
+[TestProvider("None")]
+[TestArea("Runtime")]
+public class GrainReferenceRuntimeTests
+{
+    [Fact, TestCategory("BVT")]
+    public void TypedCompletionSourceReturnsToPoolWhenSendThrows()
+    {
+        var runtimeClient = new ThrowingRuntimeClient();
+        var runtime = CreateRuntime(runtimeClient);
+        var request = new TestInvokableRequest();
+
+        Assert.Throws<InvalidOperationException>(() => runtime.InvokeMethodAsync<int>(null!, request, InvokeMethodOptions.None));
+
+        var captured = Assert.IsType<ResponseCompletionSource<int>>(runtimeClient.Context);
+        var returned = ResponseCompletionSourcePool.Get<int>();
+        Assert.Same(captured, returned);
+        Assert.Equal(1, request.DisposeCount);
+        returned.Reset();
+    }
+
+    [Fact, TestCategory("BVT")]
+    public void UntypedCompletionSourceReturnsToPoolWhenSendThrows()
+    {
+        var runtimeClient = new ThrowingRuntimeClient();
+        var runtime = CreateRuntime(runtimeClient);
+        var request = new TestInvokableRequest();
+
+        Assert.Throws<InvalidOperationException>(() => runtime.InvokeMethodAsync(null!, request, InvokeMethodOptions.None));
+
+        var captured = Assert.IsType<ResponseCompletionSource>(runtimeClient.Context);
+        var returned = ResponseCompletionSourcePool.Get();
+        Assert.Same(captured, returned);
+        Assert.Equal(1, request.DisposeCount);
+        returned.Reset();
+    }
+
+    private static GrainReferenceRuntime CreateRuntime(IRuntimeClient runtimeClient)
+        => new(runtimeClient, null!, [], null!, null!);
+
+    private sealed class ThrowingRuntimeClient : IRuntimeClient
+    {
+        public IResponseCompletionSource? Context { get; private set; }
+
+        public TimeProvider TimeProvider => TimeProvider.System;
+        public IInternalGrainFactory InternalGrainFactory => throw new NotSupportedException();
+        public string CurrentActivationIdentity => string.Empty;
+        public IServiceProvider ServiceProvider => EmptyServiceProvider.Instance;
+        public IGrainReferenceRuntime GrainReferenceRuntime => throw new NotSupportedException();
+
+        public TimeSpan GetResponseTimeout() => throw new NotSupportedException();
+
+        public void SetResponseTimeout(TimeSpan timeout) => throw new NotSupportedException();
+
+        public void SendRequest(GrainReference target, IInvokable request, IResponseCompletionSource? context, InvokeMethodOptions options)
+        {
+            Context = context;
+            throw new InvalidOperationException("Send failed.");
+        }
+
+        public void SendResponse(Message request, Response response) => throw new NotSupportedException();
+
+        public void ReceiveResponse(Message message) => throw new NotSupportedException();
+
+        public IAddressable CreateObjectReference(IAddressable obj) => throw new NotSupportedException();
+
+        public void DeleteObjectReference(IAddressable obj) => throw new NotSupportedException();
+
+        public void BreakOutstandingMessagesToSilo(SiloAddress deadSilo) => throw new NotSupportedException();
+
+        public int GetRunningRequestsCount(GrainInterfaceType grainInterfaceType) => 0;
+    }
+
+    private sealed class EmptyServiceProvider : IServiceProvider
+    {
+        public static EmptyServiceProvider Instance { get; } = new();
+
+        public object? GetService(Type serviceType) => null;
+    }
+}