From 842dfa49b17085280b7d2f56b392281c319d5c3f Mon Sep 17 00:00:00 2001
From: ZHANG Dapeng <zdapeng@google.com>
Date: Wed, 11 Sep 2019 14:00:18 -0700
Subject: [PATCH] xds: Fix a bug locality reset not fully clean up

There is a bug in LocalityStore.reset() found during import.
---
 xds/src/main/java/io/grpc/xds/LocalityStore.java     | 1 +
 xds/src/test/java/io/grpc/xds/LocalityStoreTest.java | 9 +++++++++
 2 files changed, 10 insertions(+)

diff --git a/xds/src/main/java/io/grpc/xds/LocalityStore.java b/xds/src/main/java/io/grpc/xds/LocalityStore.java
index 0e2382758c..def3b99a06 100644
--- a/xds/src/main/java/io/grpc/xds/LocalityStore.java
+++ b/xds/src/main/java/io/grpc/xds/LocalityStore.java
@@ -192,6 +192,7 @@ interface LocalityStore {
       for (XdsLocality locality : edsResponsLocalityInfo.keySet()) {
         loadStatsStore.removeLocality(locality);
       }
+      edsResponsLocalityInfo = ImmutableMap.of();
     }
 
     // This is triggered by EDS response.
diff --git a/xds/src/test/java/io/grpc/xds/LocalityStoreTest.java b/xds/src/test/java/io/grpc/xds/LocalityStoreTest.java
index 4e78905b4b..73adcc1be5 100644
--- a/xds/src/test/java/io/grpc/xds/LocalityStoreTest.java
+++ b/xds/src/test/java/io/grpc/xds/LocalityStoreTest.java
@@ -831,6 +831,15 @@ public class LocalityStoreTest {
     verify(loadBalancers.get("sz2")).shutdown();
     verify(loadStatsStore).removeLocality(locality1);
     verify(loadStatsStore).removeLocality(locality2);
+
+    // Regression test for same locality added back.
+    localityStore.updateLocalityStore(localityInfoMap);
+    assertThat(loadBalancers).hasSize(2);
+    localityStore.reset();
+    verify(loadBalancers.get("sz1")).shutdown();
+    verify(loadBalancers.get("sz2")).shutdown();
+    verify(loadStatsStore, times(2)).removeLocality(locality1);
+    verify(loadStatsStore, times(2)).removeLocality(locality2);
   }
 
   private static final class FakeLoadStatsStore implements LoadStatsStore {