From c92c9e71bbc7fd436d604c67672ca107cbcb22ef Mon Sep 17 00:00:00 2001
From: Kun Zhang <zhangkun83@users.noreply.github.com>
Date: Fri, 14 Jun 2019 14:43:54 -0700
Subject: [PATCH]  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           |  6 +++---
 .../java/io/grpc/internal/ManagedChannelImpl.java     | 11 ++++++-----
 .../java/io/grpc/internal/ManagedChannelImplTest.java |  6 +++++-
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/api/src/main/java/io/grpc/LoadBalancer.java b/api/src/main/java/io/grpc/LoadBalancer.java
index 1b390bf268..612d25151f 100644
--- a/api/src/main/java/io/grpc/LoadBalancer.java
+++ b/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
-   * method.  The implementation should shutdown all Subchannels and OOB channels, and do any other
-   * cleanup as necessary.
+   * The channel asks the load-balancer to shutdown.  No more methods on this class will be called
+   * after this method.  The implementation should shutdown all Subchannels and OOB channels, and do
+   * any other cleanup as necessary.
    *
    * @since 1.2.0
    */
diff --git a/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java b/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
index f95e938a69..eea108e0d9 100644
--- a/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
+++ b/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,11 +1463,8 @@ 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);
-          }
+          checkState(listener != null, "listener is null");
+          listener.onSubchannelState(newState);
         }
 
         @Override
diff --git a/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java b/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java
index f133fba939..b9b5ac3c5b 100644
--- a/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java
+++ b/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);