From 81d8259aaba97402f27df7d842cca219feb9a3a0 Mon Sep 17 00:00:00 2001
From: Kun Zhang <zhangkun@google.com>
Date: Thu, 15 Aug 2019 16:51:43 -0700
Subject: [PATCH 1/2] grpclb: fix pick_first mode shutdown without subchannels.

---
 .../main/java/io/grpc/grpclb/GrpclbState.java |  6 +++--
 .../grpc/grpclb/GrpclbLoadBalancerTest.java   | 22 +++++++++++++++++++
 2 files changed, 26 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..34b4b207709 100644
--- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
+++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
@@ -1813,6 +1813,28 @@ public void grpclbWorking_pickFirstMode() throws Exception {
         .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class));
   }
 
+  @Test
+  public void shutdownWithoutSubchannel_roundRobin() throws Exception {
+    InOrder inOrder = inOrder(helper);
+    String lbConfig = "{\"childPolicy\" : [ {\"round_robin\" : {}} ]}";
+    List<EquivalentAddressGroup> grpclbResolutionList = createResolvedServerAddresses(true);
+    Attributes grpclbResolutionAttrs = Attributes.newBuilder().set(
+        LoadBalancer.ATTR_LOAD_BALANCING_CONFIG, parseJsonObject(lbConfig)).build();
+    deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
+    balancer.shutdown();
+  }
+
+  @Test
+  public void shutdownWithoutSubchannel_pickFirst() throws Exception {
+    InOrder inOrder = inOrder(helper);
+    String lbConfig = "{\"childPolicy\" : [ {\"pick_first\" : {}} ]}";
+    List<EquivalentAddressGroup> grpclbResolutionList = createResolvedServerAddresses(true);
+    Attributes grpclbResolutionAttrs = Attributes.newBuilder().set(
+        LoadBalancer.ATTR_LOAD_BALANCING_CONFIG, parseJsonObject(lbConfig)).build();
+    deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
+    balancer.shutdown();
+  }
+
   @SuppressWarnings("deprecation")
   @Test
   public void grpclbWorking_pickFirstMode_expectBackendsShuffled() throws Exception {

From ef6963868008439c170bc6c9efe24fe8e55f4358 Mon Sep 17 00:00:00 2001
From: Kun Zhang <zhangkun@google.com>
Date: Thu, 15 Aug 2019 17:41:10 -0700
Subject: [PATCH 2/2] Improve test

---
 .../grpc/grpclb/GrpclbLoadBalancerTest.java   | 21 +++++++++++--------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
index 34b4b207709..4f01c89cfdc 100644
--- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
+++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
@@ -1815,24 +1815,27 @@ public void grpclbWorking_pickFirstMode() throws Exception {
 
   @Test
   public void shutdownWithoutSubchannel_roundRobin() throws Exception {
-    InOrder inOrder = inOrder(helper);
-    String lbConfig = "{\"childPolicy\" : [ {\"round_robin\" : {}} ]}";
-    List<EquivalentAddressGroup> grpclbResolutionList = createResolvedServerAddresses(true);
-    Attributes grpclbResolutionAttrs = Attributes.newBuilder().set(
-        LoadBalancer.ATTR_LOAD_BALANCING_CONFIG, parseJsonObject(lbConfig)).build();
-    deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
-    balancer.shutdown();
+    subtestShutdownWithoutSubchannel("round_robin");
   }
 
   @Test
   public void shutdownWithoutSubchannel_pickFirst() throws Exception {
-    InOrder inOrder = inOrder(helper);
-    String lbConfig = "{\"childPolicy\" : [ {\"pick_first\" : {}} ]}";
+    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")