From 6eb24e85afa0b04cafa9e7bd0acc5ad006660092 Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Tue, 10 Mar 2020 09:37:52 -0700
Subject: [PATCH] xds: handle removing drops (#3439)

---
 xds/internal/balancer/edsbalancer/eds_impl.go      |  3 +++
 xds/internal/balancer/edsbalancer/eds_impl_test.go | 14 ++++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/xds/internal/balancer/edsbalancer/eds_impl.go b/xds/internal/balancer/edsbalancer/eds_impl.go
index 1926e38f206..e6dec39571d 100644
--- a/xds/internal/balancer/edsbalancer/eds_impl.go
+++ b/xds/internal/balancer/edsbalancer/eds_impl.go
@@ -169,6 +169,9 @@ func (edsImpl *edsBalancerImpl) updateDrops(dropPolicies []xdsclient.OverloadDro
 			dropsChanged = true
 		}
 	}
+	if len(edsImpl.drops) != len(newDrops) {
+		dropsChanged = true
+	}
 	if dropsChanged {
 		edsImpl.pickerMu.Lock()
 		edsImpl.drops = newDrops
diff --git a/xds/internal/balancer/edsbalancer/eds_impl_test.go b/xds/internal/balancer/edsbalancer/eds_impl_test.go
index c26b857c9ae..60b6f23befb 100644
--- a/xds/internal/balancer/edsbalancer/eds_impl_test.go
+++ b/xds/internal/balancer/edsbalancer/eds_impl_test.go
@@ -148,6 +148,20 @@ func (s) TestEDS_OneLocality(t *testing.T) {
 			t.Errorf("The second 50%% picks should be non-drops, got error %v", err)
 		}
 	}
+
+	// The same locality, remove drops.
+	clab6 := xdsclient.NewClusterLoadAssignmentBuilder(testClusterNames[0], nil)
+	clab6.AddLocality(testSubZones[0], 1, 0, testEndpointAddrs[2:3], nil)
+	edsb.HandleEDSResponse(xdsclient.ParseEDSRespProtoForTesting(clab6.Build()))
+
+	// Pick without drops.
+	p6 := <-cc.newPickerCh
+	for i := 0; i < 5; i++ {
+		gotSCSt, _ := p6.Pick(balancer.PickInfo{})
+		if !cmp.Equal(gotSCSt.SubConn, sc3, cmp.AllowUnexported(testSubConn{})) {
+			t.Fatalf("picker.Pick, got %v, want SubConn=%v", gotSCSt, sc3)
+		}
+	}
 }
 
 // 2 locality