core: SubchannelStateListener continues to receive updates after LoadBalancer is shutdown. (#5883)

No more methods on the `LoadBalancer` will be called after
`LoadBalancer#shutdown()` is called.  This includes
`LoadBalancer#handleSubchannelState()` too.  `SubchannelStateListener`
inherited this restriction.  However, this special case makes
`onSubchannelState(SHUTDOWN)` an unreliable way of being notified
about `Subchannel` SHUTDOWN, and may confuse/complicate a
wrapping `LoadBalancer` that expects the full notification (e.g., #5875).

The javadoc isn't clear whether this restriction applies.  I think
it's more useful to make it no apply.

api/src/main/java/io/grpc/LoadBalancer.java

@@ -349,9 +349,9 @@ public abstract class LoadBalancer {
   }
 
   /**
-   * The channel asks the load-balancer to shutdown.  No more callbacks will be called after this
+   * The channel asks the load-balancer to shutdown.  No more methods on this class will be called
-   * method.  The implementation should shutdown all Subchannels and OOB channels, and do any other
+   * after this method.  The implementation should shutdown all Subchannels and OOB channels, and do
-   * cleanup as necessary.
+   * any other cleanup as necessary.
    *
    * @since 1.2.0
    */

core/src/main/java/io/grpc/internal/ManagedChannelImpl.java

@@ -1076,6 +1076,10 @@ final class ManagedChannelImpl extends ManagedChannel implements
           new LoadBalancer.SubchannelStateListener() {
             @Override
             public void onSubchannelState(ConnectivityStateInfo newState) {
+              // Call LB only if it's not shutdown.  If LB is shutdown, lbHelper won't match.
+              if (LbHelperImpl.this != ManagedChannelImpl.this.lbHelper) {
+                return;
+              }
               lb.handleSubchannelState(subchannel, newState);
             }
           };
@@ -1459,12 +1463,9 @@ final class ManagedChannelImpl extends ManagedChannel implements
         @Override
         void onStateChange(InternalSubchannel is, ConnectivityStateInfo newState) {
           handleInternalSubchannelState(newState);
-          // Call LB only if it's not shutdown.  If LB is shutdown, lbHelper won't match.
-          if (helper == ManagedChannelImpl.this.lbHelper) {
           checkState(listener != null, "listener is null");
           listener.onSubchannelState(newState);
         }
-        }
 
         @Override
         void onInUse(InternalSubchannel is) {

core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java

@@ -784,11 +784,15 @@ public class ManagedChannelImplTest {
     subchannel1.shutdown();
     subchannel2.shutdown();
 
-    // No more callback are delivered to LoadBalancer or the state listeners after it's shut down
+    // Since subchannels are shutdown, SubchannelStateListeners will only get SHUTDOWN regardless of
+    // the transport states.
     transportInfo1.listener.transportShutdown(Status.UNAVAILABLE);
     transportInfo2.listener.transportReady();
+    verify(stateListener1).onSubchannelState(ConnectivityStateInfo.forNonError(SHUTDOWN));
+    verify(stateListener2).onSubchannelState(ConnectivityStateInfo.forNonError(SHUTDOWN));
     verifyNoMoreInteractions(stateListener1, stateListener2);
 
+    // No more callback should be delivered to LoadBalancer after it's shut down
     resolver.listener.onError(resolutionError);
     resolver.resolved();
     verifyNoMoreInteractions(mockLoadBalancer);