From 5227db884d77d575158adc3bcafed6a423c5b8c4 Mon Sep 17 00:00:00 2001
From: apolcyn <apolcyn@google.com>
Date: Fri, 16 Feb 2024 17:22:47 -0800
Subject: [PATCH] [pick first] explicitly unset selected subchannel when
 promoting pending subchannel (#35865)

This change is needed in order to get certain LB-policy-unit-test-framework-based tests to pass, which involve updates on PF policies.

Without this change, in such tests, health watchers on the original subchannel can be left hanging around. The underlying reason is related to https://github.com/grpc/grpc/blob/842057d8d5a79d48538e4b276f46d56b4dfcdb16/test/core/client_channel/lb_policy/lb_policy_test_lib.h#L203.

Related: cl/563857636

Closes #35865

COPYBARA_INTEGRATE_REVIEW=https://github.com/grpc/grpc/pull/35865 from apolcyn:update_pf_Test 1a5e6b5591e9e210d74b533d08950271b4518107
PiperOrigin-RevId: 607844309
---
 .../load_balancing/pick_first/pick_first.cc   |  1 +
 .../lb_policy/lb_policy_test_lib.h            |  9 ++--
 .../lb_policy/pick_first_test.cc              | 48 ++++++++++++++++++-
 3 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/src/core/load_balancing/pick_first/pick_first.cc b/src/core/load_balancing/pick_first/pick_first.cc
index 9d3733ddbed..406f388ed34 100644
--- a/src/core/load_balancing/pick_first/pick_first.cc
+++ b/src/core/load_balancing/pick_first/pick_first.cc
@@ -1022,6 +1022,7 @@ void PickFirst::SubchannelList::SubchannelData::ProcessUnselectedReadyLocked() {
               p, p->latest_pending_subchannel_list_.get(),
               p->subchannel_list_.get());
     }
+    p->UnsetSelectedSubchannel();
     p->subchannel_list_ = std::move(p->latest_pending_subchannel_list_);
   }
   // Cases 1 and 2.
diff --git a/test/core/client_channel/lb_policy/lb_policy_test_lib.h b/test/core/client_channel/lb_policy/lb_policy_test_lib.h
index 3b75537461e..e7c49e794d0 100644
--- a/test/core/client_channel/lb_policy/lb_policy_test_lib.h
+++ b/test/core/client_channel/lb_policy/lb_policy_test_lib.h
@@ -698,8 +698,10 @@ class LoadBalancingPolicyTest : public ::testing::Test {
     const absl::optional<BackendMetricData> backend_metric_data_;
   };
 
-  explicit LoadBalancingPolicyTest(absl::string_view lb_policy_name)
-      : lb_policy_name_(lb_policy_name) {}
+  explicit LoadBalancingPolicyTest(absl::string_view lb_policy_name,
+                                   ChannelArgs channel_args = ChannelArgs())
+      : lb_policy_name_(lb_policy_name),
+        channel_args_(std::move(channel_args)) {}
 
   void SetUp() override {
     // Order is important here: Fuzzing EE needs to be created before
@@ -715,7 +717,7 @@ class LoadBalancingPolicyTest : public ::testing::Test {
     auto helper = std::make_unique<FakeHelper>(this);
     helper_ = helper.get();
     LoadBalancingPolicy::Args args = {work_serializer_, std::move(helper),
-                                      ChannelArgs()};
+                                      channel_args_};
     lb_policy_ =
         CoreConfiguration::Get().lb_policy_registry().CreateLoadBalancingPolicy(
             lb_policy_name_, std::move(args));
@@ -1491,6 +1493,7 @@ class LoadBalancingPolicyTest : public ::testing::Test {
   std::map<SubchannelKey, SubchannelState> subchannel_pool_;
   OrphanablePtr<LoadBalancingPolicy> lb_policy_;
   const absl::string_view lb_policy_name_;
+  const ChannelArgs channel_args_;
 };
 
 }  // namespace testing
diff --git a/test/core/client_channel/lb_policy/pick_first_test.cc b/test/core/client_channel/lb_policy/pick_first_test.cc
index fccc8af16ee..77ba738d1a0 100644
--- a/test/core/client_channel/lb_policy/pick_first_test.cc
+++ b/test/core/client_channel/lb_policy/pick_first_test.cc
@@ -14,6 +14,8 @@
 // limitations under the License.
 //
 
+#include "src/core/load_balancing/pick_first/pick_first.h"
+
 #include <stddef.h>
 
 #include <algorithm>
@@ -54,7 +56,8 @@ namespace {
 
 class PickFirstTest : public LoadBalancingPolicyTest {
  protected:
-  PickFirstTest() : LoadBalancingPolicyTest("pick_first") {}
+  explicit PickFirstTest(ChannelArgs channel_args = ChannelArgs())
+      : LoadBalancingPolicyTest("pick_first", channel_args) {}
 
   void SetUp() override {
     LoadBalancingPolicyTest::SetUp();
@@ -1103,6 +1106,49 @@ TEST_F(PickFirstTest, ShufflingDisabled) {
     EXPECT_THAT(address_order, ::testing::ElementsAreArray(kAddresses));
   }
 }
+
+class PickFirstHealthCheckingEnabledTest : public PickFirstTest {
+ protected:
+  PickFirstHealthCheckingEnabledTest()
+      : PickFirstTest(ChannelArgs().Set(
+            GRPC_ARG_INTERNAL_PICK_FIRST_ENABLE_HEALTH_CHECKING, true)) {}
+};
+
+TEST_F(PickFirstHealthCheckingEnabledTest, UpdateWithReadyChannel) {
+  constexpr absl::string_view kAddress = "ipv4:127.0.0.1:443";
+  LoadBalancingPolicy::UpdateArgs update =
+      BuildUpdate({kAddress}, MakePickFirstConfig());
+  absl::Status status = ApplyUpdate(update, lb_policy());
+  EXPECT_TRUE(status.ok()) << status;
+  // LB policy should have created a subchannel for the address.
+  auto* subchannel = FindSubchannel(kAddress);
+  ASSERT_NE(subchannel, nullptr);
+  // When the LB policy receives the first subchannel's initial connectivity
+  // state notification (IDLE), it will request a connection.
+  EXPECT_TRUE(subchannel->ConnectionRequested());
+  // This causes the subchannel to start to connect, so it reports CONNECTING.
+  subchannel->SetConnectivityState(GRPC_CHANNEL_CONNECTING);
+  // LB policy should have reported CONNECTING state.
+  ExpectConnectingUpdate();
+  // When the subchannel becomes connected, it reports READY.
+  subchannel->SetConnectivityState(GRPC_CHANNEL_READY);
+  // The LB policy will report CONNECTING some number of times (doesn't
+  // matter how many) and then report READY.
+  auto picker = WaitForConnected();
+  ASSERT_NE(picker, nullptr);
+  EXPECT_EQ(ExpectPickComplete(picker.get()), kAddress);
+  // Reapply the same update we did before. The the underlying
+  // subchannel will immediately become ready.
+  status =
+      ApplyUpdate(BuildUpdate({kAddress}, MakePickFirstConfig()), lb_policy());
+  EXPECT_TRUE(status.ok()) << status;
+  picker = ExpectState(GRPC_CHANNEL_READY);
+  EXPECT_EQ(ExpectPickComplete(picker.get()), kAddress);
+  // At this point, NumWatchers() should account for our
+  // subchannel connectivity watcher and our health watcher.
+  EXPECT_EQ(subchannel->NumWatchers(), 2);
+}
+
 }  // namespace
 }  // namespace testing
 }  // namespace grpc_core