Reduce per-call allocations in ChannelRecreatingCallInvoker outcome tracking

src/Client/Grpc/ChannelRecreatingCallInvoker.cs

@@ -38,6 +38,13 @@ sealed class ChannelRecreatingCallInvoker : CallInvoker, IAsyncDisposable
     readonly bool ownsChannel;
     readonly ILogger logger;
 
+    // Cached once per invoker instance (instead of once per RPC) so ObserveOutcome's ContinueWith call
+    // does not allocate a new delegate for every unary call. The method-name diagnostic is threaded
+    // through as the ContinueWith `state` argument (a string, already a reference type) instead of a
+    // boxed (self, methodName) tuple, eliminating that per-call delegate and boxed-tuple allocation
+    // (Task.ContinueWith's own internal continuation/Task bookkeeping is unaffected).
+    readonly Action<Task, object?> onUnaryCallCompleted;
+
     // Cancelled in DisposeAsync so an in-flight RecreateAsync stops promptly and does not leak the
     // freshly created channel back into our state after we've disposed.
     readonly CancellationTokenSource disposalCts = new();
@@ -68,6 +75,7 @@ public ChannelRecreatingCallInvoker(
         this.minRecreateInterval = minRecreateInterval;
         this.ownsChannel = ownsChannel;
         this.logger = logger;
+        this.onUnaryCallCompleted = this.OnUnaryCallCompleted;
         this.state = new TransportState(initialChannel, initialChannel.CreateCallInvoker());
 
         // Backdate the initial timestamp so the first recreate is never blocked by the cooldown.
@@ -195,28 +203,33 @@ static TimeSpan ElapsedSince(long previousTimestamp, long nowTimestamp)
         return TimeSpan.FromSeconds((double)elapsedTicks / Stopwatch.Frequency);
     }
 
-    void ObserveOutcome<TResponse>(Task<TResponse> responseAsync, string methodFullName)
+    void ObserveOutcome(Task responseAsync, string methodFullName)
     {
-        // Use ContinueWith with TaskScheduler.Default so we don't capture sync context.
+        // Use ContinueWith with TaskScheduler.Default so we don't capture sync context. Both the
+        // continuation delegate (this.onUnaryCallCompleted, cached once per invoker instance) and the
+        // continuation state (methodFullName, already a reference-typed string) are allocation-free
+        // per call: no boxed tuple and no per-call delegate closure.
         responseAsync.ContinueWith(
-            (t, state) =>
-            {
-                var (self, name) = ((ChannelRecreatingCallInvoker, string))state!;
-                if (t.Status == TaskStatus.RanToCompletion)
-                {
-                    self.RecordSuccess();
-                }
-                else if (t.Exception?.InnerException is RpcException rpcEx)
-                {
-                    self.RecordFailure(rpcEx.StatusCode, name);
-                }
-            },
-            (this, methodFullName),
+            this.onUnaryCallCompleted,
+            methodFullName,
             CancellationToken.None,
             TaskContinuationOptions.ExecuteSynchronously,
             TaskScheduler.Default);
     }
 
+    void OnUnaryCallCompleted(Task task, object? continuationState)
+    {
+        string methodFullName = (string)continuationState!;
+        if (task.Status == TaskStatus.RanToCompletion)
+        {
+            this.RecordSuccess();
+        }
+        else if (task.Exception?.InnerException is RpcException rpcEx)
+        {
+            this.RecordFailure(rpcEx.StatusCode, methodFullName);
+        }
+    }
+
     void RecordSuccess()
     {
         Volatile.Write(ref this.consecutiveFailures, 0);

test/Client/Grpc.Tests/GrpcDurableTaskClientChannelRecreationTests.cs

@@ -22,6 +22,12 @@ public class GrpcDurableTaskClientChannelRecreationTests
         "TestMethod",
         StringMarshaller,
         StringMarshaller);
+    static readonly Method<string, string> WaitForInstanceCompletionMethod = new(
+        MethodType.Unary,
+        "TaskHubSidecarService",
+        "WaitForInstanceCompletion",
+        StringMarshaller,
+        StringMarshaller);
     static readonly MethodInfo GetCallInvokerMethod = typeof(GrpcDurableTaskClient)
         .GetMethod("GetCallInvoker", BindingFlags.Static | BindingFlags.NonPublic)!;
     static readonly MethodInfo ToStopwatchTicksMethod = typeof(ChannelRecreatingCallInvoker)
