lb: unified the lifetime mechanism of worker load lb 2 (#45170)

Commit Message: lb: unified the lifetime mechanism of worker load lb 2
Additional Description:

Similar to previous #45100 but
this PR have updated all left LB.

Risk Level:
Testing:
Docs Changes:
Release Notes:
Platform Specific Features:
[Optional Runtime guard:]
[Optional Fixes #Issue]
[Optional Fixes commit #PR or SHA]
[Optional Deprecated:]
[Optional [API
Considerations](https://github.com/envoyproxy/envoy/blob/main/api/review_checklist.md):]

---------

Signed-off-by: wbpcode/wangbaiping <wbphub@gmail.com>
Signed-off-by: code <wbphub@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: wbpcode

source/extensions/clusters/aggregate/cluster.h

@@ -153,6 +153,7 @@ class AggregateLoadBalancerFactory : public Upstream::LoadBalancerFactory {
         cluster_.info(), cluster_.cluster_manager_, cluster_.runtime(), cluster_.random(),
         cluster_.clusters_);
   }
+  bool recreateOnHostChange() const override { return false; }
 
   const Cluster& cluster_;
 };

source/extensions/clusters/composite/cluster.h

@@ -95,6 +95,7 @@ class CompositeLoadBalancerFactory : public Upstream::LoadBalancerFactory {
     return std::make_unique<CompositeClusterLoadBalancer>(
         cluster_.info(), cluster_.cluster_manager_, cluster_.clusters_);
   }
+  bool recreateOnHostChange() const override { return false; }
 
   const Cluster& cluster_;
 };

source/extensions/clusters/dynamic_forward_proxy/cluster.h

@@ -187,6 +187,7 @@ class Cluster : public Upstream::BaseDynamicClusterImpl,
     Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams) override {
       return std::make_unique<LoadBalancer>(cluster_);
     }
+    bool recreateOnHostChange() const override { return false; }
 
   private:
     Cluster& cluster_;

source/extensions/clusters/dynamic_modules/cluster.cc

@@ -35,6 +35,7 @@ struct DynamicModuleThreadAwareLoadBalancer : public Upstream::ThreadAwareLoadBa
     Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams params) override {
       return std::make_unique<DynamicModuleLoadBalancer>(handle_, params.priority_set);
     }
+    bool recreateOnHostChange() const override { return false; }
 
     DynamicModuleClusterHandleSharedPtr handle_;
   };

source/extensions/clusters/original_dst/original_dst_cluster.h

@@ -88,11 +88,16 @@ class OriginalDstCluster : public ClusterImplBase {
    */
   class LoadBalancer : public Upstream::LoadBalancer {
   public:
-    LoadBalancer(const OriginalDstClusterHandleSharedPtr& parent)
+    LoadBalancer(const OriginalDstClusterHandleSharedPtr& parent, const PrioritySet& priority_set)
         : parent_(parent), http_header_name_(parent->cluster_->httpHeaderName()),
           metadata_key_(parent->cluster_->metadataKey()),
           port_override_(parent->cluster_->portOverride()),
-          host_map_(parent->cluster_->getCurrentHostMap()) {}
+          host_map_(parent->cluster_->getCurrentHostMap()) {
+      member_update_cb_ =
+          priority_set.addMemberUpdateCb([this](const HostVector&, const HostVector&) {
+            host_map_ = parent_->cluster_->getCurrentHostMap();
+          });
+    }
 
     // Upstream::LoadBalancer
     HostSelectionResponse chooseHost(LoadBalancerContext* context) override;
@@ -121,6 +126,7 @@ class OriginalDstCluster : public ClusterImplBase {
     const absl::optional<Config::MetadataKey>& metadata_key_;
     const absl::optional<uint32_t> port_override_;
     HostMultiMapConstSharedPtr host_map_;
+    Common::CallbackHandlePtr member_update_cb_;
   };
 
   const absl::optional<Http::LowerCaseString>& httpHeaderName() { return http_header_name_; }
@@ -139,9 +145,10 @@ class OriginalDstCluster : public ClusterImplBase {
     LoadBalancerFactory(const OriginalDstClusterHandleSharedPtr& cluster) : cluster_(cluster) {}
 
     // Upstream::LoadBalancerFactory
-    Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams) override {
-      return std::make_unique<LoadBalancer>(cluster_);
+    Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams params) override {
+      return std::make_unique<LoadBalancer>(cluster_, params.priority_set);
     }
+    bool recreateOnHostChange() const override { return false; }
 
     const OriginalDstClusterHandleSharedPtr cluster_;
   };

source/extensions/clusters/redis/redis_cluster_lb.cc

@@ -101,9 +101,24 @@ void RedisClusterLoadBalancerFactory::onHostHealthUpdate() {
   }
 }
 
