Load balancing fix to clear disposed subchannels on resolver error (#1853)

Co-authored-by: Francois Nel <francois.nel@iotnxt.com>

Author: JamesNK

src/Grpc.Net.Client/Balancer/BalancerAddress.cs

@@ -17,6 +17,7 @@
 #endregion
 
 #if SUPPORT_LOAD_BALANCING
+using System.Diagnostics;
 using System.Net;
 
 namespace Grpc.Net.Client.Balancer
@@ -35,6 +36,7 @@ public sealed class BalancerAddress
         /// Initializes a new instance of the <see cref="BalancerAddress"/> class with the specified <see cref="DnsEndPoint"/>.
         /// </summary>
         /// <param name="endPoint">The end point.</param>
+        [DebuggerStepThrough]
         public BalancerAddress(DnsEndPoint endPoint)
         {
             EndPoint = endPoint;
@@ -45,6 +47,7 @@ public BalancerAddress(DnsEndPoint endPoint)
         /// </summary>
         /// <param name="host">The host.</param>
         /// <param name="port">The port.</param>
+        [DebuggerStepThrough]
         public BalancerAddress(string host, int port) : this(new DnsEndPoint(host, port))
         {
         }

src/Grpc.Net.Client/Balancer/SubchannelsLoadBalancer.cs

@@ -76,6 +76,7 @@ private void ResolverError(Status status)
                     {
                         RemoveSubchannel(addressSubchannel.Subchannel);
                     }
+                    _addressSubchannels.Clear();
                     Controller.UpdateState(new BalancerState(ConnectivityState.TransientFailure, new ErrorPicker(status)));
                     break;
             }

test/Grpc.Net.Client.Tests/Balancer/ConnectionManagerTests.cs

@@ -327,18 +327,131 @@ public async Task UpdateAddresses_ConnectIsInProgress_InProgressConnectIsCancele
             Assert.AreEqual(81, connectAddress2.Port);
         }
 
+        [Test]
+        public async Task PickAsync_DoesNotDeadlockAfterReconnect_WithResolverError()
+        {
+            // Arrange
+            var services = new ServiceCollection();
+            services.AddNUnitLogger();
+            await using var serviceProvider = services.BuildServiceProvider();
+            var loggerFactory = serviceProvider.GetRequiredService<ILoggerFactory>();
+
+            var resolver = new TestResolver(loggerFactory);
+
+            GrpcChannelOptions channelOptions = new GrpcChannelOptions();
+            channelOptions.ServiceConfig = new ServiceConfig()
+            {
+                LoadBalancingConfigs = { new RoundRobinConfig() }
+            };
+
+            var transportFactory = new TestSubchannelTransportFactory();
+            var clientChannel = CreateConnectionManager(loggerFactory, resolver, transportFactory, new[] { new RoundRobinBalancerFactory() });
+            // Configure balancer similar to how GrpcChannel constructor does it
+            clientChannel.ConfigureBalancer(c => new ChildHandlerLoadBalancer(
+                c,
+                channelOptions.ServiceConfig,
+                clientChannel));
+
+            // Act
+            var connectTask = clientChannel.ConnectAsync(waitForReady: true, cancellationToken: CancellationToken.None);
+            var pickTask = clientChannel.PickAsync(
+                new PickContext { Request = new HttpRequestMessage() },
+                waitForReady: true,
+                CancellationToken.None).AsTask();
+
+            resolver.UpdateAddresses(new List<BalancerAddress>
+            {
+                new BalancerAddress("localhost", 80)
+            });
+            await Task.WhenAll(connectTask, pickTask).DefaultTimeout();
+
+            // Simulate transport/network issue
+            transportFactory.Transports.ForEach(t => t.Disconnect());
+            resolver.UpdateError(new Status(StatusCode.Unavailable, "Test error"));
+
+            pickTask = clientChannel.PickAsync(
+                new PickContext { Request = new HttpRequestMessage() },
+                waitForReady: true,
+                CancellationToken.None).AsTask();
+            resolver.UpdateAddresses(new List<BalancerAddress>
+            {
+                new BalancerAddress("localhost", 80)
+            });
+
+            // Assert
+            // Should not timeout (deadlock)
+            await pickTask.DefaultTimeout();
+        }
+
+        [Test]
+        public async Task PickAsync_DoesNotDeadlockAfterReconnect_WithZeroAddressResolved()
+        {
+            // Arrange
+            var services = new ServiceCollection();
+            services.AddNUnitLogger();
+            await using var serviceProvider = services.BuildServiceProvider();
+            var loggerFactory = serviceProvider.GetRequiredService<ILoggerFactory>();
+
+            var resolver = new TestResolver(loggerFactory);
+
+            GrpcChannelOptions channelOptions = new GrpcChannelOptions();
+            channelOptions.ServiceConfig = new ServiceConfig()
+            {
+                LoadBalancingConfigs = { new RoundRobinConfig() }
+            };
+
+            var transportFactory = new TestSubchannelTransportFactory();
+            var clientChannel = CreateConnectionManager(loggerFactory, resolver, transportFactory, new[] { new RoundRobinBalancerFactory() });
+            // Configure balancer similar to how GrpcChannel constructor does it
+            clientChannel.ConfigureBalancer(c => new ChildHandlerLoadBalancer(
+                c,
+                channelOptions.ServiceConfig,
+                clientChannel));
+
+            // Act
+            var connectTask = clientChannel.ConnectAsync(waitForReady: true, cancellationToken: CancellationToken.None);
+            var pickTask = clientChannel.PickAsync(
+                new PickContext { Request = new HttpRequestMessage() },
+                waitForReady: true,
+                CancellationToken.None).AsTask();
+
+            resolver.UpdateAddresses(new List<BalancerAddress>
+            {
+                new BalancerAddress("localhost", 80)
+            });
+            await Task.WhenAll(connectTask, pickTask).DefaultTimeout();
+
+            // Simulate transport/network issue (with resolver reporting no addresses)
+            transportFactory.Transports.ForEach(t => t.Disconnect());
+            resolver.UpdateAddresses(new List<BalancerAddress>());
+
+            pickTask = clientChannel.PickAsync(
+                new PickContext { Request = new HttpRequestMessage() },
+                waitForReady: true,
+                CancellationToken.None).AsTask();
+            resolver.UpdateAddresses(new List<BalancerAddress>
+            {
+                new BalancerAddress("localhost", 80)
+            });
+
+            // Assert
+            // Should not timeout (deadlock)
+            await pickTask.DefaultTimeout();
+        }
+
         private static ConnectionManager CreateConnectionManager(
             ILoggerFactory loggerFactory,
             Resolver resolver,
-            TestSubchannelTransportFactory transportFactory)
+            TestSubchannelTransportFactory transportFactory,
+            LoadBalancerFactory[]? loadBalancerFactories = null)
         {
             return new ConnectionManager(
                 resolver,
                 disableResolverServiceConfig: false,
                 loggerFactory,
                 new TestBackoffPolicyFactory(),
                 transportFactory,
-                Array.Empty<LoadBalancerFactory>());
+                loadBalancerFactories ?? Array.Empty<LoadBalancerFactory>());
         }
 
         private class DropLoadBalancer : LoadBalancer

