From 9aa5e9cf30e578a55da54c1eacee2c02bca64a5f Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 15:05:21 -0700 Subject: [PATCH 1/4] Remove unnecessary Task.Run hop in PayloadInterceptor.AsyncUnaryCall Fixes #774 Replace the Task.Run(async () => {...}) wrapper around request externalization + continuation invocation with a directly-invoked local async function. This avoids an unconditional thread-pool hop on every unary call when externalization completes synchronously (the common case), while preserving identical externalization ordering, exception propagation, cancellation, response headers, status/trailers, and disposal behavior. Add focused unit tests in Worker.Grpc.Tests covering the success path (ordering + response/headers/status/trailers), externalization failure, cancellation, and disposal (before completion and after failure), plus a regression test asserting externalization and the continuation run inline on the calling thread instead of being dispatched to the thread pool. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 42263041-0f76-4a63-baa9-e8fa73c35466 --- .../Interceptors/PayloadInterceptor.cs | 9 +- .../Grpc.Tests/PayloadInterceptorTests.cs | 249 ++++++++++++++++++ 2 files changed, 255 insertions(+), 3 deletions(-) create mode 100644 test/Worker/Grpc.Tests/PayloadInterceptorTests.cs diff --git a/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs b/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs index 8d94fbff..5aad1cac 100644 --- a/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs +++ b/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs @@ -35,15 +35,18 @@ public override AsyncUnaryCall AsyncUnaryCall( ClientInterceptorContext context, AsyncUnaryCallContinuation continuation) { - // Build the underlying call lazily after async externalization - Task> startCallTask = Task.Run(async () => + // Build the underlying call lazily after async externalization. This runs inline (no + // thread-pool hop); the continuation typically starts synchronously anyway once invoked. + Task> startCallTask = StartCallAsync(); + + async Task> StartCallAsync() { // Externalize first; if this fails, do not proceed to send the gRPC call await this.ExternalizeRequestPayloadsAsync(request, context.Options.CancellationToken); // Only if externalization succeeds, proceed with the continuation return continuation(request, context); - }); + } async Task ResponseAsync() { diff --git a/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs b/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs new file mode 100644 index 00000000..e727be7e --- /dev/null +++ b/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs @@ -0,0 +1,249 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Text; +using Grpc.Core; +using Grpc.Core.Interceptors; + +namespace Microsoft.DurableTask.Worker.Grpc.Tests; + +public class PayloadInterceptorTests +{ + static readonly Marshaller StringMarshaller = Marshallers.Create( + value => Encoding.UTF8.GetBytes(value), + bytes => Encoding.UTF8.GetString(bytes)); + static readonly Method TestMethod = new( + MethodType.Unary, + "TestService", + "TestMethod", + StringMarshaller, + StringMarshaller); + + [Fact] + public void AsyncUnaryCall_SynchronousExternalization_RunsInlineWithoutThreadHop() + { + // Arrange: externalization completes synchronously (no real async work), so there is no + // await point for the state machine to yield on. Capture the managed thread ID that both + // callbacks observe and compare it to the calling thread's ID. This is deterministic + // (unlike checking "was it invoked yet", which races against thread-pool scheduling): the + // calling thread is busy executing this synchronous test method the whole time, so it is + // physically impossible for a Task.Run-dispatched worker item to run on this same managed + // thread while we're still inside the call. If the Task.Run thread-pool hop were still + // present, the observed thread ID would either be null (not yet run) or a different thread. + int callingThreadId = Environment.CurrentManagedThreadId; + int? externalizeThreadId = null; + int? continuationThreadId = null; + TestPayloadInterceptor interceptor = CreateInterceptor( + onExternalize: (_, _) => + { + externalizeThreadId = Environment.CurrentManagedThreadId; + return Task.CompletedTask; + }); + ClientInterceptorContext context = CreateContext(); + + // Act + using AsyncUnaryCall call = interceptor.AsyncUnaryCall( + "request", + context, + (req, ctx) => + { + continuationThreadId = Environment.CurrentManagedThreadId; + return CreateFakeCall("response"); + }); + + // Assert: both callbacks already ran, synchronously, on the calling thread. + externalizeThreadId.Should().Be(callingThreadId); + continuationThreadId.Should().Be(callingThreadId); + } + + [Fact] + public async Task AsyncUnaryCall_Success_ExternalizesThenInvokesContinuationThenResolves() + { + // Arrange + List order = []; + bool disposeInvoked = false; + TestPayloadInterceptor interceptor = CreateInterceptor( + onExternalize: (_, _) => + { + order.Add("externalize"); + return Task.CompletedTask; + }, + onResolve: (_, _) => + { + order.Add("resolve"); + return Task.CompletedTask; + }); + ClientInterceptorContext context = CreateContext(); + + // Act + using AsyncUnaryCall call = interceptor.AsyncUnaryCall( + "request", + context, + (req, ctx) => + { + order.Add("continuation"); + return CreateFakeCall("response", disposeAction: () => disposeInvoked = true); + }); + string response = await call.ResponseAsync; + Metadata headers = await call.ResponseHeadersAsync; + + // Assert: ordering, response payload, and pass-through metadata are all preserved. + order.Should().Equal("externalize", "continuation", "resolve"); + response.Should().Be("response"); + headers.Should().NotBeNull(); + call.GetStatus().StatusCode.Should().Be(StatusCode.OK); + call.GetTrailers().Should().NotBeNull(); + + call.Dispose(); + disposeInvoked.Should().BeTrue(); + } + + [Fact] + public async Task AsyncUnaryCall_ExternalizeFails_PropagatesFailureAndSkipsContinuation() + { + // Arrange + InvalidOperationException failure = new("externalization failed"); + bool continuationInvoked = false; + TestPayloadInterceptor interceptor = CreateInterceptor( + onExternalize: (_, _) => throw failure); + ClientInterceptorContext context = CreateContext(); + + // Act + using AsyncUnaryCall call = interceptor.AsyncUnaryCall( + "request", + context, + (req, ctx) => + { + continuationInvoked = true; + return CreateFakeCall("response"); + }); + Func act = async () => await call.ResponseAsync; + + // Assert: the gRPC call is never sent, and the original exception surfaces to the caller. + continuationInvoked.Should().BeFalse(); + (await act.Should().ThrowAsync()).Which.Should().Be(failure); + call.GetStatus().StatusCode.Should().Be(StatusCode.Internal); + } + + [Fact] + public async Task AsyncUnaryCall_CancellationRequested_PropagatesCancellationAndSkipsContinuation() + { + // Arrange + using CancellationTokenSource cts = new(); + cts.Cancel(); + bool continuationInvoked = false; + TestPayloadInterceptor interceptor = CreateInterceptor( + onExternalize: (_, ct) => + { + ct.ThrowIfCancellationRequested(); + return Task.CompletedTask; + }); + ClientInterceptorContext context = CreateContext(cts.Token); + + // Act + using AsyncUnaryCall call = interceptor.AsyncUnaryCall( + "request", + context, + (req, ctx) => + { + continuationInvoked = true; + return CreateFakeCall("response"); + }); + Func act = async () => await call.ResponseAsync; + + // Assert + continuationInvoked.Should().BeFalse(); + await act.Should().ThrowAsync(); + call.GetStatus().StatusCode.Should().Be(StatusCode.Cancelled); + } + + [Fact] + public async Task AsyncUnaryCall_DisposeBeforeCompletion_DisposesInnerCallOnceStarted() + { + // Arrange: hold externalization open so Dispose() races ahead of call construction. + TaskCompletionSource externalizeGate = new(TaskCreationOptions.RunContinuationsAsynchronously); + bool disposeInvoked = false; + TestPayloadInterceptor interceptor = CreateInterceptor( + onExternalize: async (_, _) => await externalizeGate.Task); + ClientInterceptorContext context = CreateContext(); + + // Act + AsyncUnaryCall call = interceptor.AsyncUnaryCall( + "request", + context, + (req, ctx) => CreateFakeCall("response", disposeAction: () => disposeInvoked = true)); + call.Dispose(); + + // Assert: not disposed yet because the inner call has not been created. + disposeInvoked.Should().BeFalse(); + + externalizeGate.SetResult(); + await call.ResponseAsync; + + // Assert: once the inner call exists, disposal is applied to it. + disposeInvoked.Should().BeTrue(); + } + + [Fact] + public void AsyncUnaryCall_DisposeAfterExternalizeFailure_DoesNotThrowOrDisposeMissingCall() + { + // Arrange + bool disposeInvoked = false; + TestPayloadInterceptor interceptor = CreateInterceptor( + onExternalize: (_, _) => throw new InvalidOperationException("boom")); + ClientInterceptorContext context = CreateContext(); + + // Act + AsyncUnaryCall call = interceptor.AsyncUnaryCall( + "request", + context, + (req, ctx) => CreateFakeCall("response", disposeAction: () => disposeInvoked = true)); + Action act = call.Dispose; + + // Assert: no inner call was ever created, so Dispose is a safe no-op. + act.Should().NotThrow(); + disposeInvoked.Should().BeFalse(); + } + + static TestPayloadInterceptor CreateInterceptor( + Func? onExternalize = null, + Func? onResolve = null) + { + Mock payloadStore = new(); + return new TestPayloadInterceptor(payloadStore.Object, new LargePayloadStorageOptions()) + { + OnExternalize = onExternalize, + OnResolve = onResolve, + }; + } + + static ClientInterceptorContext CreateContext(CancellationToken cancellationToken = default) + { + CallOptions options = new(cancellationToken: cancellationToken); + return new ClientInterceptorContext(TestMethod, null, options); + } + + static AsyncUnaryCall CreateFakeCall(string response, Action? disposeAction = null) + { + return new AsyncUnaryCall( + Task.FromResult(response), + Task.FromResult(new Metadata()), + () => new Status(StatusCode.OK, string.Empty), + () => new Metadata(), + disposeAction ?? (() => { })); + } + + sealed class TestPayloadInterceptor(PayloadStore payloadStore, LargePayloadStorageOptions options) + : PayloadInterceptor(payloadStore, options) + { + public Func? OnExternalize { get; set; } + + public Func? OnResolve { get; set; } + + protected override Task ExternalizeRequestPayloadsAsync(TRequest request, CancellationToken cancellation) + => this.OnExternalize?.Invoke(request, cancellation) ?? Task.CompletedTask; + + protected override Task ResolveResponsePayloadsAsync(TResponse response, CancellationToken cancellation) + => this.OnResolve?.Invoke(response, cancellation) ?? Task.CompletedTask; + } +} From 534f88f6dcdd18d75ad70a16625a54990c350fe3 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 16:16:24 -0700 Subject: [PATCH 2/4] Fix flaky dispose-ordering assertion in PayloadInterceptorTests Dispose() schedules the inner call's disposal via a ContinueWith on startCallTask, which is a separate continuation from the one driving ResponseAsync(). The two continuations are not ordered relative to each other, so asserting disposeInvoked immediately after awaiting ResponseAsync could race and fail intermittently. Signal a TaskCompletionSource from the fake call's dispose action and await it (with a bounded 5s timeout) before asserting, so the test deterministically waits for Dispose's continuation instead of relying on incidental ordering. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 42263041-0f76-4a63-baa9-e8fa73c35466 --- .../Grpc.Tests/PayloadInterceptorTests.cs | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs b/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs index e727be7e..dba03a47 100644 --- a/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs +++ b/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs @@ -163,6 +163,13 @@ public async Task AsyncUnaryCall_DisposeBeforeCompletion_DisposesInnerCallOnceSt // Arrange: hold externalization open so Dispose() races ahead of call construction. TaskCompletionSource externalizeGate = new(TaskCreationOptions.RunContinuationsAsynchronously); bool disposeInvoked = false; + + // Dispose() schedules its inner-call disposal via a ContinueWith on startCallTask, which is + // a separate continuation from the one driving ResponseAsync(); the two are not ordered + // relative to each other. Signal a TCS from the disposal callback and await it (with a + // bounded timeout) instead of asserting immediately after ResponseAsync completes, so the + // test doesn't race against Dispose's continuation. + TaskCompletionSource disposeSignal = new(TaskCreationOptions.RunContinuationsAsynchronously); TestPayloadInterceptor interceptor = CreateInterceptor( onExternalize: async (_, _) => await externalizeGate.Task); ClientInterceptorContext context = CreateContext(); @@ -171,7 +178,13 @@ public async Task AsyncUnaryCall_DisposeBeforeCompletion_DisposesInnerCallOnceSt AsyncUnaryCall call = interceptor.AsyncUnaryCall( "request", context, - (req, ctx) => CreateFakeCall("response", disposeAction: () => disposeInvoked = true)); + (req, ctx) => CreateFakeCall( + "response", + disposeAction: () => + { + disposeInvoked = true; + disposeSignal.TrySetResult(); + })); call.Dispose(); // Assert: not disposed yet because the inner call has not been created. @@ -180,6 +193,9 @@ public async Task AsyncUnaryCall_DisposeBeforeCompletion_DisposesInnerCallOnceSt externalizeGate.SetResult(); await call.ResponseAsync; + Task completed = await Task.WhenAny(disposeSignal.Task, Task.Delay(TimeSpan.FromSeconds(5))); + completed.Should().Be(disposeSignal.Task, "the inner call's dispose action should be invoked once it exists"); + // Assert: once the inner call exists, disposal is applied to it. disposeInvoked.Should().BeTrue(); } From 0fcd0dd7718ffb7fda36fb7f1f53ed21e2b5f76f Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Mon, 27 Jul 2026 11:37:54 -0700 Subject: [PATCH 3/4] Add ConfigureAwait(false) to unary awaits; fix double-dispose test - PayloadInterceptor.AsyncUnaryCall's local functions (StartCallAsync, ResponseAsync, ResponseHeadersAsync) now use ConfigureAwait(false) on every await. Since Task.Run no longer isolates this chain onto a context-free thread-pool thread, continuations could otherwise capture and resume on the caller's SynchronizationContext, risking deadlocks in contexts like legacy ASP.NET or WPF/WinForms UI threads. - Fixed AsyncUnaryCall_Success_ExternalizesThenInvokesContinuationThenResolves which declared the call with a using declaration while also calling Dispose() explicitly, causing Dispose() to run twice. Removed the using declaration so disposal happens exactly once via the explicit call, matching the existing assertion. - Streaming path (AsyncServerStreamingCall/TransformingStreamReader) and the Maybe* externalize/resolve helpers are intentionally left unchanged (out of scope for this fix). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 42263041-0f76-4a63-baa9-e8fa73c35466 --- .../Interceptors/PayloadInterceptor.cs | 18 +++++++++++------- .../Grpc.Tests/PayloadInterceptorTests.cs | 5 +++-- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs b/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs index 5aad1cac..e67e6297 100644 --- a/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs +++ b/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs @@ -41,8 +41,12 @@ public override AsyncUnaryCall AsyncUnaryCall( async Task> StartCallAsync() { - // Externalize first; if this fails, do not proceed to send the gRPC call - await this.ExternalizeRequestPayloadsAsync(request, context.Options.CancellationToken); + // Externalize first; if this fails, do not proceed to send the gRPC call. This now runs + // inline on the caller's thread rather than on a dedicated thread-pool thread (via + // Task.Run), so ConfigureAwait(false) is required here to avoid capturing and resuming + // on the caller's SynchronizationContext/TaskScheduler, matching library-safe behavior. + await this.ExternalizeRequestPayloadsAsync(request, context.Options.CancellationToken) + .ConfigureAwait(false); // Only if externalization succeeds, proceed with the continuation return continuation(request, context); @@ -50,16 +54,16 @@ async Task> StartCallAsync() async Task ResponseAsync() { - AsyncUnaryCall innerCall = await startCallTask; - TResponse response = await innerCall.ResponseAsync; - await this.ResolveResponsePayloadsAsync(response, context.Options.CancellationToken); + AsyncUnaryCall innerCall = await startCallTask.ConfigureAwait(false); + TResponse response = await innerCall.ResponseAsync.ConfigureAwait(false); + await this.ResolveResponsePayloadsAsync(response, context.Options.CancellationToken).ConfigureAwait(false); return response; } async Task ResponseHeadersAsync() { - AsyncUnaryCall innerCall = await startCallTask; - return await innerCall.ResponseHeadersAsync; + AsyncUnaryCall innerCall = await startCallTask.ConfigureAwait(false); + return await innerCall.ResponseHeadersAsync.ConfigureAwait(false); } Status GetStatus() diff --git a/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs b/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs index dba03a47..9572431f 100644 --- a/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs +++ b/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs @@ -75,8 +75,9 @@ public async Task AsyncUnaryCall_Success_ExternalizesThenInvokesContinuationThen }); ClientInterceptorContext context = CreateContext(); - // Act - using AsyncUnaryCall call = interceptor.AsyncUnaryCall( + // Act: no using declaration here — Dispose() is called explicitly exactly once below so + // the disposal assertion isn't muddied by a second, implicit Dispose() at scope exit. + AsyncUnaryCall call = interceptor.AsyncUnaryCall( "request", context, (req, ctx) => From b5ffc29587c0fa91675a15e17984785fb21b6c4e Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Mon, 27 Jul 2026 16:32:46 -0700 Subject: [PATCH 4/4] Preserve payload interceptor context isolation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01dc5b7c-c742-45b7-ae8c-5cda9bcfbd2c --- .../Interceptors/PayloadInterceptor.cs | 22 +-- .../Grpc.Tests/PayloadInterceptorTests.cs | 135 ++++++++++++++++-- 2 files changed, 139 insertions(+), 18 deletions(-) diff --git a/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs b/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs index e67e6297..9c494a30 100644 --- a/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs +++ b/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs @@ -35,16 +35,19 @@ public override AsyncUnaryCall AsyncUnaryCall( ClientInterceptorContext context, AsyncUnaryCallContinuation continuation) { - // Build the underlying call lazily after async externalization. This runs inline (no - // thread-pool hop); the continuation typically starts synchronously anyway once invoked. - Task> startCallTask = StartCallAsync(); + // Build the underlying call lazily after async externalization. Avoid a thread-pool hop in + // the normal server path, but preserve isolation for callers with a custom context/scheduler: + // public interceptor and payload-store implementations may capture either before returning. + Task> startCallTask = RequiresSchedulerIsolation() + ? Task.Run(StartCallAsync) + : StartCallAsync(); + + static bool RequiresSchedulerIsolation() + => SynchronizationContext.Current is not null || TaskScheduler.Current != TaskScheduler.Default; async Task> StartCallAsync() { - // Externalize first; if this fails, do not proceed to send the gRPC call. This now runs - // inline on the caller's thread rather than on a dedicated thread-pool thread (via - // Task.Run), so ConfigureAwait(false) is required here to avoid capturing and resuming - // on the caller's SynchronizationContext/TaskScheduler, matching library-safe behavior. + // Externalize first; if this fails, do not proceed to send the gRPC call. await this.ExternalizeRequestPayloadsAsync(request, context.Options.CancellationToken) .ConfigureAwait(false); @@ -56,7 +59,10 @@ async Task ResponseAsync() { AsyncUnaryCall innerCall = await startCallTask.ConfigureAwait(false); TResponse response = await innerCall.ResponseAsync.ConfigureAwait(false); - await this.ResolveResponsePayloadsAsync(response, context.Options.CancellationToken).ConfigureAwait(false); + Task resolveTask = RequiresSchedulerIsolation() + ? Task.Run(() => this.ResolveResponsePayloadsAsync(response, context.Options.CancellationToken)) + : this.ResolveResponsePayloadsAsync(response, context.Options.CancellationToken); + await resolveTask.ConfigureAwait(false); return response; } diff --git a/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs b/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs index 9572431f..a092f145 100644 --- a/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs +++ b/test/Worker/Grpc.Tests/PayloadInterceptorTests.cs @@ -41,19 +41,31 @@ public void AsyncUnaryCall_SynchronousExternalization_RunsInlineWithoutThreadHop }); ClientInterceptorContext context = CreateContext(); - // Act - using AsyncUnaryCall call = interceptor.AsyncUnaryCall( - "request", - context, - (req, ctx) => - { - continuationThreadId = Environment.CurrentManagedThreadId; - return CreateFakeCall("response"); - }); + SynchronizationContext? previous = SynchronizationContext.Current; + AsyncUnaryCall call; + try + { + SynchronizationContext.SetSynchronizationContext(null); + + // Act + call = interceptor.AsyncUnaryCall( + "request", + context, + (req, ctx) => + { + continuationThreadId = Environment.CurrentManagedThreadId; + return CreateFakeCall("response"); + }); + } + finally + { + SynchronizationContext.SetSynchronizationContext(previous); + } // Assert: both callbacks already ran, synchronously, on the calling thread. externalizeThreadId.Should().Be(callingThreadId); continuationThreadId.Should().Be(callingThreadId); + call.Dispose(); } [Fact] @@ -222,6 +234,62 @@ public void AsyncUnaryCall_DisposeAfterExternalizeFailure_DoesNotThrowOrDisposeM disposeInvoked.Should().BeFalse(); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task AsyncUnaryCall_IncompleteInterceptorAwait_DoesNotRequireCallerContextPump(bool resolve) + { + // Arrange + TaskCompletionSource gate = new(TaskCreationOptions.RunContinuationsAsynchronously); + TaskCompletionSource entered = new(TaskCreationOptions.RunContinuationsAsynchronously); + QueuedSynchronizationContext blockedContext = new(); + Func asynchronousStep = async (_, _) => + { + entered.SetResult(); + await gate.Task; + }; + TestPayloadInterceptor interceptor = resolve + ? CreateInterceptor(onResolve: asynchronousStep) + : CreateInterceptor(onExternalize: asynchronousStep); + ClientInterceptorContext context = CreateContext(); + + SynchronizationContext? previous = SynchronizationContext.Current; + AsyncUnaryCall call; + try + { + SynchronizationContext.SetSynchronizationContext(blockedContext); + + // Act + call = interceptor.AsyncUnaryCall( + "request", + context, + (_, _) => CreateFakeCall("response")); + } + finally + { + SynchronizationContext.SetSynchronizationContext(previous); + } + + try + { + await entered.Task.WaitAsync(TimeSpan.FromSeconds(5)); + gate.SetResult(); + + Task winner = await Task.WhenAny(call.ResponseAsync, blockedContext.CallbackPosted) + .WaitAsync(TimeSpan.FromSeconds(5)); + + // Assert + winner.Should().Be(call.ResponseAsync, "the response must complete without pumping the caller's context"); + await call.ResponseAsync; + } + finally + { + gate.TrySetResult(); + blockedContext.Drain(); + call.Dispose(); + } + } + static TestPayloadInterceptor CreateInterceptor( Func? onExternalize = null, Func? onResolve = null) @@ -237,7 +305,7 @@ static TestPayloadInterceptor CreateInterceptor( static ClientInterceptorContext CreateContext(CancellationToken cancellationToken = default) { CallOptions options = new(cancellationToken: cancellationToken); - return new ClientInterceptorContext(TestMethod, null, options); + return new ClientInterceptorContext(TestMethod, "localhost", options); } static AsyncUnaryCall CreateFakeCall(string response, Action? disposeAction = null) @@ -263,4 +331,51 @@ protected override Task ExternalizeRequestPayloadsAsync(TRequest reque protected override Task ResolveResponsePayloadsAsync(TResponse response, CancellationToken cancellation) => this.OnResolve?.Invoke(response, cancellation) ?? Task.CompletedTask; } + + sealed class QueuedSynchronizationContext : SynchronizationContext + { + readonly Queue<(SendOrPostCallback Callback, object? State)> callbacks = []; + readonly TaskCompletionSource callbackPosted = + new(TaskCreationOptions.RunContinuationsAsynchronously); + + public Task CallbackPosted => this.callbackPosted.Task; + + public override void Post(SendOrPostCallback callback, object? state) + { + lock (this.callbacks) + { + this.callbacks.Enqueue((callback, state)); + } + + this.callbackPosted.TrySetResult(); + } + + public void Drain() + { + SynchronizationContext? previous = Current; + SetSynchronizationContext(this); + try + { + while (true) + { + (SendOrPostCallback Callback, object? State) work; + lock (this.callbacks) + { + if (this.callbacks.Count == 0) + { + return; + } + + work = this.callbacks.Dequeue(); + } + + work.Callback(work.State); + } + } + finally + { + SetSynchronizationContext(previous); + } + } + } }