-Upstream::LoadBalancerPtr RedisClusterLoadBalancerFactory::create(Upstream::LoadBalancerParams) {
-  absl::ReaderMutexLock lock(mutex_);
-  return std::make_unique<RedisClusterLoadBalancer>(slot_array_, shard_vector_, random_);
+Upstream::LoadBalancerPtr
+RedisClusterLoadBalancerFactory::create(Upstream::LoadBalancerParams params) {
+  return std::make_unique<RedisClusterLoadBalancer>(shared_from_this(), params.priority_set);
+}
+
+RedisClusterLoadBalancerFactory::RedisClusterLoadBalancer::RedisClusterLoadBalancer(
+    std::shared_ptr<RedisClusterLoadBalancerFactory> factory,
+    const Upstream::PrioritySet& priority_set)
+    : factory_(factory), random_(factory->random_) {
+  refresh();
+  member_update_cb_ = priority_set.addMemberUpdateCb(
+      [this](const Upstream::HostVector&, const Upstream::HostVector&) { refresh(); });
+}
+
+void RedisClusterLoadBalancerFactory::RedisClusterLoadBalancer::refresh() {
+  absl::ReaderMutexLock lock(factory_->mutex_);
+  slot_array_ = factory_->slot_array_;
+  shard_vector_ = factory_->shard_vector_;
 }
 
 namespace {

source/extensions/clusters/redis/redis_cluster_lb.h

@@ -175,8 +175,10 @@ using ClusterSlotUpdateCallBackSharedPtr = std::shared_ptr<ClusterSlotUpdateCall
  * This factory is created and returned by RedisCluster's factory() method, the create() method will
  * be called on each thread to create a thread local RedisClusterLoadBalancer.
  */
-class RedisClusterLoadBalancerFactory : public ClusterSlotUpdateCallBack,
-                                        public Upstream::LoadBalancerFactory {
+class RedisClusterLoadBalancerFactory
+    : public ClusterSlotUpdateCallBack,
+      public Upstream::LoadBalancerFactory,
+      public std::enable_shared_from_this<RedisClusterLoadBalancerFactory> {
 public:
   RedisClusterLoadBalancerFactory(Random::RandomGenerator& random) : random_(random) {}
 
@@ -186,7 +188,8 @@ class RedisClusterLoadBalancerFactory : public ClusterSlotUpdateCallBack,
   void onHostHealthUpdate() override;
 
   // Upstream::LoadBalancerFactory
-  Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams) override;
+  Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams params) override;
+  bool recreateOnHostChange() const override { return false; }
 
 private:
   class RedisShard {
@@ -267,10 +270,8 @@ class RedisClusterLoadBalancerFactory : public ClusterSlotUpdateCallBack,
    */
   class RedisClusterLoadBalancer : public Upstream::LoadBalancer {
   public:
-    RedisClusterLoadBalancer(SlotArraySharedPtr slot_array, ShardVectorSharedPtr shard_vector,
-                             Random::RandomGenerator& random)
-        : slot_array_(std::move(slot_array)), shard_vector_(std::move(shard_vector)),
-          random_(random) {}
+    RedisClusterLoadBalancer(std::shared_ptr<RedisClusterLoadBalancerFactory> factory,
+                             const Upstream::PrioritySet& priority_set);
 
     // Upstream::LoadBalancerBase
     Upstream::HostSelectionResponse chooseHost(Upstream::LoadBalancerContext*) override;
@@ -290,9 +291,14 @@ class RedisClusterLoadBalancerFactory : public ClusterSlotUpdateCallBack,
     }
 
   private:
-    const SlotArraySharedPtr slot_array_;
-    const ShardVectorSharedPtr shard_vector_;
+    // Re-snapshots the topology from the parent factory under its mutex.
+    void refresh();
+
+    const std::shared_ptr<RedisClusterLoadBalancerFactory> factory_;
+    SlotArraySharedPtr slot_array_;
+    ShardVectorSharedPtr shard_vector_;
     Random::RandomGenerator& random_;
+    ::Envoy::Common::CallbackHandlePtr member_update_cb_;
   };
 
   absl::Mutex mutex_;

source/extensions/clusters/reverse_connection/reverse_connection.h

@@ -183,6 +183,7 @@ class RevConCluster : public Upstream::ClusterImplBase {
     // Upstream::LoadBalancerFactory.
     Upstream::LoadBalancerPtr create() { return std::make_unique<LoadBalancer>(cluster_); }
     Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams) override { return create(); }
+    bool recreateOnHostChange() const override { return false; }
 
     const std::shared_ptr<RevConCluster> cluster_;
   };

source/extensions/load_balancing_policies/override_host/load_balancer.h

@@ -189,6 +189,8 @@ class OverrideHostLoadBalancer : public Upstream::ThreadAwareLoadBalancer,
     // Called by worker threads to create a thread-local load balancer.
     LoadBalancerPtr create(LoadBalancerParams params) override;
 
+    bool recreateOnHostChange() const override { return false; }
+
   private:
     // Hosts in the load balancer. Owned by the cluster manager.
     const OverrideHostLbConfig& config_;