Remove Subchannel.State and track state inside load balancer (#1590)

Author: JamesNK

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

@@ -83,7 +83,7 @@ public void Complete(CompletionContext context)
         /// A result created with a <see cref="Balancer.Subchannel"/> won't necessarily be used by a gRPC call.
         /// The subchannel's state may change at the same time the picker is making a decision. That means the
         /// decision may be made with outdated information. For example, a picker may return a subchannel
-        /// with a <see cref="Subchannel.State"/> that is <see cref="ConnectivityState.Ready"/>, but
+        /// with a state that is <see cref="ConnectivityState.Ready"/>, but
         /// becomes <see cref="ConnectivityState.Idle"/> when the subchannel is about to be used. In that situation
         /// the gRPC call waits for the load balancer to react to the new state and create a new picker.
         /// </para>

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

@@ -48,6 +48,14 @@ public sealed class Subchannel : IDisposable
         internal ISubchannelTransport Transport { get; set; } = default!;
         internal int Id { get; }
 
+        /// <summary>
+        /// Connectivity state is internal rather than public because it can be updated by multiple threads while
+        /// a load balancer is building the picker.
+        /// Load balancers that care about multiple subchannels should track state by subscribing to
+        /// Subchannel.OnStateChanged and storing results.
+        /// </summary>
+        internal ConnectivityState State => _state;
+
         private readonly ConnectionManager _manager;
         private readonly ILogger _logger;
 
@@ -60,11 +68,6 @@ public sealed class Subchannel : IDisposable
         /// </summary>
         public BalancerAddress? CurrentAddress => Transport.CurrentAddress;
 
-        /// <summary>
-        /// Gets the connectivity state.
-        /// </summary>
-        public ConnectivityState State => _state;
-
         /// <summary>
         /// Gets the metadata attributes.
         /// </summary>
@@ -327,7 +330,7 @@ public override string ToString()
         {
             lock (Lock)
             {
-                return $"Id: {Id}, Addresses: {string.Join(", ", _addresses)}, State: {State}, Current address: {CurrentAddress}";
+                return $"Id: {Id}, Addresses: {string.Join(", ", _addresses)}, State: {_state}, Current address: {CurrentAddress}";
             }
         }
 
@@ -345,7 +348,7 @@ public IReadOnlyList<BalancerAddress> GetAddresses()
 
         /// <summary>
         /// Disposes the <see cref="Subchannel"/>.
-        /// The subchannel <see cref="State"/> is updated to <see cref="ConnectivityState.Shutdown"/>.
+        /// The subchannel state is updated to <see cref="ConnectivityState.Shutdown"/>.
         /// After dispose the subchannel should no longer be returned by the latest <see cref="SubchannelPicker"/>.
         /// </summary>
         public void Dispose()

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

@@ -94,14 +94,14 @@ private void ResolverError(Status status)
             return null;
         }
 
-        private int? FindSubchannel(List<AddressSubchannel> addressSubchannels, Subchannel subchannel)
+        private AddressSubchannel? FindSubchannel(List<AddressSubchannel> addressSubchannels, Subchannel subchannel)
         {
             for (var i = 0; i < addressSubchannels.Count; i++)
             {
                 var s = addressSubchannels[i];
                 if (Equals(s.Subchannel, subchannel))
                 {
-                    return i;
+                    return s;
                 }
             }
 
@@ -189,7 +189,7 @@ private void UpdateBalancingState(Status status)
             for (var i = 0; i < _addressSubchannels.Count; i++)
             {
                 var addressSubchannel = _addressSubchannels[i];
-                if (addressSubchannel.Subchannel.State == ConnectivityState.Ready)
+                if (addressSubchannel.LastKnownState == ConnectivityState.Ready)
                 {
                     readySubchannels.Add(addressSubchannel.Subchannel);
                 }
@@ -201,7 +201,7 @@ private void UpdateBalancingState(Status status)
                 var isConnecting = false;
                 foreach (var subchannel in _addressSubchannels)
                 {
-                    var state = subchannel.Subchannel.State;
+                    var state = subchannel.LastKnownState;
 
                     if (state == ConnectivityState.Connecting || state == ConnectivityState.Idle)
                     {
@@ -239,13 +239,14 @@ private void UpdateChannelState(ConnectivityState state, SubchannelPicker subcha
 
         private void UpdateSubchannelState(Subchannel subchannel, SubchannelState state)
         {
-            var index = FindSubchannel(_addressSubchannels, subchannel);
-            if (index == null)
+            var addressSubchannel = FindSubchannel(_addressSubchannels, subchannel);
+            if (addressSubchannel == null)
             {
                 SubchannelsLoadBalancerLog.IgnoredSubchannelStateChange(_logger, subchannel.Id);
                 return;
             }
 
+            addressSubchannel.UpdateKnownState(state.State);
             SubchannelsLoadBalancerLog.ProcessingSubchannelStateChanged(_logger, subchannel.Id, state.State, state.Status);
 
             UpdateBalancingState(state.Status);
@@ -297,7 +298,29 @@ protected override void Dispose(bool disposing)
         /// <returns>A subchannel picker.</returns>
         protected abstract SubchannelPicker CreatePicker(IReadOnlyList<Subchannel> readySubchannels);
 
-        private record AddressSubchannel(Subchannel Subchannel, BalancerAddress Address);
+        private class AddressSubchannel
+        {
+            private ConnectivityState _lastKnownState;
+
+            public AddressSubchannel(Subchannel subchannel, BalancerAddress address)
+            {
+                Subchannel = subchannel;
+                Address = address;
+                _lastKnownState = ConnectivityState.Idle;
+            }
+
+            // Track connectivity state that has been updated to load balancer.
+            // This is used instead of state on subchannel because subchannel state
+            // can be updated from other threads while load balancer is running.
+            public ConnectivityState LastKnownState => _lastKnownState;
+            public Subchannel Subchannel { get; }
+            public BalancerAddress Address { get; }
+
+            public void UpdateKnownState(ConnectivityState knownState)
+            {
+                _lastKnownState = knownState;
+            }
+        }
     }
 
     internal static class SubchannelsLoadBalancerLog