From 28587b449bfcb1aaee0032ba60d273d3f2ef9d78 Mon Sep 17 00:00:00 2001
From: Kun Zhang <zhangkun83@users.noreply.github.com>
Date: Wed, 5 Dec 2018 16:35:42 -0800
Subject: [PATCH] grpclb: refresh name resolution when balancer stream closed.
 (#5127)

---
 grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java   |  2 ++
 .../java/io/grpc/grpclb/GrpclbLoadBalancerTest.java    | 10 ++++++++--
 2 files changed, 10 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 c2732466bc..62c4cf26c9 100644
--- a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java
+++ b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java
@@ -583,6 +583,8 @@ final class GrpclbState {
             syncContext.schedule(new LbRpcRetryTask(), delayNanos, TimeUnit.NANOSECONDS,
                 timerService);
       }
+
+      helper.refreshNameResolution();
     }
 
     void close(@Nullable Exception error) {
diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
index cc4bffdadd..675fdc4315 100644
--- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
+++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
@@ -1467,14 +1467,14 @@ public class GrpclbLoadBalancerTest {
   }
 
   @Test
-  public void grpclbBalancerStreamRetry() throws Exception {
+  public void grpclbBalancerStreamClosedAndRetried() throws Exception {
     LoadBalanceRequest expectedInitialRequest =
         LoadBalanceRequest.newBuilder()
             .setInitialRequest(
                 InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build())
             .build();
     InOrder inOrder =
-        inOrder(mockLbService, backoffPolicyProvider, backoffPolicy1, backoffPolicy2);
+        inOrder(mockLbService, backoffPolicyProvider, backoffPolicy1, backoffPolicy2, helper);
     List<EquivalentAddressGroup> grpclbResolutionList = createResolvedServerAddresses(true);
     Attributes grpclbResolutionAttrs = Attributes.EMPTY;
     deliverResolvedAddresses(grpclbResolutionList, grpclbResolutionAttrs);
@@ -1493,10 +1493,12 @@ public class GrpclbLoadBalancerTest {
 
     // Balancer closes it immediately (erroneously)
     lbResponseObserver.onCompleted();
+
     // Will start backoff sequence 1 (10ns)
     inOrder.verify(backoffPolicyProvider).get();
     inOrder.verify(backoffPolicy1).nextBackoffNanos();
     assertEquals(1, fakeClock.numPendingTasks(LB_RPC_RETRY_TASK_FILTER));
+    inOrder.verify(helper).refreshNameResolution();
 
     // Fast-forward to a moment before the retry
     fakeClock.forwardNanos(9);
@@ -1516,6 +1518,7 @@ public class GrpclbLoadBalancerTest {
     verifyNoMoreInteractions(backoffPolicyProvider);
     inOrder.verify(backoffPolicy1).nextBackoffNanos();
     assertEquals(1, fakeClock.numPendingTasks(LB_RPC_RETRY_TASK_FILTER));
+    inOrder.verify(helper).refreshNameResolution();
 
     // Fast-forward to a moment before the retry
     fakeClock.forwardNanos(100 - 1);
@@ -1542,6 +1545,7 @@ public class GrpclbLoadBalancerTest {
     assertEquals(1, lbRequestObservers.size());
     lbRequestObserver = lbRequestObservers.poll();
     verify(lbRequestObserver).onNext(eq(expectedInitialRequest));
+    inOrder.verify(helper).refreshNameResolution();
 
     // Fail the retry after spending 4ns
     fakeClock.forwardNanos(4);
@@ -1550,6 +1554,7 @@ public class GrpclbLoadBalancerTest {
     // Will be on the first retry (10ns) of backoff sequence 2.
     inOrder.verify(backoffPolicy2).nextBackoffNanos();
     assertEquals(1, fakeClock.numPendingTasks(LB_RPC_RETRY_TASK_FILTER));
+    inOrder.verify(helper).refreshNameResolution();
 
     // Fast-forward to a moment before the retry, the time spent in the last try is deducted.
     fakeClock.forwardNanos(10 - 4 - 1);
@@ -1566,6 +1571,7 @@ public class GrpclbLoadBalancerTest {
     verify(backoffPolicyProvider, times(2)).get();
     verify(backoffPolicy1, times(2)).nextBackoffNanos();
     verify(backoffPolicy2, times(1)).nextBackoffNanos();
+    verify(helper, times(4)).refreshNameResolution();
   }
 
   private void deliverSubchannelState(