Force yielding after awaiting CallTask to avoid holding onto locks

src/Grpc.Net.Client/Internal/Retry/HedgingCall.cs

@@ -129,7 +129,8 @@ private async Task StartCall(Action<GrpcCall<TRequest, TResponse>> startCallFunc
 
                 // Wait until the call has finished and then check its status code
                 // to update retry throttling tokens.
-                var status = await call.CallTask.ConfigureAwait(false);
+                // Force yield here to prevent continuation running with any locks.
+                var status = await CompatibilityHelpers.AwaitWithYieldAsync(call.CallTask).ConfigureAwait(false);
                 if (status.StatusCode == StatusCode.OK)
                 {
                     RetryAttemptCallSuccess();
@@ -202,7 +203,8 @@ private async Task StartCall(Action<GrpcCall<TRequest, TResponse>> startCallFunc
             if (CommitedCallTask.IsCompletedSuccessfully() && CommitedCallTask.Result == call)
             {
                 // Wait until the commited call is finished and then clean up hedging call.
-                await call.CallTask.ConfigureAwait(false);
+                // Force yield here to prevent continuation running with any locks.
+                await CompatibilityHelpers.AwaitWithYieldAsync(call.CallTask).ConfigureAwait(false);
                 Cleanup();
             }
         }

src/Grpc.Net.Client/Internal/Retry/RetryCall.cs

@@ -162,7 +162,8 @@ private async Task StartRetry(Action<GrpcCall<TRequest, TResponse>> startCallFun
                     // Headers were returned. We're commited.
                     CommitCall(currentCall, CommitReason.ResponseHeadersReceived);
 
-                    responseStatus = await currentCall.CallTask.ConfigureAwait(false);
+                    // Force yield here to prevent continuation running with any locks.
+                    responseStatus = await CompatibilityHelpers.AwaitWithYieldAsync(currentCall.CallTask).ConfigureAwait(false);
                     if (responseStatus.Value.StatusCode == StatusCode.OK)
                     {
                         RetryAttemptCallSuccess();
@@ -252,7 +253,8 @@ private async Task StartRetry(Action<GrpcCall<TRequest, TResponse>> startCallFun
                 if (CommitedCallTask.Result is GrpcCall<TRequest, TResponse> call)
                 {
                     // Wait until the commited call is finished and then clean up retry call.
-                    await call.CallTask.ConfigureAwait(false);
+                    // Force yield here to prevent continuation running with any locks.
+                    await CompatibilityHelpers.AwaitWithYieldAsync(call.CallTask).ConfigureAwait(false);
                     Cleanup();
                 }
             }

src/Shared/CompatibilityHelpers.cs

@@ -46,7 +46,7 @@ public static int IndexOf(string s, char value, StringComparison comparisonType)
     }
 
 #if !NET6_0_OR_GREATER
-    public async static Task<T> WaitAsync<T>(this Task<T> task, CancellationToken cancellationToken)
+    public static async Task<T> WaitAsync<T>(this Task<T> task, CancellationToken cancellationToken)
     {
         var tcs = new TaskCompletionSource<T>();
         using (cancellationToken.Register(static s => ((TaskCompletionSource<T>)s!).TrySetCanceled(), tcs))
@@ -66,4 +66,26 @@ public static CancellationTokenRegistration RegisterWithCancellationTokenCallbac
         return cancellationToken.Register((state) => callback(state, cancellationToken), state);
 #endif
     }
+
+    public static Task<T> AwaitWithYieldAsync<T>(Task<T> callTask)
+    {
+        // A completed task doesn't need to yield because code after it isn't run in a continuation.
+        if (callTask.IsCompleted)
+        {
+            return callTask;
+        }
+
+        return AwaitWithYieldAsyncCore(callTask);
+
+        static async Task<T> AwaitWithYieldAsyncCore(Task<T> callTask)
+        {
+#if NET8_0_OR_GREATER
+            return await callTask.ConfigureAwait(ConfigureAwaitOptions.ForceYielding);
+#else
+            var status = await callTask.ConfigureAwait(false);
+            await Task.Yield();
+            return status;
+#endif
+        }
+    }
 }