test/Grpc.Net.Client.Tests/Balancer/RoundRobinBalancerTests.cs

@@ -214,25 +214,34 @@ public async Task HasSubchannels_SubchannelStatusChanges_RefreshResolver()
             var services = new ServiceCollection();
             services.AddNUnitLogger();
 
+            ILogger? logger = null;
             SyncPoint? syncPoint = new SyncPoint(runContinuationsAsynchronously: true);
 
             var connectState = ConnectivityState.Ready;
 
-            var transportFactory = new TestSubchannelTransportFactory((s, c) => Task.FromResult(connectState));
+            var transportFactory = new TestSubchannelTransportFactory((s, c) =>
+            {
+                logger.LogInformation($"Transport factory returning state: {connectState}");
+                return Task.FromResult(connectState);
+            });
             services.AddSingleton<TestResolver>(s =>
             {
                 return new TestResolver(
                     s.GetRequiredService<ILoggerFactory>(),
                     async () =>
                     {
+                        logger.LogInformation("Resolver waiting to continue.");
                         await syncPoint.WaitToContinue().DefaultTimeout();
+
+                        logger.LogInformation("Resolver creating new sync point.");
                         syncPoint = new SyncPoint(runContinuationsAsynchronously: true);
                     });
             });
             services.AddSingleton<ResolverFactory, TestResolverFactory>();
             services.AddSingleton<ISubchannelTransportFactory>(transportFactory);
             var serviceProvider = services.BuildServiceProvider();
 
+            logger = serviceProvider.GetRequiredService<ILogger<RoundRobinBalancerTests>>();
             var handler = new TestHttpMessageHandler((r, ct) => default!);
             var channelOptions = new GrpcChannelOptions
             {
@@ -250,10 +259,13 @@ public async Task HasSubchannels_SubchannelStatusChanges_RefreshResolver()
 
             // Act
             var channel = GrpcChannel.ForAddress("test:///localhost", channelOptions);
+
+            logger.LogInformation("Client connecting");
             var connectTask = channel.ConnectAsync();
 
             // Assert
             syncPoint!.Continue();
+            logger.LogInformation("Client waiting for connect to complete.");
             await connectTask.DefaultTimeout();
 
             var subchannels = channel.ConnectionManager.GetSubchannels();
@@ -265,10 +277,12 @@ public async Task HasSubchannels_SubchannelStatusChanges_RefreshResolver()
             await transportFactory.Transports.Single().TryConnectTask.DefaultTimeout();
             Assert.AreEqual(ConnectivityState.Ready, subchannels[0].State);
 
+            logger.LogInformation("Transport factory updating state.");
             connectState = ConnectivityState.TransientFailure;
             transportFactory.Transports.Single().UpdateState(ConnectivityState.Idle);
 
             // Transport will refresh resolver after some failures
+            logger.LogInformation("Waiting for sync point in resolver.");
             await syncPoint!.WaitForSyncPoint().DefaultTimeout();
             syncPoint.Continue();
         }