@@ -93,14 +99,215 @@ public async Task GetCallInvoker_WithProvidedChannel_RecreatesTransportAfterUnar
     }
 
     [Fact]
-    public async Task GetCallInvoker_WithAddressAndRecreator_UsesWrapperThatOwnsCreatedChannel()
+    public async Task AsyncUnaryCall_SuccessAfterFailure_ResetsConsecutiveFailureCountAndSuppressesRecreate()
+    {
+        // Arrange: threshold of 2 consecutive failures, but a success is interleaved so the counter
+        // should reset and the recreator should never be invoked.
+        int callIndex = 0;
+        CallbackHttpMessageHandler handler = new((request, cancellationToken) =>
+        {
+            int index = Interlocked.Increment(ref callIndex);
+            return Task.FromResult(index == 2
+                ? CreateSuccessResponse("pong")
+                : CreateFailureResponse(StatusCode.Unavailable, "transient failure"));
+        });
+
+        GrpcChannel channel = CreateChannel("http://success-reset.client.test", handler);
+        GrpcDurableTaskClientOptions options = new() { Channel = channel };
+        options.Internal.ChannelRecreateFailureThreshold = 2;
+        options.Internal.MinRecreateInterval = TimeSpan.Zero;
+
+        int recreatorCalls = 0;
+        options.SetChannelRecreator((existingChannel, ct) =>
+        {
+            Interlocked.Increment(ref recreatorCalls);
+            return Task.FromResult(existingChannel);
+        });
+
+        try
+        {
+            (AsyncDisposable disposable, CallInvoker callInvoker) = InvokeGetCallInvoker(options);
+            try
+            {
+                // Act: failure, success, failure -- the success in the middle must reset the streak so
+                // the trailing failure alone never reaches the threshold of 2.
+                await AssertRpcFailureAsync(callInvoker);
+                await AssertRpcSuccessAsync(callInvoker);
+                await AssertRpcFailureAsync(callInvoker);
+
+                await Task.Delay(TimeSpan.FromMilliseconds(200));
+
+                // Assert
+                GetConsecutiveFailures(callInvoker).Should().Be(1);
+                recreatorCalls.Should().Be(0);
+            }
+            finally
+            {
+                await disposable.DisposeAsync();
+            }
+        }
+        finally
+        {
+            await DisposeChannelAsync(channel);
+        }
+    }
+
+    [Fact]
+    public async Task AsyncUnaryCall_DeadlineExceededOnAllowedLongPollMethod_DoesNotCountTowardThreshold()
+    {
+        // Arrange: WaitForInstanceCompletion is long-poll and DeadlineExceeded is expected behavior
+        // there, so repeated timeouts on it must never count toward the recreate threshold.
+        CallbackHttpMessageHandler handler = new((request, cancellationToken) =>
+            Task.FromResult(CreateFailureResponse(StatusCode.DeadlineExceeded, "long-poll wait elapsed")));
+
+        GrpcChannel channel = CreateChannel("http://deadline-allowed.client.test", handler);
+        GrpcDurableTaskClientOptions options = new() { Channel = channel };
+        options.Internal.ChannelRecreateFailureThreshold = 1;
+        options.Internal.MinRecreateInterval = TimeSpan.Zero;
+
+        int recreatorCalls = 0;
+        options.SetChannelRecreator((existingChannel, ct) =>
+        {
+            Interlocked.Increment(ref recreatorCalls);
+            return Task.FromResult(existingChannel);
+        });
+
+        try
+        {
+            (AsyncDisposable disposable, CallInvoker callInvoker) = InvokeGetCallInvoker(options);
+            try
+            {
+                // Act
+                await AssertRpcFailureAsync(
+                    callInvoker, WaitForInstanceCompletionMethod, StatusCode.DeadlineExceeded);
+                await AssertRpcFailureAsync(
+                    callInvoker, WaitForInstanceCompletionMethod, StatusCode.DeadlineExceeded);
+
+                await Task.Delay(TimeSpan.FromMilliseconds(200));
+
+                // Assert
+                GetConsecutiveFailures(callInvoker).Should().Be(0);
+                recreatorCalls.Should().Be(0);
+            }
+            finally
+            {
+                await disposable.DisposeAsync();
+            }
+        }
+        finally
+        {
+            await DisposeChannelAsync(channel);
+        }
+    }
+
+    [Fact]
+    public async Task AsyncUnaryCall_DeadlineExceededOnRegularMethod_CountsTowardThresholdAndTriggersRecreate()
+    {
+        // Arrange: same DeadlineExceeded status as above, but on a method that is NOT in the long-poll
+        // allow-list, so it must count toward the threshold and trigger a recreate.
+        CallbackHttpMessageHandler handler = new((request, cancellationToken) =>
+            Task.FromResult(CreateFailureResponse(StatusCode.DeadlineExceeded, "unexpected timeout")));
+
+        GrpcChannel channel = CreateChannel("http://deadline-counts.client.test", handler);
+        GrpcDurableTaskClientOptions options = new() { Channel = channel };
+        options.Internal.ChannelRecreateFailureThreshold = 1;
+        options.Internal.MinRecreateInterval = TimeSpan.Zero;
+
+        TaskCompletionSource<GrpcChannel> recreateRequested = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        options.SetChannelRecreator((existingChannel, ct) =>
         {
-            // Arrange
-            GrpcDurableTaskClientOptions options = new()
+            recreateRequested.TrySetResult(existingChannel);
+            return Task.FromResult(existingChannel);
+        });
+
+        try
+        {
+            (AsyncDisposable disposable, CallInvoker callInvoker) = InvokeGetCallInvoker(options);
+            try
             {
-                Address = "http://owned.client.test",
-            };
-            options.SetChannelRecreator((existingChannel, ct) => Task.FromResult(existingChannel));
+                // Act
+                await AssertRpcFailureAsync(callInvoker, TestMethod, StatusCode.DeadlineExceeded);
+                GrpcChannel recreatorChannel = await recreateRequested.Task.WaitAsync(TimeSpan.FromSeconds(5));
+
+                // Assert
+                recreatorChannel.Should().BeSameAs(channel);
+            }
+            finally
+            {
+                await disposable.DisposeAsync();
+            }
+        }
+        finally
+        {
+            await DisposeChannelAsync(channel);
+        }
+    }
+
+    [Fact]
+    public async Task AsyncUnaryCall_MultipleCalls_ReuseSameCachedOutcomeDelegateInstance()
+    {
+        // Arrange/Act: the outcome-observation delegate must be created exactly once per invoker
+        // instance (in the constructor) and reused for every call, rather than allocated per RPC.
+        // Use the deterministic fake-handler pattern (as in the other tests in this file) instead of
+        // a real address, so the test never depends on DNS/proxy/network and cannot hang in CI.
+        CallbackHttpMessageHandler handler = new((request, cancellationToken) =>
+            Task.FromResult(CreateFailureResponse(StatusCode.Unavailable, "deterministic failure")));
+
+        GrpcChannel channel = CreateChannel("http://cached-delegate.client.test", handler);
+        GrpcDurableTaskClientOptions options = new() { Channel = channel };
+        options.SetChannelRecreator((existingChannel, ct) => Task.FromResult(existingChannel));
+
+        try
+        {
+            (AsyncDisposable disposable, CallInvoker callInvoker) = InvokeGetCallInvoker(options);
+            try
+            {
+                object? delegateBeforeCalls = GetOnUnaryCallCompletedDelegate(callInvoker);
+                delegateBeforeCalls.Should().NotBeNull();
+
+                for (int i = 0; i < 5; i++)
+                {
+                    try
+                    {
+                        using AsyncUnaryCall<string> call = callInvoker.AsyncUnaryCall(
+                            TestMethod,
+                            host: null,
+                            new CallOptions(deadline: DateTime.UtcNow.AddSeconds(1)),
+                            request: "ping");
+                        await call.ResponseAsync;
+                    }
+                    catch (RpcException)
+                    {
+                        // Each call deterministically fails via the fake handler above; only the
+                        // delegate identity across calls is under test here.
+                    }
+                }
+
+                // Assert: still the exact same delegate instance (a readonly field set once in the
+                // constructor), proving no per-call delegate is allocated.
+                object? delegateAfterCalls = GetOnUnaryCallCompletedDelegate(callInvoker);
+                delegateAfterCalls.Should().BeSameAs(delegateBeforeCalls);
+            }
+            finally
+            {
+                await disposable.DisposeAsync();
+            }
+        }
+        finally
+        {
+            await DisposeChannelAsync(channel);
+        }
+    }
+
+    [Fact]
+    public async Task GetCallInvoker_WithAddressAndRecreator_UsesWrapperThatOwnsCreatedChannel()
+    {
+        // Arrange
+        GrpcDurableTaskClientOptions options = new()
+        {
+            Address = "http://owned.client.test",
+        };
+        options.SetChannelRecreator((existingChannel, ct) => Task.FromResult(existingChannel));
 
         // Act
         (AsyncDisposable disposable, CallInvoker callInvoker) = InvokeGetCallInvoker(options);
@@ -218,17 +425,34 @@ static void SetDisposed(CallInvoker callInvoker, int value)
             .SetValue(callInvoker, value);
     }
 
