lb: unified the lifetime mechanism of worker local lb (envoyproxy#45100)

Commit Message: lb: unify the lifetime mechanism of worker local lb
Additional Description:

In the previous implementation, the thread aware lb will recreate the
worker local lb when endpoint set changes. This PR updated the
implementation of thread aware lb to update the worker local lb in place
rather than to create a new one. The new behavior is same with all other
LB implementation.

Risk Level: mid.
Testing: unit.
Docs Changes: n/a.
Release Notes: n/a.
Platform Specific Features: n/a.

---------

Signed-off-by: wbpcode/wangbaiping <wbphub@gmail.com>
Signed-off-by: code <wbphub@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: wbpcode

source/extensions/load_balancing_policies/common/thread_aware_lb_impl.cc

@@ -236,16 +236,30 @@ ThreadAwareLoadBalancerBase::LoadBalancerImpl::chooseHost(LoadBalancerContext* c
   return host;
 }
 
-LoadBalancerPtr ThreadAwareLoadBalancerBase::LoadBalancerFactoryImpl::create(LoadBalancerParams) {
-  auto lb = std::make_unique<LoadBalancerImpl>(stats_, random_, hash_policy_);
-
-  // We must protect current_lb_ via a RW lock since it is accessed and written to by multiple
-  // threads. All complex processing has already been precalculated however.
-  absl::ReaderMutexLock lock(mutex_);
-  lb->healthy_per_priority_load_ = healthy_per_priority_load_;
-  lb->degraded_per_priority_load_ = degraded_per_priority_load_;
-  lb->per_priority_state_ = per_priority_state_;
-  return lb;
+void ThreadAwareLoadBalancerBase::LoadBalancerImpl::refresh() {
+  // The per priority state is shared across all threads and refreshed on main thread. We need to
+  // copy the latest per priority state to the worker thread load balancer instance under lock.
+  absl::ReaderMutexLock lock(factory_->mutex_);
+  healthy_per_priority_load_ = factory_->healthy_per_priority_load_;
+  degraded_per_priority_load_ = factory_->degraded_per_priority_load_;
+  per_priority_state_ = factory_->per_priority_state_;
+}
+
+ThreadAwareLoadBalancerBase::LoadBalancerImpl::LoadBalancerImpl(
+    std::shared_ptr<LoadBalancerFactoryImpl> factory, ClusterLbStats& stats,
+    Random::RandomGenerator& random, HashPolicySharedPtr hash_policy,
+    const Upstream::PrioritySet& priority_set)
+    : factory_(std::move(factory)), stats_(stats), random_(random),
+      hash_policy_(std::move(hash_policy)) {
+  member_update_cb_ =
+      priority_set.addMemberUpdateCb([this](const HostVector&, const HostVector&) { refresh(); });
+  refresh();
+}
+
+LoadBalancerPtr
+ThreadAwareLoadBalancerBase::LoadBalancerFactoryImpl::create(LoadBalancerParams params) {
+  return std::make_unique<LoadBalancerImpl>(shared_from_this(), stats_, random_, hash_policy_,
+                                            params.priority_set);
 }
 
 double ThreadAwareLoadBalancerBase::BoundedLoadHashingLoadBalancer::hostOverloadFactor(

source/extensions/load_balancing_policies/common/thread_aware_lb_impl.h

@@ -126,10 +126,11 @@ class ThreadAwareLoadBalancerBase : public LoadBalancerBase, public ThreadAwareL
   };
   using PerPriorityStatePtr = std::unique_ptr<PerPriorityState>;
 
+  struct LoadBalancerFactoryImpl;
   struct LoadBalancerImpl : public LoadBalancer {
-    LoadBalancerImpl(ClusterLbStats& stats, Random::RandomGenerator& random,
-                     HashPolicySharedPtr hash_policy)
-        : stats_(stats), random_(random), hash_policy_(std::move(hash_policy)) {}
+    LoadBalancerImpl(std::shared_ptr<LoadBalancerFactoryImpl> factory, ClusterLbStats& stats,
+                     Random::RandomGenerator& random, HashPolicySharedPtr hash_policy,
+                     const Upstream::PrioritySet& priority_set);
 
     // Upstream::LoadBalancer
     HostSelectionResponse chooseHost(LoadBalancerContext* context) override;
@@ -145,23 +146,32 @@ class ThreadAwareLoadBalancerBase : public LoadBalancerBase, public ThreadAwareL
       return {};
     }
 
+    void refresh();
+
+    std::shared_ptr<LoadBalancerFactoryImpl> factory_;
+
     ClusterLbStats& stats_;
     Random::RandomGenerator& random_;
     HashPolicySharedPtr hash_policy_;
 
     std::shared_ptr<std::vector<PerPriorityStatePtr>> per_priority_state_;
     std::shared_ptr<HealthyLoad> healthy_per_priority_load_;
     std::shared_ptr<DegradedLoad> degraded_per_priority_load_;
+
+    Common::CallbackHandlePtr member_update_cb_;
   };
 
-  struct LoadBalancerFactoryImpl : public LoadBalancerFactory {
+  struct LoadBalancerFactoryImpl : public LoadBalancerFactory,
+                                   public std::enable_shared_from_this<LoadBalancerFactoryImpl> {
     LoadBalancerFactoryImpl(ClusterLbStats& stats, Random::RandomGenerator& random,
                             std::shared_ptr<Http::HashPolicy> hash_policy)
         : stats_(stats), random_(random), hash_policy_(std::move(hash_policy)) {}
 
     // Upstream::LoadBalancerFactory
-    // Ignore the params for the thread-aware LB.
+    // Uses the per-worker params to create the thread-aware LB instance, including the worker
+    // priority_set used to register member-update callbacks.
     LoadBalancerPtr create(LoadBalancerParams) override;
+    bool recreateOnHostChange() const override { return false; }
 
     ClusterLbStats& stats_;
     Random::RandomGenerator& random_;

test/common/upstream/deferred_cluster_initialization_test.cc

@@ -546,11 +546,9 @@ TEST_P(EdsTest, ShouldNotMergeAddingHostsForDifferentClustersWithSameName) {
   addEndpoint(cluster_load_assignment, 1000);
   doOnConfigUpdateVerifyNoThrow(cluster_load_assignment);
 
-  auto initiailization_instance =
+  auto initialization_instance =
       cluster_manager_->clusterInitializationMap().find("cluster_1")->second;
-  EXPECT_NE(nullptr, initiailization_instance->load_balancer_factory_);
-  // RING_HASH lb policy requires Envoy re-create the load balancer when the cluster is updated.
-  EXPECT_TRUE(initiailization_instance->load_balancer_factory_->recreateOnHostChange());
+  EXPECT_NE(nullptr, initialization_instance->load_balancer_factory_);
 
   // Update the cluster with a different lb policy. Now it's a different cluster and should
   // not be merged.
@@ -565,7 +563,7 @@ TEST_P(EdsTest, ShouldNotMergeAddingHostsForDifferentClustersWithSameName) {
 
   auto new_initialization_instance =
       cluster_manager_->clusterInitializationMap().find("cluster_1")->second;
-  EXPECT_NE(initiailization_instance.get(), new_initialization_instance.get());
+  EXPECT_NE(initialization_instance.get(), new_initialization_instance.get());
 
   EXPECT_EQ(1, new_initialization_instance->per_priority_state_.at(1).hosts_added_.size());
   // Ensure the hosts_added_ is empty for priority 0. Because if unexpected merge happens,

test/extensions/load_balancing_policies/maglev/maglev_lb_test.cc

@@ -135,6 +135,93 @@ TEST_F(MaglevLoadBalancerTest, LbDestructedBeforeFactory) {
   EXPECT_NE(nullptr, factory->create(lb_params_));
 }
 
+// The thread-aware factory updates the worker LB in place via a member update callback on the
+// worker priority set, so the cluster manager should not recreate the worker LB on host changes.
+TEST_F(MaglevLoadBalancerTest, FactoryDoesNotRecreateOnHostChange) {
+  init(7);
+  EXPECT_FALSE(lb_->factory()->recreateOnHostChange());
+}
+
+// Worker LB instances pick up new factory state when the worker priority set fires a member
+// update, without needing to be recreated.
+TEST_F(MaglevLoadBalancerTest, WorkerLbRefreshesOnMemberUpdate) {
+  host_set_.hosts_ = {makeTestHost(info_, "tcp://127.0.0.1:90")};
+  host_set_.healthy_hosts_ = host_set_.hosts_;
+  host_set_.runCallbacks({}, {});
+  init(7);
+
+  // The worker LB is created once and keeps a reference to the factory.
+  LoadBalancerPtr lb = lb_->factory()->create(lb_params_);
+  EXPECT_EQ(host_set_.hosts_[0], lb->chooseHost(nullptr).host);
+
+  // Replace the host set. host_set_.runCallbacks fires the main priority set's callbacks, which
+  // recomputes the factory state on the main thread.
+  auto new_host = makeTestHost(info_, "tcp://127.0.0.1:91");
+  host_set_.hosts_ = {new_host};
+  host_set_.healthy_hosts_ = host_set_.hosts_;
+  host_set_.runCallbacks({}, {});
+
+  // The worker LB still references the previous per-priority state until its own priority set
+  // fires the member update.
+  EXPECT_EQ("127.0.0.1:90", lb->chooseHost(nullptr).host->address()->asString());
+
+  // Simulate the worker priority set firing the member update on the worker. The worker LB must
+  // refresh its cached state from the factory.
+  worker_priority_set_.member_update_cb_helper_.runCallbacks({}, {});
+
+  EXPECT_EQ(new_host, lb->chooseHost(nullptr).host);
+}
+
+// Multiple worker LBs share the same factory; each refreshes independently when its own priority
+// set fires a member update.
+TEST_F(MaglevLoadBalancerTest, MultipleWorkerLbsRefreshIndependently) {
+  host_set_.hosts_ = {makeTestHost(info_, "tcp://127.0.0.1:90")};
+  host_set_.healthy_hosts_ = host_set_.hosts_;
+  host_set_.runCallbacks({}, {});
+  init(7);
+
+  NiceMock<MockPrioritySet> worker_priority_set_a;
+  NiceMock<MockPrioritySet> worker_priority_set_b;
+  LoadBalancerParams params_a{worker_priority_set_a, {}};
+  LoadBalancerParams params_b{worker_priority_set_b, {}};
+
+  LoadBalancerPtr lb_a = lb_->factory()->create(params_a);
+  LoadBalancerPtr lb_b = lb_->factory()->create(params_b);
+
+  EXPECT_EQ(host_set_.hosts_[0], lb_a->chooseHost(nullptr).host);
+  EXPECT_EQ(host_set_.hosts_[0], lb_b->chooseHost(nullptr).host);
+
+  auto new_host = makeTestHost(info_, "tcp://127.0.0.1:91");
+  host_set_.hosts_ = {new_host};
+  host_set_.healthy_hosts_ = host_set_.hosts_;
+  host_set_.runCallbacks({}, {});
+
+  // Only refresh worker A; worker B keeps the previous state until its own priority set fires.
+  worker_priority_set_a.member_update_cb_helper_.runCallbacks({}, {});
+  EXPECT_EQ(new_host, lb_a->chooseHost(nullptr).host);
+  EXPECT_EQ("127.0.0.1:90", lb_b->chooseHost(nullptr).host->address()->asString());
+
+  worker_priority_set_b.member_update_cb_helper_.runCallbacks({}, {});
+  EXPECT_EQ(new_host, lb_b->chooseHost(nullptr).host);
+}
+
+// The worker LB unregisters its member update callback on destruction; firing the worker priority
+// set's callback after the LB is gone must not touch freed memory.
+TEST_F(MaglevLoadBalancerTest, WorkerLbCallbackUnregisteredOnDestruction) {
+  host_set_.hosts_ = {makeTestHost(info_, "tcp://127.0.0.1:90")};
+  host_set_.healthy_hosts_ = host_set_.hosts_;
+  host_set_.runCallbacks({}, {});
+  init(7);
+
+  LoadBalancerPtr lb = lb_->factory()->create(lb_params_);
+  EXPECT_EQ(host_set_.hosts_[0], lb->chooseHost(nullptr).host);
+
+  lb.reset();
+
+  // Must be a no-op rather than calling into freed memory.
+  worker_priority_set_.member_update_cb_helper_.runCallbacks({}, {});
+}
+
 // Throws an exception if table size is not a prime number.
 TEST_F(MaglevLoadBalancerTest, NoPrimeNumber) {
   EXPECT_THROW_WITH_MESSAGE(init(8), EnvoyException,