extend load balancer context interface to provide override host

envoy/upstream/load_balancer.h

@@ -83,6 +83,16 @@ class LoadBalancerContext {
    * Returns the transport socket options which should be applied on upstream connections
    */
   virtual Network::TransportSocketOptionsConstSharedPtr upstreamTransportSocketOptions() const PURE;
+
+  using OverrideHostStatus = uint32_t;
+  using OverrideHost = std::pair<std::string, OverrideHostStatus>;
+
+  /**
+   * Returns the host the load balancer should select directly. If the expected host exists and
+   * the health status of the host matches the expectation, the load balancer can bypass the load
+   * balancing algorithm and return the corresponding host directly.
+   */
+  virtual absl::optional<OverrideHost> overrideHostToSelect() const PURE;
 };
 
 /**

source/common/upstream/load_balancer_impl.cc

@@ -1,5 +1,6 @@
 #include "source/common/upstream/load_balancer_impl.h"
 
+#include <atomic>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -22,6 +23,10 @@ static const std::string RuntimeZoneEnabled = "upstream.zone_routing.enabled";
 static const std::string RuntimeMinClusterSize = "upstream.zone_routing.min_cluster_size";
 static const std::string RuntimePanicThreshold = "upstream.healthy_panic_threshold";
 
+static constexpr uint32_t UnhealthyStatus = 1u << static_cast<size_t>(Host::Health::Unhealthy);
+static constexpr uint32_t DegradedStatus = 1u << static_cast<size_t>(Host::Health::Degraded);
+static constexpr uint32_t HealthyStatus = 1u << static_cast<size_t>(Host::Health::Healthy);
+
 bool tooManyPreconnects(size_t num_preconnect_picks, uint32_t healthy_hosts) {
   // Currently we only allow the number of preconnected connections to equal the
   // number of healthy hosts.
@@ -361,6 +366,9 @@ ZoneAwareLoadBalancerBase::ZoneAwareLoadBalancerBase(
   resizePerPriorityState();
   priority_update_cb_ = priority_set_.addPriorityUpdateCb(
       [this](uint32_t priority, const HostVector&, const HostVector&) -> void {
+        // Update cross priority host map for fast host searching.
+        cross_priority_host_map_ = priority_set_.crossPriorityHostMap();
+
         // Make sure per_priority_state_ is as large as priority_set_.hostSetsPerPriority()
         resizePerPriorityState();
         // If P=0 changes, regenerate locality routing structures. Locality based routing is
@@ -507,8 +515,52 @@ bool ZoneAwareLoadBalancerBase::earlyExitNonLocalityRouting() {
   return false;
 }
 
-HostConstSharedPtr LoadBalancerBase::chooseHost(LoadBalancerContext* context) {
+bool LoadBalancerContextBase::validateOverrideHostStatus(Host::Health health,
+                                                         OverrideHostStatus status) {
+  switch (health) {
+  case Host::Health::Unhealthy:
+    return status & UnhealthyStatus;
+  case Host::Health::Degraded:
+    return status & DegradedStatus;
+  case Host::Health::Healthy:
+    return status & HealthyStatus;
+  }
+  return false;
+}
+
+HostConstSharedPtr LoadBalancerContextBase::selectOverrideHost(const HostMap* host_map,
+                                                               LoadBalancerContext* context) {
+  if (context == nullptr) {
+    return nullptr;
+  }
+
+  auto override_host = context->overrideHostToSelect();
+  if (!override_host.has_value()) {
+    return nullptr;
+  }
+
+  if (host_map == nullptr) {
+    return nullptr;
+  }
+
+  auto host_iter = host_map->find(override_host.value().first);
+  HostConstSharedPtr host = host_iter != host_map->end() ? host_iter->second : nullptr;
+
+  if (host != nullptr && LoadBalancerContextBase::validateOverrideHostStatus(
+                             host->health(), override_host.value().second)) {
+    return host;
+  }
+  return nullptr;
+}
+
+HostConstSharedPtr ZoneAwareLoadBalancerBase::chooseHost(LoadBalancerContext* context) {
   HostConstSharedPtr host;
+
+  if (host = LoadBalancerContextBase::selectOverrideHost(cross_priority_host_map_.get(), context);
+      host != nullptr) {
+    return host;
+  }
+
   const size_t max_attempts = context ? context->hostSelectionRetryCount() + 1 : 1;
   for (size_t i = 0; i < max_attempts; ++i) {
     host = chooseHostOnce(context);

source/common/upstream/load_balancer_impl.h

@@ -40,24 +40,7 @@ class LoadBalancerBase : public LoadBalancer {
   choosePriority(uint64_t hash, const HealthyLoad& healthy_per_priority_load,
                  const DegradedLoad& degraded_per_priority_load);
 
-  HostConstSharedPtr chooseHost(LoadBalancerContext* context) override;
-
 protected:
-  /**
-   * By implementing this method instead of chooseHost, host selection will
-   * be subject to host filters specified by LoadBalancerContext.
-   *
-   * Host selection will be retried up to the number specified by
-   * hostSelectionRetryCount on LoadBalancerContext, and if no hosts are found
-   * within the allowed attempts, the host that was selected during the last
-   * attempt will be returned.
-   *
-   * If host selection is idempotent (i.e. retrying will not change the outcome),
-   * sub classes should override chooseHost to avoid the unnecessary overhead of
-   * retrying host selection.
-   */
-  virtual HostConstSharedPtr chooseHostOnce(LoadBalancerContext* context) PURE;
-
   /**
    * For the given host_set @return if we should be in a panic mode or not. For example, if the
    * majority of hosts are unhealthy we'll be likely in a panic mode. In this case we'll route
@@ -147,6 +130,7 @@ class LoadBalancerBase : public LoadBalancer {
 
     return std::min<uint32_t>(health + degraded, 100);
   }
+
   // The percentage load (0-100) for each priority level when targeting healthy hosts and
   // the percentage load (0-100) for each priority level when targeting degraded hosts.
   HealthyAndDegradedLoad per_priority_load_;
@@ -165,6 +149,11 @@ class LoadBalancerBase : public LoadBalancer {
 
 class LoadBalancerContextBase : public LoadBalancerContext {
 public:
+  static bool validateOverrideHostStatus(Host::Health health, OverrideHostStatus status);
+
+  static HostConstSharedPtr selectOverrideHost(const HostMap* host_map,
+                                               LoadBalancerContext* context);
+
   absl::optional<uint64_t> computeHashKey() override { return {}; }
 
   const Network::Connection* downstreamConnection() const override { return nullptr; }
@@ -188,12 +177,17 @@ class LoadBalancerContextBase : public LoadBalancerContext {
   Network::TransportSocketOptionsConstSharedPtr upstreamTransportSocketOptions() const override {
     return nullptr;
   }
+
+  absl::optional<OverrideHost> overrideHostToSelect() const override { return {}; }
 };
 
 /**
  * Base class for zone aware load balancers
  */
 class ZoneAwareLoadBalancerBase : public LoadBalancerBase {
+public:
+  HostConstSharedPtr chooseHost(LoadBalancerContext* context) override;
+
 protected:
   // Both priority_set and local_priority_set if non-null must have at least one host set.
   ZoneAwareLoadBalancerBase(
@@ -258,6 +252,21 @@ class ZoneAwareLoadBalancerBase : public LoadBalancerBase {
     }
   };
 
+  /**
+   * By implementing this method instead of chooseHost, host selection will
+   * be subject to host filters specified by LoadBalancerContext.
+   *
+   * Host selection will be retried up to the number specified by
+   * hostSelectionRetryCount on LoadBalancerContext, and if no hosts are found
+   * within the allowed attempts, the host that was selected during the last
+   * attempt will be returned.
+   *
+   * If host selection is idempotent (i.e. retrying will not change the outcome),
+   * sub classes should override chooseHost to avoid the unnecessary overhead of
+   * retrying host selection.
+   */
+  virtual HostConstSharedPtr chooseHostOnce(LoadBalancerContext* context) PURE;
+
   /**
    * Pick the host source to use, doing zone aware routing when the hosts are sufficiently healthy.
    * If no host is chosen (due to fail_traffic_on_panic being set), return absl::nullopt.
@@ -269,6 +278,10 @@ class ZoneAwareLoadBalancerBase : public LoadBalancerBase {
    */
   const HostVector& hostSourceToHosts(HostsSource hosts_source) const;
 
+  // Cross priority host map for fast cross priority host searching. When the priority update
+  // callback is executed, the host map will also be updated.
+  HostMapConstSharedPtr cross_priority_host_map_;
+
 private:
   enum class LocalityRoutingState {
     // Locality based routing is off.
@@ -381,7 +394,7 @@ class EdfLoadBalancerBase : public ZoneAwareLoadBalancerBase {
                       Random::RandomGenerator& random,
                       const envoy::config::cluster::v3::Cluster::CommonLbConfig& common_config);
 
-  // Upstream::LoadBalancerBase
+  // Upstream::ZoneAwareLoadBalancerBase
   HostConstSharedPtr peekAnotherHost(LoadBalancerContext* context) override;
   HostConstSharedPtr chooseHostOnce(LoadBalancerContext* context) override;
 
@@ -579,7 +592,7 @@ class RandomLoadBalancer : public ZoneAwareLoadBalancerBase {
       : ZoneAwareLoadBalancerBase(priority_set, local_priority_set, stats, runtime, random,
                                   common_config) {}
 
-  // Upstream::LoadBalancerBase
+  // Upstream::ZoneAwareLoadBalancerBase
   HostConstSharedPtr chooseHostOnce(LoadBalancerContext* context) override;
   HostConstSharedPtr peekAnotherHost(LoadBalancerContext* context) override;
 

source/common/upstream/subset_lb.cc

@@ -82,6 +82,9 @@ SubsetLoadBalancer::SubsetLoadBalancer(
         // performed.
         rebuildSingle();
 
+        // Update cross priority host map.
+        cross_priority_host_map_ = original_priority_set_.crossPriorityHostMap();
+
         if (hosts_added.empty() && hosts_removed.empty()) {
           // It's possible that metadata changed, without hosts being added nor removed.
           // If so we need to add any new subsets, remove unused ones, and regroup hosts into
@@ -276,6 +279,12 @@ void SubsetLoadBalancer::initSelectorFallbackSubset(
 }
 
 HostConstSharedPtr SubsetLoadBalancer::chooseHost(LoadBalancerContext* context) {
+  HostConstSharedPtr override_host =
+      LoadBalancerContextBase::selectOverrideHost(cross_priority_host_map_.get(), context);
+  if (override_host != nullptr) {
+    return override_host;
+  }
+
   if (context) {
     bool host_chosen;
     HostConstSharedPtr host = tryChooseHostFromContext(context, host_chosen);

source/common/upstream/subset_lb.h

@@ -162,6 +162,10 @@ class SubsetLoadBalancer : public LoadBalancer, Logger::Loggable<Logger::Id::ups
       return wrapped_->upstreamTransportSocketOptions();
     }
 
+    absl::optional<OverrideHost> overrideHostToSelect() const override {
+      return wrapped_->overrideHostToSelect();
+    }
+
   private:
     LoadBalancerContext* wrapped_;
     Router::MetadataMatchCriteriaConstPtr metadata_match_;
@@ -268,6 +272,10 @@ class SubsetLoadBalancer : public LoadBalancer, Logger::Loggable<Logger::Id::ups
   absl::flat_hash_map<HashedValue, HostConstSharedPtr> single_host_per_subset_map_;
   Stats::Gauge* single_duplicate_stat_{};
 
+  // Cross priority host map for fast cross priority host searching. When the priority update
+  // callback is executed, the host map will also be updated.
+  HostMapConstSharedPtr cross_priority_host_map_;
+
   const bool locality_weight_aware_;
   const bool scale_locality_weight_;
   const bool list_as_any_;

source/common/upstream/thread_aware_lb_impl.cc

@@ -95,7 +95,10 @@ void ThreadAwareLoadBalancerBase::initialize() {
   // complicated initialization as the load balancer would need its own initialized callback. I
   // think the synchronous/asynchronous split is probably the best option.
   priority_update_cb_ = priority_set_.addPriorityUpdateCb(
-      [this](uint32_t, const HostVector&, const HostVector&) -> void { refresh(); });
+      [this](uint32_t, const HostVector&, const HostVector&) -> void {
+        refresh();
+        threadSafeSetCrossPriorityHostMap(priority_set_.crossPriorityHostMap());
+      });
 
   refresh();
 }
@@ -141,6 +144,12 @@ ThreadAwareLoadBalancerBase::LoadBalancerImpl::chooseHost(LoadBalancerContext* c
     return nullptr;
   }
 
+  HostConstSharedPtr host;
+  if (host = LoadBalancerContextBase::selectOverrideHost(cross_priority_host_map_.get(), context);
+      host != nullptr) {
+    return host;
+  }
+
   // If there is no hash in the context, just choose a random value (this effectively becomes
   // the random LB but it won't crash if someone configures it this way).
   // computeHashKey() may be computed on demand, so get it only once.
@@ -158,7 +167,6 @@ ThreadAwareLoadBalancerBase::LoadBalancerImpl::chooseHost(LoadBalancerContext* c
     stats_.lb_healthy_panic_.inc();
   }
 
-  HostConstSharedPtr host;
   const uint32_t max_attempts = context ? context->hostSelectionRetryCount() + 1 : 1;
   for (uint32_t i = 0; i < max_attempts; ++i) {
     host = per_priority_state->current_lb_->chooseHost(h, i);
@@ -173,7 +181,8 @@ ThreadAwareLoadBalancerBase::LoadBalancerImpl::chooseHost(LoadBalancerContext* c
 }
 
 LoadBalancerPtr ThreadAwareLoadBalancerBase::LoadBalancerFactoryImpl::create() {
-  auto lb = std::make_unique<LoadBalancerImpl>(stats_, random_);
+  auto lb = std::make_unique<LoadBalancerImpl>(
+      stats_, random_, thread_aware_lb_.threadSafeGetCrossPriorityHostMap());
 
   // 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.

source/common/upstream/thread_aware_lb_impl.h

@@ -87,10 +87,8 @@ class ThreadAwareLoadBalancerBase : public LoadBalancerBase, public ThreadAwareL
   LoadBalancerFactorySharedPtr factory() override { return factory_; }
   void initialize() override;
 
-  // Upstream::LoadBalancerBase
-  HostConstSharedPtr chooseHostOnce(LoadBalancerContext*) override {
-    NOT_IMPLEMENTED_GCOVR_EXCL_LINE;
-  }
+  // Upstream::LoadBalancer
+  HostConstSharedPtr chooseHost(LoadBalancerContext*) override { NOT_IMPLEMENTED_GCOVR_EXCL_LINE; }
   // Preconnect not implemented for hash based load balancing
   HostConstSharedPtr peekAnotherHost(LoadBalancerContext*) override { return nullptr; }
 
@@ -100,7 +98,7 @@ class ThreadAwareLoadBalancerBase : public LoadBalancerBase, public ThreadAwareL
       Random::RandomGenerator& random,
       const envoy::config::cluster::v3::Cluster::CommonLbConfig& common_config)
       : LoadBalancerBase(priority_set, stats, runtime, random, common_config),
-        factory_(new LoadBalancerFactoryImpl(stats, random)) {}
+        factory_(new LoadBalancerFactoryImpl(stats, random, *this)) {}
 
 private:
   struct PerPriorityState {
@@ -110,8 +108,9 @@ class ThreadAwareLoadBalancerBase : public LoadBalancerBase, public ThreadAwareL
   using PerPriorityStatePtr = std::unique_ptr<PerPriorityState>;
 
   struct LoadBalancerImpl : public LoadBalancer {
-    LoadBalancerImpl(ClusterStats& stats, Random::RandomGenerator& random)
-        : stats_(stats), random_(random) {}
+    LoadBalancerImpl(ClusterStats& stats, Random::RandomGenerator& random,
+                     HostMapConstSharedPtr host_map)
+        : stats_(stats), random_(random), cross_priority_host_map_(std::move(host_map)) {}
 
     // Upstream::LoadBalancer
     HostConstSharedPtr chooseHost(LoadBalancerContext* context) override;
@@ -123,15 +122,21 @@ class ThreadAwareLoadBalancerBase : public LoadBalancerBase, public ThreadAwareL
     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_;
+
+    // Cross priority host map.
+    HostMapConstSharedPtr cross_priority_host_map_;
   };
 
   struct LoadBalancerFactoryImpl : public LoadBalancerFactory {
-    LoadBalancerFactoryImpl(ClusterStats& stats, Random::RandomGenerator& random)
-        : stats_(stats), random_(random) {}
+    LoadBalancerFactoryImpl(ClusterStats& stats, Random::RandomGenerator& random,
+                            ThreadAwareLoadBalancerBase& thread_aware_lb)
+        : thread_aware_lb_(thread_aware_lb), stats_(stats), random_(random) {}
 
     // Upstream::LoadBalancerFactory
     LoadBalancerPtr create() override;
 
+    ThreadAwareLoadBalancerBase& thread_aware_lb_;
+
     ClusterStats& stats_;
     Random::RandomGenerator& random_;
     absl::Mutex mutex_;
@@ -146,8 +151,29 @@ class ThreadAwareLoadBalancerBase : public LoadBalancerBase, public ThreadAwareL
                      double min_normalized_weight, double max_normalized_weight) PURE;
   void refresh();
 
+  void threadSafeSetCrossPriorityHostMap(HostMapConstSharedPtr host_map) {
+    absl::MutexLock ml(&cross_priority_host_map_mutex_);
+    cross_priority_host_map_ = std::move(host_map);
+  }
+  HostMapConstSharedPtr threadSafeGetCrossPriorityHostMap() {
+    absl::MutexLock ml(&cross_priority_host_map_mutex_);
+    return cross_priority_host_map_;
+  }
+
   std::shared_ptr<LoadBalancerFactoryImpl> factory_;
   Common::CallbackHandlePtr priority_update_cb_;
+
+  // Whenever the membership changes, the cross_priority_host_map_ will be updated automatically.
+  // And all workers will create a new worker local load balancer and copy the
+  // cross_priority_host_map_.
+  //
+  // This leads to the possibility of simultaneous reading and writing of cross_priority_host_map_
+  // in different threads. For this reason, an additional mutex is necessary to guard
+  // cross_priority_host_map_.
+  absl::Mutex cross_priority_host_map_mutex_;
+  // Cross priority host map for fast cross priority host searching. When the priority update
+  // callback is executed, the host map will also be updated.
+  HostMapConstSharedPtr cross_priority_host_map_ ABSL_GUARDED_BY(cross_priority_host_map_mutex_);
 };
 
 } // namespace Upstream