+    static int GetConsecutiveFailures(CallInvoker callInvoker)
+    {
+        return (int)typeof(ChannelRecreatingCallInvoker)
+            .GetField("consecutiveFailures", BindingFlags.Instance | BindingFlags.NonPublic)!
+            .GetValue(callInvoker)!;
+    }
+
+    static object GetOnUnaryCallCompletedDelegate(CallInvoker callInvoker)
+    {
+        return typeof(ChannelRecreatingCallInvoker)
+            .GetField("onUnaryCallCompleted", BindingFlags.Instance | BindingFlags.NonPublic)!
+            .GetValue(callInvoker)!;
+    }
+
     static long InvokeToStopwatchTicks(TimeSpan interval)
     {
         return (long)ToStopwatchTicksMethod.Invoke(null, new object?[] { interval })!;
     }
 
-    static async Task AssertRpcFailureAsync(CallInvoker callInvoker)
+    static async Task AssertRpcFailureAsync(CallInvoker callInvoker) =>
+        await AssertRpcFailureAsync(callInvoker, TestMethod, StatusCode.Unavailable);
+
+    static async Task AssertRpcFailureAsync(CallInvoker callInvoker, Method<string, string> method, StatusCode expectedStatus)
     {
         Func<Task> act = async () =>
         {
             using AsyncUnaryCall<string> call = callInvoker.AsyncUnaryCall(
-                TestMethod,
+                method,
                 host: null,
                 new CallOptions(deadline: DateTime.UtcNow.AddSeconds(1)),
                 request: "ping");
@@ -237,7 +461,19 @@ static async Task AssertRpcFailureAsync(CallInvoker callInvoker)
         };
 
         RpcException rpcException = (await act.Should().ThrowAsync<RpcException>()).Which;
-        rpcException.StatusCode.Should().Be(StatusCode.Unavailable);
+        rpcException.StatusCode.Should().Be(expectedStatus);
+    }
+
+    static async Task AssertRpcSuccessAsync(CallInvoker callInvoker)
+    {
+        using AsyncUnaryCall<string> call = callInvoker.AsyncUnaryCall(
+            TestMethod,
+            host: null,
+            new CallOptions(deadline: DateTime.UtcNow.AddSeconds(5)),
+            request: "ping");
+
+        string response = await call.ResponseAsync;
+        response.Should().Be("pong");
     }
 
     static GrpcChannel CreateChannel(string address, HttpMessageHandler handler)
@@ -268,6 +504,28 @@ static HttpResponseMessage CreateFailureResponse(StatusCode statusCode, string d
         return response;
     }
 
+    static HttpResponseMessage CreateSuccessResponse(string payload)
+    {
+        byte[] messageBytes = StringMarshaller.Serializer(payload);
+        byte[] frame = new byte[5 + messageBytes.Length];
+        uint length = (uint)messageBytes.Length;
+        frame[1] = (byte)(length >> 24);
+        frame[2] = (byte)(length >> 16);
+        frame[3] = (byte)(length >> 8);
+        frame[4] = (byte)length;
+        Buffer.BlockCopy(messageBytes, 0, frame, 5, messageBytes.Length);
+
+        HttpResponseMessage response = new(System.Net.HttpStatusCode.OK)
+        {
+            Version = new Version(2, 0),
+            Content = new ByteArrayContent(frame),
+        };
+
+        response.Content.Headers.ContentType = new System.Net.Http.Headers.MediaTypeHeaderValue("application/grpc");
+        response.TrailingHeaders.Add("grpc-status", "0");
+        return response;
+    }
+
     sealed class CallbackHttpMessageHandler : HttpMessageHandler
     {
         readonly Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> callback;