Fix call cancellation impacting other calls (#1688) (#1694)

Author: JamesNK

src/Grpc.Net.Client/Balancer/Internal/ConnectionManager.cs

@@ -36,7 +36,6 @@ internal sealed class ConnectionManager : IDisposable, IChannelControlHelper
     {
         private static readonly ServiceConfig DefaultServiceConfig = new ServiceConfig();
 
-        private readonly SemaphoreSlim _nextPickerLock;
         private readonly object _lock;
         internal readonly Resolver _resolver;
         private readonly ISubchannelTransportFactory _subchannelTransportFactory;
@@ -47,6 +46,8 @@ internal sealed class ConnectionManager : IDisposable, IChannelControlHelper
         // Internal for testing
         internal LoadBalancer? _balancer;
         internal SubchannelPicker? _picker;
+        // Cache picker wrapped in task once and reuse.
+        private Task<SubchannelPicker>? _pickerTask;
         private bool _resolverStarted;
         private TaskCompletionSource<SubchannelPicker> _nextPickerTcs;
         private int _currentSubchannelId;
@@ -61,7 +62,6 @@ internal ConnectionManager(
             LoadBalancerFactory[] loadBalancerFactories)
         {
             _lock = new object();
-            _nextPickerLock = new SemaphoreSlim(1);
             _nextPickerTcs = new TaskCompletionSource<SubchannelPicker>(TaskCreationOptions.RunContinuationsAsynchronously);
             _resolverStartedTcs = new TaskCompletionSource<object?>(TaskCreationOptions.RunContinuationsAsynchronously);
 
@@ -200,7 +200,6 @@ private void OnResolverResult(ResolverResult result)
         public void Dispose()
         {
             _resolver.Dispose();
-            _nextPickerLock.Dispose();
             lock (_lock)
             {
                 _balancer?.Dispose();
@@ -301,6 +300,7 @@ public void UpdateState(BalancerState state)
                 {
                     ConnectionManagerLog.ChannelPickerUpdated(Logger);
                     _picker = state.Picker;
+                    _pickerTask = Task.FromResult(state.Picker);
                     _nextPickerTcs.SetResult(state.Picker);
                     _nextPickerTcs = new TaskCompletionSource<SubchannelPicker>(TaskCreationOptions.RunContinuationsAsynchronously);
                 }
@@ -365,70 +365,22 @@ public void UpdateState(BalancerState state)
             }
         }
 
-        private
-#if !NETSTANDARD2_0
-            ValueTask<SubchannelPicker>
-#else
-            Task<SubchannelPicker>
-#endif
-            GetPickerAsync(SubchannelPicker? currentPicker, CancellationToken cancellationToken)
+        private Task<SubchannelPicker> GetPickerAsync(SubchannelPicker? currentPicker, CancellationToken cancellationToken)
         {
             lock (_lock)
             {
                 if (_picker != null && _picker != currentPicker)
                 {
-#if !NETSTANDARD2_0
-                    return new ValueTask<SubchannelPicker>(_picker);
-#else
-                    return Task.FromResult<SubchannelPicker>(_picker);
-#endif
+                    Debug.Assert(_pickerTask != null);
+                    return _pickerTask;
                 }
                 else
                 {
-                    return GetNextPickerAsync(cancellationToken);
-                }
-            }
-        }
-
-        private async
-#if !NETSTANDARD2_0
-            ValueTask<SubchannelPicker>
-#else
-            Task<SubchannelPicker>
-#endif
-            GetNextPickerAsync(CancellationToken cancellationToken)
-        {
-            ConnectionManagerLog.PickWaiting(Logger);
-
-            Debug.Assert(Monitor.IsEntered(_lock));
-
-            var nextPickerTcs = _nextPickerTcs;
+                    ConnectionManagerLog.PickWaiting(Logger);
 
-            await _nextPickerLock.WaitAsync(cancellationToken).ConfigureAwait(false);
-            try
-            {
-                using (cancellationToken.Register(
-                    static s => ((TaskCompletionSource<SubchannelPicker?>)s!).TrySetCanceled(),
-                    nextPickerTcs))
-                {
-                    try
-                    {
-                        return await nextPickerTcs.Task.ConfigureAwait(false);
-                    }
-                    finally
-                    {
-                        // Picker can throw when canceled so reset picker in finally block.
-                        lock (_lock)
-                        {
-                            _nextPickerTcs = new TaskCompletionSource<SubchannelPicker>(TaskCreationOptions.RunContinuationsAsynchronously);
-                        }
-                    }
+                    return _nextPickerTcs.Task.WaitAsync(cancellationToken);
                 }
             }
-            finally
-            {
-                _nextPickerLock.Release();
-            }
         }
 
         internal Task WaitForStateChangedAsync(ConnectivityState lastObservedState, ConnectivityState? waitForState, CancellationToken cancellationToken)

src/Shared/CompatibilityHelpers.cs

@@ -44,5 +44,16 @@ public static int IndexOf(string s, char value, StringComparison comparisonType)
             return s.IndexOf(value, comparisonType);
 #endif
         }
+
+#if !NET6_0_OR_GREATER
+        public async static 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))
+            {
+                return await (await Task.WhenAny(task, tcs.Task).ConfigureAwait(false)).ConfigureAwait(false);
+            }
+        }
+#endif
     }
 }

test/FunctionalTests/Balancer/PickFirstBalancerTests.cs

@@ -61,6 +61,50 @@ private GrpcChannel CreateGrpcWebChannel(TestServerEndpointName endpointName, Ve
             return channel;
         }
 
+        [Test]
+        public async Task UnaryCall_CallAfterCancellation_Success()
+        {
+            // Ignore errors
+            SetExpectedErrorsFilter(writeContext =>
+            {
+                return true;
+            });
+
+            string? host = null;
+            Task<HelloReply> UnaryMethod(HelloRequest request, ServerCallContext context)
+            {
+                host = context.Host;
+                return Task.FromResult(new HelloReply { Message = request.Name });
+            }
+
+            // Arrange
+            using var endpoint = BalancerHelpers.CreateGrpcEndpoint<HelloRequest, HelloReply>(50051, UnaryMethod, nameof(UnaryMethod));
+
+            var channel = await BalancerHelpers.CreateChannel(LoggerFactory, new PickFirstConfig(), new[] { endpoint.Address }).DefaultTimeout();
+            var client = TestClientFactory.Create(channel, endpoint.Method);
+            
+            // Kill endpoint so client can't connect.
+            endpoint.Dispose();
+
+            var cts = new CancellationTokenSource(TimeSpan.FromSeconds(0.5));
+            cts.Token.Register(() => Logger.LogInformation("Cancellation token raised"));
+
+            // Start call that is canceled while getting picker.
+            await ExceptionAssert.ThrowsAsync<RpcException>(() => client.UnaryCall(
+                new HelloRequest { Name = "Balancer" },
+                new CallOptions(cancellationToken: cts.Token)).ResponseAsync).DefaultTimeout();
+
+            // Restart endpoint.
+            using var endpoint1 = BalancerHelpers.CreateGrpcEndpoint<HelloRequest, HelloReply>(50051, UnaryMethod, nameof(UnaryMethod));
+
+            // Act
+            var reply = await client.UnaryCall(new HelloRequest { Name = "Balancer" }, new CallOptions().WithWaitForReady(true)).ResponseAsync.DefaultTimeout();
+
+            // Assert
+            Assert.AreEqual("Balancer", reply.Message);
+            Assert.AreEqual("127.0.0.1:50051", host);
+        }
+
         [Test]
         public async Task UnaryCall_ReconnectBetweenCalls_Success()
         {