Subchannel sockets transport cleanup and tests (#1562)

Author: JamesNK

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

@@ -58,7 +58,7 @@ internal record struct ActiveStream(BalancerAddress Address, Socket Socket, Stre
         private readonly Timer _socketConnectedTimer;
 
         private int _lastEndPointIndex;
-        private Socket? _initialSocket;
+        internal Socket? _initialSocket;
         private BalancerAddress? _initialSocketAddress;
         private bool _disposed;
         private BalancerAddress? _currentAddress;
@@ -89,19 +89,29 @@ public void Disconnect()
         {
             lock (Lock)
             {
-                _initialSocket?.Dispose();
-                _initialSocket = null;
-                _initialSocketAddress = null;
-                _lastEndPointIndex = 0;
-                if (!_disposed)
+                if (_disposed)
                 {
-                    _socketConnectedTimer.Change(TimeSpan.Zero, TimeSpan.Zero);
+                    return;
                 }
-                _currentAddress = null;
+
+                DisconnectUnsynchronized();
+                _socketConnectedTimer.Change(TimeSpan.Zero, TimeSpan.Zero);
             }
             _subchannel.UpdateConnectivityState(ConnectivityState.Idle, "Disconnected.");
         }
 
+        private void DisconnectUnsynchronized()
+        {
+            Debug.Assert(Monitor.IsEntered(Lock));
+            Debug.Assert(!_disposed);
+
+            _initialSocket?.Dispose();
+            _initialSocket = null;
+            _initialSocketAddress = null;
+            _lastEndPointIndex = 0;
+            _currentAddress = null;
+        }
+
         public async ValueTask<bool> TryConnectAsync(CancellationToken cancellationToken)
         {
             Debug.Assert(CurrentAddress == null);
@@ -196,13 +206,14 @@ private async void OnCheckSocketConnection(object? state)
                     {
                         lock (Lock)
                         {
+                            if (_disposed)
+                            {
+                                return;
+                            }
+
                             if (_initialSocket == socket)
                             {
-                                _initialSocket.Dispose();
-                                _initialSocket = null;
-                                _initialSocketAddress = null;
-                                _currentAddress = null;
-                                _lastEndPointIndex = 0;
+                                DisconnectUnsynchronized();
                             }
                         }
                         _subchannel.UpdateConnectivityState(ConnectivityState.Idle, new Status(StatusCode.Unavailable, "Lost connection to socket.", sendException));
@@ -278,45 +289,58 @@ private static bool IsSocketInBadState(Socket socket)
 
         private void OnStreamDisposed(Stream streamWrapper)
         {
-            var disconnect = false;
-            lock (Lock)
+            try
             {
-                for (var i = _activeStreams.Count - 1; i >= 0; i--)
+                var disconnect = false;
+                lock (Lock)
                 {
-                    var t = _activeStreams[i];
-                    if (t.Stream == streamWrapper)
+                    for (var i = _activeStreams.Count - 1; i >= 0; i--)
                     {
-                        _activeStreams.RemoveAt(i);
-                        SocketConnectivitySubchannelTransportLog.DisposingStream(_logger, _subchannel.Id, t.Address);
+                        var t = _activeStreams[i];
+                        if (t.Stream == streamWrapper)
+                        {
+                            _activeStreams.RemoveAt(i);
+                            SocketConnectivitySubchannelTransportLog.DisposingStream(_logger, _subchannel.Id, t.Address);
 
-                        // If the last active streams is removed then there is no active connection.
-                        disconnect = _activeStreams.Count == 0;
+                            // If the last active streams is removed then there is no active connection.
+                            disconnect = _activeStreams.Count == 0;
 
-                        break;
+                            break;
+                        }
                     }
                 }
-            }
 
-            if (disconnect)
+                if (disconnect)
+                {
+                    // What happens after disconnect depends if the load balancer requests a new connection.
+                    // For example:
+                    // - Pick first will go into an idle state.
+                    // - Round-robin will reconnect to get back to a ready state.
+                    Disconnect();
+                }
+            }
+            catch (Exception ex)
             {
-                // What happens after disconnect depends if the load balancer requests a new connection.
-                // For example:
-                // - Pick first will go into an idle state.
-                // - Round-robin will reconnect to get back to a ready state.
-                Disconnect();
+                // Don't throw error to Stream.Dispose() caller.
+                SocketConnectivitySubchannelTransportLog.ErrorOnDisposingStream(_logger, _subchannel.Id, ex);
             }
         }
 
         public void Dispose()
         {
             lock (Lock)
             {
-                if (!_disposed)
+                if (_disposed)
                 {
-                    SocketConnectivitySubchannelTransportLog.DisposingTransport(_logger, _subchannel.Id);
-                    _socketConnectedTimer.Dispose();
-                    _disposed = true;
+                    return;
                 }
+
+                SocketConnectivitySubchannelTransportLog.DisposingTransport(_logger, _subchannel.Id);
+
+                DisconnectUnsynchronized();
+
+                _socketConnectedTimer.Dispose();
+                _disposed = true;
             }
         }
 
@@ -354,6 +378,9 @@ internal static class SocketConnectivitySubchannelTransportLog
         private static readonly Action<ILogger, int, Exception?> _disposingTransport =
             LoggerMessage.Define<int>(LogLevel.Trace, new EventId(9, "DisposingTransport"), "Subchannel id '{SubchannelId}' disposing transport.");
 
+        private static readonly Action<ILogger, int, Exception> _errorOnDisposingStream =
+            LoggerMessage.Define<int>(LogLevel.Error, new EventId(10, "ErrorOnDisposingStream"), "Subchannel id '{SubchannelId}' unexpected error when reacting to transport stream dispose.");
+
         public static void ConnectingSocket(ILogger logger, int subchannelId, BalancerAddress address)
         {
             _connectingSocket(logger, subchannelId, address, null);
@@ -398,6 +425,11 @@ public static void DisposingTransport(ILogger logger, int subchannelId)
         {
             _disposingTransport(logger, subchannelId, null);
         }
+
+        public static void ErrorOnDisposingStream(ILogger logger, int subchannelId, Exception ex)
+        {
+            _errorOnDisposingStream(logger, subchannelId, ex);
+        }
     }
 #endif
 }

test/FunctionalTests/Balancer/BalancerHelpers.cs

@@ -76,6 +76,7 @@ public EndpointContext(GrpcTestFixture<Startup> server, Method<TRequest, TRespon
             public Method<TRequest, TResponse> Method { get; }
             public Uri Address { get; }
             public ILoggerFactory LoggerFactory => _server.LoggerFactory;
+            public EndPoint EndPoint => new DnsEndPoint(Address.Host, Address.Port);
 
             public void Dispose()
             {
@@ -172,7 +173,13 @@ public static async Task WaitForChannelStatesAsync(ILogger logger, GrpcChannel c
             logger.LogInformation($"Channel id {channelId}: Current channel state '{currentState}' matches expected states {statesText}.");
         }
 
-        public static async Task<Subchannel[]> WaitForSubChannelsToBeReadyAsync(ILogger logger, GrpcChannel channel, int expectedCount, Func<SubchannelPicker?, Subchannel[]>? getPickerSubchannels = null)
+        public static async Task<Subchannel> WaitForSubchannelToBeReadyAsync(ILogger logger, GrpcChannel channel, Func<SubchannelPicker?, Subchannel[]>? getPickerSubchannels = null)
+        {
+            var subChannel = (await WaitForSubchannelsToBeReadyAsync(logger, channel, 1)).Single();
+            return subChannel;
+        }
+
+        public static async Task<Subchannel[]> WaitForSubchannelsToBeReadyAsync(ILogger logger, GrpcChannel channel, int expectedCount, Func<SubchannelPicker?, Subchannel[]>? getPickerSubchannels = null)
         {
             if (getPickerSubchannels == null)
             {

test/FunctionalTests/Balancer/LeastUsedBalancerTests.cs

@@ -67,7 +67,7 @@ async Task<HelloReply> UnaryMethod(HelloRequest request, ServerCallContext conte
 
             var channel = await BalancerHelpers.CreateChannel(LoggerFactory, new LoadBalancingConfig("least_used"), new[] { endpoint1.Address, endpoint2.Address }, connect: true);
 
-            await BalancerHelpers.WaitForSubChannelsToBeReadyAsync(
+            await BalancerHelpers.WaitForSubchannelsToBeReadyAsync(
                 Logger,
                 channel,
                 expectedCount: 2,

test/FunctionalTests/Balancer/PickFirstBalancerTests.cs

@@ -194,7 +194,7 @@ Task<HelloReply> UnaryMethod(HelloRequest request, ServerCallContext context)
             Logger.LogInformation("Ending " + endpoint1.Address);
             endpoint1.Dispose();
 
-            await BalancerHelpers.WaitForSubChannelsToBeReadyAsync(Logger, channel, expectedCount: 1,
+            await BalancerHelpers.WaitForSubchannelsToBeReadyAsync(Logger, channel, expectedCount: 1,
                 getPickerSubchannels: picker=>
                 {
                     // We want a subchannel that has no current address