From cf032b300e509b9ec4398333749710f8597cc8fc Mon Sep 17 00:00:00 2001
From: "Mark D. Roth" <roth@google.com>
Date: Tue, 4 Feb 2020 15:13:53 -0800
Subject: [PATCH] xds: Fix crash when moving all localities from a priority to
 a higher priority.

---
 .../client_channel/lb_policy/xds/xds.cc       |  2 +
 .../filters/client_channel/xds/xds_client.cc  |  4 +-
 test/cpp/end2end/xds_end2end_test.cc          | 40 ++++++++++++++++++-
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/src/core/ext/filters/client_channel/lb_policy/xds/xds.cc b/src/core/ext/filters/client_channel/lb_policy/xds/xds.cc
index 1c8c46eaace..d735adfbc0f 100644
--- a/src/core/ext/filters/client_channel/lb_policy/xds/xds.cc
+++ b/src/core/ext/filters/client_channel/lb_policy/xds/xds.cc
@@ -1032,6 +1032,8 @@ void XdsLb::UpdatePrioritiesLocked() {
   for (uint32_t priority = 0; priority < priorities_.size(); ++priority) {
     LocalityMap* locality_map = priorities_[priority].get();
     const auto* locality_map_update = priority_list_update_.Find(priority);
+    // If we have more current priorities than exist in the update, stop here.
+    if (locality_map_update == nullptr) break;
     // Propagate locality_map_update.
     // TODO(juanlishen): Find a clean way to skip duplicate update for a
     // priority.
diff --git a/src/core/ext/filters/client_channel/xds/xds_client.cc b/src/core/ext/filters/client_channel/xds/xds_client.cc
index ee965c01ee5..e157a4305e4 100644
--- a/src/core/ext/filters/client_channel/xds/xds_client.cc
+++ b/src/core/ext/filters/client_channel/xds/xds_client.cc
@@ -1015,9 +1015,9 @@ void XdsClient::ChannelState::AdsCallState::AcceptEdsUpdate(
           const auto& locality = p.second;
           gpr_log(GPR_INFO,
                   "[xds_client %p] Priority %" PRIuPTR ", locality %" PRIuPTR
-                  " %s contains %" PRIuPTR " server addresses",
+                  " %s has weight %d, contains %" PRIuPTR " server addresses",
                   xds_client(), priority, locality_count,
-                  locality.name->AsHumanReadableString(),
+                  locality.name->AsHumanReadableString(), locality.lb_weight,
                   locality.serverlist.size());
           for (size_t i = 0; i < locality.serverlist.size(); ++i) {
             char* ipport;
diff --git a/test/cpp/end2end/xds_end2end_test.cc b/test/cpp/end2end/xds_end2end_test.cc
index f86d032a878..b3b0867fc1b 100644
--- a/test/cpp/end2end/xds_end2end_test.cc
+++ b/test/cpp/end2end/xds_end2end_test.cc
@@ -991,7 +991,8 @@ class XdsEnd2endTest : public ::testing::TestWithParam<TestType> {
   }
 
   std::tuple<int, int, int> WaitForAllBackends(size_t start_index = 0,
-                                               size_t stop_index = 0) {
+                                               size_t stop_index = 0,
+                                               bool reset_counters = true) {
     int num_ok = 0;
     int num_failure = 0;
     int num_drops = 0;
@@ -999,7 +1000,7 @@ class XdsEnd2endTest : public ::testing::TestWithParam<TestType> {
     while (!SeenAllBackends(start_index, stop_index)) {
       SendRpcAndCount(&num_total, &num_ok, &num_failure, &num_drops);
     }
-    ResetBackendCounters();
+    if (reset_counters) ResetBackendCounters();
     gpr_log(GPR_INFO,
             "Performed %d warm up requests against the backends. "
             "%d succeeded, %d failed, %d dropped.",
@@ -2202,6 +2203,41 @@ TEST_P(FailoverTest, UpdatePriority) {
   EXPECT_EQ(2U, balancers_[0]->ads_service()->response_count());
 }
 
+// Moves all localities in the current priority to a higher priority.
+TEST_P(FailoverTest, MoveAllLocalitiesInCurrentPriorityToHigherPriority) {
+  SetNextResolution({});
+  SetNextResolutionForLbChannelAllBalancers();
+  // First update:
+  // - Priority 0 is locality 0, containing backend 0, which is down.
+  // - Priority 1 is locality 1, containing backends 1 and 2, which are up.
+  ShutdownBackend(0);
+  AdsServiceImpl::ResponseArgs args({
+      {"locality0", GetBackendPorts(0, 1), kDefaultLocalityWeight, 0},
+      {"locality1", GetBackendPorts(1, 3), kDefaultLocalityWeight, 1},
+  });
+  ScheduleResponseForBalancer(0, AdsServiceImpl::BuildResponse(args), 0);
+  // Second update:
+  // - Priority 0 contains both localities 0 and 1.
+  // - Priority 1 is not present.
+  // - We add backend 3 to locality 1, just so we have a way to know
+  //   when the update has been seen by the client.
+  args = AdsServiceImpl::ResponseArgs({
+      {"locality0", GetBackendPorts(0, 1), kDefaultLocalityWeight, 0},
+      {"locality1", GetBackendPorts(1, 4), kDefaultLocalityWeight, 0},
+  });
+  ScheduleResponseForBalancer(0, AdsServiceImpl::BuildResponse(args), 1000);
+  // When we get the first update, all backends in priority 0 are down,
+  // so we will create priority 1.  Backends 1 and 2 should have traffic,
+  // but backend 3 should not.
+  WaitForAllBackends(1, 3, false);
+  EXPECT_EQ(0UL, backends_[3]->backend_service()->request_count());
+  // When backend 3 gets traffic, we know the second update has been seen.
+  WaitForBackend(3);
+  // The ADS service got a single request, and sent a single response.
+  EXPECT_EQ(1U, balancers_[0]->ads_service()->request_count());
+  EXPECT_EQ(2U, balancers_[0]->ads_service()->response_count());
+}
+
 using DropTest = BasicTest;
 
 // Tests that RPCs are dropped according to the drop config.