From 2c0b2de8628e62b9c2dc0a23752f7c50322a3f87 Mon Sep 17 00:00:00 2001
From: Kun Zhang <zhangkun83@users.noreply.github.com>
Date: Fri, 16 Aug 2019 10:21:25 -0700
Subject: [PATCH] grpclb: fix pick_first mode shutdown without subchannels.
 (#6072)

---
 .../main/java/io/grpc/grpclb/GrpclbState.java |  6 +++--
 .../grpc/grpclb/GrpclbLoadBalancerTest.java   | 25 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java
index 125875d2455..7c024843773 100644
--- a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java
+++ b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java
@@ -341,8 +341,10 @@ void shutdown() {
         subchannelPool.clear();
         break;
       case PICK_FIRST:
-        checkState(subchannels.size() == 1, "Excessive Subchannels: %s", subchannels);
-        subchannels.values().iterator().next().shutdown();
+        if (!subchannels.isEmpty()) {
+          checkState(subchannels.size() == 1, "Excessive Subchannels: %s", subchannels);
+          subchannels.values().iterator().next().shutdown();
+        }
         break;
       default:
         throw new AssertionError("Missing case for " + mode);
diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
index a3f84cbbe6f..4f01c89cfdc 100644
--- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
+++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
@@ -1813,6 +1813,31 @@ public void grpclbWorking_pickFirstMode() throws Exception {
         .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class));
   }
 
+  @Test
+  public void shutdownWithoutSubchannel_roundRobin() throws Exception {
+    subtestShutdownWithoutSubchannel("round_robin");
+  }
+
+  @Test
+  public void shutdownWithoutSubchannel_pickFirst() throws Exception {
+    subtestShutdownWithoutSubchannel("pick_first");
+  }
+
+  private void subtestShutdownWithoutSubchannel(String childPolicy) throws Exception {
+    String lbConfig = "{\"childPolicy\" : [ {\"" + childPolicy + "\" : {}} ]}";
+    List<EquivalentAddressGroup> grpclbResolutionList = createResolvedServerAddresses(true);
+    Attributes grpclbResolutionAttrs = Attributes.newBuilder().set(
+        LoadBalancer.ATTR_LOAD_BALANCING_CONFIG, parseJsonObject(lbConfig)).build();
+    deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
+    verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture());
+    assertEquals(1, lbRequestObservers.size());
+    StreamObserver<LoadBalanceRequest> requestObserver = lbRequestObservers.poll();
+
+    verify(requestObserver, never()).onCompleted();
+    balancer.shutdown();
+    verify(requestObserver).onCompleted();
+  }
+
   @SuppressWarnings("deprecation")
   @Test
   public void grpclbWorking_pickFirstMode_expectBackendsShuffled() throws Exception {