From f60ed8aa8ee1ebbcfebd33851941784d8820b32d Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Mon, 7 Dec 2020 15:10:33 -0800
Subject: [PATCH] eds/lrs: handle nil when LRS is disabled (#4086)

---
 .../balancer/edsbalancer/eds_impl_test.go     | 30 +++++++++++++++++++
 .../edsbalancer/xds_client_wrapper.go         | 16 +++++++---
 xds/internal/balancer/lrs/balancer.go         | 16 +++++++---
 3 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/xds/internal/balancer/edsbalancer/eds_impl_test.go b/xds/internal/balancer/edsbalancer/eds_impl_test.go
index 292ea4f80a70..159f7b2aa805 100644
--- a/xds/internal/balancer/edsbalancer/eds_impl_test.go
+++ b/xds/internal/balancer/edsbalancer/eds_impl_test.go
@@ -746,3 +746,33 @@ func (s) TestEDS_LoadReport(t *testing.T) {
 		t.Errorf("store.stats() returned unexpected diff (-want +got):\n%s", diff)
 	}
 }
+
+// TestEDS_LoadReportDisabled covers the case that LRS is disabled. It makes
+// sure the EDS implementation isn't broken (doesn't panic).
+func (s) TestEDS_LoadReportDisabled(t *testing.T) {
+	lsWrapper := &loadStoreWrapper{}
+	lsWrapper.updateServiceName(testClusterNames[0])
+	// Not calling lsWrapper.updateLoadStore(loadStore) because LRS is disabled.
+	cw := &xdsClientWrapper{
+		loadWrapper: lsWrapper,
+	}
+
+	cc := testutils.NewTestClientConn(t)
+	edsb := newEDSBalancerImpl(cc, nil, cw, nil)
+	edsb.enqueueChildBalancerStateUpdate = edsb.updateState
+
+	// One localities, with one backend.
+	clab1 := testutils.NewClusterLoadAssignmentBuilder(testClusterNames[0], nil)
+	clab1.AddLocality(testSubZones[0], 1, 0, testEndpointAddrs[:1], nil)
+	edsb.handleEDSResponse(parseEDSRespProtoForTesting(clab1.Build()))
+	sc1 := <-cc.NewSubConnCh
+	edsb.handleSubConnStateChange(sc1, connectivity.Connecting)
+	edsb.handleSubConnStateChange(sc1, connectivity.Ready)
+
+	// Test roundrobin with two subconns.
+	p1 := <-cc.NewPickerCh
+	// We call picks to make sure they don't panic.
+	for i := 0; i < 10; i++ {
+		p1.Pick(balancer.PickInfo{})
+	}
+}
diff --git a/xds/internal/balancer/edsbalancer/xds_client_wrapper.go b/xds/internal/balancer/edsbalancer/xds_client_wrapper.go
index 16466016bbcd..39dcb7aab768 100644
--- a/xds/internal/balancer/edsbalancer/xds_client_wrapper.go
+++ b/xds/internal/balancer/edsbalancer/xds_client_wrapper.go
@@ -72,25 +72,33 @@ func (lsw *loadStoreWrapper) updateLoadStore(store *load.Store) {
 func (lsw *loadStoreWrapper) CallStarted(locality string) {
 	lsw.mu.RLock()
 	defer lsw.mu.RUnlock()
-	lsw.perCluster.CallStarted(locality)
+	if lsw.perCluster != nil {
+		lsw.perCluster.CallStarted(locality)
+	}
 }
 
 func (lsw *loadStoreWrapper) CallFinished(locality string, err error) {
 	lsw.mu.RLock()
 	defer lsw.mu.RUnlock()
-	lsw.perCluster.CallFinished(locality, err)
+	if lsw.perCluster != nil {
+		lsw.perCluster.CallFinished(locality, err)
+	}
 }
 
 func (lsw *loadStoreWrapper) CallServerLoad(locality, name string, val float64) {
 	lsw.mu.RLock()
 	defer lsw.mu.RUnlock()
-	lsw.perCluster.CallServerLoad(locality, name, val)
+	if lsw.perCluster != nil {
+		lsw.perCluster.CallServerLoad(locality, name, val)
+	}
 }
 
 func (lsw *loadStoreWrapper) CallDropped(category string) {
 	lsw.mu.RLock()
 	defer lsw.mu.RUnlock()
-	lsw.perCluster.CallDropped(category)
+	if lsw.perCluster != nil {
+		lsw.perCluster.CallDropped(category)
+	}
 }
 
 // xdsclientWrapper is responsible for getting the xds client from attributes or
diff --git a/xds/internal/balancer/lrs/balancer.go b/xds/internal/balancer/lrs/balancer.go
index f8e7673f7d8e..4ed1d61c39ff 100644
--- a/xds/internal/balancer/lrs/balancer.go
+++ b/xds/internal/balancer/lrs/balancer.go
@@ -187,25 +187,33 @@ func (lsw *loadStoreWrapper) updateLoadStore(store *load.Store) {
 func (lsw *loadStoreWrapper) CallStarted(locality string) {
 	lsw.mu.RLock()
 	defer lsw.mu.RUnlock()
-	lsw.perCluster.CallStarted(locality)
+	if lsw.perCluster != nil {
+		lsw.perCluster.CallStarted(locality)
+	}
 }
 
 func (lsw *loadStoreWrapper) CallFinished(locality string, err error) {
 	lsw.mu.RLock()
 	defer lsw.mu.RUnlock()
-	lsw.perCluster.CallFinished(locality, err)
+	if lsw.perCluster != nil {
+		lsw.perCluster.CallFinished(locality, err)
+	}
 }
 
 func (lsw *loadStoreWrapper) CallServerLoad(locality, name string, val float64) {
 	lsw.mu.RLock()
 	defer lsw.mu.RUnlock()
-	lsw.perCluster.CallServerLoad(locality, name, val)
+	if lsw.perCluster != nil {
+		lsw.perCluster.CallServerLoad(locality, name, val)
+	}
 }
 
 func (lsw *loadStoreWrapper) CallDropped(category string) {
 	lsw.mu.RLock()
 	defer lsw.mu.RUnlock()
-	lsw.perCluster.CallDropped(category)
+	if lsw.perCluster != nil {
+		lsw.perCluster.CallDropped(category)
+	}
 }
 
 type xdsClientWrapper struct {