From ff75116e44f54798d729adfbd6d466cf30af8670 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Fri, 3 Jul 2026 12:31:45 -0600 Subject: [PATCH 01/16] load_aware_locality: config, factory registration, metadata Resolve the endpoint-picking child policy, validate policy-level knobs (update period, smoothing constant, variance and probe fractions), reject unsupported enable_oob_load_report, and register the LoadAwareLocality factory. Promote the extension from wip to alpha. Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- source/extensions/extensions_metadata.yaml | 2 +- .../load_aware_locality/BUILD | 30 +++++- .../load_aware_locality/config.cc | 95 +++++++++++++++++-- .../load_aware_locality/config.h | 11 ++- 4 files changed, 122 insertions(+), 16 deletions(-) diff --git a/source/extensions/extensions_metadata.yaml b/source/extensions/extensions_metadata.yaml index 3c71e27bee679..dd77435f1a098 100644 --- a/source/extensions/extensions_metadata.yaml +++ b/source/extensions/extensions_metadata.yaml @@ -2131,7 +2131,7 @@ envoy.load_balancing_policies.load_aware_locality: categories: - envoy.load_balancing_policies security_posture: unknown - status: wip + status: alpha type_urls: - envoy.extensions.load_balancing_policies.load_aware_locality.v3.LoadAwareLocality envoy.load_balancing_policies.random: diff --git a/source/extensions/load_balancing_policies/load_aware_locality/BUILD b/source/extensions/load_balancing_policies/load_aware_locality/BUILD index 890af1964ef23..61866ccb9142f 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/BUILD +++ b/source/extensions/load_balancing_policies/load_aware_locality/BUILD @@ -1,6 +1,7 @@ load( "//bazel:envoy_build_system.bzl", "envoy_cc_extension", + "envoy_cc_library", "envoy_extension_package", ) @@ -13,7 +14,34 @@ envoy_cc_extension( srcs = ["config.cc"], hdrs = ["config.h"], deps = [ - "//source/common/upstream:load_balancer_factory_base_lib", + ":load_aware_locality_lb_lib", + "//source/common/common:minimal_logger_lib", + "//source/common/config:utility_lib", + "//source/extensions/load_balancing_policies/common:factory_base", + "@abseil-cpp//absl/strings", "@envoy_api//envoy/extensions/load_balancing_policies/load_aware_locality/v3:pkg_cc_proto", ], ) + +envoy_cc_library( + name = "load_aware_locality_lb_lib", + srcs = ["load_aware_locality_lb.cc"], + hdrs = ["load_aware_locality_lb.h"], + deps = [ + "//envoy/common:random_generator_interface", + "//envoy/event:dispatcher_interface", + "//envoy/stats:stats_macros", + "//envoy/thread_local:thread_local_interface", + "//envoy/upstream:load_balancer_interface", + "//envoy/upstream:locality_lib", + "//envoy/upstream:upstream_interface", + "//source/common/common:minimal_logger_lib", + "//source/common/protobuf:utility_lib", + "//source/common/upstream:upstream_includes", + "//source/extensions/load_balancing_policies/common:load_balancer_lib", + "//source/extensions/load_balancing_policies/common:orca_weight_manager_lib", + "@abseil-cpp//absl/container:flat_hash_map", + "@abseil-cpp//absl/container:flat_hash_set", + "@abseil-cpp//absl/strings", + ], +) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/config.cc b/source/extensions/load_balancing_policies/load_aware_locality/config.cc index a9ca80a1165bf..0cfbf87bbd900 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/config.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/config.cc @@ -1,25 +1,100 @@ #include "source/extensions/load_balancing_policies/load_aware_locality/config.h" +#include + +#include "source/common/config/utility.h" +#include "source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h" + namespace Envoy { namespace Extensions { namespace LoadBalancingPolicies { namespace LoadAwareLocality { -Upstream::ThreadAwareLoadBalancerPtr Factory::create(OptRef, - const Upstream::ClusterInfo&, - const Upstream::PrioritySet&, Runtime::Loader&, - Envoy::Random::RandomGenerator&, TimeSource&) { - // TODO(jukie): Implement load-aware locality load balancer. - return nullptr; +Upstream::ThreadAwareLoadBalancerPtr +Factory::create(OptRef lb_config, + const Upstream::ClusterInfo& cluster_info, + const Upstream::PrioritySet& priority_set, Runtime::Loader& runtime, + Envoy::Random::RandomGenerator& random, TimeSource& time_source) { + return std::make_unique(lb_config, cluster_info, priority_set, + runtime, random, time_source); } absl::StatusOr -Factory::loadConfig(Server::Configuration::ServerFactoryContext&, const Protobuf::Message&) { - // TODO(jukie): Implement load-aware locality config loading. - return absl::UnimplementedError( - "envoy.load_balancing_policies.load_aware_locality is not yet implemented"); +Factory::loadConfig(Server::Configuration::ServerFactoryContext& context, + const Protobuf::Message& config) { + const auto& lb_config = dynamic_cast(config); + + // Validate policy-level knobs before resolving the child policy so a config error is reported + // for what it is rather than masked by an unresolvable child. + if (lb_config.has_enable_oob_load_report() && lb_config.enable_oob_load_report().value()) { + return absl::InvalidArgumentError( + "load_aware_locality: enable_oob_load_report is not yet supported"); + } + const std::chrono::milliseconds weight_update_period = + std::chrono::milliseconds(PROTOBUF_GET_MS_OR_DEFAULT(lb_config, weight_update_period, 1000)); + if (weight_update_period <= std::chrono::milliseconds(0)) { + return absl::InvalidArgumentError("weight_update_period must be positive"); + } + const std::chrono::milliseconds smoothing_time_constant = std::chrono::milliseconds( + PROTOBUF_GET_MS_OR_DEFAULT(lb_config, smoothing_time_constant, 5000)); + if (smoothing_time_constant <= std::chrono::milliseconds(0)) { + return absl::InvalidArgumentError("smoothing_time_constant must be positive"); + } + // Derive the per-tick EWMA factor so settling time is independent of tick rate. + const double ewma_alpha = + 1.0 - std::exp(-std::chrono::duration(weight_update_period).count() / + std::chrono::duration(smoothing_time_constant).count()); + const double utilization_variance_threshold = + lb_config.has_utilization_variance_threshold() + ? lb_config.utilization_variance_threshold().value() + : 0.1; + if (!(utilization_variance_threshold >= 0.0 && utilization_variance_threshold <= 1.0)) { + return absl::InvalidArgumentError("utilization_variance_threshold must be in [0, 1]"); + } + const double remote_probe_fraction = + lb_config.has_remote_probe_fraction() ? lb_config.remote_probe_fraction().value() : 0.03; + if (!(remote_probe_fraction >= 0.0 && remote_probe_fraction < 1.0)) { + return absl::InvalidArgumentError("remote_probe_fraction must be in [0, 1)"); + } + const std::chrono::milliseconds weight_expiration_period = std::chrono::milliseconds( + PROTOBUF_GET_MS_OR_DEFAULT(lb_config, weight_expiration_period, 180000)); + std::vector metric_names(lb_config.metric_names_for_computing_utilization().begin(), + lb_config.metric_names_for_computing_utilization().end()); + + // Resolve the endpoint-picking child policy. + Upstream::TypedLoadBalancerFactory* endpoint_picking_policy_factory = nullptr; + for (const auto& endpoint_picking_policy : lb_config.endpoint_picking_policy().policies()) { + endpoint_picking_policy_factory = + Config::Utility::getAndCheckFactory( + endpoint_picking_policy.typed_extension_config(), + /*is_optional=*/true); + + if (endpoint_picking_policy_factory != nullptr) { + // Load and validate the child policy configuration. + auto sub_lb_proto_message = endpoint_picking_policy_factory->createEmptyConfigProto(); + RETURN_IF_NOT_OK(Config::Utility::translateOpaqueConfig( + endpoint_picking_policy.typed_extension_config().typed_config(), + context.messageValidationVisitor(), *sub_lb_proto_message)); + + auto lb_config_or_error = + endpoint_picking_policy_factory->loadConfig(context, *sub_lb_proto_message); + RETURN_IF_NOT_OK(lb_config_or_error.status()); + + return std::make_unique( + *endpoint_picking_policy_factory, endpoint_picking_policy_factory->name(), + LoadBalancerConfigSharedPtr(std::move(lb_config_or_error.value())), weight_update_period, + utilization_variance_threshold, ewma_alpha, remote_probe_fraction, + weight_expiration_period, std::move(metric_names), context.mainThreadDispatcher(), + context.threadLocal()); + } + } + + return absl::InvalidArgumentError("No supported endpoint picking policy."); } +/** + * Static registration for the Factory. @see RegisterFactory. + */ REGISTER_FACTORY(Factory, Upstream::TypedLoadBalancerFactory); } // namespace LoadAwareLocality diff --git a/source/extensions/load_balancing_policies/load_aware_locality/config.h b/source/extensions/load_balancing_policies/load_aware_locality/config.h index 0140691580782..c52a7b703ef99 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/config.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/config.h @@ -1,21 +1,24 @@ #pragma once #include "envoy/extensions/load_balancing_policies/load_aware_locality/v3/load_aware_locality.pb.h" +#include "envoy/server/factory_context.h" #include "envoy/upstream/load_balancer.h" -#include "source/common/upstream/load_balancer_factory_base.h" +#include "source/extensions/load_balancing_policies/common/factory_base.h" namespace Envoy { namespace Extensions { namespace LoadBalancingPolicies { namespace LoadAwareLocality { -using LoadAwareLocalityProto = +using LoadAwareLocalityLbProto = envoy::extensions::load_balancing_policies::load_aware_locality::v3::LoadAwareLocality; -class Factory : public Upstream::TypedLoadBalancerFactoryBase { +class Factory : public Upstream::TypedLoadBalancerFactoryBase { public: - Factory() : TypedLoadBalancerFactoryBase("envoy.load_balancing_policies.load_aware_locality") {} + Factory() + : Upstream::TypedLoadBalancerFactoryBase( + "envoy.load_balancing_policies.load_aware_locality") {} Upstream::ThreadAwareLoadBalancerPtr create(OptRef lb_config, const Upstream::ClusterInfo& cluster_info, From 97c29407cd8d8047f125f612ce176780fe9fb3f0 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 14:04:43 -0600 Subject: [PATCH 02/16] load_aware_locality: per-host ORCA load report ingestion Add LocalityLbHostData, the per-host LB policy data that receives ORCA load reports via the multi-slot HostLbPolicyData mechanism and stores a lock-free utilization value behind a release/acquire freshness timestamp for the main-thread weight computation. Reports without a positive utilization signal are dropped (matching gRFC A58 / CSWRR) so signal-less hosts age out rather than pinning their locality as fresh-and-idle. Also lands two foundational types used by the config factory: the LoadAwareLocalityStats counter struct and the LoadAwareLocalityLbConfig holder. Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality_lb.cc | 44 +++++ .../load_aware_locality_lb.h | 160 ++++++++++++++++++ 2 files changed, 204 insertions(+) create mode 100644 source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc create mode 100644 source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc new file mode 100644 index 0000000000000..f49fbbc487dfc --- /dev/null +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -0,0 +1,44 @@ +#include "source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h" + +#include +#include +#include +#include +#include +#include + +#include "envoy/stats/stats_macros.h" + +#include "source/common/protobuf/utility.h" +#include "source/extensions/load_balancing_policies/common/load_balancer_impl.h" +#include "source/extensions/load_balancing_policies/common/orca_weight_manager.h" + +#include "absl/container/flat_hash_set.h" +#include "absl/strings/str_cat.h" + +namespace Envoy { +namespace Extensions { +namespace LoadBalancingPolicies { +namespace LoadAwareLocality { + +absl::Status LocalityLbHostData::onOrcaLoadReport(const Upstream::OrcaLoadReport& report, + const StreamInfo::StreamInfo&) { + const double util = + Common::OrcaLoadReportHandler::getUtilizationFromOrcaReport(report, *metric_names_); + // A report without a positive utilization signal is not a valid sample (same freshness rule as + // gRFC A58 / CSWRR): don't store it or refresh freshness, so signal-less hosts age out via + // weight_expiration_period instead of pinning their locality as fresh-and-idle. + if (util <= 0.0) { + return absl::OkStatus(); + } + const int64_t now_ms = std::chrono::duration_cast( + time_source_.monotonicTime().time_since_epoch()) + .count(); + storeUtilization(util, now_ms); + return absl::OkStatus(); +} + +} // namespace LoadAwareLocality +} // namespace LoadBalancingPolicies +} // namespace Extensions +} // namespace Envoy diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h new file mode 100644 index 0000000000000..da151c6c748e7 --- /dev/null +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -0,0 +1,160 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "envoy/common/random_generator.h" +#include "envoy/event/dispatcher.h" +#include "envoy/stats/stats_macros.h" +#include "envoy/thread_local/thread_local.h" +#include "envoy/thread_local/thread_local_object.h" +#include "envoy/upstream/load_balancer.h" +#include "envoy/upstream/locality.h" +#include "envoy/upstream/upstream.h" + +#include "source/common/common/logger.h" +#include "source/common/upstream/upstream_impl.h" +#include "source/extensions/load_balancing_policies/common/load_balancer_impl.h" + +#include "absl/container/flat_hash_map.h" + +namespace Envoy { +namespace Extensions { +namespace LoadBalancingPolicies { +namespace LoadAwareLocality { + +#define ALL_LOAD_AWARE_LOCALITY_STATS(COUNTER) \ + COUNTER(all_overloaded_total) \ + COUNTER(local_preferred_total) \ + COUNTER(probe_active_total) \ + COUNTER(stale_locality_total) + +struct LoadAwareLocalityStats { + ALL_LOAD_AWARE_LOCALITY_STATS(GENERATE_COUNTER_STRUCT) +}; + +/** + * Per-host LB policy data for the load-aware locality policy. + * Receives ORCA load reports via the multi-slot HostLbPolicyData mechanism + * and stores a lock-free utilization value for the main-thread weight computation. + * + * Written by worker threads via onOrcaLoadReport(); read by the main thread + * during locality weight computation. The release/acquire on last_update_time_ms_ + * ensures the reader sees the utilization that was stored before it. + * + * lastUpdateTimeMs() == kNeverReported distinguishes "never reported" from any real + * report time (monotonic 0 ms is a legitimate timestamp, e.g. under simulated time). + */ +class LocalityLbHostData : public Upstream::HostLbPolicyData { +public: + // Sentinel: no report has been stored yet. Out-of-band (int64 min) so a real monotonic + // timestamp of 0 is not mistaken for "never reported". + static constexpr int64_t kNeverReported = std::numeric_limits::min(); + + LocalityLbHostData(TimeSource& time_source, + std::shared_ptr> metric_names) + : time_source_(time_source), metric_names_(std::move(metric_names)) {} + + bool receivesOrcaLoadReport() const override { return true; } + + absl::Status onOrcaLoadReport(const Upstream::OrcaLoadReport& report, + const StreamInfo::StreamInfo&) override; + + double utilization() const { return utilization_.load(std::memory_order_relaxed); } + int64_t lastUpdateTimeMs() const { return last_update_time_ms_.load(std::memory_order_acquire); } + +private: + void storeUtilization(double util, int64_t monotonic_time_ms) { + if (!std::isfinite(util)) { + return; + } + utilization_.store(std::clamp(util, 0.0, 1.0), std::memory_order_relaxed); + last_update_time_ms_.store(monotonic_time_ms, std::memory_order_release); + } + + static_assert(std::atomic::is_always_lock_free, + "std::atomic must be lock-free for safe cross-thread utilization updates"); + std::atomic utilization_{0.0}; + std::atomic last_update_time_ms_{kNeverReported}; + TimeSource& time_source_; + const std::shared_ptr> metric_names_; +}; + +// Shared ownership wrapper for the child LB config. The config is created during loadConfig() +// and shared between LoadAwareLocalityLbConfig (which is const when accessed) and +// WorkerLocalLbFactory (which needs it for the lifetime of the cluster). +using LoadBalancerConfigSharedPtr = std::shared_ptr; + +/** + * Load balancer config for the load-aware locality policy. + */ +class LoadAwareLocalityLbConfig : public Upstream::LoadBalancerConfig { +public: + LoadAwareLocalityLbConfig(Upstream::TypedLoadBalancerFactory& endpoint_picking_policy_factory, + std::string endpoint_picking_policy_name, + LoadBalancerConfigSharedPtr endpoint_picking_policy_config, + std::chrono::milliseconds weight_update_period, + double utilization_variance_threshold, double ewma_alpha, + double remote_probe_fraction, + std::chrono::milliseconds weight_expiration_period, + std::vector metric_names_for_computing_utilization, + Event::Dispatcher& main_thread_dispatcher, + ThreadLocal::SlotAllocator& tls_slot_allocator) + : endpoint_picking_policy_factory_(endpoint_picking_policy_factory), + endpoint_picking_policy_name_(std::move(endpoint_picking_policy_name)), + endpoint_picking_policy_config_(std::move(endpoint_picking_policy_config)), + weight_update_period_(weight_update_period), + utilization_variance_threshold_(utilization_variance_threshold), ewma_alpha_(ewma_alpha), + remote_probe_fraction_(remote_probe_fraction), + weight_expiration_period_(weight_expiration_period), + metric_names_for_computing_utilization_(std::move(metric_names_for_computing_utilization)), + main_thread_dispatcher_(main_thread_dispatcher), tls_slot_allocator_(tls_slot_allocator) {} + + Upstream::TypedLoadBalancerFactory& endpointPickingPolicyFactory() const { + return endpoint_picking_policy_factory_; + } + const std::string& endpointPickingPolicyName() const { return endpoint_picking_policy_name_; } + const LoadBalancerConfigSharedPtr& endpointPickingPolicyConfig() const { + return endpoint_picking_policy_config_; + } + std::chrono::milliseconds weightUpdatePeriod() const { return weight_update_period_; } + double utilizationVarianceThreshold() const { return utilization_variance_threshold_; } + double ewmaAlpha() const { return ewma_alpha_; } + double remoteProbeFraction() const { return remote_probe_fraction_; } + std::chrono::milliseconds weightExpirationPeriod() const { return weight_expiration_period_; } + const std::vector& metricNamesForComputingUtilization() const { + return metric_names_for_computing_utilization_; + } + Event::Dispatcher& mainThreadDispatcher() const { return main_thread_dispatcher_; } + ThreadLocal::SlotAllocator& tlsSlotAllocator() const { return tls_slot_allocator_; } + absl::Status validateEndpoints(const Upstream::PriorityState& priorities) const override { + return endpoint_picking_policy_config_ != nullptr + ? endpoint_picking_policy_config_->validateEndpoints(priorities) + : absl::OkStatus(); + } + +private: + Upstream::TypedLoadBalancerFactory& endpoint_picking_policy_factory_; + const std::string endpoint_picking_policy_name_; + const LoadBalancerConfigSharedPtr endpoint_picking_policy_config_; + const std::chrono::milliseconds weight_update_period_; + const double utilization_variance_threshold_; + const double ewma_alpha_; + const double remote_probe_fraction_; + const std::chrono::milliseconds weight_expiration_period_; + const std::vector metric_names_for_computing_utilization_; + Event::Dispatcher& main_thread_dispatcher_; + ThreadLocal::SlotAllocator& tls_slot_allocator_; +}; + +} // namespace LoadAwareLocality +} // namespace LoadBalancingPolicies +} // namespace Extensions +} // namespace Envoy From 5eda981011fca1c5b0a2c9367f87a3ff39c664fe Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 14:04:43 -0600 Subject: [PATCH 03/16] load_aware_locality: routing weights snapshot and worker factory Add the immutable RoutingWeightsSnapshot (per-priority, per-source locality weights keyed by Locality identity) and the WorkerLocalLbFactory that owns the shared child ThreadAwareLoadBalancer, publishes snapshots to workers through a thread-local slot, and instantiates the worker LB and its per-locality child LBs. Keying weights by Locality identity keeps the main-thread-to-worker mapping stable when localities are added or removed between snapshot ticks. Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality_lb.cc | 48 +++++++ .../load_aware_locality_lb.h | 133 ++++++++++++++++++ 2 files changed, 181 insertions(+) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc index f49fbbc487dfc..0d536b6262c48 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -38,6 +38,54 @@ absl::Status LocalityLbHostData::onOrcaLoadReport(const Upstream::OrcaLoadReport return absl::OkStatus(); } +// --- WorkerLocalLbFactory --- + +WorkerLocalLbFactory::WorkerLocalLbFactory( + Upstream::TypedLoadBalancerFactory& child_factory, std::string child_factory_name, + LoadBalancerConfigSharedPtr child_config, const Upstream::ClusterInfo& cluster_info, + const Upstream::PrioritySet& cluster_priority_set, Runtime::Loader& runtime, + Envoy::Random::RandomGenerator& random, TimeSource& time_source, + ThreadLocal::SlotAllocator& tls_slot_allocator) + : child_factory_name_(std::move(child_factory_name)), child_config_(std::move(child_config)), + cluster_info_(cluster_info), random_(random), runtime_(runtime), + healthy_panic_threshold_(PROTOBUF_PERCENT_TO_ROUNDED_INTEGER_OR_DEFAULT( + cluster_info.lbConfig(), healthy_panic_threshold, 100, 50)), + fail_traffic_on_panic_( + cluster_info.lbConfig().zone_aware_lb_config().fail_traffic_on_panic()) { + auto child_config_ref = + makeOptRefFromPtr(child_config_.get()); + child_thread_aware_lb_ = child_factory.create( + child_config_ref, cluster_info_, cluster_priority_set, runtime, random_, time_source); + tls_ = ThreadLocal::TypedSlot::makeUnique(tls_slot_allocator); + tls_->set([](Event::Dispatcher&) { return std::make_shared(); }); +} + +absl::Status WorkerLocalLbFactory::initializeChildLb() { + if (child_thread_aware_lb_ == nullptr) { + return absl::InvalidArgumentError(absl::StrCat( + "Unsupported endpoint picking policy for load_aware_locality: ", child_factory_name_, + ". Child load balancer could not be instantiated per locality.")); + } + return child_thread_aware_lb_->initialize(); +} + +Upstream::LoadBalancerPtr +WorkerLocalLbFactory::createWorkerChildLb(Upstream::PrioritySetImpl& per_locality_priority_set) { + // initializeChildLb() must have been called on the main thread before workers call this. + ASSERT(child_thread_aware_lb_ != nullptr); + Upstream::LoadBalancerParams child_params{per_locality_priority_set, nullptr}; + return child_thread_aware_lb_->factory()->create(child_params); +} + +bool WorkerLocalLbFactory::recreateChildOnHostChange() const { + ASSERT(child_thread_aware_lb_ != nullptr); + return child_thread_aware_lb_->factory()->recreateOnHostChangeDeprecated(); +} + +Upstream::LoadBalancerPtr WorkerLocalLbFactory::create(Upstream::LoadBalancerParams params) { + return std::make_unique(*this, params.priority_set); +} + } // namespace LoadAwareLocality } // namespace LoadBalancingPolicies } // namespace Extensions diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h index da151c6c748e7..1570f312d51cf 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -154,6 +154,139 @@ class LoadAwareLocalityLbConfig : public Upstream::LoadBalancerConfig { ThreadLocal::SlotAllocator& tls_slot_allocator_; }; +// Advisory locality weights keyed by Locality identity (not HostsPerLocality position). Keying by +// identity makes the main-thread→worker mapping robust to a locality being added/removed between +// the snapshot tick and the worker's live membership: the lexicographic index of later localities +// shifts on add/remove, but their identity does not. Uses the same comparators as +// load_balancer_impl.cc. +using LocalityWeightsMap = absl::flat_hash_map; + +// Per-locality EWMA-smoothed utilization keyed by Locality identity. An entry exists iff the +// locality has contributed at least one valid sample: a present entry whose locality has zero +// valid hosts this tick is stale (carry the prior value, count the stat); an absent entry is +// cold (host-count baseline, no stat). Identity keying survives the index shifts caused by +// locality add/remove, unlike position-based state. +using LocalityEwmaMap = absl::flat_hash_map; + +/** + * Immutable snapshot of per-locality routing weights for a single priority. + */ +struct PriorityRoutingWeights { + enum class SelectionSource : uint8_t { Healthy = 0, Degraded = 1, AllHosts = 2 }; + + struct SourceWeights { + // Per-locality routing weights (host_count * headroom), keyed by Locality identity. + LocalityWeightsMap weights; + // Informational flag: true if local preference was triggered (variance below threshold). + // Not consulted on the hot path — the routing decision is fully encoded in weights. + // With remote_probe_fraction > 0, weights still include a small remote share even when true. + bool all_local{false}; + }; + + // Per-source routing weights: [Healthy=0, Degraded=1, AllHosts=2]. + std::array by_source; + // True when the hosts-per-locality has a local locality (index 0 is "local"). + // Workers use this to pick the correct zone routing stat counter. + bool has_local_locality{false}; + + const LocalityWeightsMap& weightsFor(SelectionSource s) const { + return by_source[static_cast(s)].weights; + } + bool allLocalFor(SelectionSource s) const { return by_source[static_cast(s)].all_local; } +}; + +/** + * Immutable snapshot of advisory per-priority locality weights shared between the main thread and + * workers. Priority/health/panic selection is live worker state (LoadBalancerBase), not + * snapshotted. + */ +struct RoutingWeightsSnapshot { + // Per-priority routing weights, indexed by cluster priority. + std::vector priority_weights; +}; + +using RoutingWeightsSnapshotConstSharedPtr = std::shared_ptr; + +/** + * Thread-local shim holding the routing weights pointer, pushed from the main thread via TLS. + */ +struct ThreadLocalShim : public ThreadLocal::ThreadLocalObject { + RoutingWeightsSnapshotConstSharedPtr routing_weights; +}; + +class WorkerLocalLb; + +/** + * Factory shared across workers. Owns the shared child ThreadAwareLoadBalancer and publishes + * routing-weight snapshots to workers via a thread-local slot; workers create WorkerLocalLbs. + */ +class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { +public: + WorkerLocalLbFactory(Upstream::TypedLoadBalancerFactory& child_factory, + std::string child_factory_name, LoadBalancerConfigSharedPtr child_config, + const Upstream::ClusterInfo& cluster_info, + const Upstream::PrioritySet& cluster_priority_set, Runtime::Loader& runtime, + Envoy::Random::RandomGenerator& random, TimeSource& time_source, + ThreadLocal::SlotAllocator& tls_slot_allocator); + + // Upstream::LoadBalancerFactory + Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams params) override; + bool recreateOnHostChangeDeprecated() const override { return false; } + + // Called by the main thread (from LoadAwareLocalityLoadBalancer::initialize()) to initialize + // the shared child ThreadAwareLoadBalancer. Must be called before any worker creates a LB. + absl::Status initializeChildLb(); + + // Called by the main thread to publish new routing weights via TLS. + void updateRoutingWeights(RoutingWeightsSnapshotConstSharedPtr snapshot) { + tls_->runOnAllThreads([snapshot](OptRef shim) { + if (shim.has_value()) { + shim->routing_weights = snapshot; + } + }); + } + + // Called by workers to get the current routing weights (lock-free TLS read). + // SAFETY: Must only be called on a thread that owns a TLS slot instance (worker or main + // thread). The returned pointer is valid only for the duration of the current task; do not + // store it across yield points or after the TLS slot could be updated. + const RoutingWeightsSnapshot* routingWeights() const { + auto shim = tls_->get(); + return shim.has_value() ? shim->routing_weights.get() : nullptr; + } + + // Called by workers to create a per-locality worker LB from the shared child factory. + // Must only be called after initializeChildLb() has returned on the main thread. + Upstream::LoadBalancerPtr + createWorkerChildLb(Upstream::PrioritySetImpl& per_locality_priority_set); + + // Whether the child policy requires the worker LB to be recreated on host membership changes. + bool recreateChildOnHostChange() const; + + Envoy::Random::RandomGenerator& random() const { return random_; } + Runtime::Loader& runtime() const { return runtime_; } + uint32_t healthyPanicThreshold() const { return healthy_panic_threshold_; } + bool failTrafficOnPanic() const { return fail_traffic_on_panic_; } + Upstream::ClusterLbStats& lbStats() const { return cluster_info_.lbStats(); } + +private: + std::string child_factory_name_; + LoadBalancerConfigSharedPtr child_config_; + const Upstream::ClusterInfo& cluster_info_; + Envoy::Random::RandomGenerator& random_; + Runtime::Loader& runtime_; + uint32_t healthy_panic_threshold_; + const bool fail_traffic_on_panic_; + + // Single child ThreadAwareLoadBalancer created and initialized on the main thread. + // Workers only call factory()->create() from it. + Upstream::ThreadAwareLoadBalancerPtr child_thread_aware_lb_; + + std::unique_ptr> tls_; +}; + } // namespace LoadAwareLocality } // namespace LoadBalancingPolicies } // namespace Extensions From 563d682512abc3b916eda8d0de10bd7bbd235359 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 14:04:43 -0600 Subject: [PATCH 04/16] load_aware_locality: worker-local load balancer Add WorkerLocalLb, the per-worker data-path load balancer. It derives LoadBalancerBase for live priority, health, and panic selection, chooses a locality within the selected priority/source by capacity-weighted random over the routing snapshot (looked up by Locality identity), and delegates endpoint selection to a per-locality child LB. Per-source child LBs (healthy/degraded/all-hosts) are built per locality and re-synced on membership changes, with topology rebuilds gated on an actual membership delta. Endpoint retry stays with the child LB; a stack-scoped context wrapper suppresses re-running priority and hash policies, and async child selection is rejected to protect that wrapper. Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality_lb.cc | 428 ++++++++++++++++++ .../load_aware_locality_lb.h | 102 +++++ 2 files changed, 530 insertions(+) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc index 0d536b6262c48..2a37d02389b49 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -86,6 +86,434 @@ Upstream::LoadBalancerPtr WorkerLocalLbFactory::create(Upstream::LoadBalancerPar return std::make_unique(*this, params.priority_set); } +namespace { + +// Context handed to per-locality child LBs. Overrides only determinePriorityLoad, which returns the +// child's own default: the parent already consulted the real context's (possibly stateful) +// retry-priority plugin against the cluster-wide priority set; re-invoking it against the +// single-priority per-locality set would corrupt its state. Everything else (hash key, retry hooks +// shouldSelectAnotherHost/hostSelectionRetryCount, metadata, headers) passes through to the wrapped +// context; endpoint retry is owned by the child within the chosen locality. Stack-allocated per +// pick; children must not retain it beyond the call (same constraint as the subset LB wrapper). +class ChildLoadBalancerContext : public Upstream::LoadBalancerContext { +public: + explicit ChildLoadBalancerContext(Upstream::LoadBalancerContext& wrapped) : wrapped_(wrapped) {} + + std::optional computeHashKey() override { return wrapped_.computeHashKey(); } + const Upstream::HealthyAndDegradedLoad& + determinePriorityLoad(const Upstream::PrioritySet&, + const Upstream::HealthyAndDegradedLoad& default_priority_load, + const Upstream::RetryPriority::PriorityMappingFunc&) override { + return default_priority_load; + } + const Router::MetadataMatchCriteria* metadataMatchCriteria() override { + return wrapped_.metadataMatchCriteria(); + } + const Network::Connection* downstreamConnection() const override { + return wrapped_.downstreamConnection(); + } + StreamInfo::StreamInfo* requestStreamInfo() const override { + return wrapped_.requestStreamInfo(); + } + const Http::RequestHeaderMap* downstreamHeaders() const override { + return wrapped_.downstreamHeaders(); + } + bool shouldSelectAnotherHost(const Upstream::Host& host) override { + return wrapped_.shouldSelectAnotherHost(host); + } + uint32_t hostSelectionRetryCount() const override { return wrapped_.hostSelectionRetryCount(); } + Network::Socket::OptionsSharedPtr upstreamSocketOptions() const override { + return wrapped_.upstreamSocketOptions(); + } + Network::TransportSocketOptionsConstSharedPtr upstreamTransportSocketOptions() const override { + return wrapped_.upstreamTransportSocketOptions(); + } + OptRef overrideHostToSelect() const override { + return wrapped_.overrideHostToSelect(); + } + // Async child selection is unsupported (the async handle is cancelled in chooseHost), so this is + // never invoked; a no-op matches the subset LB and avoids implying otherwise. + void onAsyncHostSelection(Upstream::HostConstSharedPtr&&, std::string&&) override {} + void setHeadersModifier(std::function modifier) override { + wrapped_.setHeadersModifier(std::move(modifier)); + } + +private: + Upstream::LoadBalancerContext& wrapped_; +}; + +// Async child host selection is unsupported: the child would retain the stack-scoped +// ChildLoadBalancerContext and call onAsyncHostSelection after chooseHost returns (use-after-free). +// Cancel any async handle (HostSelectionResponse guarantees no callback after cancel) and +// fail synchronously, matching the subset LB. Synchronous responses pass through unchanged. +Upstream::HostSelectionResponse failOnAsyncSelection(Upstream::HostSelectionResponse response) { + if (response.cancelable != nullptr) { + response.cancelable->cancel(); + return {nullptr}; + } + return response; +} + +} // namespace + +// --- WorkerLocalLb (per-worker) --- + +WorkerLocalLb::WorkerLocalLb(WorkerLocalLbFactory& factory, + const Upstream::PrioritySet& priority_set) + : Upstream::LoadBalancerBase(priority_set, factory.lbStats(), factory.runtime(), + factory.random(), factory.healthyPanicThreshold()), + factory_(factory) { + buildPerPriorityLocalities(); + // Register AFTER initial build so callback doesn't fire during construction. + priority_sync_cb_ = priority_set_.addPriorityUpdateCb( + [this](uint32_t priority, const Upstream::HostVector& hosts_added, + const Upstream::HostVector& hosts_removed) { + // A membership delta permits a topology rebuild. An empty-delta update still matters for + // child policies (in-place host attribute changes such as health, weight, or metadata) but + // never rebuilds — a real topology change always arrives with a delta. + syncPriority(priority, /*allow_rebuild=*/!hosts_added.empty() || !hosts_removed.empty()); + }); +} + +WorkerLocalLb::~WorkerLocalLb() { + // Reset callback handle before other members are destroyed, so the callback + // doesn't fire during destruction and access freed per-locality state. + priority_sync_cb_.reset(); +} + +void WorkerLocalLb::updateLocalityHosts(PerSourceLocalityState& state, + const Upstream::HostVector& hosts, bool is_local, + const Upstream::HostVector& hosts_added, + const Upstream::HostVector& hosts_removed) { + auto hosts_shared = std::make_shared(hosts); + auto per_locality = std::make_shared(hosts, is_local); + // All passed-in hosts are marked as "healthy" in the child priority set regardless of their + // actual health status. The caller (syncLocalityState) already partitions hosts by source + // (healthy/degraded/all), so this child LB only sees hosts that can be selected. Marking + // them all healthy ensures the child policy (e.g. RoundRobin) considers every host eligible. + auto healthy_hosts = std::make_shared(hosts); + auto update_params = Upstream::HostSetImpl::updateHostsParams( + hosts_shared, per_locality, healthy_hosts, per_locality, + std::make_shared(), + Upstream::HostsPerLocalityImpl::empty(), + std::make_shared(), + Upstream::HostsPerLocalityImpl::empty()); + state.priority_set->updateHosts(0, std::move(update_params), nullptr, hosts_added, hosts_removed, + absl::nullopt, absl::nullopt); +} + +void WorkerLocalLb::buildPerPriorityLocalities() { + const auto& host_sets = priority_set_.hostSetsPerPriority(); + per_priority_locality_.clear(); + per_priority_locality_.resize(host_sets.size()); + + for (size_t priority = 0; priority < host_sets.size(); ++priority) { + buildPerLocality(priority, *host_sets[priority]); + } +} + +void WorkerLocalLb::syncLocalityState(PerLocalityState& state, const Upstream::HostSet& host_set, + size_t locality_index, bool recreate_child) { + static const Upstream::HostVector empty_hosts; + const auto& all_localities = host_set.hostsPerLocality().get(); + const auto& healthy_localities = host_set.healthyHostsPerLocality().get(); + const auto& degraded_localities = host_set.degradedHostsPerLocality().get(); + const bool is_local = (locality_index == 0 && host_set.hostsPerLocality().hasLocalLocality()); + + const auto& all_hosts = all_localities[locality_index]; + const auto& healthy_hosts = + locality_index < healthy_localities.size() ? healthy_localities[locality_index] : empty_hosts; + const auto& degraded_hosts = locality_index < degraded_localities.size() + ? degraded_localities[locality_index] + : empty_hosts; + + // Record this locality's identity for chooseLocality's weight lookup. Empty group → no identity. + state.locality = + all_hosts.empty() ? envoy::config::core::v3::Locality() : all_hosts[0]->locality(); + + const auto sync_source = [this, is_local, recreate_child](PerSourceLocalityState& source_state, + const Upstream::HostVector& new_hosts) { + if (new_hosts.empty()) { + if (source_state.priority_set != nullptr) { + source_state.lb.reset(); + source_state.priority_set.reset(); + } + return; + } + + if (source_state.priority_set == nullptr) { + source_state.priority_set = std::make_unique(); + updateLocalityHosts(source_state, new_hosts, is_local, new_hosts, {}); + source_state.lb = factory_.createWorkerChildLb(*source_state.priority_set); + return; + } + + const auto& old_hosts = source_state.priority_set->hostSetsPerPriority()[0]->hosts(); + absl::flat_hash_set old_set(old_hosts.begin(), old_hosts.end()); + + Upstream::HostVector hosts_added; + Upstream::HostVector hosts_removed; + for (const auto& host : new_hosts) { + if (!old_set.erase(host)) { + hosts_added.push_back(host); + } + } + hosts_removed.reserve(old_set.size()); + for (const auto& host : old_set) { + hosts_removed.push_back(std::const_pointer_cast(host)); + } + + const bool membership_changed = !hosts_added.empty() || !hosts_removed.empty(); + if (!membership_changed) { + // Host identity is unchanged, but child policies still need an update to observe in-place + // host attribute changes such as weight or metadata. + updateLocalityHosts(source_state, new_hosts, is_local, {}, {}); + return; + } + + updateLocalityHosts(source_state, new_hosts, is_local, hosts_added, hosts_removed); + if (recreate_child) { + source_state.lb = factory_.createWorkerChildLb(*source_state.priority_set); + } + }; + + sync_source(state.stateFor(PriorityRoutingWeights::SelectionSource::Healthy), healthy_hosts); + sync_source(state.stateFor(PriorityRoutingWeights::SelectionSource::Degraded), degraded_hosts); + sync_source(state.stateFor(PriorityRoutingWeights::SelectionSource::AllHosts), all_hosts); +} + +void WorkerLocalLb::buildPerLocality(uint32_t priority, const Upstream::HostSet& host_set) { + // Locality routing structures for this priority are (re)generated on this worker — the same + // event the zone-aware counter tracks. + stats_.lb_recalculate_zone_structures_.inc(); + const auto& locality_hosts = host_set.hostsPerLocality().get(); + auto& per_locality = per_priority_locality_[priority].localities; + per_locality.clear(); + per_locality.resize(locality_hosts.size()); + + for (size_t i = 0; i < locality_hosts.size(); ++i) { + syncLocalityState(per_locality[i], host_set, i, /*recreate_child=*/false); + } +} + +void WorkerLocalLb::syncPriority(uint32_t priority, bool allow_rebuild) { + const auto& host_sets = priority_set_.hostSetsPerPriority(); + if (priority >= host_sets.size()) { + return; + } + + // Priority count changed — full rebuild when permitted, else wait for a membership-delta update. + if (host_sets.size() != per_priority_locality_.size()) { + if (allow_rebuild) { + buildPerPriorityLocalities(); + } + return; + } + + const auto& host_set = host_sets[priority]; + const auto& locality_hosts = host_set->hostsPerLocality().get(); + auto& per_locality = per_priority_locality_[priority].localities; + + // Topology change (locality added/removed) — rebuild this priority when permitted, else wait. + if (locality_hosts.size() != per_locality.size()) { + if (allow_rebuild) { + buildPerLocality(priority, *host_set); + } + return; + } + + const bool recreate_child = factory_.recreateChildOnHostChange(); + for (size_t i = 0; i < locality_hosts.size(); ++i) { + syncLocalityState(per_locality[i], *host_set, i, recreate_child); + } +} + +Upstream::LoadBalancer* +WorkerLocalLb::pickLocalityLb(const std::vector& per_locality, + PriorityRoutingWeights::SelectionSource source, size_t preferred_idx, + size_t& actual_idx) const { + auto* lb = per_locality[preferred_idx].stateFor(source).lb.get(); + if (lb != nullptr) { + actual_idx = preferred_idx; + return lb; + } + // Stale routing snapshot — the preferred locality's child LB was torn down after a host + // change but before the routing weights were recomputed. Scan for any locality with a + // usable LB. + for (size_t i = 0; i < per_locality.size(); ++i) { + lb = per_locality[i].stateFor(source).lb.get(); + if (lb != nullptr) { + actual_idx = i; + return lb; + } + } + return nullptr; +} + +WorkerLocalLb::PrioritySourcePick +WorkerLocalLb::resolvePrioritySource(Upstream::LoadBalancerContext* context, bool peeking) { + using SelectionSource = PriorityRoutingWeights::SelectionSource; + if (priority_set_.hostSetsPerPriority().empty()) { + return {0, SelectionSource::Healthy, /*in_panic=*/false, /*fail=*/true}; + } + + const uint64_t hash = random(peeking); + + // Live priority/health selection from worker-local LoadBalancerBase state. + const auto host_set_and_availability = chooseHostSet(context, hash); + const uint32_t priority = host_set_and_availability.first.priority(); + + if (isInPanic(priority)) { + return {priority, SelectionSource::AllHosts, /*in_panic=*/true, + /*fail=*/factory_.failTrafficOnPanic()}; + } + const SelectionSource source = + host_set_and_availability.second == Upstream::LoadBalancerBase::HostAvailability::Healthy + ? SelectionSource::Healthy + : SelectionSource::Degraded; + return {priority, source, /*in_panic=*/false, /*fail=*/false}; +} + +Upstream::HostSelectionResponse WorkerLocalLb::chooseHost(Upstream::LoadBalancerContext* context) { + const auto pick = resolvePrioritySource(context, /*peeking=*/false); + // Panic stat is counted here only; peekAnotherHost deliberately does not double-count it. + if (pick.in_panic) { + stats_.lb_healthy_panic_.inc(); + } + if (pick.fail) { + return {nullptr}; + } + const uint32_t priority = pick.priority; + const PriorityRoutingWeights::SelectionSource source = pick.source; + + // A custom retry-priority plugin can route to a just-created priority that + // per_priority_locality_ doesn't cover yet (empty-delta update, rebuild pending). + if (priority >= per_priority_locality_.size()) { + return {nullptr}; + } + auto& per_locality = per_priority_locality_[priority].localities; + if (per_locality.empty()) { + return {nullptr}; + } + + // Fetch the routing snapshot once and reuse it for locality selection and the zone-routing stat + // below (avoids a second TLS lookup per request). + const auto* snapshot = factory_.routingWeights(); + const size_t preferred_idx = chooseLocality(snapshot, priority, source, /*peeking=*/false); + size_t locality_idx = preferred_idx; + auto* lb = pickLocalityLb(per_locality, source, preferred_idx, locality_idx); + if (lb == nullptr) { + return {nullptr}; + } + + // Endpoint retry is owned by the child LB; delegate through the child context and return its + // response (failOnAsyncSelection rejects async children to protect the stack-scoped context). + // Record the zone-routing stat once for the selected locality. + if (!pick.in_panic && snapshot != nullptr && priority < snapshot->priority_weights.size()) { + const auto& priority_snapshot = snapshot->priority_weights[priority]; + if (priority_snapshot.has_local_locality) { + if (locality_idx == 0) { + if (priority_snapshot.allLocalFor(source)) { + stats_.lb_zone_routing_all_directly_.inc(); + } else { + stats_.lb_zone_routing_sampled_.inc(); + } + } else { + stats_.lb_zone_routing_cross_zone_.inc(); + } + } + } + + if (context == nullptr) { + return failOnAsyncSelection(lb->chooseHost(nullptr)); + } + ChildLoadBalancerContext child_context(*context); + return failOnAsyncSelection(lb->chooseHost(&child_context)); +} + +size_t WorkerLocalLb::chooseLocality(const RoutingWeightsSnapshot* snapshot, uint32_t priority, + PriorityRoutingWeights::SelectionSource source, bool peeking) { + const auto& per_locality = per_priority_locality_[priority].localities; + if (per_locality.size() <= 1) { + return 0; + } + if (snapshot == nullptr || priority >= snapshot->priority_weights.size()) { + return 0; + } + const auto& weights = snapshot->priority_weights[priority].weightsFor(source); + if (weights.empty()) { + return 0; + } + + // Look up each LIVE locality's advisory weight by identity. A locality present on the worker but + // missing from the snapshot (e.g. just-added, snapshot lagging by ≤1 tick) gets weight 0.0 and is + // excluded from selection until the next recompute — matching the prior count-mismatch behavior. + const auto weight_at = [&weights, &per_locality](size_t i) -> double { + auto it = weights.find(per_locality[i].locality); + return it != weights.end() ? it->second : 0.0; + }; + + double effective_total = 0.0; + for (size_t i = 0; i < per_locality.size(); ++i) { + effective_total += weight_at(i); + } + if (effective_total <= 0.0) { + return 0; + } + + // This draw is skipped by the early returns above, so a snapshot published between a paired peek + // and choose can offset the stashed-draw replay by one (harmless; membership changes clear it). + const double target = (static_cast(random(peeking)) / + static_cast(std::numeric_limits::max())) * + effective_total; + double cumulative = 0.0; + for (size_t i = 0; i < per_locality.size(); ++i) { + cumulative += weight_at(i); + if (target < cumulative) { + return i; + } + } + + // Floating-point rounding guard: return the last locality, not an arbitrary one. + return per_locality.size() - 1; +} + +Upstream::HostConstSharedPtr +WorkerLocalLb::peekAnotherHost(Upstream::LoadBalancerContext* context) { + // Cap preconnect look-ahead as the sibling LBs do, so stashed_random_ cannot grow unbounded and + // the peek/choose replay stream stays bounded. + if (Upstream::tooManyPreconnects(stashed_random_.size(), total_healthy_hosts_)) { + return nullptr; + } + const auto pick = resolvePrioritySource(context, /*peeking=*/true); + if (pick.fail) { + return nullptr; + } + const uint32_t priority = pick.priority; + const PriorityRoutingWeights::SelectionSource source = pick.source; + + if (priority >= per_priority_locality_.size()) { + return nullptr; + } + auto& per_locality = per_priority_locality_[priority].localities; + if (per_locality.empty()) { + return nullptr; + } + + const size_t preferred_idx = + chooseLocality(factory_.routingWeights(), priority, source, /*peeking=*/true); + size_t actual_idx = preferred_idx; + auto* lb = pickLocalityLb(per_locality, source, preferred_idx, actual_idx); + if (lb == nullptr) { + return nullptr; + } + if (context == nullptr) { + return lb->peekAnotherHost(nullptr); + } + ChildLoadBalancerContext child_context(*context); + return lb->peekAnotherHost(&child_context); +} + } // namespace LoadAwareLocality } // namespace LoadBalancingPolicies } // namespace Extensions diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h index 1570f312d51cf..ad2394689c88f 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -287,6 +287,108 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { std::unique_ptr> tls_; }; +/** + * Per-locality state in the worker-local LB. Holds a child PrioritySet and child LB for + * one locality. + */ +struct PerSourceLocalityState { + // PrioritySet containing only the hosts that can be selected for one source in one locality. + std::unique_ptr priority_set; + // The worker-local LB for this source/locality pair, created from the shared child factory. + Upstream::LoadBalancerPtr lb; +}; + +struct PerLocalityState { + // Per-source child LB state: [Healthy=0, Degraded=1, AllHosts=2]. + std::array by_source; + // Identity of this worker locality, used to look up its advisory weight in the snapshot map. + // Empty for a locality with no hosts (which carries no weight regardless). + envoy::config::core::v3::Locality locality; + + PerSourceLocalityState& stateFor(PriorityRoutingWeights::SelectionSource s) { + return by_source[static_cast(s)]; + } + const PerSourceLocalityState& stateFor(PriorityRoutingWeights::SelectionSource s) const { + return by_source[static_cast(s)]; + } +}; + +struct PerPriorityLocalityState { + std::vector localities; +}; + +/** + * Worker-local load balancer. Derives LoadBalancerBase for live priority/health/panic selection, + * then selects a locality by capacity-weighted random and delegates to the per-locality child LB + * for endpoint selection. + */ +class WorkerLocalLb : public Upstream::LoadBalancerBase { +public: + WorkerLocalLb(WorkerLocalLbFactory& factory, const Upstream::PrioritySet& priority_set); + ~WorkerLocalLb() override; + + // Upstream::LoadBalancer + Upstream::HostSelectionResponse chooseHost(Upstream::LoadBalancerContext* context) override; + Upstream::HostConstSharedPtr peekAnotherHost(Upstream::LoadBalancerContext* context) override; + +private: + // Build per-locality child LBs for all priorities. + void buildPerPriorityLocalities(); + + // Build or rebuild the per-locality child LBs for a single priority. + void buildPerLocality(uint32_t priority, const Upstream::HostSet& host_set); + + // Re-sync a priority's per-locality child LBs after a priority update. allow_rebuild permits + // tearing down and rebuilding on a topology change (priority/locality count mismatch); when false + // (in-place attribute update with no membership delta) a mismatch is left for the next + // membership-change callback to rebuild. + void syncPriority(uint32_t priority, bool allow_rebuild); + + // Update a locality/source PrioritySet with a pre-selected host subset. + void updateLocalityHosts(PerSourceLocalityState& state, const Upstream::HostVector& hosts, + bool is_local, const Upstream::HostVector& hosts_added, + const Upstream::HostVector& hosts_removed); + + // Update the per-source child LBs for one locality from the cluster's current host set. + void syncLocalityState(PerLocalityState& state, const Upstream::HostSet& host_set, + size_t locality_index, bool recreate_child); + + // Live priority + selection source for a request, resolved from worker-local LoadBalancerBase + // state. fail=true means the request must be failed (empty priority set, or panic + + // fail_traffic_on_panic). in_panic lets callers count the panic stat themselves (peek does not). + struct PrioritySourcePick { + uint32_t priority; + PriorityRoutingWeights::SelectionSource source; + bool in_panic; + bool fail; + }; + // peeking routes random draws through LoadBalancerBase::random(peeking) so peekAnotherHost and + // the subsequent chooseHost replay the same sequence. + PrioritySourcePick resolvePrioritySource(Upstream::LoadBalancerContext* context, bool peeking); + + // Choose a locality index within the selected priority/source by weighted-random over the routing + // snapshot. Weights are looked up by identity (see LocalityWeightsMap; missing → 0.0). Returns 0 + // when there is no usable weight (single locality, no/stale snapshot, or zero effective total). + // FUTURE: for very high locality counts, cache a per-worker index-aligned weight vector rebuilt + // on snapshot publish + membership change to restore pure index reads (avoids per-pick hashing). + size_t chooseLocality(const RoutingWeightsSnapshot* snapshot, uint32_t priority, + PriorityRoutingWeights::SelectionSource source, bool peeking); + + // Find a usable child LB at the preferred locality, falling back to any locality with a + // non-null LB when the routing snapshot is stale (e.g. a locality's hosts were removed but + // the routing weights haven't been recomputed yet). Returns nullptr if no locality has a + // usable LB. Sets actual_idx to the index of the locality whose LB was returned. + Upstream::LoadBalancer* pickLocalityLb(const std::vector& per_locality, + PriorityRoutingWeights::SelectionSource source, + size_t preferred_idx, size_t& actual_idx) const; + + WorkerLocalLbFactory& factory_; + std::vector per_priority_locality_; + // Destroyed explicitly in the destructor before other members so the callback doesn't fire + // during destruction and access freed per-locality state. + Envoy::Common::CallbackHandlePtr priority_sync_cb_; +}; + } // namespace LoadAwareLocality } // namespace LoadBalancingPolicies } // namespace Extensions From 7ee349969f998d4303e8bbbe6cb5cf4762670347 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 14:04:43 -0600 Subject: [PATCH 05/16] load_aware_locality: main-thread weight computation Add LoadAwareLocalityLoadBalancer, the main-thread ThreadAwareLoadBalancer. On a periodic timer it reads each host's ORCA utilization and computes per-locality routing weights for the healthy, degraded, and all-host source slices: EWMA-smoothed headroom (host_count * (1 - smoothed_util)), local preference when cross-locality variance is below threshold, and a small remote probe fraction to keep remote signal fresh. The result is published to workers as an immutable snapshot. Attaches LocalityLbHostData to hosts on membership changes and emits the all_overloaded, local_preferred, probe_active, and stale_locality counters. Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality_lb.cc | 289 ++++++++++++++++++ .../load_aware_locality_lb.h | 57 ++++ 2 files changed, 346 insertions(+) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc index 2a37d02389b49..8ed394e14bd0e 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -38,6 +38,295 @@ absl::Status LocalityLbHostData::onOrcaLoadReport(const Upstream::OrcaLoadReport return absl::OkStatus(); } +// --- LoadAwareLocalityLoadBalancer (main thread) --- + +LoadAwareLocalityLoadBalancer::LoadAwareLocalityLoadBalancer( + OptRef lb_config, const Upstream::ClusterInfo& cluster_info, + const Upstream::PrioritySet& priority_set, Runtime::Loader& runtime, + Envoy::Random::RandomGenerator& random, TimeSource& time_source) + : priority_set_(priority_set), lb_stats_{ALL_LOAD_AWARE_LOCALITY_STATS(POOL_COUNTER_PREFIX( + cluster_info.statsScope(), "load_aware_locality"))}, + time_source_(time_source) { + const auto* typed_config = dynamic_cast(lb_config.ptr()); + ASSERT(typed_config != nullptr); + + utilization_variance_threshold_ = typed_config->utilizationVarianceThreshold(); + ewma_alpha_ = typed_config->ewmaAlpha(); + remote_probe_fraction_ = typed_config->remoteProbeFraction(); + weight_expiration_period_ = typed_config->weightExpirationPeriod(); + weight_update_period_ = typed_config->weightUpdatePeriod(); + metric_names_ = std::make_shared>( + typed_config->metricNamesForComputingUtilization()); + + factory_ = std::make_shared( + typed_config->endpointPickingPolicyFactory(), typed_config->endpointPickingPolicyName(), + typed_config->endpointPickingPolicyConfig(), cluster_info, priority_set, runtime, random, + time_source, typed_config->tlsSlotAllocator()); + + weight_update_timer_ = typed_config->mainThreadDispatcher().createTimer( + [this]() { computeLocalityRoutingWeights(); }); +} + +LoadAwareLocalityLoadBalancer::~LoadAwareLocalityLoadBalancer() = default; + +void LoadAwareLocalityLoadBalancer::addLbPolicyDataToHosts(const Upstream::HostVector& hosts) { + for (const auto& host_ptr : hosts) { + if (!host_ptr->typedLbPolicyData().has_value()) { + auto data = std::make_unique(time_source_, metric_names_); + host_ptr->addLbPolicyData(std::move(data)); + } + } +} + +absl::Status LoadAwareLocalityLoadBalancer::initialize() { + RETURN_IF_NOT_OK(factory_->initializeChildLb()); + + for (const auto& host_set : priority_set_.hostSetsPerPriority()) { + addLbPolicyDataToHosts(host_set->hosts()); + } + + priority_update_cb_ = priority_set_.addPriorityUpdateCb( + [this](uint32_t, const Upstream::HostVector& hosts_added, const Upstream::HostVector&) { + addLbPolicyDataToHosts(hosts_added); + }); + + computeLocalityRoutingWeights(); + return absl::OkStatus(); +} + +void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { + // Re-arm first so the timer always fires on schedule regardless of early returns below. + weight_update_timer_->enableTimer(weight_update_period_); + + auto snapshot = std::make_shared(); + const auto& host_sets = priority_set_.hostSetsPerPriority(); + snapshot->priority_weights.resize(host_sets.size()); + + for (auto& per_source_state : ewma_state_) { + per_source_state.resize(host_sets.size()); + } + + // Current monotonic time in milliseconds, used for weight expiration checks. + const int64_t now_ms = std::chrono::duration_cast( + time_source_.monotonicTime().time_since_epoch()) + .count(); + + // Tick-scoped stat accumulators: OR'd across every (source × priority) pass below. + bool tick_all_overloaded = false; + bool tick_local_preferred = false; + bool tick_probe_active = false; + uint32_t tick_stale_localities = 0; // summed only from the all_hosts pass + + for (size_t priority = 0; priority < host_sets.size(); ++priority) { + const auto& host_set = host_sets[priority]; + const auto& hosts_per_locality = host_set->hostsPerLocality(); + const auto& locality_hosts = hosts_per_locality.get(); + auto& priority_snapshot = snapshot->priority_weights[priority]; + + if (locality_hosts.empty()) { + for (auto& per_source_state : ewma_state_) { + per_source_state[priority].clear(); + } + continue; + } + + priority_snapshot.has_local_locality = hosts_per_locality.hasLocalLocality(); + using SelectionSource = PriorityRoutingWeights::SelectionSource; + const auto run_source = [&](SelectionSource src, + const std::vector& eligible_hosts) { + const size_t s = static_cast(src); + auto& sw = priority_snapshot.by_source[s]; + return computeSourceWeights(hosts_per_locality, eligible_hosts, now_ms, sw.weights, + sw.all_local, ewma_state_[s][priority]); + }; + const auto healthy_result = + run_source(SelectionSource::Healthy, host_set->healthyHostsPerLocality().get()); + const auto degraded_result = + run_source(SelectionSource::Degraded, host_set->degradedHostsPerLocality().get()); + // all_hosts pass: run last so its staleness reflects the canonical "zero fresh hosts". + const auto all_hosts_result = run_source(SelectionSource::AllHosts, locality_hosts); + + tick_all_overloaded |= healthy_result.all_overloaded || degraded_result.all_overloaded || + all_hosts_result.all_overloaded; + tick_local_preferred |= healthy_result.local_preferred || degraded_result.local_preferred || + all_hosts_result.local_preferred; + tick_probe_active |= healthy_result.probe_active || degraded_result.probe_active || + all_hosts_result.probe_active; + // Count stale localities from the all_hosts pass only (canonical; avoids triple-counting). + tick_stale_localities += all_hosts_result.stale_localities; + } + + if (tick_all_overloaded) { + lb_stats_.all_overloaded_total_.inc(); + } + if (tick_local_preferred) { + lb_stats_.local_preferred_total_.inc(); + } + if (tick_probe_active) { + lb_stats_.probe_active_total_.inc(); + } + if (tick_stale_localities > 0) { + lb_stats_.stale_locality_total_.add(tick_stale_localities); + } + + ENVOY_LOG(trace, "computeLocalityRoutingWeights: {} priorities", + snapshot->priority_weights.size()); + factory_->updateRoutingWeights(std::move(snapshot)); +} + +LoadAwareLocalityLoadBalancer::SourceComputeResult +LoadAwareLocalityLoadBalancer::computeSourceWeights( + const Upstream::HostsPerLocality& all_hosts_per_locality, + const std::vector& eligible_hosts_per_locality, int64_t now_ms, + LocalityWeightsMap& weights_map, bool& all_local, LocalityEwmaMap& ewma_state) { + SourceComputeResult result; + const auto& locality_hosts = all_hosts_per_locality.get(); + const size_t locality_count = locality_hosts.size(); + // Compute weights index-based (local locality is index 0), then key by identity at the end. + std::vector weights(locality_count, 0.0); + all_local = false; + + std::vector avg_utils(locality_count, 0.0); + std::vector valid_counts(locality_count, 0); + std::vector host_counts(locality_count, 0); + + for (size_t i = 0; i < locality_count; ++i) { + const bool has_eligible = i < eligible_hosts_per_locality.size(); + host_counts[i] = + has_eligible ? static_cast(eligible_hosts_per_locality[i].size()) : 0u; + + double util_sum = 0.0; + uint32_t valid_count = 0; + if (has_eligible) { + for (const auto& host : eligible_hosts_per_locality[i]) { + auto host_data = host->typedLbPolicyData(); + if (!host_data.has_value()) { + continue; + } + const int64_t last_update_ms = host_data->lastUpdateTimeMs(); + if (last_update_ms == LocalityLbHostData::kNeverReported) { + continue; + } + + if (weight_expiration_period_.count() > 0 && + (now_ms - last_update_ms) > weight_expiration_period_.count()) { + continue; + } + + util_sum += host_data->utilization(); + valid_count++; + } + } + + avg_utils[i] = valid_count > 0 ? util_sum / valid_count : 0.0; + valid_counts[i] = valid_count; + } + + // Advance the identity-keyed EWMA state. Rebuilding the map each tick carries state for + // localities that persist, applies the first valid sample raw, and drops departed localities. + LocalityEwmaMap new_state; + new_state.reserve(locality_count); + std::vector utilizations(locality_count, 0.0); + std::vector stale(locality_count, false); + for (size_t i = 0; i < locality_count; ++i) { + if (locality_hosts[i].empty()) { + continue; // No identity and no hosts: weight below is 0 regardless. + } + const auto prev = ewma_state.find(locality_hosts[i][0]->locality()); + if (valid_counts[i] > 0) { + const double smoothed = prev == ewma_state.end() + ? avg_utils[i] + : ewma_alpha_ * avg_utils[i] + (1.0 - ewma_alpha_) * prev->second; + new_state[locality_hosts[i][0]->locality()] = smoothed; + utilizations[i] = smoothed; + } else if (prev != ewma_state.end()) { + // Stale: had valid data before, none now. Carry the prior smoothed value so a locality + // whose load we no longer know is not mistaken for idle; host-count baseline below. + new_state[locality_hosts[i][0]->locality()] = prev->second; + utilizations[i] = prev->second; + stale[i] = true; + ++result.stale_localities; + } + // else: cold (never had a valid sample) — utilization stays 0 (full headroom), so the + // weight below reduces to the host-count baseline without counting the stale stat. + } + ewma_state = std::move(new_state); + + uint32_t total_hosts = 0; + for (size_t i = 0; i < locality_count; ++i) { + weights[i] = stale[i] ? static_cast(host_counts[i]) + : host_counts[i] * std::max(0.0, 1.0 - utilizations[i]); + total_hosts += host_counts[i]; + } + + const auto set_all_local = [&weights, &all_local]() { + all_local = true; + std::fill(weights.begin(), weights.end(), 0.0); + if (!weights.empty()) { + weights[0] = 1.0; + } + }; + + // All-overloaded fallback, checked first per the rst: when no locality has headroom, spread by + // host count and skip local preference / probe. Otherwise apply stage-4 local preference and + // stage-5 remote probing. + const double total_base_weight = std::accumulate(weights.begin(), weights.end(), 0.0); + if (total_base_weight == 0.0 && total_hosts > 0) { + for (size_t i = 0; i < locality_count; ++i) { + weights[i] = static_cast(host_counts[i]); + } + result.all_overloaded = true; + } else { + // Remote host count, shared by the local-preference target and the probe redistribution below. + uint32_t remote_hosts = 0; + if (all_hosts_per_locality.hasLocalLocality()) { + double remote_util_sum = 0.0; + for (size_t i = 1; i < locality_count; ++i) { + remote_util_sum += utilizations[i] * host_counts[i]; + remote_hosts += host_counts[i]; + } + + if (total_hosts > 0 && !host_counts.empty() && host_counts[0] > 0 && remote_hosts > 0) { + const double target_util = remote_util_sum / remote_hosts; + if (utilizations[0] <= target_util + utilization_variance_threshold_) { + set_all_local(); + result.local_preferred = true; + } + } else if (total_hosts == 0) { + set_all_local(); + } + } + + if (remote_hosts > 0 && remote_probe_fraction_ > 0.0) { + const double total = std::accumulate(weights.begin(), weights.end(), 0.0); + const double remote_target = total * remote_probe_fraction_; + const double remote_sum = total - weights[0]; + if (total > 0.0 && remote_sum < remote_target) { + const double take_from_local = std::min(remote_target - remote_sum, weights[0]); + weights[0] -= take_from_local; + for (size_t i = 1; i < locality_count; ++i) { + weights[i] += take_from_local * static_cast(host_counts[i]) / remote_hosts; + } + result.probe_active = true; + } + } + } + + // Key the computed weights by Locality identity. An empty locality group carries no entry + // (it has 0 weight anyway); the local locality is stored under its own identity like the + // rest (local-vs-remote preference is already baked into the weight values). + weights_map.clear(); + weights_map.reserve(locality_count); + for (size_t i = 0; i < locality_count; ++i) { + if (locality_hosts[i].empty()) { + continue; + } + weights_map[locality_hosts[i][0]->locality()] = weights[i]; + } + + return result; +} + // --- WorkerLocalLbFactory --- WorkerLocalLbFactory::WorkerLocalLbFactory( diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h index ad2394689c88f..3497d86a83765 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -389,6 +389,63 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { Envoy::Common::CallbackHandlePtr priority_sync_cb_; }; +/** + * Main thread load balancer. Reads host-level utilization data and computes locality routing + * weights on a periodic timer. + */ +class LoadAwareLocalityLoadBalancer : public Upstream::ThreadAwareLoadBalancer, + protected Logger::Loggable { +public: + LoadAwareLocalityLoadBalancer(OptRef lb_config, + const Upstream::ClusterInfo& cluster_info, + const Upstream::PrioritySet& priority_set, Runtime::Loader& runtime, + Envoy::Random::RandomGenerator& random, TimeSource& time_source); + ~LoadAwareLocalityLoadBalancer() override; + + // Upstream::ThreadAwareLoadBalancer + Upstream::LoadBalancerFactorySharedPtr factory() override { return factory_; } + absl::Status initialize() override; + +private: + // Compute per-locality routing weights from host-level utilization and publish to factory. + void computeLocalityRoutingWeights(); + + // Attach LocalityLbHostData to hosts that don't already have it. + void addLbPolicyDataToHosts(const Upstream::HostVector& hosts); + + // Per-source outcome flags from computeSourceWeights, OR'd into the per-tick stats by the caller. + // stale_localities is only consumed from the all-hosts pass (the canonical staleness source). + struct SourceComputeResult { + bool all_overloaded{false}; + bool local_preferred{false}; + bool probe_active{false}; + uint32_t stale_localities{0}; + }; + + // Compute a source's per-locality weights into weights_map/all_local, advancing the caller-owned + // EWMA state. Returns this pass's per-tick stat signals. + SourceComputeResult + computeSourceWeights(const Upstream::HostsPerLocality& all_hosts_per_locality, + const std::vector& eligible_hosts_per_locality, + int64_t now_ms, LocalityWeightsMap& weights_map, bool& all_local, + LocalityEwmaMap& ewma_state); + + const Upstream::PrioritySet& priority_set_; + LoadAwareLocalityStats lb_stats_; + TimeSource& time_source_; + double utilization_variance_threshold_; + double ewma_alpha_; + double remote_probe_fraction_; + std::chrono::milliseconds weight_expiration_period_; + std::shared_ptr> metric_names_; + // Per-source, per-priority EWMA state keyed by locality identity (main thread only). + std::array, 3> ewma_state_; + Event::TimerPtr weight_update_timer_; + std::chrono::milliseconds weight_update_period_; + std::shared_ptr factory_; + Envoy::Common::CallbackHandlePtr priority_update_cb_; +}; + } // namespace LoadAwareLocality } // namespace LoadBalancingPolicies } // namespace Extensions From ff96479cfadf4d55a73aab60bd54da918ccba8fd Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Fri, 3 Jul 2026 23:29:23 -0600 Subject: [PATCH 06/16] load_aware_locality: architecture doc and changelog Bring the load_aware_locality architecture doc in line with the implementation, un-orphan it into the load balancing toctree, and add the new-feature changelog entry. The doc covers the five-stage locality weight computation, the child endpoint-picking policy, the emitted stats, and migration notes from zone-aware routing. Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- ...lancing__load-aware-locality-lb-policy.rst | 5 + .../load_balancing/load_aware_locality.rst | 202 ++++++++++-------- .../load_balancing/load_balancing.rst | 1 + 3 files changed, 113 insertions(+), 95 deletions(-) create mode 100644 changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst diff --git a/changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst b/changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst new file mode 100644 index 0000000000000..74680f49ad817 --- /dev/null +++ b/changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst @@ -0,0 +1,5 @@ +load_balancing: implemented the +:ref:`envoy.load_balancing_policies.load_aware_locality +` +locality-picking load balancer. It weights localities by ORCA-derived utilization headroom, +consumes in-band ORCA reporting, and applies at all priority levels. diff --git a/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst b/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst index 055658a40f421..92e6c9dc6853e 100644 --- a/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst +++ b/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst @@ -1,14 +1,8 @@ -:orphan: - .. _arch_overview_load_balancing_load_aware_locality: Load-aware locality load balancing ----------------------------------- -.. attention:: - - This extension is **work-in-progress** and is not yet implemented. - The load-aware locality LB policy (:ref:`envoy.load_balancing_policies.load_aware_locality `) @@ -42,7 +36,14 @@ imbalance. - **Pick zone-aware routing when** only local-zone preference is needed and traffic is already balanced by other means. It has no ORCA dependency and is simpler to operate, but it only applies at priority - 0 and does not react to backend load. + 0 and does not react to backend load. Zone-aware routing is not + extended here because it is priority-0 only, recomputes synchronously + per-worker inside membership callbacks (no timer, EWMA, or + cross-thread snapshot), and is built around a local-cluster percentage + model that would require rewriting the base class shared by + round-robin, least-request, and random. This policy requires + all-priority support, async main-thread ORCA recompute, and no + local-cluster dependency. - **Pick** :ref:`WrrLocality ` @@ -84,45 +85,34 @@ The policy is implemented as a ``ThreadAwareLoadBalancer``: - A main-thread timer recomputes per-locality weights from ORCA data and publishes an immutable snapshot to worker threads via a thread-local slot. - The snapshot carries a generation counter; workers rebuild per-locality - child LB instances when membership changes bump the generation. -- Worker threads read the latest snapshot lock-free on the request path, - pick a locality, and delegate endpoint selection to the child LB for that - locality. + The snapshot carries only advisory per-priority locality weights; it does + not carry priority loads, panic state, or health information. +- Workers rebuild per-locality child LB instances via a membership-update + callback registered on the cluster priority set, independent of the + weight snapshot and ``weight_update_period``. +- Priority selection, health/degraded mode, and panic detection use standard + ``LoadBalancerBase`` machinery over live worker-thread host-set state, + recomputed on membership and health callbacks. Failover between priorities + is therefore callback-fresh, not timer-driven. +- Worker threads read the latest locality-weight snapshot lock-free on the + request path, pick a locality, and delegate endpoint selection to the + child LB for that locality. - ORCA reports flow through per-host ``HostLbPolicyData`` slots shared with other ORCA consumers; see :ref:`ORCA data flow ` below for coexistence details. -- When ``enable_oob_load_report`` is set, a cluster-level OOB manager runs - on the main-thread dispatcher and owns one ORCA gRPC streaming session - per host, reacting to membership updates to add and remove sessions. It - decodes reports into the same shared report handler that backs the - in-band path, so workers see OOB and in-band samples through the same - ``HostLbPolicyData`` slots. .. _load_aware_locality_orca_data_flow: ORCA data flow """""""""""""" -Upstream endpoints must report ORCA utilization. The policy supports both -in-band and out-of-band reporting modes: - -- **In-band (default).** ORCA reports returned on the response headers or - trailers of upstream responses. Sample rate is tied to the request rate - to each host, so probing (``remote_probe_fraction``) is required to keep - remote-locality data fresh. -- **Out-of-band (OOB).** When ``enable_oob_load_report`` is set, the policy - opens a per-host ORCA gRPC stream and the endpoint pushes reports every - ``oob_reporting_period`` independent of request traffic. OOB reuses the - same central ORCA client as - :ref:`CSWRR `: - a cluster-level OOB manager owns one streaming session per host, reacts to - membership changes to add and remove sessions, and feeds every decoded - report into the shared ORCA report handler. Because OOB decouples sample - rate from request rate, ``remote_probe_fraction`` may safely be set to 0. - -Either way reports land in the same per-host ``HostLbPolicyData`` slots, so -weight computation is identical regardless of reporting mode. +Upstream endpoints must report ORCA utilization in-band: ORCA reports are +returned on the response headers or trailers of upstream responses. Sample +rate is tied to the request rate to each host, so probing +(``remote_probe_fraction``) is required to keep remote-locality data fresh. + +Reports land in per-host ``HostLbPolicyData`` slots, which feed weight +computation. Pairing this policy with :ref:`CSWRR ` @@ -132,16 +122,23 @@ per-endpoint capacity. Each consumer attaches independent ``HostLbPolicyData`` entries, so the two policies do not interfere. Utilization is derived from each host's ORCA report using the same -extraction as CSWRR (precedence may be flipped by the +extraction as CSWRR. The runtime flag ``envoy.reloadable_features.orca_weight_manager_use_named_metrics_first`` -runtime feature). By default: +defaults to ``true``, so the default precedence is: -1. ``application_utilization`` -- value in [0, 1], used when reported and - greater than 0. -2. Named metrics via ``metric_names_for_computing_utilization`` -- max of - present values, used when ``application_utilization`` is not reported. +1. Named metrics via ``metric_names_for_computing_utilization`` -- max of + present values, used when the result is greater than 0. +2. ``application_utilization`` -- used when greater than 0. 3. ``cpu_utilization`` -- final fallback. +Disabling the runtime flag reverses steps 1 and 2 so +``application_utilization`` is preferred over named metrics. + +A report whose derived utilization is not greater than 0 does not count as a +fresh sample: the host is treated as non-reporting and ages out per +``weight_expiration_period``. Note that a configured named metric that +legitimately reads exactly 0 is therefore treated as no signal. + .. _load_aware_locality_weight_computation: Weight computation @@ -160,9 +157,10 @@ locality routing weights in five stages: locality is applied raw (no blending) so the policy begins differentiating within a single tick after cold start; subsequent samples blend with the prior smoothed value. At startup with no ORCA data every locality - defaults to utilization 0 (full headroom), so weights reduce to host - counts -- equivalent to round-robin locality selection until the first - reports arrive. + defaults to utilization 0 (full headroom), so the stage-4 local-preference + check applies and routing snaps to the local locality (minus the probe + floor) until the first reports arrive. See + :ref:`Caveats ` for cold-start implications. 3. **Headroom weight.** Compute each locality's base weight as ``host_count * (1 - smoothed_util)`` -- capacity-weighted headroom. Stale localities fall back to ``host_count`` so traffic keeps flowing without @@ -307,9 +305,9 @@ Configuration parameters * - ``endpoint_picking_policy`` - (required) - Child LB policy for selecting an endpoint within the chosen locality. - Any LB policy may be configured here, including ring hash and Maglev, - though policies that build cluster-wide structures will operate over - only the chosen locality's host slice. See + Any LB policy may be configured here, but hash-based policies (ring + hash, Maglev) build their structures cluster-wide and are not + constrained by the locality pick. See :ref:`Caveats `. * - ``weight_update_period`` - 1 s @@ -317,8 +315,8 @@ Configuration parameters least 100 ms. * - ``metric_names_for_computing_utilization`` - (unset) - - Named ORCA metrics used to compute utilization when - ``application_utilization`` is not reported. The max of matching + - Named ORCA metrics used to compute utilization (by default preferred + over ``application_utilization``). The max of matching values is taken. Map entries use ``.`` (e.g. ``named_metrics.foo``). See :ref:`Weight computation ` for @@ -342,9 +340,8 @@ Configuration parameters - 0.03 - Minimum fraction of traffic sent to non-local localities to keep ORCA data fresh in all-local mode. The deficit is redistributed - proportionally to host count. Set to 0 to disable (safe only with - out-of-band ORCA reporting or when cross-zone traffic must be strictly - avoided). Range: [0, 1). See + proportionally to host count. Set to 0 to disable (safe only when + cross-zone traffic must be strictly avoided). Range: [0, 1). See :ref:`Caveats ` for scaling notes. * - ``weight_expiration_period`` - 3 minutes @@ -355,19 +352,6 @@ Configuration parameters falls back to host-count-proportional weighting. Tune higher to tolerate longer reporting gaps; tune lower to prune draining backends faster. Set to 0 s to disable expiration. - * - ``enable_oob_load_report`` - - (unset / false) - - Enables out-of-band (OOB) ORCA utilization reporting. When set, the - policy opens a per-host ORCA gRPC stream and the endpoint pushes - reports on its own schedule rather than piggybacking on responses. - When unset, only in-band reports on response headers/trailers are - consumed. See :ref:`ORCA data flow - `. - * - ``oob_reporting_period`` - - 10 s - - Requested load-reporting interval, used only when - ``enable_oob_load_report`` is true. The upstream may report less - frequently than requested. Priority support ^^^^^^^^^^^^^^^^ @@ -379,6 +363,11 @@ priority load calculation; locality selection then applies within the chosen priority. Unlike zone-aware routing (priority 0 only), this policy applies at all priority levels. +The policy honors ``common_lb_config.healthy_panic_threshold`` and +``common_lb_config.zone_aware_lb_config.fail_traffic_on_panic``. When +in panic mode the policy routes to all hosts; when ``fail_traffic_on_panic`` +is set it returns no host instead. + Three independent weight sets are maintained per priority: - **Healthy** -- common case, healthy hosts only. @@ -395,17 +384,34 @@ weight, computed from the same per-host ORCA data in a single tick pass. Caveats and known limitations ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -- **Probing is required with in-band reporting.** In the default in-band - mode, a locality only produces fresh ORCA samples when it receives - traffic, so ``remote_probe_fraction`` must stay above 0 to keep remote - localities reporting. Set it to 0 only when ``enable_oob_load_report`` - is on (OOB streams report independently of traffic) or when cross-zone +- **Probing is required.** A locality only produces fresh ORCA samples + when it receives traffic, so ``remote_probe_fraction`` must stay above 0 + to keep remote localities reporting. Set it to 0 only when cross-zone traffic must be strictly avoided. +- **Cold start snaps local.** Until the first ORCA reports arrive, every + locality reads as utilization 0, so the stage-4 local-preference check + routes ~100% of traffic to the local locality (minus + ``remote_probe_fraction``) rather than a host-count split. Convergence is + fast -- local hosts receive traffic immediately and report within about one + tick -- but a local locality much smaller than its share of cold-start + traffic can be briefly overloaded after a mass restart. Backends that never + report ORCA leave the policy in this all-local mode permanently: the policy + assumes ORCA-reporting backends. - **Hash-based child policies.** Ring hash and Maglev build their hash - structures from the host set they are given. With this policy, that set - is the chosen locality's hosts, not the full cluster. The same request - hash will not necessarily map to the same endpoint cluster-wide, so - consistency guarantees apply only within a locality. + structures once over the full cluster host set and select from those + structures directly, ignoring the per-locality host slice this policy hands + them. The locality pick therefore does not constrain them: locality + weighting (local preference, variance threshold, ``remote_probe_fraction``) + has no effect on traffic, and the zone-routing stats record locality + decisions that traffic does not follow. The resulting behavior is that of + the plain hash policy. Support for per-locality hash structures may be + added in the future. Endpoint-selection retry is owned by the child LB. +- **Asynchronous child policies.** Child policies that select hosts + asynchronously (for example dynamic modules or dynamic forward proxy) are + not supported: this policy wraps the request context in a short-lived + per-pick object, so an asynchronous selection is cancelled and the pick + fails synchronously rather than calling back after the context is gone. + Use a synchronous child policy. - **Probe-fraction scaling.** ``remote_probe_fraction`` is a global value divided across all remote localities, then again across each locality's hosts. The per-host probe rate is therefore approximately @@ -434,15 +440,18 @@ Caveats and known limitations row, samples expire before the next probe arrives, so remote localities will alternate between fresh data and host-count fallback every few ticks. To avoid this, either reduce locality count, raise - ``remote_probe_fraction``, raise ``weight_expiration_period`` to - tolerate longer gaps, or enable OOB ORCA reporting via - ``enable_oob_load_report`` (which decouples sample rate from request - rate entirely). + ``remote_probe_fraction``, or raise ``weight_expiration_period`` to + tolerate longer gaps. - **Variance-threshold oscillation.** Workloads sitting near the ``utilization_variance_threshold`` boundary can theoretically oscillate between snap-to-local and spillover modes across consecutive ticks. EWMA smoothing dampens this in practice; tune ``smoothing_time_constant`` higher if oscillation is observed. +- **Child LB health view.** The child endpoint-picking LB sees only the + hosts in its chosen locality, all marked healthy (a flattened health + view). This is benign for ORCA-weighting children like CSWRR but + means a child that relies on Envoy health flags will not see degraded + or unhealthy markers within the locality slice. - **Subsetting.** Load balancer :ref:`subsetting ` partitions hosts orthogonally to locality boundaries. The policy will operate @@ -461,8 +470,6 @@ The policy emits stats under ``cluster..load_aware_locality.*``: * - Counter - Increments when - * - ``recompute_total`` - - Per main-thread tick that recomputes weights. * - ``all_overloaded_total`` - Per tick where every locality's headroom was 0 (fallback to host-count weighting). @@ -472,16 +479,21 @@ The policy emits stats under ``cluster..load_aware_locality.*``: * - ``probe_active_total`` - Per tick where ``remote_probe_fraction`` redistribution kicked in. * - ``stale_locality_total`` - - Incremented once per stale locality per tick (a locality whose - hosts were all stale and fell back to host-count baseline). A - 5-locality cluster with 2 stale localities adds 2 each tick. - -Migrating from zone-aware routing? The closest counter mappings: - -+--------------------------------------+----------------------------------------------+ -| Zone-aware counter | Load-aware-locality equivalent | -+======================================+==============================================+ -| ``lb_zone_routing_all_directly`` | ``load_aware_locality.local_preferred_total``| -+--------------------------------------+----------------------------------------------+ -| ``lb_recalculate_zone_structures`` | ``load_aware_locality.recompute_total`` | -+--------------------------------------+----------------------------------------------+ + - Incremented once per stale locality per priority per tick: a locality + that previously had fresh reports but whose hosts' samples have all + expired (fell back to host-count baseline). A locality stale in + multiple priorities is counted once for each. Localities that have + never reported -- e.g. at cold start -- are not counted. A + single-priority 5-locality cluster with 2 stale localities adds 2 + each tick. + +Migrating from zone-aware routing? The per-request zone routing counters +(``lb_zone_routing_all_directly``, ``lb_zone_routing_sampled``, +``lb_zone_routing_cross_zone``) are still emitted with equivalent semantics, +recorded for the locality this policy selects. They are not incremented +while the priority is in panic mode. +``lb_recalculate_zone_structures`` is still emitted at cluster scope when +per-locality routing structures are rebuilt on membership changes (as in +zone-aware routing), not per weight-update tick. The per-tick counter closest +to ``lb_zone_routing_all_directly`` is +``load_aware_locality.local_preferred_total``. diff --git a/docs/root/intro/arch_overview/upstream/load_balancing/load_balancing.rst b/docs/root/intro/arch_overview/upstream/load_balancing/load_balancing.rst index be284b6e99eaf..e531e559697fe 100644 --- a/docs/root/intro/arch_overview/upstream/load_balancing/load_balancing.rst +++ b/docs/root/intro/arch_overview/upstream/load_balancing/load_balancing.rst @@ -14,6 +14,7 @@ Load Balancing excluded original_dst zone_aware + load_aware_locality subsets slow_start override_host From e09d8a2507350f54d40d44f787d3cdf4af5eaa8e Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 13:45:04 -0600 Subject: [PATCH 07/16] load_aware_locality: tests Add config, unit, and integration tests. config_test covers policy-knob validation and child-policy resolution; the unit test covers the main-thread weight computation (headroom, EWMA stale-carry, local preference, remote probe, all-overloaded host-count fallback) and the worker-local LB (live priority/panic selection, capacity-weighted locality selection, child delegation, and membership churn); the integration test exercises end-to-end ORCA-driven locality routing. Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality/BUILD | 45 + .../load_aware_locality/config_test.cc | 338 ++- .../load_aware_locality/integration_test.cc | 278 +++ .../load_aware_locality_lb_test.cc | 1927 +++++++++++++++++ 4 files changed, 2567 insertions(+), 21 deletions(-) create mode 100644 test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc create mode 100644 test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc diff --git a/test/extensions/load_balancing_policies/load_aware_locality/BUILD b/test/extensions/load_balancing_policies/load_aware_locality/BUILD index f60293a7e7baa..10604f86bb1d8 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/BUILD +++ b/test/extensions/load_balancing_policies/load_aware_locality/BUILD @@ -18,9 +18,54 @@ envoy_extension_cc_test( rbe_pool = "6gig", deps = [ "//source/extensions/load_balancing_policies/load_aware_locality:config", + "//source/extensions/load_balancing_policies/round_robin:config", "//test/mocks/server:factory_context_mocks", "//test/mocks/upstream:cluster_info_mocks", "//test/mocks/upstream:priority_set_mocks", + "//test/mocks/upstream:typed_load_balancer_factory_mocks", + "//test/test_common:registry_lib", "@envoy_api//envoy/config/core/v3:pkg_cc_proto", + "@envoy_api//envoy/extensions/load_balancing_policies/round_robin/v3:pkg_cc_proto", + ], +) + +envoy_extension_cc_test( + name = "load_aware_locality_lb_test", + srcs = ["load_aware_locality_lb_test.cc"], + extension_names = ["envoy.load_balancing_policies.load_aware_locality"], + rbe_pool = "6gig", + deps = [ + "//source/extensions/load_balancing_policies/load_aware_locality:config", + "//source/extensions/load_balancing_policies/load_aware_locality:load_aware_locality_lb_lib", + "//source/extensions/load_balancing_policies/round_robin:config", + "//test/common/upstream:utility_lib", + "//test/mocks/server:factory_context_mocks", + "//test/mocks/stream_info:stream_info_mocks", + "//test/mocks/upstream:cluster_info_mocks", + "//test/mocks/upstream:load_balancer_context_mock", + "//test/mocks/upstream:priority_set_mocks", + "//test/mocks/upstream:typed_load_balancer_factory_mocks", + "//test/test_common:simulated_time_system_lib", + "//test/test_common:utility_lib", + "@abseil-cpp//absl/container:flat_hash_map", + "@abseil-cpp//absl/strings", + "@envoy_api//envoy/extensions/load_balancing_policies/round_robin/v3:pkg_cc_proto", + "@xds//xds/data/orca/v3:pkg_cc_proto", + ], +) + +envoy_extension_cc_test( + name = "integration_test", + size = "large", + srcs = ["integration_test.cc"], + extension_names = ["envoy.load_balancing_policies.load_aware_locality"], + rbe_pool = "6gig", + deps = [ + "//source/common/protobuf", + "//source/extensions/load_balancing_policies/load_aware_locality:config", + "//source/extensions/load_balancing_policies/round_robin:config", + "//test/integration:http_integration_lib", + "//test/test_common:utility_lib", + "@xds//xds/data/orca/v3:pkg_cc_proto", ], ) diff --git a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc index ffad5d6d5e148..42c4a3955ad79 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc @@ -1,10 +1,20 @@ +#include + #include "envoy/config/core/v3/extension.pb.h" +#include "envoy/extensions/load_balancing_policies/round_robin/v3/round_robin.pb.h" +#include "source/common/protobuf/protobuf.h" #include "source/extensions/load_balancing_policies/load_aware_locality/config.h" +#include "source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h" #include "test/mocks/server/factory_context.h" #include "test/mocks/upstream/cluster_info.h" #include "test/mocks/upstream/priority_set.h" +#include "test/mocks/upstream/typed_load_balancer_factory.h" +#include "test/test_common/registry.h" + +#include "gmock/gmock.h" +#include "gtest/gtest.h" namespace Envoy { namespace Extensions { @@ -12,29 +22,315 @@ namespace LoadBalancingPolicies { namespace LoadAwareLocality { namespace { -TEST(LoadAwareLocalityConfigTest, CreateFactory) { - NiceMock context; +envoy::config::core::v3::TypedExtensionConfig +typedExtensionConfig(const std::string& name, const Protobuf::Message& message) { + envoy::config::core::v3::TypedExtensionConfig config; + config.set_name(name); + std::ignore = config.mutable_typed_config()->PackFrom(message); + return config; +} + +envoy::config::core::v3::TypedExtensionConfig roundRobinEndpointPolicy() { + envoy::extensions::load_balancing_policies::round_robin::v3::RoundRobin round_robin; + return typedExtensionConfig("envoy.load_balancing_policies.round_robin", round_robin); +} + +LoadAwareLocalityLbProto +loadAwareConfig(std::initializer_list children) { + LoadAwareLocalityLbProto config; + for (const auto& child : children) { + *config.mutable_endpoint_picking_policy()->add_policies()->mutable_typed_extension_config() = + child; + } + return config; +} + +const LoadAwareLocalityLbConfig& typedConfig(const Upstream::LoadBalancerConfigPtr& config) { + const auto* typed = dynamic_cast(config.get()); + EXPECT_NE(nullptr, typed); + return *typed; +} + +void expectFactoryCreateSucceeds(const Upstream::LoadBalancerConfigPtr& config, + Server::Configuration::MockServerFactoryContext& context) { NiceMock cluster_info; - NiceMock main_thread_priority_set; + NiceMock priority_set; + NiceMock dispatcher; + ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); - envoy::config::core::v3::TypedExtensionConfig config; - config.set_name("envoy.load_balancing_policies.load_aware_locality"); - LoadAwareLocalityProto config_msg; - std::ignore = config.mutable_typed_config()->PackFrom(config_msg); - - auto& factory = Config::Utility::getAndCheckFactory(config); - EXPECT_EQ("envoy.load_balancing_policies.load_aware_locality", factory.name()); - - // loadConfig returns UnimplementedError until the policy is functional. - auto lb_config = factory.loadConfig(context, *factory.createEmptyConfigProto()); - EXPECT_FALSE(lb_config.ok()); - EXPECT_EQ(lb_config.status().code(), absl::StatusCode::kUnimplemented); - - // create is stubbed to return nullptr for now. - auto thread_aware_lb = - factory.create({}, cluster_info, main_thread_priority_set, context.runtime_loader_, - context.api_.random_, context.time_system_); - EXPECT_EQ(nullptr, thread_aware_lb); + Factory factory; + auto lb = factory.create(*config, cluster_info, priority_set, context.runtime_loader_, + context.api_.random_, context.time_system_); + ASSERT_NE(nullptr, lb); + ASSERT_TRUE(lb->initialize().ok()); + + auto worker_factory = lb->factory(); + ASSERT_NE(nullptr, worker_factory); + EXPECT_NE(nullptr, worker_factory->create({priority_set, nullptr})); +} + +class FailingLoadConfigFactory : public Upstream::MockTypedLoadBalancerFactory { +public: + std::string name() const override { + return "envoy.load_balancing_policies.load_aware_locality_test_error"; + } + + absl::StatusOr + loadConfig(Server::Configuration::ServerFactoryContext&, const Protobuf::Message&) override { + return absl::InvalidArgumentError("failing child loadConfig"); + } +}; + +class RejectingEndpointValidationConfig : public Upstream::LoadBalancerConfig { +public: + absl::Status validateEndpoints(const Upstream::PriorityState&) const override { + return absl::InvalidArgumentError("child endpoint validation failed"); + } +}; + +TEST(LoadAwareLocalityConfigTest, Defaults) { + NiceMock context; + NiceMock dispatcher; + ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); + + Factory factory; + auto result = factory.loadConfig(context, loadAwareConfig({roundRobinEndpointPolicy()})); + ASSERT_TRUE(result.ok()); + + const auto& config = typedConfig(result.value()); + EXPECT_EQ(std::chrono::milliseconds(1000), config.weightUpdatePeriod()); + EXPECT_DOUBLE_EQ(0.1, config.utilizationVarianceThreshold()); + // Default alpha derived from weight_update_period=1s and smoothing_time_constant=5s. + EXPECT_DOUBLE_EQ(1.0 - std::exp(-1.0 / 5.0), config.ewmaAlpha()); + EXPECT_DOUBLE_EQ(0.03, config.remoteProbeFraction()); + EXPECT_EQ(std::chrono::milliseconds(180000), config.weightExpirationPeriod()); + + expectFactoryCreateSucceeds(result.value(), context); +} + +TEST(LoadAwareLocalityConfigTest, Overrides) { + NiceMock context; + NiceMock dispatcher; + ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); + + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_weight_update_period()->set_seconds(2); + config_proto.mutable_utilization_variance_threshold()->set_value(0.05); + config_proto.mutable_smoothing_time_constant()->set_seconds(4); + config_proto.mutable_remote_probe_fraction()->set_value(0.07); + config_proto.mutable_weight_expiration_period()->set_seconds(15); + + Factory factory; + auto result = factory.loadConfig(context, config_proto); + ASSERT_TRUE(result.ok()); + + const auto& config = typedConfig(result.value()); + EXPECT_EQ(std::chrono::milliseconds(2000), config.weightUpdatePeriod()); + EXPECT_DOUBLE_EQ(0.05, config.utilizationVarianceThreshold()); + // Alpha derived from weight_update_period=2s and smoothing_time_constant=4s. + EXPECT_DOUBLE_EQ(1.0 - std::exp(-2.0 / 4.0), config.ewmaAlpha()); + EXPECT_DOUBLE_EQ(0.07, config.remoteProbeFraction()); + EXPECT_EQ(std::chrono::milliseconds(15000), config.weightExpirationPeriod()); + + expectFactoryCreateSucceeds(result.value(), context); +} + +TEST(LoadAwareLocalityConfigTest, MissingChildPolicy) { + NiceMock context; + + Factory factory; + auto result = factory.loadConfig(context, loadAwareConfig({})); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_EQ(result.status(), absl::InvalidArgumentError("No supported endpoint picking policy.")); +} + +TEST(LoadAwareLocalityConfigTest, FirstSupportedChildWins) { + NiceMock context; + NiceMock dispatcher; + ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); + + envoy::config::core::v3::TypedExtensionConfig unknown_policy; + unknown_policy.set_name("envoy.load_balancing_policies.does_not_exist"); + std::ignore = unknown_policy.mutable_typed_config()->PackFrom(Protobuf::Struct{}); + + Factory factory; + auto result = + factory.loadConfig(context, loadAwareConfig({unknown_policy, roundRobinEndpointPolicy()})); + ASSERT_TRUE(result.ok()); + + const auto& config = typedConfig(result.value()); + EXPECT_EQ("envoy.load_balancing_policies.round_robin", config.endpointPickingPolicyName()); + expectFactoryCreateSucceeds(result.value(), context); +} + +TEST(LoadAwareLocalityConfigTest, InvalidWeightUpdatePeriod) { + NiceMock context; + + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_weight_update_period()->set_seconds(0); + + Factory factory; + auto result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("weight_update_period")); +} + +TEST(LoadAwareLocalityConfigTest, ZeroSmoothingTimeConstantRejected) { + NiceMock context; + + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_smoothing_time_constant()->set_seconds(0); + + Factory factory; + auto result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("smoothing_time_constant")); +} + +TEST(LoadAwareLocalityConfigTest, OutOfRangeVarianceThresholdRejected) { + NiceMock context; + Factory factory; + + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_utilization_variance_threshold()->set_value(-0.1); + auto result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("utilization_variance_threshold")); + + config_proto.mutable_utilization_variance_threshold()->set_value(1.5); + result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); +} + +TEST(LoadAwareLocalityConfigTest, OutOfRangeRemoteProbeFractionRejected) { + NiceMock context; + Factory factory; + + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_remote_probe_fraction()->set_value(1.0); + auto result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("remote_probe_fraction")); + + config_proto.mutable_remote_probe_fraction()->set_value(-0.01); + result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); +} + +TEST(LoadAwareLocalityConfigTest, PolicyValidationPrecedesChildResolution) { + NiceMock context; + + // With an unresolvable child AND an invalid policy knob, the policy-knob error must win so the + // user isn't sent chasing the wrong problem. + auto config_proto = loadAwareConfig({}); + config_proto.mutable_enable_oob_load_report()->set_value(true); + + Factory factory; + auto result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("enable_oob_load_report")); +} + +TEST(LoadAwareLocalityConfigTest, MalformedSupportedChildConfigIsRejected) { + NiceMock context; + + envoy::config::core::v3::TypedExtensionConfig malformed_round_robin; + malformed_round_robin.set_name("envoy.load_balancing_policies.round_robin"); + malformed_round_robin.mutable_typed_config()->set_type_url( + "type.googleapis.com/" + "envoy.extensions.load_balancing_policies.round_robin.v3.RoundRobin"); + malformed_round_robin.mutable_typed_config()->set_value("not-a-valid-proto"); + + Factory factory; + auto result = factory.loadConfig(context, loadAwareConfig({malformed_round_robin})); + EXPECT_FALSE(result.ok()); + EXPECT_THAT(result.status().code(), + testing::AnyOf(absl::StatusCode::kInvalidArgument, absl::StatusCode::kInternal)); + EXPECT_THAT(result.status().message(), + testing::AnyOf(testing::HasSubstr("round_robin.v3.RoundRobin"), + testing::HasSubstr("Unable to unpack"))); +} + +TEST(LoadAwareLocalityConfigTest, SupportedChildLoadConfigErrorIsPropagated) { + NiceMock context; + FailingLoadConfigFactory child_factory; + Registry::InjectFactory registered_factory(child_factory); + + Factory factory; + auto result = factory.loadConfig( + context, loadAwareConfig({typedExtensionConfig(child_factory.name(), Protobuf::Struct{})})); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("failing child loadConfig")); +} + +TEST(LoadAwareLocalityConfigTest, EndpointValidationDelegatesToChildConfig) { + NiceMock context; + NiceMock child_factory; + NiceMock dispatcher; + + LoadAwareLocalityLbConfig config( + child_factory, "mock_child", std::make_shared(), + std::chrono::milliseconds(1000), 0.1, 0.3, 0.03, std::chrono::milliseconds(180000), + std::vector{}, dispatcher, context.thread_local_); + + const Upstream::PriorityState priorities; + auto status = config.validateEndpoints(priorities); + EXPECT_FALSE(status.ok()); + EXPECT_THAT(status.message(), testing::HasSubstr("child endpoint validation failed")); +} + +TEST(LoadAwareLocalityConfigTest, MetricNamesReachConfig) { + NiceMock context; + NiceMock dispatcher; + ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); + + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.add_metric_names_for_computing_utilization("named_metrics.foo"); + config_proto.add_metric_names_for_computing_utilization("named_metrics.bar"); + + Factory factory; + auto result = factory.loadConfig(context, config_proto); + ASSERT_TRUE(result.ok()); + + const auto& config = typedConfig(result.value()); + ASSERT_EQ(2u, config.metricNamesForComputingUtilization().size()); + EXPECT_EQ("named_metrics.foo", config.metricNamesForComputingUtilization()[0]); + EXPECT_EQ("named_metrics.bar", config.metricNamesForComputingUtilization()[1]); + + expectFactoryCreateSucceeds(result.value(), context); +} + +TEST(LoadAwareLocalityConfigTest, OobLoadReportRejected) { + NiceMock context; + + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_enable_oob_load_report()->set_value(true); + + Factory factory; + auto result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("enable_oob_load_report")); +} + +TEST(LoadAwareLocalityConfigTest, NegativeWeightExpirationRejected) { + NiceMock context; + + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_weight_expiration_period()->set_seconds(-1); + + Factory factory; + EXPECT_THROW( + { auto result = factory.loadConfig(context, config_proto); }, ProtoValidationException); } } // namespace diff --git a/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc new file mode 100644 index 0000000000000..4237b781a5aae --- /dev/null +++ b/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc @@ -0,0 +1,278 @@ +#include +#include +#include +#include + +#include "source/common/common/base64.h" +#include "source/common/protobuf/protobuf.h" + +#include "test/integration/http_integration.h" +#include "test/test_common/utility.h" + +#include "gtest/gtest.h" +#include "xds/data/orca/v3/orca_load_report.pb.h" + +namespace Envoy { +namespace Extensions { +namespace LoadBalancingPolicies { +namespace LoadAwareLocality { +namespace { + +class LoadAwareLocalityIntegrationTest : public testing::TestWithParam, + public HttpIntegrationTest { +public: + LoadAwareLocalityIntegrationTest() : HttpIntegrationTest(Http::CodecType::HTTP1, GetParam()) { + setUpstreamCount(num_upstreams_); + use_bootstrap_node_metadata_ = true; + } + + void TearDown() override { HttpIntegrationTest::cleanupUpstreamAndDownstream(); } + + void initializeConfig(double variance_threshold = 0.1, double remote_probe_fraction = 0.1, + int weight_update_period_seconds = 10, + int smoothing_time_constant_seconds = 1, + std::vector remote_zones = {"zone-b"}) { + num_upstreams_ = static_cast(2 * (1 + remote_zones.size())); + setUpstreamCount(num_upstreams_); + + const auto ip_version = GetParam(); + config_helper_.addConfigModifier( + [ip_version, variance_threshold, remote_probe_fraction, weight_update_period_seconds, + smoothing_time_constant_seconds, + remote_zones](envoy::config::bootstrap::v3::Bootstrap& bootstrap) { + auto* node = bootstrap.mutable_node(); + node->set_id("node_name"); + node->set_cluster("cluster_name"); + auto* locality = node->mutable_locality(); + locality->set_region("test-region"); + locality->set_zone("zone-a"); + + auto* cluster = bootstrap.mutable_static_resources()->mutable_clusters()->Mutable(0); + ASSERT(cluster->name() == "cluster_0"); + cluster->mutable_load_assignment()->clear_endpoints(); + cluster->mutable_load_assignment()->set_cluster_name("cluster_0"); + + const std::string local_address = Network::Test::getLoopbackAddressString(ip_version); + std::vector all_zones = {"zone-a"}; + all_zones.insert(all_zones.end(), remote_zones.begin(), remote_zones.end()); + for (const auto& zone : all_zones) { + auto* locality_pb = cluster->mutable_load_assignment()->add_endpoints(); + locality_pb->mutable_locality()->set_region("test-region"); + locality_pb->mutable_locality()->set_zone(zone); + for (int i = 0; i < 2; ++i) { + auto* addr = locality_pb->add_lb_endpoints() + ->mutable_endpoint() + ->mutable_address() + ->mutable_socket_address(); + addr->set_address(local_address); + addr->set_port_value(0); + } + } + + const std::string policy_yaml = + fmt::format(R"EOF( + policies: + - typed_extension_config: + name: envoy.load_balancing_policies.load_aware_locality + typed_config: + "@type": type.googleapis.com/envoy.extensions.load_balancing_policies.load_aware_locality.v3.LoadAwareLocality + endpoint_picking_policy: + policies: + - typed_extension_config: + name: envoy.load_balancing_policies.round_robin + typed_config: + "@type": type.googleapis.com/envoy.extensions.load_balancing_policies.round_robin.v3.RoundRobin + weight_update_period: + seconds: {} + utilization_variance_threshold: + value: {} + remote_probe_fraction: + value: {} + smoothing_time_constant: + seconds: {} + )EOF", + weight_update_period_seconds, variance_threshold, remote_probe_fraction, + smoothing_time_constant_seconds); + TestUtility::loadFromYaml(policy_yaml, *cluster->mutable_load_balancing_policy()); + }); + + HttpIntegrationTest::initialize(); + } + + Http::TestResponseHeaderMapImpl + responseHeadersWithOrcaUtilization(double application_utilization) { + xds::data::orca::v3::OrcaLoadReport report; + report.set_application_utilization(application_utilization); + const std::string proto_string = TestUtility::getProtobufBinaryStringFromMessage(report); + Http::TestResponseHeaderMapImpl response_headers{{":status", "200"}}; + response_headers.addCopy("endpoint-load-metrics-bin", + Envoy::Base64::encode(proto_string.c_str(), proto_string.length())); + return response_headers; + } + + uint64_t sendRequestWithOrcaResponse(const std::vector& utilizations) { + codec_client_ = makeHttpConnection(lookupPort("http")); + Http::TestRequestHeaderMapImpl request_headers{ + {":method", "GET"}, {":path", "/"}, {":scheme", "http"}, {":authority", "example.com"}}; + auto response = codec_client_->makeRequestWithBody(request_headers, 0); + + std::vector upstream_indices(num_upstreams_); + std::iota(upstream_indices.begin(), upstream_indices.end(), 0); + auto upstream_index = waitForNextUpstreamRequest(upstream_indices); + RELEASE_ASSERT(upstream_index.has_value(), "Expected upstream request"); + + upstream_request_->encodeHeaders( + responseHeadersWithOrcaUtilization(utilizations[*upstream_index]), true); + RELEASE_ASSERT(response->waitForEndStream(), "Expected response"); + cleanupUpstreamAndDownstream(); + return *upstream_index; + } + + std::vector sendRequestsAndTrack(uint64_t count, + const std::vector& utilizations) { + std::vector usage(num_upstreams_, 0); + for (uint64_t i = 0; i < count; ++i) { + usage[sendRequestWithOrcaResponse(utilizations)]++; + } + return usage; + } + + void advanceWeightTick(uint64_t expected_counter) { + timeSystem().advanceTimeWait(std::chrono::seconds(11)); + test_server_->waitForCounter("cluster.cluster_0.lb_recalculate_zone_structures", + testing::Ge(expected_counter)); + } + + void seedWithTwoCycles(const std::vector& utilizations, uint64_t starting_count = 1) { + sendRequestsAndTrack(30, utilizations); + advanceWeightTick(starting_count + 1); + sendRequestsAndTrack(30, utilizations); + advanceWeightTick(starting_count + 2); + } + + uint64_t zoneTraffic(const std::vector& usage, size_t zone_index) const { + const size_t start = zone_index * 2; + return usage[start] + usage[start + 1]; + } + + uint64_t localTraffic(const std::vector& usage) const { return zoneTraffic(usage, 0); } + + uint64_t remoteTraffic(const std::vector& usage) const { + uint64_t total = 0; + for (size_t zone = 1; zone < num_upstreams_ / 2; ++zone) { + total += zoneTraffic(usage, zone); + } + return total; + } + +protected: + uint32_t num_upstreams_{4}; +}; + +INSTANTIATE_TEST_SUITE_P(IpVersions, LoadAwareLocalityIntegrationTest, + testing::ValuesIn(TestEnvironment::getIpVersionsForTest()), + TestUtility::ipTestParamsToString); + +TEST_P(LoadAwareLocalityIntegrationTest, StrictLocalWithRoundRobinWithinLocality) { + initializeConfig(/*variance_threshold=*/0.1, /*remote_probe_fraction=*/0.0, + /*weight_update_period_seconds=*/10); + + const std::vector utilizations = {0.05, 0.05, 0.05, 0.05}; + sendRequestsAndTrack(30, utilizations); + advanceWeightTick(2); + + const auto usage = sendRequestsAndTrack(100, utilizations); + + EXPECT_EQ(100u, localTraffic(usage)); + EXPECT_EQ(0u, remoteTraffic(usage)); + EXPECT_GE(usage[0], 30u); + EXPECT_LE(usage[0], 70u); + EXPECT_GE(usage[1], 30u); + EXPECT_LE(usage[1], 70u); +} + +TEST_P(LoadAwareLocalityIntegrationTest, ProbeFloorKeepsRemoteSampledInAllLocalMode) { + // Balanced utilization keeps the variance snap in all-local mode; the probe floor must still + // route ~10% of traffic to the remote locality end-to-end so its ORCA data stays fresh. + initializeConfig(/*variance_threshold=*/0.1, /*remote_probe_fraction=*/0.1, + /*weight_update_period_seconds=*/10); + + const std::vector balanced = {0.4, 0.4, 0.4, 0.4}; + seedWithTwoCycles(balanced); + + const auto usage = sendRequestsAndTrack(400, balanced); + const uint64_t local = localTraffic(usage); + const uint64_t remote = remoteTraffic(usage); + // Expected remote share is the probe floor: 10% of 400 = ~40. Generous sampling slack, but + // the floor must be visibly active (not zero, not dominating). + EXPECT_GE(remote, 15u); + EXPECT_LE(remote, 90u); + EXPECT_GT(local, remote); +} + +TEST_P(LoadAwareLocalityIntegrationTest, AdaptiveSpillAndRecovery) { + initializeConfig(/*variance_threshold=*/0.1, /*remote_probe_fraction=*/0.1, + /*weight_update_period_seconds=*/10); + + const std::vector overloaded_local = {0.9, 0.9, 0.2, 0.2}; + seedWithTwoCycles(overloaded_local); + + const auto phase1_usage = sendRequestsAndTrack(100, overloaded_local); + const uint64_t phase1_local = localTraffic(phase1_usage); + const uint64_t phase1_remote = remoteTraffic(phase1_usage); + EXPECT_GE(phase1_remote, 60u); + EXPECT_GT(phase1_remote, phase1_local); + + const std::vector rebalanced = {0.3, 0.3, 0.3, 0.3}; + sendRequestsAndTrack(60, rebalanced); + advanceWeightTick(4); + + const auto phase2_usage = sendRequestsAndTrack(100, rebalanced); + const uint64_t phase2_local = localTraffic(phase2_usage); + EXPECT_GT(phase2_local, phase1_local); + EXPECT_GE(phase2_local, 35u); +} + +TEST_P(LoadAwareLocalityIntegrationTest, EwmaDampensSpike) { + // smoothing_time_constant=14s with weight_update_period=10s yields alpha ~= 0.51. + initializeConfig(/*variance_threshold=*/0.1, /*remote_probe_fraction=*/0.1, + /*weight_update_period_seconds=*/10, /*smoothing_time_constant_seconds=*/14); + + const std::vector baseline = {0.2, 0.2, 0.2, 0.2}; + seedWithTwoCycles(baseline); + + const std::vector spiked_local = {0.9, 0.9, 0.2, 0.2}; + sendRequestsAndTrack(30, spiked_local); + advanceWeightTick(4); + + const auto usage = sendRequestsAndTrack(100, spiked_local); + const uint64_t remote = remoteTraffic(usage); + EXPECT_GE(remote, 40u); + EXPECT_LE(remote, 85u); +} + +TEST_P(LoadAwareLocalityIntegrationTest, ThreeLocalityDistribution) { + initializeConfig(/*variance_threshold=*/0.1, /*remote_probe_fraction=*/0.1, + /*weight_update_period_seconds=*/10, /*smoothing_time_constant_seconds=*/1, + /*remote_zones=*/{"zone-b", "zone-c"}); + + const std::vector utilizations = {0.8, 0.8, 0.3, 0.3, 0.5, 0.5}; + seedWithTwoCycles(utilizations); + + const auto usage = sendRequestsAndTrack(400, utilizations); + const uint64_t zone_a = zoneTraffic(usage, 0); + const uint64_t zone_b = zoneTraffic(usage, 1); + const uint64_t zone_c = zoneTraffic(usage, 2); + + EXPECT_GT(zone_a, 0u); + EXPECT_GT(zone_b, 0u); + EXPECT_GT(zone_c, 0u); + EXPECT_GT(zone_b, zone_c); + EXPECT_GT(zone_c, zone_a); +} + +} // namespace +} // namespace LoadAwareLocality +} // namespace LoadBalancingPolicies +} // namespace Extensions +} // namespace Envoy diff --git a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc new file mode 100644 index 0000000000000..e7823508be4ae --- /dev/null +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -0,0 +1,1927 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "envoy/extensions/load_balancing_policies/round_robin/v3/round_robin.pb.h" + +#include "source/extensions/load_balancing_policies/load_aware_locality/config.h" +#include "source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h" + +#include "test/common/upstream/utility.h" +#include "test/mocks/server/factory_context.h" +#include "test/mocks/stream_info/mocks.h" +#include "test/mocks/upstream/cluster_info.h" +#include "test/mocks/upstream/load_balancer.h" +#include "test/mocks/upstream/load_balancer_context.h" +#include "test/mocks/upstream/priority_set.h" +#include "test/mocks/upstream/typed_load_balancer_factory.h" +#include "test/test_common/simulated_time_system.h" +#include "test/test_common/utility.h" + +#include "absl/container/flat_hash_map.h" +#include "absl/strings/str_cat.h" +#include "gmock/gmock.h" +#include "gtest/gtest.h" +#include "xds/data/orca/v3/orca_load_report.pb.h" + +namespace Envoy { +namespace Extensions { +namespace LoadBalancingPolicies { +namespace LoadAwareLocality { +namespace { + +using testing::NiceMock; +using testing::Return; + +// Distinct Locality identity for locality-group index `i`. Region strings sort lexicographically by +// index ("r00", "r01", ...) so the default setup order matches a local-first/lexicographic +// ordering. +envoy::config::core::v3::Locality localityForIndex(size_t i) { + return Upstream::Locality(absl::StrCat("r", absl::Dec(i, absl::kZeroPad2)), "z", "sz"); +} + +std::shared_ptr> +makeWeightTrackingMockHost(uint32_t initial_weight = 1) { + auto host = std::make_shared>(); + auto weight = std::make_shared(initial_weight); + ON_CALL(*host, weight()).WillByDefault([weight]() -> uint32_t { return *weight; }); + ON_CALL(*host, weight(::testing::_)).WillByDefault([weight](uint32_t new_weight) { + *weight = new_weight; + }); + return host; +} + +Upstream::HostVector makeHosts(uint32_t count, uint32_t weight = 1) { + Upstream::HostVector hosts; + hosts.reserve(count); + for (uint32_t i = 0; i < count; ++i) { + hosts.push_back(makeWeightTrackingMockHost(weight)); + } + return hosts; +} + +class RecordingWorkerChildFactory : public Upstream::LoadBalancerFactory { +public: + explicit RecordingWorkerChildFactory(bool recreate_on_host_change) + : recreate_on_host_change_(recreate_on_host_change) {} + + Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams) override { + ++create_count_; + return std::make_unique>(); + } + + bool recreateOnHostChangeDeprecated() const override { return recreate_on_host_change_; } + + uint32_t createCount() const { return create_count_; } + +private: + const bool recreate_on_host_change_; + uint32_t create_count_{0}; +}; + +// A child LB that always selects asynchronously: chooseHost returns no host plus a cancelable +// handle and (in production) would retain the context to call back later, mimicking dynamic modules +// or dynamic forward proxy. Flips `cancelled` when its handle is cancelled. +class AsyncSelectingChildLb : public Upstream::LoadBalancer { +public: + explicit AsyncSelectingChildLb(bool& cancelled) : cancelled_(cancelled) {} + + Upstream::HostSelectionResponse chooseHost(Upstream::LoadBalancerContext*) override { + return {nullptr, std::make_unique(cancelled_)}; + } + Upstream::HostConstSharedPtr peekAnotherHost(Upstream::LoadBalancerContext*) override { + return nullptr; + } + OptRef lifetimeCallbacks() override { + return {}; + } + absl::optional + selectExistingConnection(Upstream::LoadBalancerContext*, const Upstream::Host&, + std::vector&) override { + return absl::nullopt; + } + +private: + struct Handle : public Upstream::AsyncHostSelectionHandle { + explicit Handle(bool& cancelled) : cancelled_(cancelled) {} + void cancel() override { cancelled_ = true; } + bool& cancelled_; + }; + bool& cancelled_; +}; + +class AsyncSelectingChildFactory : public Upstream::LoadBalancerFactory { +public: + explicit AsyncSelectingChildFactory(bool& cancelled) : cancelled_(cancelled) {} + Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams) override { + return std::make_unique(cancelled_); + } + bool recreateOnHostChangeDeprecated() const override { return false; } + +private: + bool& cancelled_; +}; + +class StaticThreadAwareLoadBalancer : public Upstream::ThreadAwareLoadBalancer { +public: + explicit StaticThreadAwareLoadBalancer(Upstream::LoadBalancerFactorySharedPtr factory) + : factory_(std::move(factory)) {} + + Upstream::LoadBalancerFactorySharedPtr factory() override { return factory_; } + absl::Status initialize() override { return absl::OkStatus(); } + +private: + Upstream::LoadBalancerFactorySharedPtr factory_; +}; + +uint64_t counterValue(NiceMock& info, const std::string& name) { + auto c = info.stats_store_.findCounterByString(name); + return c.has_value() ? c->get().value() : 0; +} + +class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public testing::Test { +protected: + void SetUp() override { + // Tests run from the simulated-time epoch (monotonic 0): "never reported" is tracked by an + // out-of-band sentinel, so a report stored at time 0 is a valid sample. + ON_CALL(context_, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher_)); + ON_CALL(random_, random()).WillByDefault([this]() -> uint64_t { + if (forced_random_index_ < forced_random_values_.size()) { + return forced_random_values_[forced_random_index_++]; + } + constexpr uint64_t step = std::numeric_limits::max() / 100; + return (random_call_count_++ % 100) * step; + }); + } + + void forceRandomValues(std::vector values) { + forced_random_values_ = std::move(values); + forced_random_index_ = 0; + } + + void clearForcedRandomValues() { + forced_random_values_.clear(); + forced_random_index_ = 0; + } + + void createLb(double variance_threshold = 0.1, double ewma_alpha = 1.0, + double remote_probe_fraction = 0.0, + std::chrono::milliseconds weight_expiration_period = std::chrono::milliseconds(0)) { + auto weight_update_period = std::chrono::milliseconds(1000); + + envoy::extensions::load_balancing_policies::round_robin::v3::RoundRobin round_robin; + envoy::config::core::v3::TypedExtensionConfig round_robin_config; + round_robin_config.set_name("envoy.load_balancing_policies.round_robin"); + std::ignore = round_robin_config.mutable_typed_config()->PackFrom(round_robin); + + auto& round_robin_factory = + Config::Utility::getAndCheckFactory(round_robin_config); + + auto round_robin_proto = round_robin_factory.createEmptyConfigProto(); + ASSERT_TRUE(Config::Utility::translateOpaqueConfig(round_robin_config.typed_config(), + context_.messageValidationVisitor(), + *round_robin_proto) + .ok()); + auto round_robin_lb_config = + round_robin_factory.loadConfig(context_, *round_robin_proto).value(); + + auto lb_config = std::make_unique( + round_robin_factory, round_robin_factory.name(), std::move(round_robin_lb_config), + weight_update_period, variance_threshold, ewma_alpha, remote_probe_fraction, + weight_expiration_period, std::vector{}, dispatcher_, context_.thread_local_); + + timer_ = new NiceMock(&dispatcher_); + + thread_aware_lb_ = std::make_unique( + *lb_config, cluster_info_, priority_set_, context_.runtime_loader_, random_, + context_.time_system_); + + ASSERT_TRUE(thread_aware_lb_->initialize().ok()); + factory_ = thread_aware_lb_->factory(); + } + + std::shared_ptr createRoundRobinWorkerFactory() { + envoy::extensions::load_balancing_policies::round_robin::v3::RoundRobin round_robin; + envoy::config::core::v3::TypedExtensionConfig round_robin_config; + round_robin_config.set_name("envoy.load_balancing_policies.round_robin"); + std::ignore = round_robin_config.mutable_typed_config()->PackFrom(round_robin); + + auto& round_robin_factory = + Config::Utility::getAndCheckFactory(round_robin_config); + + auto round_robin_proto = round_robin_factory.createEmptyConfigProto(); + EXPECT_TRUE(Config::Utility::translateOpaqueConfig(round_robin_config.typed_config(), + context_.messageValidationVisitor(), + *round_robin_proto) + .ok()); + auto round_robin_lb_config = + round_robin_factory.loadConfig(context_, *round_robin_proto).value(); + + auto factory = std::make_shared( + round_robin_factory, round_robin_factory.name(), + LoadBalancerConfigSharedPtr(std::move(round_robin_lb_config)), cluster_info_, priority_set_, + context_.runtime_loader_, random_, context_.time_system_, context_.thread_local_); + EXPECT_TRUE(factory->initializeChildLb().ok()); + return factory; + } + + // Builds a WorkerLocalLbFactory whose child create() yields a StaticThreadAwareLoadBalancer + // wrapping `child`. + std::shared_ptr + makeWorkerFactoryWithChild(Upstream::LoadBalancerFactorySharedPtr child, + const std::string& name) { + auto typed_child_factory = std::make_shared>(); + ON_CALL(*typed_child_factory, name()).WillByDefault(Return(name)); + ON_CALL(*typed_child_factory, + create(testing::_, testing::_, testing::_, testing::_, testing::_, testing::_)) + .WillByDefault( + [child](auto&&...) { return std::make_unique(child); }); + recording_typed_child_factory_ = typed_child_factory; + + auto factory = std::make_shared( + *typed_child_factory, name, + std::make_shared(), + cluster_info_, priority_set_, context_.runtime_loader_, random_, context_.time_system_, + context_.thread_local_); + EXPECT_TRUE(factory->initializeChildLb().ok()); + return factory; + } + + void setupPriorityLocalities( + uint32_t priority, std::vector localities, + bool has_local_locality = false, + absl::optional> healthy_localities = absl::nullopt, + absl::optional> degraded_localities = absl::nullopt) { + auto* host_set = priority_set_.getMockHostSet(priority); + // Stamp each host with its locality-group identity so the policy can key weights by Locality. + for (size_t i = 0; i < localities.size(); ++i) { + for (const auto& host : localities[i]) { + setHostLocality(host, i); + } + } + Upstream::HostVector all_hosts; + for (const auto& locality : localities) { + all_hosts.insert(all_hosts.end(), locality.begin(), locality.end()); + } + host_set->hosts_ = all_hosts; + host_set->hosts_per_locality_ = + Upstream::makeHostsPerLocality(std::move(localities), !has_local_locality); + + if (healthy_localities.has_value()) { + Upstream::HostVector healthy_hosts; + for (const auto& locality : *healthy_localities) { + healthy_hosts.insert(healthy_hosts.end(), locality.begin(), locality.end()); + } + host_set->healthy_hosts_ = healthy_hosts; + host_set->healthy_hosts_per_locality_ = + Upstream::makeHostsPerLocality(std::move(*healthy_localities), !has_local_locality); + } else { + host_set->healthy_hosts_ = all_hosts; + host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + } + + if (degraded_localities.has_value()) { + Upstream::HostVector degraded_hosts; + for (const auto& locality : *degraded_localities) { + degraded_hosts.insert(degraded_hosts.end(), locality.begin(), locality.end()); + } + host_set->degraded_hosts_ = degraded_hosts; + host_set->degraded_hosts_per_locality_ = + Upstream::makeHostsPerLocality(std::move(*degraded_localities), !has_local_locality); + } else { + host_set->degraded_hosts_.clear(); + host_set->degraded_hosts_per_locality_ = + Upstream::makeHostsPerLocality({}, !has_local_locality); + } + } + + void setupLocalities( + std::vector localities, bool has_local_locality = false, + absl::optional> healthy_localities = absl::nullopt, + absl::optional> degraded_localities = absl::nullopt) { + setupPriorityLocalities(0, std::move(localities), has_local_locality, + std::move(healthy_localities), std::move(degraded_localities)); + } + + // Stamps `host` with the Locality identity for locality-group index `i`. + void setHostLocality(const Upstream::HostSharedPtr& host, size_t i) { + setHostLocalityExplicit(host, localityForIndex(i)); + } + + // Stamps `host` with an explicit Locality identity. The Locality is owned by the fixture (stable + // address) so the mock's locality() can return a reference to it for the test's lifetime. + void setHostLocalityExplicit(const Upstream::HostSharedPtr& host, + const envoy::config::core::v3::Locality& locality) { + auto* mock = dynamic_cast(host.get()); + if (mock == nullptr) { + return; + } + localities_.push_back(locality); + ON_CALL(*mock, locality()).WillByDefault(testing::ReturnRef(localities_.back())); + } + + // Builds a LocalityWeightsMap keyed by the canonical per-index Locality identities, matching the + // setup order used by setupLocalities/reshapeLocalities (locality i ↔ localityForIndex(i)). + static LocalityWeightsMap makeWeightsMap(const std::vector& weights) { + LocalityWeightsMap map; + for (size_t i = 0; i < weights.size(); ++i) { + map[localityForIndex(i)] = weights[i]; + } + return map; + } + + // Triggers a routing-weight recompute by re-running the thread-aware initialize. + void recomputeWeights() { ASSERT_TRUE(thread_aware_lb_->initialize().ok()); } + + // Reshapes a priority's host set, then fires the membership update callback. Two modes: + // - full membership reshape: sets hosts_/healthy_hosts_ and both per-locality vectors from + // `localities` (healthy mirrors all hosts). + // - health-only flip (when `healthy_override` is set): leaves hosts_/hosts_per_locality_ intact + // and only sets healthy_hosts_/healthy_hosts_per_locality_ from `healthy_override`. + void reshapeLocalities( + uint32_t priority, std::vector localities, + Upstream::HostVector added = {}, Upstream::HostVector removed = {}, bool has_local = false, + absl::optional> healthy_override = absl::nullopt) { + auto* host_set = priority_set_.getMockHostSet(priority); + if (healthy_override.has_value()) { + Upstream::HostVector healthy_hosts; + for (const auto& locality : *healthy_override) { + healthy_hosts.insert(healthy_hosts.end(), locality.begin(), locality.end()); + } + host_set->healthy_hosts_ = healthy_hosts; + host_set->healthy_hosts_per_locality_ = + Upstream::makeHostsPerLocality(std::move(*healthy_override), !has_local); + } else { + for (size_t i = 0; i < localities.size(); ++i) { + for (const auto& host : localities[i]) { + setHostLocality(host, i); + } + } + Upstream::HostVector all_hosts; + for (const auto& locality : localities) { + all_hosts.insert(all_hosts.end(), locality.begin(), locality.end()); + } + host_set->hosts_ = all_hosts; + host_set->healthy_hosts_ = all_hosts; + host_set->hosts_per_locality_ = + Upstream::makeHostsPerLocality(std::move(localities), !has_local); + host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + } + priority_set_.runUpdateCallbacks(priority, added, removed); + } + + static xds::data::orca::v3::OrcaLoadReport makeOrcaReport(double app_utilization) { + xds::data::orca::v3::OrcaLoadReport report; + report.set_application_utilization(app_utilization); + return report; + } + + void setHostUtilization(Upstream::HostSharedPtr host, double utilization) { + auto host_data = host->typedLbPolicyData(); + ASSERT(host_data.has_value()); + EXPECT_TRUE(host_data->onOrcaLoadReport(makeOrcaReport(utilization), stream_info_).ok()); + } + + void setUtilizationForHosts(const Upstream::HostVector& hosts, double utilization) { + for (const auto& host : hosts) { + setHostUtilization(host, utilization); + } + } + + int64_t nowMs() { + return std::chrono::duration_cast( + context_.time_system_.monotonicTime().time_since_epoch()) + .count(); + } + + Upstream::LoadBalancerPtr createWorkerLb() { + EXPECT_NE(nullptr, factory_); + auto worker_lb = factory_->create({priority_set_, nullptr}); + EXPECT_NE(nullptr, worker_lb); + return worker_lb; + } + + const RoutingWeightsSnapshot* routingSnapshot() const { + auto* typed_factory = dynamic_cast(factory_.get()); + EXPECT_NE(nullptr, typed_factory); + return typed_factory != nullptr ? typed_factory->routingWeights() : nullptr; + } + + void publishSnapshot(std::shared_ptr snapshot) { + auto* typed_factory = dynamic_cast(factory_.get()); + ASSERT_NE(nullptr, typed_factory); + typed_factory->updateRoutingWeights(std::move(snapshot)); + } + + // Builds an advisory locality-weights snapshot. Priority/health/panic selection is now live + // worker state (LoadBalancerBase), so the snapshot only carries per-priority locality weights. + std::shared_ptr + makeSnapshot(std::vector priority_weights) { + auto snapshot = std::make_shared(); + snapshot->priority_weights = std::move(priority_weights); + return snapshot; + } + + // Publishes a single-priority snapshot with Healthy weights keyed by per-index Locality identity. + void publishHealthyWeights(const std::vector& weights) { + PriorityRoutingWeights pw; + auto& healthy = + pw.by_source[static_cast(PriorityRoutingWeights::SelectionSource::Healthy)]; + healthy.weights = makeWeightsMap(weights); + publishSnapshot(makeSnapshot({pw})); + } + + absl::flat_hash_map countPicks(Upstream::LoadBalancer& lb, + int num_picks) { + absl::flat_hash_map counts; + for (int i = 0; i < num_picks; ++i) { + auto result = lb.chooseHost(nullptr); + EXPECT_NE(nullptr, result.host); + if (result.host != nullptr) { + counts[result.host.get()]++; + } + } + return counts; + } + + bool hostSeen(Upstream::LoadBalancer& lb, const Upstream::HostConstSharedPtr& target, + int picks = 200) { + for (int i = 0; i < picks; ++i) { + auto result = lb.chooseHost(nullptr); + if (result.host == target) { + return true; + } + } + return false; + } + + void expectOnlyHost(Upstream::LoadBalancer& lb, const Upstream::HostConstSharedPtr& host, + int picks = 50) { + for (int i = 0; i < picks; ++i) { + auto result = lb.chooseHost(nullptr); + ASSERT_NE(nullptr, result.host); + EXPECT_EQ(host, result.host); + } + } + + // Asserts the live routing snapshot's `source` weights for `priority` equal `expected`, matching + // each entry by Locality identity (expected[i] ↔ localityForIndex(i)). A missing identity reads + // as 0.0 (a locality not in the snapshot map carries no weight). Asserts the snapshot is + // non-null. + void expectWeights(PriorityRoutingWeights::SelectionSource source, std::vector expected, + double tol = 0.01, uint32_t priority = 0) { + const auto* snapshot = routingSnapshot(); + ASSERT_NE(nullptr, snapshot); + const auto& weights = snapshot->priority_weights[priority].weightsFor(source); + for (size_t i = 0; i < expected.size(); ++i) { + SCOPED_TRACE(absl::StrCat("locality ", i)); + auto it = weights.find(localityForIndex(i)); + const double actual = it != weights.end() ? it->second : 0.0; + EXPECT_NEAR(actual, expected[i], tol); + } + } + + // Reads the snapshot weight for locality-group index `i` (by identity); 0.0 if absent. + double weightForIndex(PriorityRoutingWeights::SelectionSource source, size_t i, + uint32_t priority = 0) { + const auto* snapshot = routingSnapshot(); + EXPECT_NE(nullptr, snapshot); + if (snapshot == nullptr) { + return 0.0; + } + const auto& weights = snapshot->priority_weights[priority].weightsFor(source); + auto it = weights.find(localityForIndex(i)); + return it != weights.end() ? it->second : 0.0; + } + + // Asserts the live snapshot's all_local flag for `source`/`priority`. Asserts snapshot non-null. + void expectAllLocal(bool expected, + PriorityRoutingWeights::SelectionSource source = + PriorityRoutingWeights::SelectionSource::Healthy, + uint32_t priority = 0) { + const auto* snapshot = routingSnapshot(); + ASSERT_NE(nullptr, snapshot); + EXPECT_EQ(expected, snapshot->priority_weights[priority].allLocalFor(source)); + } + + // Asserts a single timer tick increments counter `name` by exactly `delta`. + void expectCounterIncrementsPerTick(const std::string& name, uint64_t delta) { + const uint64_t before = counterValue(cluster_info_, name); + timer_->invokeCallback(); + EXPECT_EQ(before + delta, counterValue(cluster_info_, name)); + } + + NiceMock context_; + NiceMock dispatcher_; + NiceMock cluster_info_; + NiceMock priority_set_; + NiceMock random_; + NiceMock stream_info_; + uint64_t random_call_count_{0}; + std::vector forced_random_values_; + size_t forced_random_index_{0}; + Event::MockTimer* timer_{}; + // Stable storage for Locality identities returned by mock hosts' locality() (addresses must + // outlive the hosts; std::deque keeps element addresses stable across pushes). + std::deque localities_; + + Upstream::ThreadAwareLoadBalancerPtr thread_aware_lb_; + Upstream::LoadBalancerFactorySharedPtr factory_; + std::shared_ptr> recording_typed_child_factory_; +}; + +// Table-driven single-tick weight-computation cases: set up localities, push one uniform ORCA +// utilization per locality, recompute once, and assert the resulting Healthy snapshot weights. +// Multi-tick (EWMA), reshape, and distribution-sampling scenarios stay as dedicated TEST_Fs. +struct WeightComputationCase { + std::string test_name; + std::vector locality_sizes; // host count per locality group + bool has_local_locality; + double variance_threshold; + double remote_probe_fraction; + std::vector locality_utils; // utilization applied to every host in the locality + bool expect_all_local; + std::vector expected_healthy_weights; + double tol; +}; + +class WeightComputationTest : public LoadAwareLocalityLbTest, + public testing::WithParamInterface {}; + +TEST_P(WeightComputationTest, ComputesHealthyWeights) { + const WeightComputationCase& c = GetParam(); + std::vector localities; + for (int size : c.locality_sizes) { + localities.push_back(makeHosts(size)); + } + setupLocalities(localities, c.has_local_locality); + createLb(c.variance_threshold, /*ewma_alpha=*/1.0, c.remote_probe_fraction); + + for (size_t i = 0; i < localities.size(); ++i) { + setUtilizationForHosts(localities[i], c.locality_utils[i]); + } + recomputeWeights(); + + expectAllLocal(c.expect_all_local); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, c.expected_healthy_weights, + c.tol); +} + +INSTANTIATE_TEST_SUITE_P(WeightComputation, WeightComputationTest, + testing::ValuesIn({ + {"ScaleWithEligibleHostCounts", + {10, 2}, + false, + 0.1, + 0.0, + {0.5, 0.5}, + false, + {5.0, 1.0}, + 0.01}, + {"LocalSaturatedSpillsBeyondThreshold", + {1, 1}, + true, + 0.1, + 0.0, + {1.0, 0.3}, + false, + {0.0, 0.7}, + 0.01}, + {"SaturatedLocalWithinThresholdSnapsLocal", + {1, 1}, + true, + 0.1, + 0.0, + {1.0, 0.95}, + true, + {1.0, 0.0}, + 0.01}, + {"ProbeDistributionAcrossRemotes", + {1, 2, 4}, + true, + 1.0, + 0.06, + {0.5, 0.5, 0.5}, + true, + {0.94, 0.02, 0.04}, + 0.01}, + {"ProbeRedistributionClampsLocalAtZero", + {1, 1}, + true, + 0.0, + 2.0, + {0.3, 0.0}, + false, + {0.0, 1.7}, + 0.01}, + {"LocalPreferenceRemoteHostWeightedTarget", + {1, 9, 1}, + true, + 0.01, + 0.0, + {0.5, 0.6, 0.0}, + true, + {1.0, 0.0, 0.0}, + 0.01}, + }), + [](const testing::TestParamInfo& info) { + return info.param.test_name; + }); + +TEST_F(LoadAwareLocalityLbTest, EmptyAndSingleLocalitySmoke) { + createLb(); + + auto empty_lb = createWorkerLb(); + ASSERT_NE(nullptr, empty_lb); + EXPECT_EQ(nullptr, empty_lb->chooseHost(nullptr).host); + EXPECT_EQ(nullptr, empty_lb->peekAnotherHost(nullptr)); + + auto h1 = makeWeightTrackingMockHost(); + setupLocalities({{h1}}); + recomputeWeights(); + + auto populated_lb = createWorkerLb(); + ASSERT_NE(nullptr, populated_lb); + EXPECT_EQ(h1, populated_lb->chooseHost(nullptr).host); + EXPECT_EQ(h1, populated_lb->peekAnotherHost(nullptr)); + EXPECT_FALSE(populated_lb->lifetimeCallbacks().has_value()); + std::vector hash_key; + EXPECT_FALSE(populated_lb->selectExistingConnection(nullptr, *h1, hash_key).has_value()); +} + +TEST_F(LoadAwareLocalityLbTest, EmptyPrioritySetHasNoAvailableHost) { + // Priority 0 exists (the production invariant for LoadBalancerBase) but carries no hosts. + auto* host_set = priority_set_.getMockHostSet(0); + host_set->hosts_.clear(); + host_set->healthy_hosts_.clear(); + + auto child_factory = std::make_shared(false); + auto factory = makeWorkerFactoryWithChild(child_factory, "mock_empty_priority"); + + auto worker_lb = factory->create({priority_set_, nullptr}); + ASSERT_NE(nullptr, worker_lb); + EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); + EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); +} + +TEST_F(LoadAwareLocalityLbTest, BasicRoutingWeightsTrackNoDataFreshDataAndZeroUtilization) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(); + + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 1.0}); + + setHostUtilization(h1, 0.9); + setHostUtilization(h2, 0.1); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.9}); + + // A zero-utilization report carries no positive signal and is ignored: h1 keeps its last-known + // 0.9 rather than being recorded as a fresh, fully idle reporter. + setHostUtilization(h1, 0.0); + setHostUtilization(h2, 0.9); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.1}); +} + +TEST_F(LoadAwareLocalityLbTest, WeightExpirationIgnoresStaleData) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, + std::chrono::milliseconds(5000)); + + // No utilization reported — hosts are cold (never reported) and contribute no data. + recomputeWeights(); + + // Unreported hosts are ignored — weights are equal (host count only). + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 1.0}); + + // Fresh data at current time is used. + setHostUtilization(h1, 0.9); + setHostUtilization(h2, 0.1); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.9}); +} + +TEST_F(LoadAwareLocalityLbTest, WeightExpirationDisabledUsesStaleData) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + // Expiration disabled (0ms) — data is always used regardless of age. + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, + std::chrono::milliseconds(0)); + + setHostUtilization(h1, 0.9); + setHostUtilization(h2, 0.1); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.9}); +} + +// Regression: a report stored at monotonic time 0 (the simulated-time epoch) is a valid sample — +// "never reported" is tracked by an out-of-band sentinel, not timestamp 0. +TEST_F(LoadAwareLocalityLbTest, ReportAtMonotonicTimeZeroIsValid) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, + std::chrono::milliseconds(5000)); + + ASSERT_EQ(0, std::chrono::duration_cast( + simTime().monotonicTime().time_since_epoch()) + .count()); + setHostUtilization(h1, 0.9); + setHostUtilization(h2, 0.1); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.9}); +} + +TEST_F(LoadAwareLocalityLbTest, StaleRemoteCarriesPriorNotZero) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, + std::chrono::milliseconds(5000)); + + // Local 0.5 is within the variance threshold of remote 0.6, so local is preferred. + setHostUtilization(h_local, 0.5); + setHostUtilization(h_remote, 0.6); + recomputeWeights(); + + expectAllLocal(true); + + // Age the remote past expiration while keeping the local fresh. The remote's last-known 0.6 + // must be carried (not reset to 0); otherwise the remote would look idle and local would + // spuriously spill toward it. + context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); + setHostUtilization(h_local, 0.5); + recomputeWeights(); + + expectAllLocal(true); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 0.0}); +} + +TEST_F(LoadAwareLocalityLbTest, StaleLocalityHostCountWeighted) { + auto locality_a = makeHosts(3); + auto locality_b = makeHosts(2); + setupLocalities({locality_a, locality_b}); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, + std::chrono::milliseconds(5000)); + + setUtilizationForHosts(locality_a, 0.8); + setUtilizationForHosts(locality_b, 0.5); + recomputeWeights(); + + // Age locality_b past expiration while keeping locality_a fresh. Locality_b is now all-stale + // and falls back to its host-count baseline (2), not host_count * (1 - 0.5) = 1. + context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); + setUtilizationForHosts(locality_a, 0.8); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.6, 2.0}); +} + +TEST_F(LoadAwareLocalityLbTest, LocalPreferenceThresholdBoundaries) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0); + + setHostUtilization(h_local, 0.3); + setHostUtilization(h_remote, 0.3); + recomputeWeights(); + expectAllLocal(true); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 0.0}); + + setHostUtilization(h_local, 0.6); + setHostUtilization(h_remote, 0.5); + recomputeWeights(); + expectAllLocal(true); + + setHostUtilization(h_local, 0.8); + setHostUtilization(h_remote, 0.2); + recomputeWeights(); + expectAllLocal(false); + EXPECT_GT(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 1), + weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 0)); +} + +TEST_F(LoadAwareLocalityLbTest, NoHealthyHostsWithLocalLocalityDefaultsToLocal) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true, + std::vector{{}, {}}); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0); + + expectAllLocal(true); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 0.0}); +} + +TEST_F(LoadAwareLocalityLbTest, ProbeRedistributesToRemoteAndSkipsWhenNoRemoteHealthyHosts) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + createLb(/*variance_threshold=*/1.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.05); + + setHostUtilization(h_local, 0.3); + setHostUtilization(h_remote, 1.0); + recomputeWeights(); + + expectAllLocal(true); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.95, 0.05}); + + reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/true, + std::vector{{h_local}, {}}); + recomputeWeights(); + + // No remote healthy hosts (remote_hosts == 0): the remote-gated snap is skipped. Routing is still + // fully local (remote weight is 0), but all_local is unset and the local weight is its base + // headroom (0.7), not the snapped 1.0. + expectAllLocal(false); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.7, 0.0}); +} + +TEST_F(LoadAwareLocalityLbTest, ProbeSkipsWhenTargetMetAndSaturatedTieSpreadsByHostCount) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + createLb(/*variance_threshold=*/0.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.1); + + setHostUtilization(h_local, 0.6); + setHostUtilization(h_remote, 0.5); + recomputeWeights(); + + auto snapshot = routingSnapshot(); + ASSERT_NE(nullptr, snapshot); + ASSERT_EQ(2, snapshot->priority_weights[0].by_source[0].weights.size()); + EXPECT_FALSE(snapshot->priority_weights[0].by_source[0].all_local); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.4, 0.5}); + + // Saturated tie: every locality is at 1.0 (zero headroom), so the all-overloaded fallback spreads + // by host count and skips local preference / probe. + setHostUtilization(h_local, 1.0); + setHostUtilization(h_remote, 1.0); + recomputeWeights(); + + expectAllLocal(false); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 1.0}); +} + +TEST_F(LoadAwareLocalityLbTest, HealthyWeightsCanDivergeFromAllHostWeights) { + auto all_a = makeHosts(4); + auto all_b = makeHosts(2); + Upstream::HostVector healthy_a = {all_a[0]}; + Upstream::HostVector healthy_b = {all_b[0], all_b[1]}; + setupLocalities({all_a, all_b}, /*has_local_locality=*/false, + std::vector{healthy_a, healthy_b}); + + createLb(); + + const auto* snapshot = routingSnapshot(); + ASSERT_NE(nullptr, snapshot); + ASSERT_EQ(2, snapshot->priority_weights[0].by_source[0].weights.size()); + ASSERT_EQ(2, + snapshot->priority_weights[0] + .by_source[static_cast(PriorityRoutingWeights::SelectionSource::AllHosts)] + .weights.size()); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 2.0}); + expectWeights(PriorityRoutingWeights::SelectionSource::AllHosts, {4.0, 2.0}); +} + +TEST_F(LoadAwareLocalityLbTest, AllLocalitiesSaturatedFallBacksToHostCount) { + auto locality_a = makeHosts(3); + auto locality_b = makeHosts(1); + setupLocalities({locality_a, locality_b}); + + createLb(); + + setUtilizationForHosts(locality_a, 1.0); + setUtilizationForHosts(locality_b, 1.0); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {3.0, 1.0}); + + auto worker_lb = createWorkerLb(); + auto counts = countPicks(*worker_lb, 400); + int locality_a_count = 0; + for (const auto& host : locality_a) { + locality_a_count += counts[host.get()]; + } + EXPECT_GT(locality_a_count, counts[locality_b[0].get()] * 2); +} + +TEST_F(LoadAwareLocalityLbTest, EwmaLifecycleDampensSpikesAndClearsExpiredState) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/0.3, /*remote_probe_fraction=*/0.0, + std::chrono::milliseconds(5000)); + + setHostUtilization(h1, 0.5); + setHostUtilization(h2, 0.5); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.5, 0.5}); + + setHostUtilization(h1, 1.0); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.35, 0.5}, 0.05); + + context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); + recomputeWeights(); + + const auto* snapshot = routingSnapshot(); + ASSERT_NE(nullptr, snapshot); + EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 1), 1.0, 0.01); +} + +TEST_F(LoadAwareLocalityLbTest, EwmaNewLocalityGetsRawFirstSampleOnTopologyChange) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/0.3); + + setHostUtilization(h1, 0.5); + setHostUtilization(h2, 0.5); + recomputeWeights(); + + auto h3 = makeWeightTrackingMockHost(); + setHostLocality(h3, 2); + auto* host_set = priority_set_.getMockHostSet(0); + host_set->hosts_ = {h1, h2, h3}; + host_set->healthy_hosts_ = {h1, h2, h3}; + host_set->hosts_per_locality_ = + Upstream::makeHostsPerLocality({{h1}, {h2}, {h3}}, /*force_no_local_locality=*/true); + host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + + // Re-initialize to attach LocalityLbHostData to the new host, then set its utilization. + recomputeWeights(); + setHostUtilization(h3, 0.8); + recomputeWeights(); + + const auto* snapshot = routingSnapshot(); + ASSERT_NE(nullptr, snapshot); + ASSERT_EQ(3, snapshot->priority_weights[0].by_source[0].weights.size()); + EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 2), 0.2, 0.05); + // Existing localities keep their smoothed state across the topology change (identity-keyed). + EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 0), 0.5, 0.05); +} + +TEST_F(LoadAwareLocalityLbTest, EwmaStateFollowsLocalityIdentityOnSameCountSwap) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/0.3); + + setHostUtilization(h1, 0.5); + setHostUtilization(h2, 0.9); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.5, 0.1}); + + // Replace locality index 1 (h2's zone) with a brand-new locality (h3) — locality COUNT is + // unchanged, so index-keyed state would silently attribute h2's 0.9 EWMA to h3. + auto h3 = makeWeightTrackingMockHost(); + setHostLocality(h3, 2); + auto* host_set = priority_set_.getMockHostSet(0); + host_set->hosts_ = {h1, h3}; + host_set->healthy_hosts_ = {h1, h3}; + host_set->hosts_per_locality_ = + Upstream::makeHostsPerLocality({{h1}, {h3}}, /*force_no_local_locality=*/true); + host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + + // Re-initialize to attach LocalityLbHostData to the new host, then report. The first sample + // must be applied raw (0.9 weight), not EWMA-blended with the departed locality's 0.9 util + // (which would give 0.3*0.1 + 0.7*0.9 = 0.66 → weight 0.34). + recomputeWeights(); + setHostUtilization(h3, 0.1); + recomputeWeights(); + EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 0), 0.5, 0.05); + EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 2), 0.9, 0.05); + // The departed locality's identity carries no snapshot entry (its EWMA state was dropped). + EXPECT_EQ(0.0, weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 1)); +} + +TEST_F(LoadAwareLocalityLbTest, HealthyAndDegradedSourcesUseDifferentSnapshots) { + auto healthy_host = makeWeightTrackingMockHost(); + auto degraded_host = makeWeightTrackingMockHost(); + + setupLocalities({{healthy_host}, {degraded_host}}, + /*has_local_locality=*/false, + std::vector{{healthy_host}, {}}, + std::vector{{}, {degraded_host}}); + + createLb(); + + expectWeights(PriorityRoutingWeights::SelectionSource::Degraded, {0.0, 1.0}); + + auto worker_lb = createWorkerLb(); + auto counts = countPicks(*worker_lb, 800); + EXPECT_GT(counts[healthy_host.get()], 300); + EXPECT_GT(counts[degraded_host.get()], 100); +} + +TEST_F(LoadAwareLocalityLbTest, HealthTransitionRefreshesHealthyAndDegradedChildren) { + auto healthy_host = makeWeightTrackingMockHost(); + auto remote_host = makeWeightTrackingMockHost(); + setupLocalities({{healthy_host}, {remote_host}}, + /*has_local_locality=*/false, + std::vector{{healthy_host}, {remote_host}}, + std::vector{{}, {}}); + + createLb(); + auto worker_lb = createWorkerLb(); + EXPECT_TRUE(hostSeen(*worker_lb, healthy_host, 100)); + EXPECT_TRUE(hostSeen(*worker_lb, remote_host, 100)); + + auto* host_set = priority_set_.getMockHostSet(0); + host_set->healthy_hosts_ = {healthy_host}; + host_set->healthy_hosts_per_locality_ = + Upstream::makeHostsPerLocality(std::vector{{healthy_host}, {}}, + /*force_no_local_locality=*/true); + host_set->degraded_hosts_ = {remote_host}; + host_set->degraded_hosts_per_locality_ = + Upstream::makeHostsPerLocality(std::vector{{}, {remote_host}}, + /*force_no_local_locality=*/true); + priority_set_.runUpdateCallbacks(0, {}, {}); + recomputeWeights(); + + const auto* snapshot = routingSnapshot(); + ASSERT_NE(nullptr, snapshot); + EXPECT_GT(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 0), 0.0); + EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 1), 0.0, 0.01); + EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Degraded, 0), 0.0, 0.01); + EXPECT_GT(weightForIndex(PriorityRoutingWeights::SelectionSource::Degraded, 1), 0.0); + + EXPECT_TRUE(hostSeen(*worker_lb, healthy_host, 200)); + EXPECT_TRUE(hostSeen(*worker_lb, remote_host, 400)); +} + +TEST_F(LoadAwareLocalityLbTest, PanicUsesAllHosts) { + auto healthy_host = makeWeightTrackingMockHost(); + auto unhealthy_host_1 = makeWeightTrackingMockHost(); + auto unhealthy_host_2 = makeWeightTrackingMockHost(); + + setupLocalities({{healthy_host}, {unhealthy_host_1, unhealthy_host_2}}, + /*has_local_locality=*/false, + std::vector{{healthy_host}, {}}, + std::vector{{}, {}}); + + createLb(); + + // Only 1 of 3 hosts healthy → live LoadBalancerBase panic on priority 0; all hosts selectable. + expectWeights(PriorityRoutingWeights::SelectionSource::AllHosts, {1.0, 2.0}); + + auto worker_lb = createWorkerLb(); + auto counts = countPicks(*worker_lb, 400); + EXPECT_GT(counts[healthy_host.get()], 100); + EXPECT_GT(counts[unhealthy_host_1.get()] + counts[unhealthy_host_2.get()], 200); + EXPECT_GT(cluster_info_.lbStats().lb_healthy_panic_.value(), 0u); +} + +TEST_F(LoadAwareLocalityLbTest, FailTrafficOnPanicReturnsNoHost) { + cluster_info_.lb_config_.mutable_zone_aware_lb_config()->set_fail_traffic_on_panic(true); + + auto healthy_host = makeWeightTrackingMockHost(); + auto unhealthy_host_1 = makeWeightTrackingMockHost(); + auto unhealthy_host_2 = makeWeightTrackingMockHost(); + // Only 1 of 3 hosts healthy → live panic on priority 0. + setupLocalities({{healthy_host}, {unhealthy_host_1, unhealthy_host_2}}, + /*has_local_locality=*/false, + std::vector{{healthy_host}, {}}, + std::vector{{}, {}}); + + createLb(); + + auto worker_lb = createWorkerLb(); + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); + EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); + } + EXPECT_GT(cluster_info_.lbStats().lb_healthy_panic_.value(), 0u); +} + +TEST_F(LoadAwareLocalityLbTest, LiveSnapshotRefreshUsesLatestWeights) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(); + auto worker_lb = createWorkerLb(); + EXPECT_TRUE(hostSeen(*worker_lb, h1, 100)); + EXPECT_TRUE(hostSeen(*worker_lb, h2, 100)); + + publishHealthyWeights({0.0, 1.0}); + + auto counts = countPicks(*worker_lb, 100); + EXPECT_GE(counts[h2.get()], 98); +} + +TEST_F(LoadAwareLocalityLbTest, MembershipAndTopologyUpdatesShareTheSameWorkerLb) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(); + auto worker_lb = createWorkerLb(); + EXPECT_TRUE(hostSeen(*worker_lb, h1, 100)); + EXPECT_TRUE(hostSeen(*worker_lb, h2, 100)); + + auto h3 = makeWeightTrackingMockHost(); + reshapeLocalities(0, {{h1}, {h2, h3}}, /*added=*/{h3}); + EXPECT_TRUE(hostSeen(*worker_lb, h3, 200)); + + reshapeLocalities(0, {{h1}, {h3}}, /*added=*/{}, /*removed=*/{h2}); + EXPECT_FALSE(hostSeen(*worker_lb, h2, 200)); + EXPECT_TRUE(hostSeen(*worker_lb, h3, 200)); + + auto h4 = makeWeightTrackingMockHost(); + reshapeLocalities(0, {{h1}, {h3}, {h4}}, /*added=*/{h4}); + + publishHealthyWeights({1.0, 1.0, 1.0}); + EXPECT_TRUE(hostSeen(*worker_lb, h4, 300)); +} + +TEST_F(LoadAwareLocalityLbTest, StaleSnapshotFallbackAndAllLocalitiesRemovedAreGraceful) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + auto h3 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}, {h3}}); + + createLb(); + setHostUtilization(h1, 0.5); + setHostUtilization(h2, 0.5); + setHostUtilization(h3, 0.5); + recomputeWeights(); + + auto worker_lb = createWorkerLb(); + EXPECT_TRUE(hostSeen(*worker_lb, h1, 200)); + EXPECT_TRUE(hostSeen(*worker_lb, h2, 200)); + EXPECT_TRUE(hostSeen(*worker_lb, h3, 200)); + + reshapeLocalities(0, {{h1}, {}, {}}, /*added=*/{}, /*removed=*/{h2, h3}); + + expectOnlyHost(*worker_lb, h1, 100); + for (int i = 0; i < 50; ++i) { + auto host = worker_lb->peekAnotherHost(nullptr); + ASSERT_NE(nullptr, host); + EXPECT_EQ(h1, host); + worker_lb->chooseHost(nullptr); + } + + reshapeLocalities(0, {{}, {}, {}}, /*added=*/{}, /*removed=*/{h1}); + + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); + EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); + } +} + +TEST_F(LoadAwareLocalityLbTest, EmptyDeltaUpdatesRefreshChildWeights) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1, h2}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + auto initial_counts = countPicks(*worker_lb, 200); + EXPECT_GT(initial_counts[h1.get()], 50); + EXPECT_GT(initial_counts[h2.get()], 50); + + h1->weight(100); + priority_set_.runUpdateCallbacks(0, {}, {}); + + auto updated_counts = countPicks(*worker_lb, 400); + EXPECT_GT(updated_counts[h1.get()], updated_counts[h2.get()] * 10); + EXPECT_GT(updated_counts[h1.get()], 300); +} + +TEST_F(LoadAwareLocalityLbTest, InPlaceUpdateTopologyMismatchWaitsForMembershipChange) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + auto h3 = makeWeightTrackingMockHost(); + reshapeLocalities(0, {{h1}, {h2}, {h3}}); + EXPECT_TRUE(hostSeen(*worker_lb, h1, 100)); + EXPECT_TRUE(hostSeen(*worker_lb, h2, 100)); + EXPECT_FALSE(hostSeen(*worker_lb, h3, 100)); + + priority_set_.runUpdateCallbacks(0, {h3}, {}); + publishHealthyWeights({1.0, 1.0, 1.0}); + EXPECT_TRUE(hostSeen(*worker_lb, h3, 300)); +} + +TEST_F(LoadAwareLocalityLbTest, PriorityFailoverAndFailback) { + auto h_p0 = makeWeightTrackingMockHost(); + auto h_p1 = makeWeightTrackingMockHost(); + + setupPriorityLocalities(0, {{h_p0}}, /*has_local_locality=*/false, + std::vector{{}}); + setupPriorityLocalities(1, {{h_p1}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(h_p1, worker_lb->chooseHost(nullptr).host); + } + + reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, + std::vector{{h_p0}}); + recomputeWeights(); + + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(h_p0, worker_lb->chooseHost(nullptr).host); + } + + reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, + std::vector{{}}); + recomputeWeights(); + + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(h_p1, worker_lb->chooseHost(nullptr).host); + } +} + +TEST_F(LoadAwareLocalityLbTest, FailoverIsCallbackFreshNotTimerGated) { + auto h_p0 = makeWeightTrackingMockHost(); + auto h_p1 = makeWeightTrackingMockHost(); + + setupPriorityLocalities(0, {{h_p0}}); + setupPriorityLocalities(1, {{h_p1}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + // Priority 0 is healthy, so live failover keeps load on it. + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(h_p0, worker_lb->chooseHost(nullptr).host); + } + + // Flip priority 0 to all-unhealthy (keep hosts_, clear healthy_hosts_) and fire the membership + // update callback ONLY. The weight-update timer is NOT invoked and simulated time is NOT + // advanced past weight_update_period, so no new routing-weight snapshot is published. + reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, + std::vector{{}}); + + // Failover to priority 1 happens immediately from the live LoadBalancerBase callback, proving it + // is not gated on a timer-published weight snapshot. + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(h_p1, worker_lb->chooseHost(nullptr).host); + } +} + +TEST_F(LoadAwareLocalityLbTest, PriorityUpdateGuardPaths) { + auto h1 = makeWeightTrackingMockHost(); + setupLocalities({{h1}}); + + createLb(); + auto worker_lb = createWorkerLb(); + EXPECT_TRUE(hostSeen(*worker_lb, h1, 50)); + + // Add a healthy priority 1. Priority 0 is still healthy, so live failover keeps load there. + auto h_p1 = makeWeightTrackingMockHost(); + setupPriorityLocalities(1, {{h_p1}}); + priority_set_.runUpdateCallbacks(1, {h_p1}, {}); + PriorityRoutingWeights p0; + p0.by_source[0].weights = makeWeightsMap({1.0}); + PriorityRoutingWeights p1; + p1.by_source[0].weights = makeWeightsMap({1.0}); + publishSnapshot(makeSnapshot({p0, p1})); + expectOnlyHost(*worker_lb, h1, 50); + + // Drop priority 0 to zero availability; live failover shifts load to priority 1. + reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, + std::vector{{}}); + EXPECT_TRUE(hostSeen(*worker_lb, h_p1, 100)); + + // Growing the backing priority set without a matching snapshot is handled gracefully. + priority_set_.getMockHostSet(5); + priority_set_.runUpdateCallbacks(5, {makeWeightTrackingMockHost()}, {}); + priority_set_.runUpdateCallbacks(5, {}, {}); + EXPECT_TRUE(hostSeen(*worker_lb, h_p1, 100)); +} + +TEST_F(LoadAwareLocalityLbTest, AdvisorySnapshotSizeMismatchesAreHandledGracefully) { + auto h1 = makeWeightTrackingMockHost(); + setupLocalities({{h1}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + PriorityRoutingWeights weights; + weights.by_source[0].weights = makeWeightsMap({1.0}); + + // The advisory snapshot may carry more priorities than the live priority set; the single live + // healthy host stays selected (priority is live, not snapshot-driven). + publishSnapshot(makeSnapshot({weights, weights})); + EXPECT_TRUE(hostSeen(*worker_lb, h1, 20)); + + publishSnapshot(makeSnapshot({weights})); + EXPECT_TRUE(hostSeen(*worker_lb, h1, 20)); + EXPECT_EQ(h1, worker_lb->peekAnotherHost(nullptr)); +} + +TEST_F(LoadAwareLocalityLbTest, ShortAdvisorySnapshotStillRoutesLivePriority) { + auto h_p1 = makeWeightTrackingMockHost(); + setupPriorityLocalities(0, {}); + setupPriorityLocalities(1, {{h_p1}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + // Priority 0 has no hosts; live failover routes to the healthy priority 1 even though the + // advisory snapshot only carries weights for priority 0. + publishHealthyWeights({1.0}); + + expectOnlyHost(*worker_lb, h_p1, 20); + EXPECT_EQ(h_p1, worker_lb->peekAnotherHost(nullptr)); +} + +TEST_F(LoadAwareLocalityLbTest, EmptyHigherPriorityDoesNotBreakLiveRouting) { + auto h_p0 = makeWeightTrackingMockHost(); + setupPriorityLocalities(0, {{h_p0}}); + setupPriorityLocalities(1, {}); + + createLb(); + auto worker_lb = createWorkerLb(); + + // Priority 1 has no hosts; live failover keeps load on the healthy priority 0. + PriorityRoutingWeights p0; + p0.by_source[0].weights = makeWeightsMap({1.0}); + PriorityRoutingWeights p1; + publishSnapshot(makeSnapshot({p0, p1})); + + expectOnlyHost(*worker_lb, h_p0, 20); +} + +TEST_F(LoadAwareLocalityLbTest, SelectLocalityFallsBackForCountMismatchAndZeroTotals) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + // Snapshot has only locality 0's identity; the worker's locality 1 (missing → 0.0) is excluded, + // so all traffic stays on h1. + publishHealthyWeights({1.0}); + expectOnlyHost(*worker_lb, h1, 20); + + publishHealthyWeights({}); + expectOnlyHost(*worker_lb, h1, 20); + + // Both live localities have weight 0 (locality 2's 5.0 is not a live identity), so selection + // falls back to locality 0. + publishHealthyWeights({0.0, 0.0, 5.0}); + expectOnlyHost(*worker_lb, h1, 20); +} + +TEST_F(LoadAwareLocalityLbTest, SelectLocalityIgnoresWeightsForUnknownLocalities) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + // The snapshot carries a third locality identity (weight 100.0) that the worker does not have. + // It is ignored; the two live localities split traffic by their identity-matched weights (1:1). + publishHealthyWeights({1.0, 1.0, 100.0}); + + auto counts = countPicks(*worker_lb, 400); + EXPECT_GT(counts[h1.get()], 120); + EXPECT_GT(counts[h2.get()], 120); +} + +TEST_F(LoadAwareLocalityLbTest, SelectLocalityRoundingGuardReturnsLastLocality) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + auto h3 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}, {h3}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + publishHealthyWeights({0.2, 0.3, 0.5}); + + // The hash draw (chooseHostSet) consumes the first random value even with a single priority, + // so the second draw must drive locality selection. + forceRandomValues({0, std::numeric_limits::max()}); + EXPECT_EQ(h3, worker_lb->chooseHost(nullptr).host); +} + +// Advisory weights are keyed by Locality identity, so adding a locality that shifts another's +// lexicographic index between the snapshot tick and the worker's live membership must NOT mis-map +// the weight. Set up local A + remote C with distinct weights, publish a snapshot, then add remote +// B (sorts between A and C, shifting C from index 1 to index 2) WITHOUT recomputing weights, and +// assert C's traffic still follows C's weight (and the new B, absent from the snapshot, gets none). +TEST_F(LoadAwareLocalityLbTest, AdvisoryWeightsFollowLocalityIdentityAcrossIndexShift) { + const auto locality_a = Upstream::Locality("a", "z", "sz"); // local, index 0 + const auto locality_b = Upstream::Locality("b", "z", "sz"); // sorts between A and C + const auto locality_c = + Upstream::Locality("c", "z", "sz"); // index 1, shifts to 2 when B is added + + auto host_a = makeWeightTrackingMockHost(); + auto host_c = makeWeightTrackingMockHost(); + setHostLocalityExplicit(host_a, locality_a); + setHostLocalityExplicit(host_c, locality_c); + + auto* host_set = priority_set_.getMockHostSet(0); + host_set->hosts_ = {host_a, host_c}; + host_set->healthy_hosts_ = {host_a, host_c}; + host_set->hosts_per_locality_ = Upstream::makeHostsPerLocality({{host_a}, {host_c}}); + host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + + createLb(); + auto worker_lb = createWorkerLb(); + + // Snapshot computed for the {A, C} ordering: A gets almost no weight, C gets all of it. + PriorityRoutingWeights weights; + weights.by_source[0].weights[locality_a] = 0.0; + weights.by_source[0].weights[locality_c] = 1.0; + publishSnapshot(makeSnapshot({weights})); + expectOnlyHost(*worker_lb, host_c, 50); + + // Add remote B, whose identity sorts between A and C, so C's index shifts from 1 to 2. The + // snapshot is NOT recomputed: it still only knows A and C by identity. + auto host_b = makeWeightTrackingMockHost(); + setHostLocalityExplicit(host_b, locality_b); + host_set->hosts_ = {host_a, host_b, host_c}; + host_set->healthy_hosts_ = {host_a, host_b, host_c}; + host_set->hosts_per_locality_ = Upstream::makeHostsPerLocality({{host_a}, {host_b}, {host_c}}); + host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + priority_set_.runUpdateCallbacks(0, {host_b}, {}); + + // Identity mapping survives the index shift: C still carries all the weight (its weight was NOT + // applied to B at the now-stale index 1), and B — absent from the snapshot — receives nothing. + expectOnlyHost(*worker_lb, host_c, 50); + EXPECT_FALSE(hostSeen(*worker_lb, host_b, 200)); + EXPECT_FALSE(hostSeen(*worker_lb, host_a, 200)); +} + +TEST_F(LoadAwareLocalityLbTest, DedicatedStatsAllOverloadedTotalIncOncePerTickWhenAllZeroHeadroom) { + auto h_a = makeHosts(2); + auto h_b = makeHosts(1); + setupLocalities({h_a, h_b}); + + createLb(); + EXPECT_EQ(0, counterValue(cluster_info_, "load_aware_locality.all_overloaded_total")); + + // Make all localities fully saturated — headroom = 0. + setUtilizationForHosts(h_a, 1.0); + setUtilizationForHosts(h_b, 1.0); + recomputeWeights(); + EXPECT_EQ(1, counterValue(cluster_info_, "load_aware_locality.all_overloaded_total")); + + // Next tick, still saturated — counter goes to 2 (one per tick, not per source). + expectCounterIncrementsPerTick("load_aware_locality.all_overloaded_total", 1); +} + +TEST_F(LoadAwareLocalityLbTest, DedicatedStatsLocalPreferredTotalIncOncePerTickOnVarianceSnap) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + // High threshold so local is NOT preferred when both are saturated (utils=1.0). + createLb(/*variance_threshold=*/0.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0); + // Cold start: no data yet, both utils=0 → utils[0]=0 <= 0+0.0=0 → snap fires on the initial + // tick, deterministically. + const uint64_t after_create = + counterValue(cluster_info_, "load_aware_locality.local_preferred_total"); + EXPECT_EQ(1u, after_create); + + // Force high local utilization so snap does NOT fire. + setHostUtilization(h_local, 0.9); + setHostUtilization(h_remote, 0.1); + recomputeWeights(); + // 0.9 > 0.1 + 0.0 → no snap. + const uint64_t after_no_snap = + counterValue(cluster_info_, "load_aware_locality.local_preferred_total"); + EXPECT_EQ(after_create, after_no_snap); + + // Now snap fires (local <= remote + threshold). + setHostUtilization(h_local, 0.1); + setHostUtilization(h_remote, 0.9); + recomputeWeights(); + EXPECT_EQ(after_no_snap + 1, + counterValue(cluster_info_, "load_aware_locality.local_preferred_total")); + + // Another snap tick. + expectCounterIncrementsPerTick("load_aware_locality.local_preferred_total", 1); +} + +TEST_F(LoadAwareLocalityLbTest, DedicatedStatsProbeActiveTotalIncOncePerTickWhenProbeKicksIn) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + // remote_probe_fraction=0.05 → probe redistribution fires when remote < 5% of total. + // variance_threshold=1.0 so local is always preferred. + createLb(/*variance_threshold=*/1.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.05); + // Cold start tick: utils=0, local preferred, probe fires on the initial tick. + const uint64_t after_create = + counterValue(cluster_info_, "load_aware_locality.probe_active_total"); + EXPECT_EQ(1u, after_create); + + setHostUtilization(h_local, 0.3); + setHostUtilization(h_remote, 1.0); + recomputeWeights(); + // all_local fires (local preferred), then probe redistribution shifts weight to remote. + EXPECT_EQ(after_create + 1, + counterValue(cluster_info_, "load_aware_locality.probe_active_total")); + + expectCounterIncrementsPerTick("load_aware_locality.probe_active_total", 1); +} + +TEST_F(LoadAwareLocalityLbTest, DedicatedStatsStaleLocalityTotalCountsOncePerStalePerTick) { + auto locality_a = makeHosts(2); + auto locality_b = makeHosts(1); + auto locality_c = makeHosts(3); + setupLocalities({locality_a, locality_b, locality_c}); + + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, + std::chrono::milliseconds(5000)); + + // Never-reported localities are cold, not stale: nothing counted before first data arrives. + EXPECT_EQ(0, counterValue(cluster_info_, "load_aware_locality.stale_locality_total")); + expectCounterIncrementsPerTick("load_aware_locality.stale_locality_total", 0); + + // Give all localities fresh data and re-initialize so we have a clean baseline. + setUtilizationForHosts(locality_a, 0.5); + setUtilizationForHosts(locality_b, 0.4); + setUtilizationForHosts(locality_c, 0.6); + recomputeWeights(); + // Record counter after the tick where all data is fresh (should add 0). + const uint64_t baseline = counterValue(cluster_info_, "load_aware_locality.stale_locality_total"); + + // Age locality_b and locality_c past expiration; locality_a stays fresh. + context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); + setUtilizationForHosts(locality_a, 0.5); + recomputeWeights(); + // 2 stale localities (b and c) — counter adds 2 per tick, NOT 6 (not 3× for 3 sources). + EXPECT_EQ(baseline + 2, counterValue(cluster_info_, "load_aware_locality.stale_locality_total")); + + // Next tick, same state → adds another 2. + expectCounterIncrementsPerTick("load_aware_locality.stale_locality_total", 2); +} + +TEST_F(LoadAwareLocalityLbTest, WeightUpdateTimerRearmsOnEachTick) { + auto h = makeWeightTrackingMockHost(); + setupLocalities({{h}}); + createLb(); + // initialize() ran the first tick, which re-arms before computing. + ASSERT_TRUE(timer_->enabled_); + // Every fired tick must re-arm; a conditional or late re-arm would leave the timer disabled. + timer_->invokeCallback(); + EXPECT_TRUE(timer_->enabled_); + timer_->invokeCallback(); + EXPECT_TRUE(timer_->enabled_); +} + +TEST_F(LoadAwareLocalityLbTest, ChildContextShieldsRetryPriorityWithoutComputingHash) { + auto h_a = makeWeightTrackingMockHost(); + auto h_b = makeWeightTrackingMockHost(); + setupLocalities({{h_a}, {h_b}}); + createLb(); + auto worker_lb = createWorkerLb(); + + NiceMock ctx; + // The (possibly stateful) retry-priority hook must be consulted exactly once per pick, by the + // parent against the cluster priority set, never re-invoked by the per-locality child LB whose + // single-priority view would corrupt retry-priority state. Route-level hash policies should not + // be evaluated by non-hash children. + EXPECT_CALL(ctx, determinePriorityLoad(testing::Ref(priority_set_), testing::_, testing::_)) + .WillOnce([](const Upstream::PrioritySet&, + const Upstream::HealthyAndDegradedLoad& default_priority_load, + const Upstream::RetryPriority::PriorityMappingFunc&) + -> const Upstream::HealthyAndDegradedLoad& { return default_priority_load; }); + EXPECT_CALL(ctx, computeHashKey()).Times(0); + + EXPECT_NE(nullptr, worker_lb->chooseHost(&ctx).host); +} + +// An asynchronous child LB would retain the stack-scoped child context and call back after +// chooseHost returns (use-after-free). The policy must cancel the async selection and fail +// synchronously, never propagating the cancelable handle to its caller. +TEST_F(LoadAwareLocalityLbTest, AsyncChildSelectionIsCancelledAndFailsSynchronously) { + auto h = makeWeightTrackingMockHost(); + setupLocalities({{h}}); + + bool cancelled = false; + auto child_factory = std::make_shared(cancelled); + auto factory = makeWorkerFactoryWithChild(child_factory, "mock_async_child"); + auto worker_lb = factory->create({priority_set_, nullptr}); + ASSERT_NE(nullptr, worker_lb); + + NiceMock ctx; + auto response = worker_lb->chooseHost(&ctx); + EXPECT_EQ(nullptr, response.host); + EXPECT_EQ(nullptr, response.cancelable); + EXPECT_TRUE(cancelled); +} + +TEST_F(LoadAwareLocalityLbTest, PeekAndChooseReplayTheSameRandomSequence) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + createLb(); + + setHostUtilization(h1, 0.9); + setHostUtilization(h2, 0.1); + recomputeWeights(); + auto worker_lb = createWorkerLb(); + + // Peek draws (priority hash + locality draw) are stashed: 0 → locality 0 (weight 0.1 of 1.0). + forceRandomValues({0, 0}); + auto peeked = worker_lb->peekAnotherHost(nullptr); + ASSERT_NE(nullptr, peeked); + + // If chooseHost drew fresh randoms these values would steer it to locality 1; replaying the + // stashed peek draws must land on the same locality (and host) the peek returned. + forceRandomValues( + {std::numeric_limits::max() - 1, std::numeric_limits::max() - 1}); + auto chosen = worker_lb->chooseHost(nullptr).host; + EXPECT_EQ(peeked.get(), chosen.get()); +} + +TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverHealthyRoutingPaths) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + createLb(/*variance_threshold=*/1.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0); + setHostUtilization(h_local, 0.3); + setHostUtilization(h_remote, 0.3); + recomputeWeights(); + + auto worker_lb = createWorkerLb(); + expectOnlyHost(*worker_lb, h_local, 10); + EXPECT_EQ(10, static_cast(cluster_info_.lbStats().lb_zone_routing_all_directly_.value())); + + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + createLb(/*variance_threshold=*/0.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0); + setHostUtilization(h_local, 0.31); + setHostUtilization(h_remote, 0.3); + recomputeWeights(); + + worker_lb = createWorkerLb(); + for (int i = 0; i < 200; ++i) { + worker_lb->chooseHost(nullptr); + } + EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_sampled_.value(), 0); + EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_cross_zone_.value(), 0); +} + +TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverDegradedAndPanicSources) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + // No healthy hosts, all degraded (but enough to avoid panic): live priority targets the + // Degraded source. + setupLocalities({{h_local}, {h_remote}}, + /*has_local_locality=*/true, std::vector{{}, {}}, + std::vector{{h_local}, {h_remote}}); + + createLb(); + recomputeWeights(); + + auto worker_lb = createWorkerLb(); + auto degraded_counts = countPicks(*worker_lb, 400); + EXPECT_GT(degraded_counts[h_local.get()] + degraded_counts[h_remote.get()], 0); + + auto h_remote_2 = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote, h_remote_2}}, + /*has_local_locality=*/true, std::vector{{h_local}, {}}, + std::vector{{}, {}}); + + createLb(); + recomputeWeights(); + worker_lb = createWorkerLb(); + auto panic_counts = countPicks(*worker_lb, 400); + EXPECT_GT(panic_counts[h_local.get()] + panic_counts[h_remote.get()] + + panic_counts[h_remote_2.get()], + 0); + EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_sampled_.value() + + cluster_info_.lbStats().lb_zone_routing_cross_zone_.value() + + cluster_info_.lbStats().lb_zone_routing_all_directly_.value(), + 0u); +} + +TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsStayZeroWithoutLocalLocality) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}, /*has_local_locality=*/false); + + createLb(); + setHostUtilization(h1, 0.3); + setHostUtilization(h2, 0.3); + recomputeWeights(); + + auto worker_lb = createWorkerLb(); + for (int i = 0; i < 100; ++i) { + worker_lb->chooseHost(nullptr); + } + + EXPECT_EQ(0, cluster_info_.lbStats().lb_zone_routing_all_directly_.value()); + EXPECT_EQ(0, cluster_info_.lbStats().lb_zone_routing_cross_zone_.value()); + EXPECT_EQ(0, cluster_info_.lbStats().lb_zone_routing_sampled_.value()); +} + +TEST_F(LoadAwareLocalityLbTest, EmptyLocalityHostsDoNotCrash) { + auto* host_set = priority_set_.getMockHostSet(0); + host_set->hosts_.clear(); + host_set->healthy_hosts_.clear(); + host_set->hosts_per_locality_ = + Upstream::makeHostsPerLocality({}, /*force_no_local_locality=*/true); + host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + + createLb(); + auto worker_lb = createWorkerLb(); + EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); +} + +TEST_F(LoadAwareLocalityLbTest, WorkerFallsBackWithoutPublishedRoutingSnapshot) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + auto factory = createRoundRobinWorkerFactory(); + auto worker_lb = factory->create({priority_set_, nullptr}); + ASSERT_NE(nullptr, worker_lb); + + expectOnlyHost(*worker_lb, h1, 20); + EXPECT_EQ(h1, worker_lb->peekAnotherHost(nullptr)); +} + +TEST_F(LoadAwareLocalityLbTest, InitializePropagatesChildFactoryError) { + NiceMock null_child_factory; + ON_CALL(null_child_factory, name()).WillByDefault(Return("mock_null_factory")); + ON_CALL(null_child_factory, + create(testing::_, testing::_, testing::_, testing::_, testing::_, testing::_)) + .WillByDefault(testing::Return(testing::ByMove(nullptr))); + + auto lb_config = std::make_unique( + null_child_factory, "mock_null_factory", nullptr, std::chrono::milliseconds(1000), 0.1, 0.3, + 0.03, std::chrono::milliseconds(180000), std::vector{}, dispatcher_, + context_.thread_local_); + + LoadAwareLocalityLoadBalancer lb(*lb_config, cluster_info_, priority_set_, + context_.runtime_loader_, random_, context_.time_system_); + auto status = lb.initialize(); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(status.message(), testing::HasSubstr("mock_null_factory")); +} + +TEST_F(LoadAwareLocalityLbTest, ChildRecreatedWhenUnderlyingFactoryRequiresIt) { + auto local = makeWeightTrackingMockHost(); + auto remote = makeWeightTrackingMockHost(); + setupLocalities({{local}, {remote}}); + + auto child_factory = std::make_shared(true); + auto factory = makeWorkerFactoryWithChild(child_factory, "mock_recreate_child"); + + auto worker_lb = factory->create({priority_set_, nullptr}); + ASSERT_NE(nullptr, worker_lb); + const uint32_t initial_create_count = child_factory->createCount(); + EXPECT_GT(initial_create_count, 0); + + auto extra = makeWeightTrackingMockHost(); + reshapeLocalities(0, {{local}, {remote, extra}}, /*added=*/{extra}); + EXPECT_GT(child_factory->createCount(), initial_create_count); +} + +// Wraps metric names for LocalityLbHostData's shared_ptr constructor. +std::shared_ptr> makeMetricNames(std::vector names) { + return std::make_shared>(std::move(names)); +} + +TEST_F(LoadAwareLocalityLbTest, MetricNamesDriveUtilization) { + // LocalityLbHostData with metric_names={"named_metrics.foo"} delegates to + // OrcaLoadReportHandler::getUtilizationFromOrcaReport, which follows the runtime-feature + // controlled precedence: named_metrics → application_utilization → cpu_utilization + // (when the default orca_weight_manager_use_named_metrics_first flag is enabled). + Event::GlobalTimeSystem ts; + + { + // Named metric present → use its value. + LocalityLbHostData slot(ts, makeMetricNames({"named_metrics.foo"})); + xds::data::orca::v3::OrcaLoadReport report; + (*report.mutable_named_metrics())["foo"] = 0.6; + ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); + EXPECT_DOUBLE_EQ(0.6, slot.utilization()); + } + + { + // Named metric absent and no application_utilization → cpu_utilization fallback. + LocalityLbHostData slot(ts, makeMetricNames({"named_metrics.foo"})); + xds::data::orca::v3::OrcaLoadReport report; + report.set_cpu_utilization(0.3); + ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); + EXPECT_DOUBLE_EQ(0.3, slot.utilization()); + } + + { + // Empty metric_names → application_utilization used when set. + LocalityLbHostData slot(ts, makeMetricNames({})); + xds::data::orca::v3::OrcaLoadReport report; + report.set_application_utilization(0.7); + ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); + EXPECT_DOUBLE_EQ(0.7, slot.utilization()); + } + + { + // Empty metric_names and no application_utilization → cpu_utilization fallback. + LocalityLbHostData slot(ts, makeMetricNames({})); + xds::data::orca::v3::OrcaLoadReport report; + report.set_cpu_utilization(0.4); + ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); + EXPECT_DOUBLE_EQ(0.4, slot.utilization()); + } +} + +// Tests 1+2: chooseHost and peekAnotherHost empty-priority-set guard. +// The guards at the top of each function return early when hostSetsPerPriority() is empty. +// We construct the LB against a real priority set (so LoadBalancerBase initialises safely), then +// override the mock to present an empty set for subsequent calls. +TEST_F(LoadAwareLocalityLbTest, EmptyPrioritySetGuardInChooseHostAndPeekAnotherHost) { + auto h1 = makeWeightTrackingMockHost(); + setupLocalities({{h1}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + // Override so that hostSetsPerPriority() returns empty from this point on. + static const std::vector kEmptyHostSets; + ON_CALL(testing::Const(priority_set_), hostSetsPerPriority()) + .WillByDefault(testing::ReturnRef(kEmptyHostSets)); + + // Both guards fire: `if (priority_set_.hostSetsPerPriority().empty()) return {nullptr/nullptr}`. + EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); + EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); +} + +// Test 3: empty-delta (no-rebuild) priority-count-mismatch early-return. +// When the priority count in the live set grows but the worker hasn't rebuilt yet, an empty-delta +// callback runs syncPriority with allow_rebuild=false. The size mismatch guard returns early +// without crashing. +TEST_F(LoadAwareLocalityLbTest, InPlaceHostUpdateSizeMismatchReturnsEarly) { + auto h1 = makeWeightTrackingMockHost(); + setupLocalities({{h1}}); + + createLb(); + auto worker_lb = createWorkerLb(); + expectOnlyHost(*worker_lb, h1, 10); + + // Grow the backing priority set to 6 priorities without a membership-delta rebuild first. + priority_set_.getMockHostSet(5); + // Empty-delta update on priority 5: syncPriority(5, allow_rebuild=false). + // host_sets.size()=6, per_priority_locality_.size()=1 → size-mismatch guard returns early. + priority_set_.runUpdateCallbacks(5, {}, {}); + + // The LB is still functional for priority 0 (early-return did not corrupt state). + expectOnlyHost(*worker_lb, h1, 10); +} + +// Test 4: storeUtilization non-finite guard. +// When getUtilizationFromOrcaReport returns a non-finite value (Inf/NaN), storeUtilization +// discards it: utilization() stays 0.0 and lastUpdateTimeMs() stays 0. +TEST_F(LoadAwareLocalityLbTest, StoreUtilizationRejectsNonFiniteValue) { + Event::GlobalTimeSystem ts; + LocalityLbHostData slot(ts, makeMetricNames({})); + + // cpu_utilization is the last-resort fallback when no metric names are configured and + // application_utilization is 0. Setting it to Inf produces a non-finite util value. + xds::data::orca::v3::OrcaLoadReport report; + report.set_cpu_utilization(std::numeric_limits::infinity()); + ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); + + // Non-finite value rejected: slot remains at its initial state. + EXPECT_DOUBLE_EQ(0.0, slot.utilization()); + EXPECT_EQ(LocalityLbHostData::kNeverReported, slot.lastUpdateTimeMs()); +} + +} // namespace +} // namespace LoadAwareLocality +} // namespace LoadBalancingPolicies +} // namespace Extensions +} // namespace Envoy From 3fb4a7574534b2493ad12ce7742aec8f5fa39ec8 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 16:12:58 -0600 Subject: [PATCH 08/16] move child LB ownership to main thread, enforce 100ms tick floor, single-pass locality pick Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality/BUILD | 1 + .../load_aware_locality/config.cc | 6 +- .../load_aware_locality_lb.cc | 65 ++++++++++--------- .../load_aware_locality_lb.h | 35 +++++----- .../load_aware_locality/config_test.cc | 7 ++ .../load_aware_locality_lb_test.cc | 59 ++++++----------- 6 files changed, 81 insertions(+), 92 deletions(-) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/BUILD b/source/extensions/load_balancing_policies/load_aware_locality/BUILD index 61866ccb9142f..713cdbc82cfb2 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/BUILD +++ b/source/extensions/load_balancing_policies/load_aware_locality/BUILD @@ -42,6 +42,7 @@ envoy_cc_library( "//source/extensions/load_balancing_policies/common:orca_weight_manager_lib", "@abseil-cpp//absl/container:flat_hash_map", "@abseil-cpp//absl/container:flat_hash_set", + "@abseil-cpp//absl/container:inlined_vector", "@abseil-cpp//absl/strings", ], ) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/config.cc b/source/extensions/load_balancing_policies/load_aware_locality/config.cc index 0cfbf87bbd900..448c3923f4f16 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/config.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/config.cc @@ -32,8 +32,10 @@ Factory::loadConfig(Server::Configuration::ServerFactoryContext& context, } const std::chrono::milliseconds weight_update_period = std::chrono::milliseconds(PROTOBUF_GET_MS_OR_DEFAULT(lb_config, weight_update_period, 1000)); - if (weight_update_period <= std::chrono::milliseconds(0)) { - return absl::InvalidArgumentError("weight_update_period must be positive"); + // PGV constraints are not enforced on unpacked LB configs, so the documented floor is owned + // here. + if (weight_update_period < std::chrono::milliseconds(100)) { + return absl::InvalidArgumentError("weight_update_period must be at least 100ms"); } const std::chrono::milliseconds smoothing_time_constant = std::chrono::milliseconds( PROTOBUF_GET_MS_OR_DEFAULT(lb_config, smoothing_time_constant, 5000)); diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc index 8ed394e14bd0e..51f584ff079da 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -14,6 +14,7 @@ #include "source/extensions/load_balancing_policies/common/orca_weight_manager.h" #include "absl/container/flat_hash_set.h" +#include "absl/container/inlined_vector.h" #include "absl/strings/str_cat.h" namespace Envoy { @@ -58,10 +59,16 @@ LoadAwareLocalityLoadBalancer::LoadAwareLocalityLoadBalancer( metric_names_ = std::make_shared>( typed_config->metricNamesForComputingUtilization()); + child_factory_name_ = typed_config->endpointPickingPolicyName(); + auto child_config_ref = makeOptRefFromPtr( + typed_config->endpointPickingPolicyConfig().get()); + child_thread_aware_lb_ = typed_config->endpointPickingPolicyFactory().create( + child_config_ref, cluster_info, priority_set, runtime, random, time_source); + factory_ = std::make_shared( - typed_config->endpointPickingPolicyFactory(), typed_config->endpointPickingPolicyName(), - typed_config->endpointPickingPolicyConfig(), cluster_info, priority_set, runtime, random, - time_source, typed_config->tlsSlotAllocator()); + child_thread_aware_lb_ != nullptr ? child_thread_aware_lb_->factory() : nullptr, + typed_config->endpointPickingPolicyConfig(), cluster_info, runtime, random, + typed_config->tlsSlotAllocator()); weight_update_timer_ = typed_config->mainThreadDispatcher().createTimer( [this]() { computeLocalityRoutingWeights(); }); @@ -79,7 +86,12 @@ void LoadAwareLocalityLoadBalancer::addLbPolicyDataToHosts(const Upstream::HostV } absl::Status LoadAwareLocalityLoadBalancer::initialize() { - RETURN_IF_NOT_OK(factory_->initializeChildLb()); + if (child_thread_aware_lb_ == nullptr) { + return absl::InvalidArgumentError(absl::StrCat( + "Unsupported endpoint picking policy for load_aware_locality: ", child_factory_name_, + ". Child load balancer could not be instantiated per locality.")); + } + RETURN_IF_NOT_OK(child_thread_aware_lb_->initialize()); for (const auto& host_set : priority_set_.hostSetsPerPriority()) { addLbPolicyDataToHosts(host_set->hosts()); @@ -330,45 +342,32 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( // --- WorkerLocalLbFactory --- WorkerLocalLbFactory::WorkerLocalLbFactory( - Upstream::TypedLoadBalancerFactory& child_factory, std::string child_factory_name, + Upstream::LoadBalancerFactorySharedPtr child_worker_factory, LoadBalancerConfigSharedPtr child_config, const Upstream::ClusterInfo& cluster_info, - const Upstream::PrioritySet& cluster_priority_set, Runtime::Loader& runtime, - Envoy::Random::RandomGenerator& random, TimeSource& time_source, + Runtime::Loader& runtime, Envoy::Random::RandomGenerator& random, ThreadLocal::SlotAllocator& tls_slot_allocator) - : child_factory_name_(std::move(child_factory_name)), child_config_(std::move(child_config)), - cluster_info_(cluster_info), random_(random), runtime_(runtime), - healthy_panic_threshold_(PROTOBUF_PERCENT_TO_ROUNDED_INTEGER_OR_DEFAULT( - cluster_info.lbConfig(), healthy_panic_threshold, 100, 50)), + : child_worker_factory_(std::move(child_worker_factory)), + child_config_(std::move(child_config)), cluster_info_(cluster_info), random_(random), + runtime_(runtime), healthy_panic_threshold_(PROTOBUF_PERCENT_TO_ROUNDED_INTEGER_OR_DEFAULT( + cluster_info.lbConfig(), healthy_panic_threshold, 100, 50)), fail_traffic_on_panic_( cluster_info.lbConfig().zone_aware_lb_config().fail_traffic_on_panic()) { - auto child_config_ref = - makeOptRefFromPtr(child_config_.get()); - child_thread_aware_lb_ = child_factory.create( - child_config_ref, cluster_info_, cluster_priority_set, runtime, random_, time_source); tls_ = ThreadLocal::TypedSlot::makeUnique(tls_slot_allocator); tls_->set([](Event::Dispatcher&) { return std::make_shared(); }); } -absl::Status WorkerLocalLbFactory::initializeChildLb() { - if (child_thread_aware_lb_ == nullptr) { - return absl::InvalidArgumentError(absl::StrCat( - "Unsupported endpoint picking policy for load_aware_locality: ", child_factory_name_, - ". Child load balancer could not be instantiated per locality.")); - } - return child_thread_aware_lb_->initialize(); -} - Upstream::LoadBalancerPtr WorkerLocalLbFactory::createWorkerChildLb(Upstream::PrioritySetImpl& per_locality_priority_set) { - // initializeChildLb() must have been called on the main thread before workers call this. - ASSERT(child_thread_aware_lb_ != nullptr); + // Non-null is guaranteed: the main-thread LB fails initialize() when the child policy could not + // be instantiated, so the cluster never reaches workers. + ASSERT(child_worker_factory_ != nullptr); Upstream::LoadBalancerParams child_params{per_locality_priority_set, nullptr}; - return child_thread_aware_lb_->factory()->create(child_params); + return child_worker_factory_->create(child_params); } bool WorkerLocalLbFactory::recreateChildOnHostChange() const { - ASSERT(child_thread_aware_lb_ != nullptr); - return child_thread_aware_lb_->factory()->recreateOnHostChangeDeprecated(); + ASSERT(child_worker_factory_ != nullptr); + return child_worker_factory_->recreateOnHostChangeDeprecated(); } Upstream::LoadBalancerPtr WorkerLocalLbFactory::create(Upstream::LoadBalancerParams params) { @@ -742,9 +741,13 @@ size_t WorkerLocalLb::chooseLocality(const RoutingWeightsSnapshot* snapshot, uin return it != weights.end() ? it->second : 0.0; }; + // Materialize weights once so the cumulative walk below re-reads the stack copy instead of + // re-hashing the identity map. + absl::InlinedVector locality_weights(per_locality.size()); double effective_total = 0.0; for (size_t i = 0; i < per_locality.size(); ++i) { - effective_total += weight_at(i); + locality_weights[i] = weight_at(i); + effective_total += locality_weights[i]; } if (effective_total <= 0.0) { return 0; @@ -757,7 +760,7 @@ size_t WorkerLocalLb::chooseLocality(const RoutingWeightsSnapshot* snapshot, uin effective_total; double cumulative = 0.0; for (size_t i = 0; i < per_locality.size(); ++i) { - cumulative += weight_at(i); + cumulative += locality_weights[i]; if (target < cumulative) { return i; } diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h index 3497d86a83765..a883bc1204d27 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -179,8 +179,8 @@ struct PriorityRoutingWeights { struct SourceWeights { // Per-locality routing weights (host_count * headroom), keyed by Locality identity. LocalityWeightsMap weights; - // Informational flag: true if local preference was triggered (variance below threshold). - // Not consulted on the hot path — the routing decision is fully encoded in weights. + // True if local preference was triggered (variance below threshold). Does not affect routing + // (fully encoded in weights); read per pick only to choose the zone-routing stat counter. // With remote_probe_fraction > 0, weights still include a small remote share even when true. bool all_local{false}; }; @@ -219,26 +219,21 @@ struct ThreadLocalShim : public ThreadLocal::ThreadLocalObject { class WorkerLocalLb; /** - * Factory shared across workers. Owns the shared child ThreadAwareLoadBalancer and publishes + * Factory shared across workers. Holds the child policy's worker factory and publishes * routing-weight snapshots to workers via a thread-local slot; workers create WorkerLocalLbs. */ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { public: - WorkerLocalLbFactory(Upstream::TypedLoadBalancerFactory& child_factory, - std::string child_factory_name, LoadBalancerConfigSharedPtr child_config, - const Upstream::ClusterInfo& cluster_info, - const Upstream::PrioritySet& cluster_priority_set, Runtime::Loader& runtime, - Envoy::Random::RandomGenerator& random, TimeSource& time_source, + WorkerLocalLbFactory(Upstream::LoadBalancerFactorySharedPtr child_worker_factory, + LoadBalancerConfigSharedPtr child_config, + const Upstream::ClusterInfo& cluster_info, Runtime::Loader& runtime, + Envoy::Random::RandomGenerator& random, ThreadLocal::SlotAllocator& tls_slot_allocator); // Upstream::LoadBalancerFactory Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams params) override; bool recreateOnHostChangeDeprecated() const override { return false; } - // Called by the main thread (from LoadAwareLocalityLoadBalancer::initialize()) to initialize - // the shared child ThreadAwareLoadBalancer. Must be called before any worker creates a LB. - absl::Status initializeChildLb(); - // Called by the main thread to publish new routing weights via TLS. void updateRoutingWeights(RoutingWeightsSnapshotConstSharedPtr snapshot) { tls_->runOnAllThreads([snapshot](OptRef shim) { @@ -257,8 +252,7 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { return shim.has_value() ? shim->routing_weights.get() : nullptr; } - // Called by workers to create a per-locality worker LB from the shared child factory. - // Must only be called after initializeChildLb() has returned on the main thread. + // Called by workers to create a per-locality worker LB from the shared child worker factory. Upstream::LoadBalancerPtr createWorkerChildLb(Upstream::PrioritySetImpl& per_locality_priority_set); @@ -272,7 +266,10 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { Upstream::ClusterLbStats& lbStats() const { return cluster_info_.lbStats(); } private: - std::string child_factory_name_; + // Worker factory of the child ThreadAwareLoadBalancer (owned by the main-thread LB). Shared + // ownership keeps it valid on workers after the child LB is destroyed, per the + // ThreadAwareLoadBalancer factory contract. + Upstream::LoadBalancerFactorySharedPtr child_worker_factory_; LoadBalancerConfigSharedPtr child_config_; const Upstream::ClusterInfo& cluster_info_; Envoy::Random::RandomGenerator& random_; @@ -280,10 +277,6 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { uint32_t healthy_panic_threshold_; const bool fail_traffic_on_panic_; - // Single child ThreadAwareLoadBalancer created and initialized on the main thread. - // Workers only call factory()->create() from it. - Upstream::ThreadAwareLoadBalancerPtr child_thread_aware_lb_; - std::unique_ptr> tls_; }; @@ -442,6 +435,10 @@ class LoadAwareLocalityLoadBalancer : public Upstream::ThreadAwareLoadBalancer, std::array, 3> ewma_state_; Event::TimerPtr weight_update_timer_; std::chrono::milliseconds weight_update_period_; + std::string child_factory_name_; + // Child ThreadAwareLoadBalancer owned here so any priority-set callback handles it registers + // are created and destroyed on the main thread; workers hold only its worker factory. + Upstream::ThreadAwareLoadBalancerPtr child_thread_aware_lb_; std::shared_ptr factory_; Envoy::Common::CallbackHandlePtr priority_update_cb_; }; diff --git a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc index 42c4a3955ad79..827eed7987f07 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc @@ -175,6 +175,13 @@ TEST(LoadAwareLocalityConfigTest, InvalidWeightUpdatePeriod) { EXPECT_FALSE(result.ok()); EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); EXPECT_THAT(result.status().message(), testing::HasSubstr("weight_update_period")); + + // Below the 100ms floor: PGV does not run on unpacked LB configs, so loadConfig enforces it. + config_proto.mutable_weight_update_period()->set_nanos(50000000); + result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("at least 100ms")); } TEST(LoadAwareLocalityConfigTest, ZeroSmoothingTimeConstantRejected) { diff --git a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc index e7823508be4ae..83f2b1d0763ed 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -126,18 +126,6 @@ class AsyncSelectingChildFactory : public Upstream::LoadBalancerFactory { bool& cancelled_; }; -class StaticThreadAwareLoadBalancer : public Upstream::ThreadAwareLoadBalancer { -public: - explicit StaticThreadAwareLoadBalancer(Upstream::LoadBalancerFactorySharedPtr factory) - : factory_(std::move(factory)) {} - - Upstream::LoadBalancerFactorySharedPtr factory() override { return factory_; } - absl::Status initialize() override { return absl::OkStatus(); } - -private: - Upstream::LoadBalancerFactorySharedPtr factory_; -}; - uint64_t counterValue(NiceMock& info, const std::string& name) { auto c = info.stats_store_.findCounterByString(name); return c.has_value() ? c->get().value() : 0; @@ -221,34 +209,23 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes auto round_robin_lb_config = round_robin_factory.loadConfig(context_, *round_robin_proto).value(); - auto factory = std::make_shared( - round_robin_factory, round_robin_factory.name(), - LoadBalancerConfigSharedPtr(std::move(round_robin_lb_config)), cluster_info_, priority_set_, - context_.runtime_loader_, random_, context_.time_system_, context_.thread_local_); - EXPECT_TRUE(factory->initializeChildLb().ok()); - return factory; + auto child_config = LoadBalancerConfigSharedPtr(std::move(round_robin_lb_config)); + retained_child_talb_ = + round_robin_factory.create(*child_config, cluster_info_, priority_set_, + context_.runtime_loader_, random_, context_.time_system_); + EXPECT_TRUE(retained_child_talb_->initialize().ok()); + return std::make_shared(retained_child_talb_->factory(), child_config, + cluster_info_, context_.runtime_loader_, random_, + context_.thread_local_); } - // Builds a WorkerLocalLbFactory whose child create() yields a StaticThreadAwareLoadBalancer - // wrapping `child`. + // Builds a WorkerLocalLbFactory using `child` directly as the shared child worker factory. std::shared_ptr - makeWorkerFactoryWithChild(Upstream::LoadBalancerFactorySharedPtr child, - const std::string& name) { - auto typed_child_factory = std::make_shared>(); - ON_CALL(*typed_child_factory, name()).WillByDefault(Return(name)); - ON_CALL(*typed_child_factory, - create(testing::_, testing::_, testing::_, testing::_, testing::_, testing::_)) - .WillByDefault( - [child](auto&&...) { return std::make_unique(child); }); - recording_typed_child_factory_ = typed_child_factory; - - auto factory = std::make_shared( - *typed_child_factory, name, + makeWorkerFactoryWithChild(Upstream::LoadBalancerFactorySharedPtr child) { + return std::make_shared( + std::move(child), std::make_shared(), - cluster_info_, priority_set_, context_.runtime_loader_, random_, context_.time_system_, - context_.thread_local_); - EXPECT_TRUE(factory->initializeChildLb().ok()); - return factory; + cluster_info_, context_.runtime_loader_, random_, context_.thread_local_); } void setupPriorityLocalities( @@ -531,7 +508,9 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes Upstream::ThreadAwareLoadBalancerPtr thread_aware_lb_; Upstream::LoadBalancerFactorySharedPtr factory_; - std::shared_ptr> recording_typed_child_factory_; + // Child ThreadAwareLoadBalancer backing createRoundRobinWorkerFactory (owned by the main-thread + // LB in production). + Upstream::ThreadAwareLoadBalancerPtr retained_child_talb_; }; // Table-driven single-tick weight-computation cases: set up localities, push one uniform ORCA @@ -660,7 +639,7 @@ TEST_F(LoadAwareLocalityLbTest, EmptyPrioritySetHasNoAvailableHost) { host_set->healthy_hosts_.clear(); auto child_factory = std::make_shared(false); - auto factory = makeWorkerFactoryWithChild(child_factory, "mock_empty_priority"); + auto factory = makeWorkerFactoryWithChild(child_factory); auto worker_lb = factory->create({priority_set_, nullptr}); ASSERT_NE(nullptr, worker_lb); @@ -1629,7 +1608,7 @@ TEST_F(LoadAwareLocalityLbTest, AsyncChildSelectionIsCancelledAndFailsSynchronou bool cancelled = false; auto child_factory = std::make_shared(cancelled); - auto factory = makeWorkerFactoryWithChild(child_factory, "mock_async_child"); + auto factory = makeWorkerFactoryWithChild(child_factory); auto worker_lb = factory->create({priority_set_, nullptr}); ASSERT_NE(nullptr, worker_lb); @@ -1798,7 +1777,7 @@ TEST_F(LoadAwareLocalityLbTest, ChildRecreatedWhenUnderlyingFactoryRequiresIt) { setupLocalities({{local}, {remote}}); auto child_factory = std::make_shared(true); - auto factory = makeWorkerFactoryWithChild(child_factory, "mock_recreate_child"); + auto factory = makeWorkerFactoryWithChild(child_factory); auto worker_lb = factory->create({priority_set_, nullptr}); ASSERT_NE(nullptr, worker_lb); From 8acb258cd1e12512de58e251173ac3445dfdc166 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 17:43:42 -0600 Subject: [PATCH 09/16] absl --> std Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality/load_aware_locality_lb.cc | 2 +- .../load_aware_locality_lb_test.cc | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc index 51f584ff079da..8187131fb60cd 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -487,7 +487,7 @@ void WorkerLocalLb::updateLocalityHosts(PerSourceLocalityState& state, std::make_shared(), Upstream::HostsPerLocalityImpl::empty()); state.priority_set->updateHosts(0, std::move(update_params), nullptr, hosts_added, hosts_removed, - absl::nullopt, absl::nullopt); + std::nullopt, std::nullopt); } void WorkerLocalLb::buildPerPriorityLocalities() { diff --git a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc index 83f2b1d0763ed..954c27fdee499 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -99,10 +99,10 @@ class AsyncSelectingChildLb : public Upstream::LoadBalancer { OptRef lifetimeCallbacks() override { return {}; } - absl::optional + std::nullopt selectExistingConnection(Upstream::LoadBalancerContext*, const Upstream::Host&, std::vector&) override { - return absl::nullopt; + return std::nullopt; } private: @@ -231,8 +231,8 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes void setupPriorityLocalities( uint32_t priority, std::vector localities, bool has_local_locality = false, - absl::optional> healthy_localities = absl::nullopt, - absl::optional> degraded_localities = absl::nullopt) { + std::nullopt> healthy_localities = std::nullopt, + std::nullopt> degraded_localities = std::nullopt) { auto* host_set = priority_set_.getMockHostSet(priority); // Stamp each host with its locality-group identity so the policy can key weights by Locality. for (size_t i = 0; i < localities.size(); ++i) { @@ -278,8 +278,8 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes void setupLocalities( std::vector localities, bool has_local_locality = false, - absl::optional> healthy_localities = absl::nullopt, - absl::optional> degraded_localities = absl::nullopt) { + std::nullopt> healthy_localities = std::nullopt, + std::nullopt> degraded_localities = std::nullopt) { setupPriorityLocalities(0, std::move(localities), has_local_locality, std::move(healthy_localities), std::move(degraded_localities)); } @@ -322,7 +322,7 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes void reshapeLocalities( uint32_t priority, std::vector localities, Upstream::HostVector added = {}, Upstream::HostVector removed = {}, bool has_local = false, - absl::optional> healthy_override = absl::nullopt) { + std::nullopt> healthy_override = std::nullopt) { auto* host_set = priority_set_.getMockHostSet(priority); if (healthy_override.has_value()) { Upstream::HostVector healthy_hosts; From 2523af0b06706695df597b8b6a7985122c522993 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 18:30:20 -0600 Subject: [PATCH 10/16] Fix std typos Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality_lb_test.cc | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc index 954c27fdee499..1ebac2f887cf6 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -99,7 +100,7 @@ class AsyncSelectingChildLb : public Upstream::LoadBalancer { OptRef lifetimeCallbacks() override { return {}; } - std::nullopt + std::optional selectExistingConnection(Upstream::LoadBalancerContext*, const Upstream::Host&, std::vector&) override { return std::nullopt; @@ -231,8 +232,8 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes void setupPriorityLocalities( uint32_t priority, std::vector localities, bool has_local_locality = false, - std::nullopt> healthy_localities = std::nullopt, - std::nullopt> degraded_localities = std::nullopt) { + std::optional> healthy_localities = std::nullopt, + std::optional> degraded_localities = std::nullopt) { auto* host_set = priority_set_.getMockHostSet(priority); // Stamp each host with its locality-group identity so the policy can key weights by Locality. for (size_t i = 0; i < localities.size(); ++i) { @@ -278,8 +279,8 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes void setupLocalities( std::vector localities, bool has_local_locality = false, - std::nullopt> healthy_localities = std::nullopt, - std::nullopt> degraded_localities = std::nullopt) { + std::optional> healthy_localities = std::nullopt, + std::optional> degraded_localities = std::nullopt) { setupPriorityLocalities(0, std::move(localities), has_local_locality, std::move(healthy_localities), std::move(degraded_localities)); } @@ -322,7 +323,7 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes void reshapeLocalities( uint32_t priority, std::vector localities, Upstream::HostVector added = {}, Upstream::HostVector removed = {}, bool has_local = false, - std::nullopt> healthy_override = std::nullopt) { + std::optional> healthy_override = std::nullopt) { auto* host_set = priority_set_.getMockHostSet(priority); if (healthy_override.has_value()) { Upstream::HostVector healthy_hosts; From 71ace56fa2cd0db995352d83d73f6b10c5bd10d9 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 19:06:55 -0600 Subject: [PATCH 11/16] trim tests, comments, and use peekAnotherHost/chooseHost helper Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality/BUILD | 2 - .../load_aware_locality_lb.cc | 206 ++--- .../load_aware_locality_lb.h | 129 +-- .../load_aware_locality/config_test.cc | 166 ++-- .../load_aware_locality_lb_test.cc | 865 ++++-------------- 5 files changed, 343 insertions(+), 1025 deletions(-) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/BUILD b/source/extensions/load_balancing_policies/load_aware_locality/BUILD index 713cdbc82cfb2..e114f1ffad8f5 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/BUILD +++ b/source/extensions/load_balancing_policies/load_aware_locality/BUILD @@ -15,10 +15,8 @@ envoy_cc_extension( hdrs = ["config.h"], deps = [ ":load_aware_locality_lb_lib", - "//source/common/common:minimal_logger_lib", "//source/common/config:utility_lib", "//source/extensions/load_balancing_policies/common:factory_base", - "@abseil-cpp//absl/strings", "@envoy_api//envoy/extensions/load_balancing_policies/load_aware_locality/v3:pkg_cc_proto", ], ) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc index 8187131fb60cd..c9875365edde5 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -26,9 +26,7 @@ absl::Status LocalityLbHostData::onOrcaLoadReport(const Upstream::OrcaLoadReport const StreamInfo::StreamInfo&) { const double util = Common::OrcaLoadReportHandler::getUtilizationFromOrcaReport(report, *metric_names_); - // A report without a positive utilization signal is not a valid sample (same freshness rule as - // gRFC A58 / CSWRR): don't store it or refresh freshness, so signal-less hosts age out via - // weight_expiration_period instead of pinning their locality as fresh-and-idle. + // Signal-less reports do not refresh freshness; they age out instead of pinning idle. if (util <= 0.0) { return absl::OkStatus(); } @@ -107,7 +105,6 @@ absl::Status LoadAwareLocalityLoadBalancer::initialize() { } void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { - // Re-arm first so the timer always fires on schedule regardless of early returns below. weight_update_timer_->enableTimer(weight_update_period_); auto snapshot = std::make_shared(); @@ -118,16 +115,14 @@ void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { per_source_state.resize(host_sets.size()); } - // Current monotonic time in milliseconds, used for weight expiration checks. const int64_t now_ms = std::chrono::duration_cast( time_source_.monotonicTime().time_since_epoch()) .count(); - // Tick-scoped stat accumulators: OR'd across every (source × priority) pass below. bool tick_all_overloaded = false; bool tick_local_preferred = false; bool tick_probe_active = false; - uint32_t tick_stale_localities = 0; // summed only from the all_hosts pass + uint32_t tick_stale_localities = 0; for (size_t priority = 0; priority < host_sets.size(); ++priority) { const auto& host_set = host_sets[priority]; @@ -155,7 +150,6 @@ void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { run_source(SelectionSource::Healthy, host_set->healthyHostsPerLocality().get()); const auto degraded_result = run_source(SelectionSource::Degraded, host_set->degradedHostsPerLocality().get()); - // all_hosts pass: run last so its staleness reflects the canonical "zero fresh hosts". const auto all_hosts_result = run_source(SelectionSource::AllHosts, locality_hosts); tick_all_overloaded |= healthy_result.all_overloaded || degraded_result.all_overloaded || @@ -164,7 +158,6 @@ void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { all_hosts_result.local_preferred; tick_probe_active |= healthy_result.probe_active || degraded_result.probe_active || all_hosts_result.probe_active; - // Count stale localities from the all_hosts pass only (canonical; avoids triple-counting). tick_stale_localities += all_hosts_result.stale_localities; } @@ -194,7 +187,6 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( SourceComputeResult result; const auto& locality_hosts = all_hosts_per_locality.get(); const size_t locality_count = locality_hosts.size(); - // Compute weights index-based (local locality is index 0), then key by identity at the end. std::vector weights(locality_count, 0.0); all_local = false; @@ -234,8 +226,6 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( valid_counts[i] = valid_count; } - // Advance the identity-keyed EWMA state. Rebuilding the map each tick carries state for - // localities that persist, applies the first valid sample raw, and drops departed localities. LocalityEwmaMap new_state; new_state.reserve(locality_count); std::vector utilizations(locality_count, 0.0); @@ -252,15 +242,11 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( new_state[locality_hosts[i][0]->locality()] = smoothed; utilizations[i] = smoothed; } else if (prev != ewma_state.end()) { - // Stale: had valid data before, none now. Carry the prior smoothed value so a locality - // whose load we no longer know is not mistaken for idle; host-count baseline below. new_state[locality_hosts[i][0]->locality()] = prev->second; utilizations[i] = prev->second; stale[i] = true; ++result.stale_localities; } - // else: cold (never had a valid sample) — utilization stays 0 (full headroom), so the - // weight below reduces to the host-count baseline without counting the stale stat. } ewma_state = std::move(new_state); @@ -274,22 +260,16 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( const auto set_all_local = [&weights, &all_local]() { all_local = true; std::fill(weights.begin(), weights.end(), 0.0); - if (!weights.empty()) { - weights[0] = 1.0; - } + weights[0] = 1.0; }; - // All-overloaded fallback, checked first per the rst: when no locality has headroom, spread by - // host count and skip local preference / probe. Otherwise apply stage-4 local preference and - // stage-5 remote probing. const double total_base_weight = std::accumulate(weights.begin(), weights.end(), 0.0); if (total_base_weight == 0.0 && total_hosts > 0) { for (size_t i = 0; i < locality_count; ++i) { weights[i] = static_cast(host_counts[i]); } result.all_overloaded = true; - } else { - // Remote host count, shared by the local-preference target and the probe redistribution below. + } else if (total_base_weight > 0.0) { uint32_t remote_hosts = 0; if (all_hosts_per_locality.hasLocalLocality()) { double remote_util_sum = 0.0; @@ -298,14 +278,12 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( remote_hosts += host_counts[i]; } - if (total_hosts > 0 && !host_counts.empty() && host_counts[0] > 0 && remote_hosts > 0) { + if (host_counts[0] > 0 && remote_hosts > 0) { const double target_util = remote_util_sum / remote_hosts; if (utilizations[0] <= target_util + utilization_variance_threshold_) { set_all_local(); result.local_preferred = true; } - } else if (total_hosts == 0) { - set_all_local(); } } @@ -313,7 +291,7 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( const double total = std::accumulate(weights.begin(), weights.end(), 0.0); const double remote_target = total * remote_probe_fraction_; const double remote_sum = total - weights[0]; - if (total > 0.0 && remote_sum < remote_target) { + if (remote_sum < remote_target) { const double take_from_local = std::min(remote_target - remote_sum, weights[0]); weights[0] -= take_from_local; for (size_t i = 1; i < locality_count; ++i) { @@ -324,9 +302,6 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( } } - // Key the computed weights by Locality identity. An empty locality group carries no entry - // (it has 0 weight anyway); the local locality is stored under its own identity like the - // rest (local-vs-remote preference is already baked into the weight values). weights_map.clear(); weights_map.reserve(locality_count); for (size_t i = 0; i < locality_count; ++i) { @@ -376,13 +351,10 @@ Upstream::LoadBalancerPtr WorkerLocalLbFactory::create(Upstream::LoadBalancerPar namespace { -// Context handed to per-locality child LBs. Overrides only determinePriorityLoad, which returns the -// child's own default: the parent already consulted the real context's (possibly stateful) -// retry-priority plugin against the cluster-wide priority set; re-invoking it against the -// single-priority per-locality set would corrupt its state. Everything else (hash key, retry hooks -// shouldSelectAnotherHost/hostSelectionRetryCount, metadata, headers) passes through to the wrapped -// context; endpoint retry is owned by the child within the chosen locality. Stack-allocated per -// pick; children must not retain it beyond the call (same constraint as the subset LB wrapper). +// Context for per-locality child LBs. Priority selection stays fixed to the child's default because +// the parent already ran retry-priority against cluster-wide priorities; rerunning it against the +// single-priority child set can corrupt plugin state. Other methods pass through; children must not +// retain this stack object. class ChildLoadBalancerContext : public Upstream::LoadBalancerContext { public: explicit ChildLoadBalancerContext(Upstream::LoadBalancerContext& wrapped) : wrapped_(wrapped) {} @@ -430,10 +402,8 @@ class ChildLoadBalancerContext : public Upstream::LoadBalancerContext { Upstream::LoadBalancerContext& wrapped_; }; -// Async child host selection is unsupported: the child would retain the stack-scoped -// ChildLoadBalancerContext and call onAsyncHostSelection after chooseHost returns (use-after-free). -// Cancel any async handle (HostSelectionResponse guarantees no callback after cancel) and -// fail synchronously, matching the subset LB. Synchronous responses pass through unchanged. +// Async child selection would retain ChildLoadBalancerContext after chooseHost returns. Cancel any +// async handle and fail synchronously, matching the subset LB. Upstream::HostSelectionResponse failOnAsyncSelection(Upstream::HostSelectionResponse response) { if (response.cancelable != nullptr) { response.cancelable->cancel(); @@ -456,9 +426,7 @@ WorkerLocalLb::WorkerLocalLb(WorkerLocalLbFactory& factory, priority_sync_cb_ = priority_set_.addPriorityUpdateCb( [this](uint32_t priority, const Upstream::HostVector& hosts_added, const Upstream::HostVector& hosts_removed) { - // A membership delta permits a topology rebuild. An empty-delta update still matters for - // child policies (in-place host attribute changes such as health, weight, or metadata) but - // never rebuilds — a real topology change always arrives with a delta. + // Empty-delta updates refresh child host attributes without rebuilding topology. syncPriority(priority, /*allow_rebuild=*/!hosts_added.empty() || !hosts_removed.empty()); }); } @@ -475,10 +443,7 @@ void WorkerLocalLb::updateLocalityHosts(PerSourceLocalityState& state, const Upstream::HostVector& hosts_removed) { auto hosts_shared = std::make_shared(hosts); auto per_locality = std::make_shared(hosts, is_local); - // All passed-in hosts are marked as "healthy" in the child priority set regardless of their - // actual health status. The caller (syncLocalityState) already partitions hosts by source - // (healthy/degraded/all), so this child LB only sees hosts that can be selected. Marking - // them all healthy ensures the child policy (e.g. RoundRobin) considers every host eligible. + // The caller has already partitioned by source, so every host passed to the child is eligible. auto healthy_hosts = std::make_shared(hosts); auto update_params = Upstream::HostSetImpl::updateHostsParams( hosts_shared, per_locality, healthy_hosts, per_locality, @@ -515,7 +480,7 @@ void WorkerLocalLb::syncLocalityState(PerLocalityState& state, const Upstream::H ? degraded_localities[locality_index] : empty_hosts; - // Record this locality's identity for chooseLocality's weight lookup. Empty group → no identity. + // Record this locality's identity for chooseLocality's weight lookup. Empty group has none. state.locality = all_hosts.empty() ? envoy::config::core::v3::Locality() : all_hosts[0]->locality(); @@ -571,7 +536,7 @@ void WorkerLocalLb::syncLocalityState(PerLocalityState& state, const Upstream::H } void WorkerLocalLb::buildPerLocality(uint32_t priority, const Upstream::HostSet& host_set) { - // Locality routing structures for this priority are (re)generated on this worker — the same + // Locality routing structures for this priority are (re)generated on this worker, the same // event the zone-aware counter tracks. stats_.lb_recalculate_zone_structures_.inc(); const auto& locality_hosts = host_set.hostsPerLocality().get(); @@ -590,7 +555,7 @@ void WorkerLocalLb::syncPriority(uint32_t priority, bool allow_rebuild) { return; } - // Priority count changed — full rebuild when permitted, else wait for a membership-delta update. + // Priority count changed; full rebuild when permitted, else wait for a membership-delta update. if (host_sets.size() != per_priority_locality_.size()) { if (allow_rebuild) { buildPerPriorityLocalities(); @@ -602,7 +567,7 @@ void WorkerLocalLb::syncPriority(uint32_t priority, bool allow_rebuild) { const auto& locality_hosts = host_set->hostsPerLocality().get(); auto& per_locality = per_priority_locality_[priority].localities; - // Topology change (locality added/removed) — rebuild this priority when permitted, else wait. + // Topology changed; rebuild this priority when permitted, else wait. if (locality_hosts.size() != per_locality.size()) { if (allow_rebuild) { buildPerLocality(priority, *host_set); @@ -625,7 +590,7 @@ WorkerLocalLb::pickLocalityLb(const std::vector& per_locality, actual_idx = preferred_idx; return lb; } - // Stale routing snapshot — the preferred locality's child LB was torn down after a host + // Stale routing snapshot: the preferred locality's child LB was torn down after a host // change but before the routing weights were recomputed. Scan for any locality with a // usable LB. for (size_t i = 0; i < per_locality.size(); ++i) { @@ -662,61 +627,76 @@ WorkerLocalLb::resolvePrioritySource(Upstream::LoadBalancerContext* context, boo return {priority, source, /*in_panic=*/false, /*fail=*/false}; } -Upstream::HostSelectionResponse WorkerLocalLb::chooseHost(Upstream::LoadBalancerContext* context) { - const auto pick = resolvePrioritySource(context, /*peeking=*/false); - // Panic stat is counted here only; peekAnotherHost deliberately does not double-count it. - if (pick.in_panic) { - stats_.lb_healthy_panic_.inc(); - } +WorkerLocalLb::LocalityLbSelection WorkerLocalLb::selectLocalityLb(const PrioritySourcePick& pick, + bool peeking) { + LocalityLbSelection selection; if (pick.fail) { - return {nullptr}; + return selection; } - const uint32_t priority = pick.priority; - const PriorityRoutingWeights::SelectionSource source = pick.source; + + selection.priority = pick.priority; + selection.source = pick.source; // A custom retry-priority plugin can route to a just-created priority that // per_priority_locality_ doesn't cover yet (empty-delta update, rebuild pending). - if (priority >= per_priority_locality_.size()) { - return {nullptr}; + if (selection.priority >= per_priority_locality_.size()) { + return selection; } - auto& per_locality = per_priority_locality_[priority].localities; + auto& per_locality = per_priority_locality_[selection.priority].localities; if (per_locality.empty()) { - return {nullptr}; + return selection; } - // Fetch the routing snapshot once and reuse it for locality selection and the zone-routing stat - // below (avoids a second TLS lookup per request). - const auto* snapshot = factory_.routingWeights(); - const size_t preferred_idx = chooseLocality(snapshot, priority, source, /*peeking=*/false); - size_t locality_idx = preferred_idx; - auto* lb = pickLocalityLb(per_locality, source, preferred_idx, locality_idx); - if (lb == nullptr) { - return {nullptr}; + selection.snapshot = factory_.routingWeights(); + const size_t preferred_idx = + chooseLocality(selection.snapshot, selection.priority, selection.source, peeking); + selection.locality_idx = preferred_idx; + selection.lb = + pickLocalityLb(per_locality, selection.source, preferred_idx, selection.locality_idx); + return selection; +} + +void WorkerLocalLb::recordZoneRoutingStats(const PrioritySourcePick& pick, + const LocalityLbSelection& selection) { + if (pick.in_panic || selection.snapshot == nullptr || + selection.priority >= selection.snapshot->priority_weights.size()) { + return; } - // Endpoint retry is owned by the child LB; delegate through the child context and return its - // response (failOnAsyncSelection rejects async children to protect the stack-scoped context). - // Record the zone-routing stat once for the selected locality. - if (!pick.in_panic && snapshot != nullptr && priority < snapshot->priority_weights.size()) { - const auto& priority_snapshot = snapshot->priority_weights[priority]; - if (priority_snapshot.has_local_locality) { - if (locality_idx == 0) { - if (priority_snapshot.allLocalFor(source)) { - stats_.lb_zone_routing_all_directly_.inc(); - } else { - stats_.lb_zone_routing_sampled_.inc(); - } - } else { - stats_.lb_zone_routing_cross_zone_.inc(); - } - } + const auto& priority_snapshot = selection.snapshot->priority_weights[selection.priority]; + if (!priority_snapshot.has_local_locality) { + return; + } + if (selection.locality_idx != 0) { + stats_.lb_zone_routing_cross_zone_.inc(); + return; + } + if (priority_snapshot.allLocalFor(selection.source)) { + stats_.lb_zone_routing_all_directly_.inc(); + } else { + stats_.lb_zone_routing_sampled_.inc(); } +} + +Upstream::HostSelectionResponse WorkerLocalLb::chooseHost(Upstream::LoadBalancerContext* context) { + const auto pick = resolvePrioritySource(context, /*peeking=*/false); + // Panic stat is counted here only; peekAnotherHost deliberately does not double-count it. + if (pick.in_panic) { + stats_.lb_healthy_panic_.inc(); + } + + const auto selection = selectLocalityLb(pick, /*peeking=*/false); + if (selection.lb == nullptr) { + return {nullptr}; + } + + recordZoneRoutingStats(pick, selection); if (context == nullptr) { - return failOnAsyncSelection(lb->chooseHost(nullptr)); + return failOnAsyncSelection(selection.lb->chooseHost(nullptr)); } ChildLoadBalancerContext child_context(*context); - return failOnAsyncSelection(lb->chooseHost(&child_context)); + return failOnAsyncSelection(selection.lb->chooseHost(&child_context)); } size_t WorkerLocalLb::chooseLocality(const RoutingWeightsSnapshot* snapshot, uint32_t priority, @@ -733,20 +713,12 @@ size_t WorkerLocalLb::chooseLocality(const RoutingWeightsSnapshot* snapshot, uin return 0; } - // Look up each LIVE locality's advisory weight by identity. A locality present on the worker but - // missing from the snapshot (e.g. just-added, snapshot lagging by ≤1 tick) gets weight 0.0 and is - // excluded from selection until the next recompute — matching the prior count-mismatch behavior. - const auto weight_at = [&weights, &per_locality](size_t i) -> double { - auto it = weights.find(per_locality[i].locality); - return it != weights.end() ? it->second : 0.0; - }; - - // Materialize weights once so the cumulative walk below re-reads the stack copy instead of - // re-hashing the identity map. + // Snapshot-missing live localities get weight 0 until the next recompute. absl::InlinedVector locality_weights(per_locality.size()); double effective_total = 0.0; for (size_t i = 0; i < per_locality.size(); ++i) { - locality_weights[i] = weight_at(i); + const auto weight = weights.find(per_locality[i].locality); + locality_weights[i] = weight != weights.end() ? weight->second : 0.0; effective_total += locality_weights[i]; } if (effective_total <= 0.0) { @@ -777,33 +749,17 @@ WorkerLocalLb::peekAnotherHost(Upstream::LoadBalancerContext* context) { if (Upstream::tooManyPreconnects(stashed_random_.size(), total_healthy_hosts_)) { return nullptr; } - const auto pick = resolvePrioritySource(context, /*peeking=*/true); - if (pick.fail) { - return nullptr; - } - const uint32_t priority = pick.priority; - const PriorityRoutingWeights::SelectionSource source = pick.source; - - if (priority >= per_priority_locality_.size()) { - return nullptr; - } - auto& per_locality = per_priority_locality_[priority].localities; - if (per_locality.empty()) { - return nullptr; - } - const size_t preferred_idx = - chooseLocality(factory_.routingWeights(), priority, source, /*peeking=*/true); - size_t actual_idx = preferred_idx; - auto* lb = pickLocalityLb(per_locality, source, preferred_idx, actual_idx); - if (lb == nullptr) { + const auto selection = selectLocalityLb(resolvePrioritySource(context, /*peeking=*/true), + /*peeking=*/true); + if (selection.lb == nullptr) { return nullptr; } if (context == nullptr) { - return lb->peekAnotherHost(nullptr); + return selection.lb->peekAnotherHost(nullptr); } ChildLoadBalancerContext child_context(*context); - return lb->peekAnotherHost(&child_context); + return selection.lb->peekAnotherHost(&child_context); } } // namespace LoadAwareLocality diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h index a883bc1204d27..53f714bd4b886 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -40,22 +40,11 @@ struct LoadAwareLocalityStats { ALL_LOAD_AWARE_LOCALITY_STATS(GENERATE_COUNTER_STRUCT) }; -/** - * Per-host LB policy data for the load-aware locality policy. - * Receives ORCA load reports via the multi-slot HostLbPolicyData mechanism - * and stores a lock-free utilization value for the main-thread weight computation. - * - * Written by worker threads via onOrcaLoadReport(); read by the main thread - * during locality weight computation. The release/acquire on last_update_time_ms_ - * ensures the reader sees the utilization that was stored before it. - * - * lastUpdateTimeMs() == kNeverReported distinguishes "never reported" from any real - * report time (monotonic 0 ms is a legitimate timestamp, e.g. under simulated time). - */ +// Per-host ORCA data written by workers and read by the main-thread weight computation. +// The release/acquire timestamp store publishes the preceding utilization store. class LocalityLbHostData : public Upstream::HostLbPolicyData { public: - // Sentinel: no report has been stored yet. Out-of-band (int64 min) so a real monotonic - // timestamp of 0 is not mistaken for "never reported". + // Out-of-band sentinel so monotonic time 0 remains a valid report timestamp. static constexpr int64_t kNeverReported = std::numeric_limits::min(); LocalityLbHostData(TimeSource& time_source, @@ -87,9 +76,7 @@ class LocalityLbHostData : public Upstream::HostLbPolicyData { const std::shared_ptr> metric_names_; }; -// Shared ownership wrapper for the child LB config. The config is created during loadConfig() -// and shared between LoadAwareLocalityLbConfig (which is const when accessed) and -// WorkerLocalLbFactory (which needs it for the lifetime of the cluster). +// Shared between config and worker factory; must outlive the child LB. using LoadBalancerConfigSharedPtr = std::shared_ptr; /** @@ -154,41 +141,25 @@ class LoadAwareLocalityLbConfig : public Upstream::LoadBalancerConfig { ThreadLocal::SlotAllocator& tls_slot_allocator_; }; -// Advisory locality weights keyed by Locality identity (not HostsPerLocality position). Keying by -// identity makes the main-thread→worker mapping robust to a locality being added/removed between -// the snapshot tick and the worker's live membership: the lexicographic index of later localities -// shifts on add/remove, but their identity does not. Uses the same comparators as -// load_balancer_impl.cc. +// Key locality state by identity rather than HostsPerLocality position so add/remove index shifts +// cannot mis-map main-thread snapshots to worker-local membership. using LocalityWeightsMap = absl::flat_hash_map; -// Per-locality EWMA-smoothed utilization keyed by Locality identity. An entry exists iff the -// locality has contributed at least one valid sample: a present entry whose locality has zero -// valid hosts this tick is stale (carry the prior value, count the stat); an absent entry is -// cold (host-count baseline, no stat). Identity keying survives the index shifts caused by -// locality add/remove, unlike position-based state. +// EWMA state uses the same identity key: present-with-no-valid-hosts is stale; absent is cold. using LocalityEwmaMap = absl::flat_hash_map; -/** - * Immutable snapshot of per-locality routing weights for a single priority. - */ struct PriorityRoutingWeights { enum class SelectionSource : uint8_t { Healthy = 0, Degraded = 1, AllHosts = 2 }; struct SourceWeights { - // Per-locality routing weights (host_count * headroom), keyed by Locality identity. LocalityWeightsMap weights; - // True if local preference was triggered (variance below threshold). Does not affect routing - // (fully encoded in weights); read per pick only to choose the zone-routing stat counter. - // With remote_probe_fraction > 0, weights still include a small remote share even when true. + // Read per pick only to choose the zone-routing stat counter; routing is encoded in weights. bool all_local{false}; }; - // Per-source routing weights: [Healthy=0, Degraded=1, AllHosts=2]. std::array by_source; - // True when the hosts-per-locality has a local locality (index 0 is "local"). - // Workers use this to pick the correct zone routing stat counter. bool has_local_locality{false}; const LocalityWeightsMap& weightsFor(SelectionSource s) const { @@ -197,31 +168,20 @@ struct PriorityRoutingWeights { bool allLocalFor(SelectionSource s) const { return by_source[static_cast(s)].all_local; } }; -/** - * Immutable snapshot of advisory per-priority locality weights shared between the main thread and - * workers. Priority/health/panic selection is live worker state (LoadBalancerBase), not - * snapshotted. - */ +// Advisory per-priority locality weights. Priority/health/panic selection stays live on workers. struct RoutingWeightsSnapshot { - // Per-priority routing weights, indexed by cluster priority. std::vector priority_weights; }; using RoutingWeightsSnapshotConstSharedPtr = std::shared_ptr; -/** - * Thread-local shim holding the routing weights pointer, pushed from the main thread via TLS. - */ struct ThreadLocalShim : public ThreadLocal::ThreadLocalObject { RoutingWeightsSnapshotConstSharedPtr routing_weights; }; class WorkerLocalLb; -/** - * Factory shared across workers. Holds the child policy's worker factory and publishes - * routing-weight snapshots to workers via a thread-local slot; workers create WorkerLocalLbs. - */ +// Factory shared across workers; publishes routing-weight snapshots via TLS. class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { public: WorkerLocalLbFactory(Upstream::LoadBalancerFactorySharedPtr child_worker_factory, @@ -234,7 +194,6 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams params) override; bool recreateOnHostChangeDeprecated() const override { return false; } - // Called by the main thread to publish new routing weights via TLS. void updateRoutingWeights(RoutingWeightsSnapshotConstSharedPtr snapshot) { tls_->runOnAllThreads([snapshot](OptRef shim) { if (shim.has_value()) { @@ -243,7 +202,6 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { }); } - // Called by workers to get the current routing weights (lock-free TLS read). // SAFETY: Must only be called on a thread that owns a TLS slot instance (worker or main // thread). The returned pointer is valid only for the duration of the current task; do not // store it across yield points or after the TLS slot could be updated. @@ -252,7 +210,6 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { return shim.has_value() ? shim->routing_weights.get() : nullptr; } - // Called by workers to create a per-locality worker LB from the shared child worker factory. Upstream::LoadBalancerPtr createWorkerChildLb(Upstream::PrioritySetImpl& per_locality_priority_set); @@ -280,22 +237,14 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { std::unique_ptr> tls_; }; -/** - * Per-locality state in the worker-local LB. Holds a child PrioritySet and child LB for - * one locality. - */ struct PerSourceLocalityState { - // PrioritySet containing only the hosts that can be selected for one source in one locality. std::unique_ptr priority_set; - // The worker-local LB for this source/locality pair, created from the shared child factory. Upstream::LoadBalancerPtr lb; }; struct PerLocalityState { - // Per-source child LB state: [Healthy=0, Degraded=1, AllHosts=2]. std::array by_source; - // Identity of this worker locality, used to look up its advisory weight in the snapshot map. - // Empty for a locality with no hosts (which carries no weight regardless). + // Empty for a locality with no hosts. envoy::config::core::v3::Locality locality; PerSourceLocalityState& stateFor(PriorityRoutingWeights::SelectionSource s) { @@ -310,11 +259,7 @@ struct PerPriorityLocalityState { std::vector localities; }; -/** - * Worker-local load balancer. Derives LoadBalancerBase for live priority/health/panic selection, - * then selects a locality by capacity-weighted random and delegates to the per-locality child LB - * for endpoint selection. - */ +// Worker-local LB: picks live priority/source, then delegates to a per-locality child LB. class WorkerLocalLb : public Upstream::LoadBalancerBase { public: WorkerLocalLb(WorkerLocalLbFactory& factory, const Upstream::PrioritySet& priority_set); @@ -325,52 +270,47 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { Upstream::HostConstSharedPtr peekAnotherHost(Upstream::LoadBalancerContext* context) override; private: - // Build per-locality child LBs for all priorities. void buildPerPriorityLocalities(); - // Build or rebuild the per-locality child LBs for a single priority. void buildPerLocality(uint32_t priority, const Upstream::HostSet& host_set); - // Re-sync a priority's per-locality child LBs after a priority update. allow_rebuild permits - // tearing down and rebuilding on a topology change (priority/locality count mismatch); when false - // (in-place attribute update with no membership delta) a mismatch is left for the next - // membership-change callback to rebuild. + // allow_rebuild gates topology rebuilds for empty-delta in-place host updates. void syncPriority(uint32_t priority, bool allow_rebuild); - // Update a locality/source PrioritySet with a pre-selected host subset. void updateLocalityHosts(PerSourceLocalityState& state, const Upstream::HostVector& hosts, bool is_local, const Upstream::HostVector& hosts_added, const Upstream::HostVector& hosts_removed); - // Update the per-source child LBs for one locality from the cluster's current host set. void syncLocalityState(PerLocalityState& state, const Upstream::HostSet& host_set, size_t locality_index, bool recreate_child); - // Live priority + selection source for a request, resolved from worker-local LoadBalancerBase - // state. fail=true means the request must be failed (empty priority set, or panic + - // fail_traffic_on_panic). in_panic lets callers count the panic stat themselves (peek does not). struct PrioritySourcePick { uint32_t priority; PriorityRoutingWeights::SelectionSource source; bool in_panic; bool fail; }; - // peeking routes random draws through LoadBalancerBase::random(peeking) so peekAnotherHost and - // the subsequent chooseHost replay the same sequence. + + struct LocalityLbSelection { + Upstream::LoadBalancer* lb{nullptr}; + const RoutingWeightsSnapshot* snapshot{nullptr}; + uint32_t priority{0}; + PriorityRoutingWeights::SelectionSource source{ + PriorityRoutingWeights::SelectionSource::Healthy}; + size_t locality_idx{0}; + }; + PrioritySourcePick resolvePrioritySource(Upstream::LoadBalancerContext* context, bool peeking); - // Choose a locality index within the selected priority/source by weighted-random over the routing - // snapshot. Weights are looked up by identity (see LocalityWeightsMap; missing → 0.0). Returns 0 - // when there is no usable weight (single locality, no/stale snapshot, or zero effective total). - // FUTURE: for very high locality counts, cache a per-worker index-aligned weight vector rebuilt - // on snapshot publish + membership change to restore pure index reads (avoids per-pick hashing). + LocalityLbSelection selectLocalityLb(const PrioritySourcePick& pick, bool peeking); + + void recordZoneRoutingStats(const PrioritySourcePick& pick, const LocalityLbSelection& selection); + + // Returns 0 for single locality, no/stale snapshot, or zero effective total. size_t chooseLocality(const RoutingWeightsSnapshot* snapshot, uint32_t priority, PriorityRoutingWeights::SelectionSource source, bool peeking); - // Find a usable child LB at the preferred locality, falling back to any locality with a - // non-null LB when the routing snapshot is stale (e.g. a locality's hosts were removed but - // the routing weights haven't been recomputed yet). Returns nullptr if no locality has a - // usable LB. Sets actual_idx to the index of the locality whose LB was returned. + // Falls back to any usable child LB when the routing snapshot lags membership. Upstream::LoadBalancer* pickLocalityLb(const std::vector& per_locality, PriorityRoutingWeights::SelectionSource source, size_t preferred_idx, size_t& actual_idx) const; @@ -382,10 +322,7 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { Envoy::Common::CallbackHandlePtr priority_sync_cb_; }; -/** - * Main thread load balancer. Reads host-level utilization data and computes locality routing - * weights on a periodic timer. - */ +// Main-thread LB: computes locality weights and publishes snapshots to workers. class LoadAwareLocalityLoadBalancer : public Upstream::ThreadAwareLoadBalancer, protected Logger::Loggable { public: @@ -400,14 +337,10 @@ class LoadAwareLocalityLoadBalancer : public Upstream::ThreadAwareLoadBalancer, absl::Status initialize() override; private: - // Compute per-locality routing weights from host-level utilization and publish to factory. void computeLocalityRoutingWeights(); - // Attach LocalityLbHostData to hosts that don't already have it. void addLbPolicyDataToHosts(const Upstream::HostVector& hosts); - // Per-source outcome flags from computeSourceWeights, OR'd into the per-tick stats by the caller. - // stale_localities is only consumed from the all-hosts pass (the canonical staleness source). struct SourceComputeResult { bool all_overloaded{false}; bool local_preferred{false}; @@ -415,8 +348,6 @@ class LoadAwareLocalityLoadBalancer : public Upstream::ThreadAwareLoadBalancer, uint32_t stale_localities{0}; }; - // Compute a source's per-locality weights into weights_map/all_local, advancing the caller-owned - // EWMA state. Returns this pass's per-tick stat signals. SourceComputeResult computeSourceWeights(const Upstream::HostsPerLocality& all_hosts_per_locality, const std::vector& eligible_hosts_per_locality, diff --git a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc index 827eed7987f07..b6826cd03f9af 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc @@ -1,4 +1,7 @@ #include +#include +#include +#include #include "envoy/config/core/v3/extension.pb.h" #include "envoy/extensions/load_balancing_policies/round_robin/v3/round_robin.pb.h" @@ -164,71 +167,80 @@ TEST(LoadAwareLocalityConfigTest, FirstSupportedChildWins) { expectFactoryCreateSucceeds(result.value(), context); } -TEST(LoadAwareLocalityConfigTest, InvalidWeightUpdatePeriod) { - NiceMock context; - - auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); - config_proto.mutable_weight_update_period()->set_seconds(0); - - Factory factory; - auto result = factory.loadConfig(context, config_proto); - EXPECT_FALSE(result.ok()); - EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); - EXPECT_THAT(result.status().message(), testing::HasSubstr("weight_update_period")); - - // Below the 100ms floor: PGV does not run on unpacked LB configs, so loadConfig enforces it. - config_proto.mutable_weight_update_period()->set_nanos(50000000); - result = factory.loadConfig(context, config_proto); - EXPECT_FALSE(result.ok()); - EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); - EXPECT_THAT(result.status().message(), testing::HasSubstr("at least 100ms")); -} - -TEST(LoadAwareLocalityConfigTest, ZeroSmoothingTimeConstantRejected) { - NiceMock context; - - auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); - config_proto.mutable_smoothing_time_constant()->set_seconds(0); - - Factory factory; - auto result = factory.loadConfig(context, config_proto); - EXPECT_FALSE(result.ok()); - EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); - EXPECT_THAT(result.status().message(), testing::HasSubstr("smoothing_time_constant")); -} - -TEST(LoadAwareLocalityConfigTest, OutOfRangeVarianceThresholdRejected) { - NiceMock context; - Factory factory; - - auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); - config_proto.mutable_utilization_variance_threshold()->set_value(-0.1); - auto result = factory.loadConfig(context, config_proto); - EXPECT_FALSE(result.ok()); - EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); - EXPECT_THAT(result.status().message(), testing::HasSubstr("utilization_variance_threshold")); - - config_proto.mutable_utilization_variance_threshold()->set_value(1.5); - result = factory.loadConfig(context, config_proto); - EXPECT_FALSE(result.ok()); - EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); -} - -TEST(LoadAwareLocalityConfigTest, OutOfRangeRemoteProbeFractionRejected) { - NiceMock context; - Factory factory; - - auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); - config_proto.mutable_remote_probe_fraction()->set_value(1.0); - auto result = factory.loadConfig(context, config_proto); - EXPECT_FALSE(result.ok()); - EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); - EXPECT_THAT(result.status().message(), testing::HasSubstr("remote_probe_fraction")); - - config_proto.mutable_remote_probe_fraction()->set_value(-0.01); - result = factory.loadConfig(context, config_proto); - EXPECT_FALSE(result.ok()); - EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); +TEST(LoadAwareLocalityConfigTest, InvalidPolicyKnobsRejected) { + struct Case { + std::string name; + std::function mutate; + std::string message; + bool throws_validation_exception{}; + }; + + const std::vector cases = { + {"ZeroWeightUpdatePeriod", + [](LoadAwareLocalityLbProto& config) { + config.mutable_weight_update_period()->set_seconds(0); + }, + "weight_update_period"}, + {"WeightUpdatePeriodBelowFloor", + [](LoadAwareLocalityLbProto& config) { + config.mutable_weight_update_period()->set_nanos(50000000); + }, + "at least 100ms"}, + {"ZeroSmoothingTimeConstant", + [](LoadAwareLocalityLbProto& config) { + config.mutable_smoothing_time_constant()->set_seconds(0); + }, + "smoothing_time_constant"}, + {"VarianceThresholdBelowRange", + [](LoadAwareLocalityLbProto& config) { + config.mutable_utilization_variance_threshold()->set_value(-0.1); + }, + "utilization_variance_threshold"}, + {"VarianceThresholdAboveRange", + [](LoadAwareLocalityLbProto& config) { + config.mutable_utilization_variance_threshold()->set_value(1.5); + }, + "utilization_variance_threshold"}, + {"RemoteProbeFractionAboveRange", + [](LoadAwareLocalityLbProto& config) { + config.mutable_remote_probe_fraction()->set_value(1.0); + }, + "remote_probe_fraction"}, + {"RemoteProbeFractionBelowRange", + [](LoadAwareLocalityLbProto& config) { + config.mutable_remote_probe_fraction()->set_value(-0.01); + }, + "remote_probe_fraction"}, + {"OobLoadReport", + [](LoadAwareLocalityLbProto& config) { + config.mutable_enable_oob_load_report()->set_value(true); + }, + "enable_oob_load_report"}, + {"NegativeWeightExpiration", + [](LoadAwareLocalityLbProto& config) { + config.mutable_weight_expiration_period()->set_seconds(-1); + }, + "", true}, + }; + + for (const Case& c : cases) { + SCOPED_TRACE(c.name); + NiceMock context; + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + c.mutate(config_proto); + + Factory factory; + if (c.throws_validation_exception) { + EXPECT_THROW( + { auto result = factory.loadConfig(context, config_proto); }, ProtoValidationException); + continue; + } + + auto result = factory.loadConfig(context, config_proto); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr(c.message)); + } } TEST(LoadAwareLocalityConfigTest, PolicyValidationPrecedesChildResolution) { @@ -316,30 +328,6 @@ TEST(LoadAwareLocalityConfigTest, MetricNamesReachConfig) { expectFactoryCreateSucceeds(result.value(), context); } -TEST(LoadAwareLocalityConfigTest, OobLoadReportRejected) { - NiceMock context; - - auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); - config_proto.mutable_enable_oob_load_report()->set_value(true); - - Factory factory; - auto result = factory.loadConfig(context, config_proto); - EXPECT_FALSE(result.ok()); - EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); - EXPECT_THAT(result.status().message(), testing::HasSubstr("enable_oob_load_report")); -} - -TEST(LoadAwareLocalityConfigTest, NegativeWeightExpirationRejected) { - NiceMock context; - - auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); - config_proto.mutable_weight_expiration_period()->set_seconds(-1); - - Factory factory; - EXPECT_THROW( - { auto result = factory.loadConfig(context, config_proto); }, ProtoValidationException); -} - } // namespace } // namespace LoadAwareLocality } // namespace LoadBalancingPolicies diff --git a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc index 1ebac2f887cf6..a088ba49d2195 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -134,6 +134,12 @@ uint64_t counterValue(NiceMock& info, const std::stri class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public testing::Test { protected: + struct PanicHosts { + Upstream::HostSharedPtr healthy; + Upstream::HostSharedPtr unhealthy1; + Upstream::HostSharedPtr unhealthy2; + }; + void SetUp() override { // Tests run from the simulated-time epoch (monotonic 0): "never reported" is tracked by an // out-of-band sentinel, so a report stored at time 0 is a valid sample. @@ -152,11 +158,6 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes forced_random_index_ = 0; } - void clearForcedRandomValues() { - forced_random_values_.clear(); - forced_random_index_ = 0; - } - void createLb(double variance_threshold = 0.1, double ewma_alpha = 1.0, double remote_probe_fraction = 0.0, std::chrono::milliseconds weight_expiration_period = std::chrono::milliseconds(0)) { @@ -193,34 +194,6 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes factory_ = thread_aware_lb_->factory(); } - std::shared_ptr createRoundRobinWorkerFactory() { - envoy::extensions::load_balancing_policies::round_robin::v3::RoundRobin round_robin; - envoy::config::core::v3::TypedExtensionConfig round_robin_config; - round_robin_config.set_name("envoy.load_balancing_policies.round_robin"); - std::ignore = round_robin_config.mutable_typed_config()->PackFrom(round_robin); - - auto& round_robin_factory = - Config::Utility::getAndCheckFactory(round_robin_config); - - auto round_robin_proto = round_robin_factory.createEmptyConfigProto(); - EXPECT_TRUE(Config::Utility::translateOpaqueConfig(round_robin_config.typed_config(), - context_.messageValidationVisitor(), - *round_robin_proto) - .ok()); - auto round_robin_lb_config = - round_robin_factory.loadConfig(context_, *round_robin_proto).value(); - - auto child_config = LoadBalancerConfigSharedPtr(std::move(round_robin_lb_config)); - retained_child_talb_ = - round_robin_factory.create(*child_config, cluster_info_, priority_set_, - context_.runtime_loader_, random_, context_.time_system_); - EXPECT_TRUE(retained_child_talb_->initialize().ok()); - return std::make_shared(retained_child_talb_->factory(), child_config, - cluster_info_, context_.runtime_loader_, random_, - context_.thread_local_); - } - - // Builds a WorkerLocalLbFactory using `child` directly as the shared child worker factory. std::shared_ptr makeWorkerFactoryWithChild(Upstream::LoadBalancerFactorySharedPtr child) { return std::make_shared( @@ -285,13 +258,29 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes std::move(healthy_localities), std::move(degraded_localities)); } - // Stamps `host` with the Locality identity for locality-group index `i`. + std::pair + setupTwoLocalityLb(std::chrono::milliseconds weight_expiration_period, bool has_local = false) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}, has_local); + createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, + weight_expiration_period); + return {h1, h2}; + } + + PanicHosts setupPanicHosts(bool has_local_locality = false) { + PanicHosts hosts{makeWeightTrackingMockHost(), makeWeightTrackingMockHost(), + makeWeightTrackingMockHost()}; + setupLocalities({{hosts.healthy}, {hosts.unhealthy1, hosts.unhealthy2}}, has_local_locality, + std::vector{{hosts.healthy}, {}}, + std::vector{{}, {}}); + return hosts; + } + void setHostLocality(const Upstream::HostSharedPtr& host, size_t i) { setHostLocalityExplicit(host, localityForIndex(i)); } - // Stamps `host` with an explicit Locality identity. The Locality is owned by the fixture (stable - // address) so the mock's locality() can return a reference to it for the test's lifetime. void setHostLocalityExplicit(const Upstream::HostSharedPtr& host, const envoy::config::core::v3::Locality& locality) { auto* mock = dynamic_cast(host.get()); @@ -302,8 +291,6 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes ON_CALL(*mock, locality()).WillByDefault(testing::ReturnRef(localities_.back())); } - // Builds a LocalityWeightsMap keyed by the canonical per-index Locality identities, matching the - // setup order used by setupLocalities/reshapeLocalities (locality i ↔ localityForIndex(i)). static LocalityWeightsMap makeWeightsMap(const std::vector& weights) { LocalityWeightsMap map; for (size_t i = 0; i < weights.size(); ++i) { @@ -312,7 +299,6 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes return map; } - // Triggers a routing-weight recompute by re-running the thread-aware initialize. void recomputeWeights() { ASSERT_TRUE(thread_aware_lb_->initialize().ok()); } // Reshapes a priority's host set, then fires the membership update callback. Two modes: @@ -370,12 +356,6 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes } } - int64_t nowMs() { - return std::chrono::duration_cast( - context_.time_system_.monotonicTime().time_since_epoch()) - .count(); - } - Upstream::LoadBalancerPtr createWorkerLb() { EXPECT_NE(nullptr, factory_); auto worker_lb = factory_->create({priority_set_, nullptr}); @@ -476,7 +456,6 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes return it != weights.end() ? it->second : 0.0; } - // Asserts the live snapshot's all_local flag for `source`/`priority`. Asserts snapshot non-null. void expectAllLocal(bool expected, PriorityRoutingWeights::SelectionSource source = PriorityRoutingWeights::SelectionSource::Healthy, @@ -486,13 +465,45 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes EXPECT_EQ(expected, snapshot->priority_weights[priority].allLocalFor(source)); } - // Asserts a single timer tick increments counter `name` by exactly `delta`. void expectCounterIncrementsPerTick(const std::string& name, uint64_t delta) { const uint64_t before = counterValue(cluster_info_, name); timer_->invokeCallback(); EXPECT_EQ(before + delta, counterValue(cluster_info_, name)); } + uint64_t zoneRoutingStatTotal() const { + return cluster_info_.lbStats().lb_zone_routing_all_directly_.value() + + cluster_info_.lbStats().lb_zone_routing_sampled_.value() + + cluster_info_.lbStats().lb_zone_routing_cross_zone_.value(); + } + + void expectTwoHostSnapCounter(const std::string& counter, double variance_threshold, + double remote_probe_fraction, + std::optional> quiet_utilizations, + std::pair active_utilizations) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + createLb(variance_threshold, /*ewma_alpha=*/1.0, remote_probe_fraction); + const uint64_t after_create = counterValue(cluster_info_, counter); + EXPECT_EQ(1u, after_create); + + if (quiet_utilizations.has_value()) { + setHostUtilization(h_local, quiet_utilizations->first); + setHostUtilization(h_remote, quiet_utilizations->second); + recomputeWeights(); + EXPECT_EQ(after_create, counterValue(cluster_info_, counter)); + } + + const uint64_t before_active = counterValue(cluster_info_, counter); + setHostUtilization(h_local, active_utilizations.first); + setHostUtilization(h_remote, active_utilizations.second); + recomputeWeights(); + EXPECT_EQ(before_active + 1, counterValue(cluster_info_, counter)); + expectCounterIncrementsPerTick(counter, 1); + } + NiceMock context_; NiceMock dispatcher_; NiceMock cluster_info_; @@ -503,20 +514,14 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes std::vector forced_random_values_; size_t forced_random_index_{0}; Event::MockTimer* timer_{}; - // Stable storage for Locality identities returned by mock hosts' locality() (addresses must - // outlive the hosts; std::deque keeps element addresses stable across pushes). + // Stable storage for Locality references returned by mock hosts. std::deque localities_; Upstream::ThreadAwareLoadBalancerPtr thread_aware_lb_; Upstream::LoadBalancerFactorySharedPtr factory_; - // Child ThreadAwareLoadBalancer backing createRoundRobinWorkerFactory (owned by the main-thread - // LB in production). - Upstream::ThreadAwareLoadBalancerPtr retained_child_talb_; }; -// Table-driven single-tick weight-computation cases: set up localities, push one uniform ORCA -// utilization per locality, recompute once, and assert the resulting Healthy snapshot weights. -// Multi-tick (EWMA), reshape, and distribution-sampling scenarios stay as dedicated TEST_Fs. +// Single-tick weight-computation cases; multi-tick and routing-sampling scenarios stay dedicated. struct WeightComputationCase { std::string test_name; std::vector locality_sizes; // host count per locality group @@ -562,6 +567,15 @@ INSTANTIATE_TEST_SUITE_P(WeightComputation, WeightComputationTest, false, {5.0, 1.0}, 0.01}, + {"AllLocalitiesSaturatedFallBacksToHostCount", + {3, 1}, + false, + 0.1, + 0.0, + {1.0, 1.0}, + false, + {3.0, 1.0}, + 0.01}, {"LocalSaturatedSpillsBeyondThreshold", {1, 1}, true, @@ -618,7 +632,6 @@ TEST_F(LoadAwareLocalityLbTest, EmptyAndSingleLocalitySmoke) { auto empty_lb = createWorkerLb(); ASSERT_NE(nullptr, empty_lb); EXPECT_EQ(nullptr, empty_lb->chooseHost(nullptr).host); - EXPECT_EQ(nullptr, empty_lb->peekAnotherHost(nullptr)); auto h1 = makeWeightTrackingMockHost(); setupLocalities({{h1}}); @@ -627,27 +640,11 @@ TEST_F(LoadAwareLocalityLbTest, EmptyAndSingleLocalitySmoke) { auto populated_lb = createWorkerLb(); ASSERT_NE(nullptr, populated_lb); EXPECT_EQ(h1, populated_lb->chooseHost(nullptr).host); - EXPECT_EQ(h1, populated_lb->peekAnotherHost(nullptr)); EXPECT_FALSE(populated_lb->lifetimeCallbacks().has_value()); std::vector hash_key; EXPECT_FALSE(populated_lb->selectExistingConnection(nullptr, *h1, hash_key).has_value()); } -TEST_F(LoadAwareLocalityLbTest, EmptyPrioritySetHasNoAvailableHost) { - // Priority 0 exists (the production invariant for LoadBalancerBase) but carries no hosts. - auto* host_set = priority_set_.getMockHostSet(0); - host_set->hosts_.clear(); - host_set->healthy_hosts_.clear(); - - auto child_factory = std::make_shared(false); - auto factory = makeWorkerFactoryWithChild(child_factory); - - auto worker_lb = factory->create({priority_set_, nullptr}); - ASSERT_NE(nullptr, worker_lb); - EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); - EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); -} - TEST_F(LoadAwareLocalityLbTest, BasicRoutingWeightsTrackNoDataFreshDataAndZeroUtilization) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost(); @@ -662,8 +659,7 @@ TEST_F(LoadAwareLocalityLbTest, BasicRoutingWeightsTrackNoDataFreshDataAndZeroUt recomputeWeights(); expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.9}); - // A zero-utilization report carries no positive signal and is ignored: h1 keeps its last-known - // 0.9 rather than being recorded as a fresh, fully idle reporter. + // Zero-utilization reports carry no positive signal and do not refresh host data. setHostUtilization(h1, 0.0); setHostUtilization(h2, 0.9); recomputeWeights(); @@ -671,20 +667,12 @@ TEST_F(LoadAwareLocalityLbTest, BasicRoutingWeightsTrackNoDataFreshDataAndZeroUt } TEST_F(LoadAwareLocalityLbTest, WeightExpirationIgnoresStaleData) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); + auto [h1, h2] = setupTwoLocalityLb(std::chrono::milliseconds(5000)); - createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, - std::chrono::milliseconds(5000)); - - // No utilization reported — hosts are cold (never reported) and contribute no data. recomputeWeights(); - // Unreported hosts are ignored — weights are equal (host count only). expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 1.0}); - // Fresh data at current time is used. setHostUtilization(h1, 0.9); setHostUtilization(h2, 0.1); recomputeWeights(); @@ -692,29 +680,23 @@ TEST_F(LoadAwareLocalityLbTest, WeightExpirationIgnoresStaleData) { } TEST_F(LoadAwareLocalityLbTest, WeightExpirationDisabledUsesStaleData) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); - // Expiration disabled (0ms) — data is always used regardless of age. - createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, - std::chrono::milliseconds(0)); + auto [h1, h2] = setupTwoLocalityLb(std::chrono::milliseconds(0)); setHostUtilization(h1, 0.9); setHostUtilization(h2, 0.1); recomputeWeights(); expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.9}); + + context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.9}); } // Regression: a report stored at monotonic time 0 (the simulated-time epoch) is a valid sample — // "never reported" is tracked by an out-of-band sentinel, not timestamp 0. TEST_F(LoadAwareLocalityLbTest, ReportAtMonotonicTimeZeroIsValid) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); - - createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, - std::chrono::milliseconds(5000)); + auto [h1, h2] = setupTwoLocalityLb(std::chrono::milliseconds(5000)); ASSERT_EQ(0, std::chrono::duration_cast( simTime().monotonicTime().time_since_epoch()) @@ -726,23 +708,16 @@ TEST_F(LoadAwareLocalityLbTest, ReportAtMonotonicTimeZeroIsValid) { } TEST_F(LoadAwareLocalityLbTest, StaleRemoteCarriesPriorNotZero) { - auto h_local = makeWeightTrackingMockHost(); - auto h_remote = makeWeightTrackingMockHost(); - setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + auto [h_local, h_remote] = + setupTwoLocalityLb(std::chrono::milliseconds(5000), /*has_local=*/true); - createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, - std::chrono::milliseconds(5000)); - - // Local 0.5 is within the variance threshold of remote 0.6, so local is preferred. setHostUtilization(h_local, 0.5); setHostUtilization(h_remote, 0.6); recomputeWeights(); expectAllLocal(true); - // Age the remote past expiration while keeping the local fresh. The remote's last-known 0.6 - // must be carried (not reset to 0); otherwise the remote would look idle and local would - // spuriously spill toward it. + // The stale remote's last-known load must be carried, not reset to idle. context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); setHostUtilization(h_local, 0.5); recomputeWeights(); @@ -763,104 +738,13 @@ TEST_F(LoadAwareLocalityLbTest, StaleLocalityHostCountWeighted) { setUtilizationForHosts(locality_b, 0.5); recomputeWeights(); - // Age locality_b past expiration while keeping locality_a fresh. Locality_b is now all-stale - // and falls back to its host-count baseline (2), not host_count * (1 - 0.5) = 1. + // Stale locality_b falls back to its host-count baseline. context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); setUtilizationForHosts(locality_a, 0.8); recomputeWeights(); expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.6, 2.0}); } -TEST_F(LoadAwareLocalityLbTest, LocalPreferenceThresholdBoundaries) { - auto h_local = makeWeightTrackingMockHost(); - auto h_remote = makeWeightTrackingMockHost(); - setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); - - createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0); - - setHostUtilization(h_local, 0.3); - setHostUtilization(h_remote, 0.3); - recomputeWeights(); - expectAllLocal(true); - expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 0.0}); - - setHostUtilization(h_local, 0.6); - setHostUtilization(h_remote, 0.5); - recomputeWeights(); - expectAllLocal(true); - - setHostUtilization(h_local, 0.8); - setHostUtilization(h_remote, 0.2); - recomputeWeights(); - expectAllLocal(false); - EXPECT_GT(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 1), - weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 0)); -} - -TEST_F(LoadAwareLocalityLbTest, NoHealthyHostsWithLocalLocalityDefaultsToLocal) { - auto h_local = makeWeightTrackingMockHost(); - auto h_remote = makeWeightTrackingMockHost(); - setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true, - std::vector{{}, {}}); - - createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0); - - expectAllLocal(true); - expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 0.0}); -} - -TEST_F(LoadAwareLocalityLbTest, ProbeRedistributesToRemoteAndSkipsWhenNoRemoteHealthyHosts) { - auto h_local = makeWeightTrackingMockHost(); - auto h_remote = makeWeightTrackingMockHost(); - setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); - - createLb(/*variance_threshold=*/1.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.05); - - setHostUtilization(h_local, 0.3); - setHostUtilization(h_remote, 1.0); - recomputeWeights(); - - expectAllLocal(true); - expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.95, 0.05}); - - reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/true, - std::vector{{h_local}, {}}); - recomputeWeights(); - - // No remote healthy hosts (remote_hosts == 0): the remote-gated snap is skipped. Routing is still - // fully local (remote weight is 0), but all_local is unset and the local weight is its base - // headroom (0.7), not the snapped 1.0. - expectAllLocal(false); - expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.7, 0.0}); -} - -TEST_F(LoadAwareLocalityLbTest, ProbeSkipsWhenTargetMetAndSaturatedTieSpreadsByHostCount) { - auto h_local = makeWeightTrackingMockHost(); - auto h_remote = makeWeightTrackingMockHost(); - setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); - - createLb(/*variance_threshold=*/0.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.1); - - setHostUtilization(h_local, 0.6); - setHostUtilization(h_remote, 0.5); - recomputeWeights(); - - auto snapshot = routingSnapshot(); - ASSERT_NE(nullptr, snapshot); - ASSERT_EQ(2, snapshot->priority_weights[0].by_source[0].weights.size()); - EXPECT_FALSE(snapshot->priority_weights[0].by_source[0].all_local); - expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.4, 0.5}); - - // Saturated tie: every locality is at 1.0 (zero headroom), so the all-overloaded fallback spreads - // by host count and skips local preference / probe. - setHostUtilization(h_local, 1.0); - setHostUtilization(h_remote, 1.0); - recomputeWeights(); - - expectAllLocal(false); - expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 1.0}); -} - TEST_F(LoadAwareLocalityLbTest, HealthyWeightsCanDivergeFromAllHostWeights) { auto all_a = makeHosts(4); auto all_b = makeHosts(2); @@ -882,27 +766,6 @@ TEST_F(LoadAwareLocalityLbTest, HealthyWeightsCanDivergeFromAllHostWeights) { expectWeights(PriorityRoutingWeights::SelectionSource::AllHosts, {4.0, 2.0}); } -TEST_F(LoadAwareLocalityLbTest, AllLocalitiesSaturatedFallBacksToHostCount) { - auto locality_a = makeHosts(3); - auto locality_b = makeHosts(1); - setupLocalities({locality_a, locality_b}); - - createLb(); - - setUtilizationForHosts(locality_a, 1.0); - setUtilizationForHosts(locality_b, 1.0); - recomputeWeights(); - expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {3.0, 1.0}); - - auto worker_lb = createWorkerLb(); - auto counts = countPicks(*worker_lb, 400); - int locality_a_count = 0; - for (const auto& host : locality_a) { - locality_a_count += counts[host.get()]; - } - EXPECT_GT(locality_a_count, counts[locality_b[0].get()] * 2); -} - TEST_F(LoadAwareLocalityLbTest, EwmaLifecycleDampensSpikesAndClearsExpiredState) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost(); @@ -928,39 +791,6 @@ TEST_F(LoadAwareLocalityLbTest, EwmaLifecycleDampensSpikesAndClearsExpiredState) EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 1), 1.0, 0.01); } -TEST_F(LoadAwareLocalityLbTest, EwmaNewLocalityGetsRawFirstSampleOnTopologyChange) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); - - createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/0.3); - - setHostUtilization(h1, 0.5); - setHostUtilization(h2, 0.5); - recomputeWeights(); - - auto h3 = makeWeightTrackingMockHost(); - setHostLocality(h3, 2); - auto* host_set = priority_set_.getMockHostSet(0); - host_set->hosts_ = {h1, h2, h3}; - host_set->healthy_hosts_ = {h1, h2, h3}; - host_set->hosts_per_locality_ = - Upstream::makeHostsPerLocality({{h1}, {h2}, {h3}}, /*force_no_local_locality=*/true); - host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; - - // Re-initialize to attach LocalityLbHostData to the new host, then set its utilization. - recomputeWeights(); - setHostUtilization(h3, 0.8); - recomputeWeights(); - - const auto* snapshot = routingSnapshot(); - ASSERT_NE(nullptr, snapshot); - ASSERT_EQ(3, snapshot->priority_weights[0].by_source[0].weights.size()); - EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 2), 0.2, 0.05); - // Existing localities keep their smoothed state across the topology change (identity-keyed). - EXPECT_NEAR(weightForIndex(PriorityRoutingWeights::SelectionSource::Healthy, 0), 0.5, 0.05); -} - TEST_F(LoadAwareLocalityLbTest, EwmaStateFollowsLocalityIdentityOnSameCountSwap) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost(); @@ -1052,15 +882,7 @@ TEST_F(LoadAwareLocalityLbTest, HealthTransitionRefreshesHealthyAndDegradedChild } TEST_F(LoadAwareLocalityLbTest, PanicUsesAllHosts) { - auto healthy_host = makeWeightTrackingMockHost(); - auto unhealthy_host_1 = makeWeightTrackingMockHost(); - auto unhealthy_host_2 = makeWeightTrackingMockHost(); - - setupLocalities({{healthy_host}, {unhealthy_host_1, unhealthy_host_2}}, - /*has_local_locality=*/false, - std::vector{{healthy_host}, {}}, - std::vector{{}, {}}); - + const auto hosts = setupPanicHosts(); createLb(); // Only 1 of 3 hosts healthy → live LoadBalancerBase panic on priority 0; all hosts selectable. @@ -1068,49 +890,24 @@ TEST_F(LoadAwareLocalityLbTest, PanicUsesAllHosts) { auto worker_lb = createWorkerLb(); auto counts = countPicks(*worker_lb, 400); - EXPECT_GT(counts[healthy_host.get()], 100); - EXPECT_GT(counts[unhealthy_host_1.get()] + counts[unhealthy_host_2.get()], 200); + EXPECT_GT(counts[hosts.healthy.get()], 100); + EXPECT_GT(counts[hosts.unhealthy1.get()] + counts[hosts.unhealthy2.get()], 200); EXPECT_GT(cluster_info_.lbStats().lb_healthy_panic_.value(), 0u); } TEST_F(LoadAwareLocalityLbTest, FailTrafficOnPanicReturnsNoHost) { cluster_info_.lb_config_.mutable_zone_aware_lb_config()->set_fail_traffic_on_panic(true); - auto healthy_host = makeWeightTrackingMockHost(); - auto unhealthy_host_1 = makeWeightTrackingMockHost(); - auto unhealthy_host_2 = makeWeightTrackingMockHost(); - // Only 1 of 3 hosts healthy → live panic on priority 0. - setupLocalities({{healthy_host}, {unhealthy_host_1, unhealthy_host_2}}, - /*has_local_locality=*/false, - std::vector{{healthy_host}, {}}, - std::vector{{}, {}}); - + setupPanicHosts(); createLb(); auto worker_lb = createWorkerLb(); for (int i = 0; i < 20; ++i) { EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); - EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); } EXPECT_GT(cluster_info_.lbStats().lb_healthy_panic_.value(), 0u); } -TEST_F(LoadAwareLocalityLbTest, LiveSnapshotRefreshUsesLatestWeights) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); - - createLb(); - auto worker_lb = createWorkerLb(); - EXPECT_TRUE(hostSeen(*worker_lb, h1, 100)); - EXPECT_TRUE(hostSeen(*worker_lb, h2, 100)); - - publishHealthyWeights({0.0, 1.0}); - - auto counts = countPicks(*worker_lb, 100); - EXPECT_GE(counts[h2.get()], 98); -} - TEST_F(LoadAwareLocalityLbTest, MembershipAndTopologyUpdatesShareTheSameWorkerLb) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost(); @@ -1136,41 +933,6 @@ TEST_F(LoadAwareLocalityLbTest, MembershipAndTopologyUpdatesShareTheSameWorkerLb EXPECT_TRUE(hostSeen(*worker_lb, h4, 300)); } -TEST_F(LoadAwareLocalityLbTest, StaleSnapshotFallbackAndAllLocalitiesRemovedAreGraceful) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - auto h3 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}, {h3}}); - - createLb(); - setHostUtilization(h1, 0.5); - setHostUtilization(h2, 0.5); - setHostUtilization(h3, 0.5); - recomputeWeights(); - - auto worker_lb = createWorkerLb(); - EXPECT_TRUE(hostSeen(*worker_lb, h1, 200)); - EXPECT_TRUE(hostSeen(*worker_lb, h2, 200)); - EXPECT_TRUE(hostSeen(*worker_lb, h3, 200)); - - reshapeLocalities(0, {{h1}, {}, {}}, /*added=*/{}, /*removed=*/{h2, h3}); - - expectOnlyHost(*worker_lb, h1, 100); - for (int i = 0; i < 50; ++i) { - auto host = worker_lb->peekAnotherHost(nullptr); - ASSERT_NE(nullptr, host); - EXPECT_EQ(h1, host); - worker_lb->chooseHost(nullptr); - } - - reshapeLocalities(0, {{}, {}, {}}, /*added=*/{}, /*removed=*/{h1}); - - for (int i = 0; i < 20; ++i) { - EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); - EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); - } -} - TEST_F(LoadAwareLocalityLbTest, EmptyDeltaUpdatesRefreshChildWeights) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost(); @@ -1191,57 +953,6 @@ TEST_F(LoadAwareLocalityLbTest, EmptyDeltaUpdatesRefreshChildWeights) { EXPECT_GT(updated_counts[h1.get()], 300); } -TEST_F(LoadAwareLocalityLbTest, InPlaceUpdateTopologyMismatchWaitsForMembershipChange) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); - - createLb(); - auto worker_lb = createWorkerLb(); - - auto h3 = makeWeightTrackingMockHost(); - reshapeLocalities(0, {{h1}, {h2}, {h3}}); - EXPECT_TRUE(hostSeen(*worker_lb, h1, 100)); - EXPECT_TRUE(hostSeen(*worker_lb, h2, 100)); - EXPECT_FALSE(hostSeen(*worker_lb, h3, 100)); - - priority_set_.runUpdateCallbacks(0, {h3}, {}); - publishHealthyWeights({1.0, 1.0, 1.0}); - EXPECT_TRUE(hostSeen(*worker_lb, h3, 300)); -} - -TEST_F(LoadAwareLocalityLbTest, PriorityFailoverAndFailback) { - auto h_p0 = makeWeightTrackingMockHost(); - auto h_p1 = makeWeightTrackingMockHost(); - - setupPriorityLocalities(0, {{h_p0}}, /*has_local_locality=*/false, - std::vector{{}}); - setupPriorityLocalities(1, {{h_p1}}); - - createLb(); - auto worker_lb = createWorkerLb(); - - for (int i = 0; i < 20; ++i) { - EXPECT_EQ(h_p1, worker_lb->chooseHost(nullptr).host); - } - - reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, - std::vector{{h_p0}}); - recomputeWeights(); - - for (int i = 0; i < 20; ++i) { - EXPECT_EQ(h_p0, worker_lb->chooseHost(nullptr).host); - } - - reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, - std::vector{{}}); - recomputeWeights(); - - for (int i = 0; i < 20; ++i) { - EXPECT_EQ(h_p1, worker_lb->chooseHost(nullptr).host); - } -} - TEST_F(LoadAwareLocalityLbTest, FailoverIsCallbackFreshNotTimerGated) { auto h_p0 = makeWeightTrackingMockHost(); auto h_p1 = makeWeightTrackingMockHost(); @@ -1268,146 +979,12 @@ TEST_F(LoadAwareLocalityLbTest, FailoverIsCallbackFreshNotTimerGated) { for (int i = 0; i < 20; ++i) { EXPECT_EQ(h_p1, worker_lb->chooseHost(nullptr).host); } -} -TEST_F(LoadAwareLocalityLbTest, PriorityUpdateGuardPaths) { - auto h1 = makeWeightTrackingMockHost(); - setupLocalities({{h1}}); - - createLb(); - auto worker_lb = createWorkerLb(); - EXPECT_TRUE(hostSeen(*worker_lb, h1, 50)); - - // Add a healthy priority 1. Priority 0 is still healthy, so live failover keeps load there. - auto h_p1 = makeWeightTrackingMockHost(); - setupPriorityLocalities(1, {{h_p1}}); - priority_set_.runUpdateCallbacks(1, {h_p1}, {}); - PriorityRoutingWeights p0; - p0.by_source[0].weights = makeWeightsMap({1.0}); - PriorityRoutingWeights p1; - p1.by_source[0].weights = makeWeightsMap({1.0}); - publishSnapshot(makeSnapshot({p0, p1})); - expectOnlyHost(*worker_lb, h1, 50); - - // Drop priority 0 to zero availability; live failover shifts load to priority 1. reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, - std::vector{{}}); - EXPECT_TRUE(hostSeen(*worker_lb, h_p1, 100)); - - // Growing the backing priority set without a matching snapshot is handled gracefully. - priority_set_.getMockHostSet(5); - priority_set_.runUpdateCallbacks(5, {makeWeightTrackingMockHost()}, {}); - priority_set_.runUpdateCallbacks(5, {}, {}); - EXPECT_TRUE(hostSeen(*worker_lb, h_p1, 100)); -} - -TEST_F(LoadAwareLocalityLbTest, AdvisorySnapshotSizeMismatchesAreHandledGracefully) { - auto h1 = makeWeightTrackingMockHost(); - setupLocalities({{h1}}); - - createLb(); - auto worker_lb = createWorkerLb(); - - PriorityRoutingWeights weights; - weights.by_source[0].weights = makeWeightsMap({1.0}); - - // The advisory snapshot may carry more priorities than the live priority set; the single live - // healthy host stays selected (priority is live, not snapshot-driven). - publishSnapshot(makeSnapshot({weights, weights})); - EXPECT_TRUE(hostSeen(*worker_lb, h1, 20)); - - publishSnapshot(makeSnapshot({weights})); - EXPECT_TRUE(hostSeen(*worker_lb, h1, 20)); - EXPECT_EQ(h1, worker_lb->peekAnotherHost(nullptr)); -} - -TEST_F(LoadAwareLocalityLbTest, ShortAdvisorySnapshotStillRoutesLivePriority) { - auto h_p1 = makeWeightTrackingMockHost(); - setupPriorityLocalities(0, {}); - setupPriorityLocalities(1, {{h_p1}}); - - createLb(); - auto worker_lb = createWorkerLb(); - - // Priority 0 has no hosts; live failover routes to the healthy priority 1 even though the - // advisory snapshot only carries weights for priority 0. - publishHealthyWeights({1.0}); - - expectOnlyHost(*worker_lb, h_p1, 20); - EXPECT_EQ(h_p1, worker_lb->peekAnotherHost(nullptr)); -} - -TEST_F(LoadAwareLocalityLbTest, EmptyHigherPriorityDoesNotBreakLiveRouting) { - auto h_p0 = makeWeightTrackingMockHost(); - setupPriorityLocalities(0, {{h_p0}}); - setupPriorityLocalities(1, {}); - - createLb(); - auto worker_lb = createWorkerLb(); - - // Priority 1 has no hosts; live failover keeps load on the healthy priority 0. - PriorityRoutingWeights p0; - p0.by_source[0].weights = makeWeightsMap({1.0}); - PriorityRoutingWeights p1; - publishSnapshot(makeSnapshot({p0, p1})); - - expectOnlyHost(*worker_lb, h_p0, 20); -} - -TEST_F(LoadAwareLocalityLbTest, SelectLocalityFallsBackForCountMismatchAndZeroTotals) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); - - createLb(); - auto worker_lb = createWorkerLb(); - - // Snapshot has only locality 0's identity; the worker's locality 1 (missing → 0.0) is excluded, - // so all traffic stays on h1. - publishHealthyWeights({1.0}); - expectOnlyHost(*worker_lb, h1, 20); - - publishHealthyWeights({}); - expectOnlyHost(*worker_lb, h1, 20); - - // Both live localities have weight 0 (locality 2's 5.0 is not a live identity), so selection - // falls back to locality 0. - publishHealthyWeights({0.0, 0.0, 5.0}); - expectOnlyHost(*worker_lb, h1, 20); -} - -TEST_F(LoadAwareLocalityLbTest, SelectLocalityIgnoresWeightsForUnknownLocalities) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); - - createLb(); - auto worker_lb = createWorkerLb(); - - // The snapshot carries a third locality identity (weight 100.0) that the worker does not have. - // It is ignored; the two live localities split traffic by their identity-matched weights (1:1). - publishHealthyWeights({1.0, 1.0, 100.0}); - - auto counts = countPicks(*worker_lb, 400); - EXPECT_GT(counts[h1.get()], 120); - EXPECT_GT(counts[h2.get()], 120); -} - -TEST_F(LoadAwareLocalityLbTest, SelectLocalityRoundingGuardReturnsLastLocality) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - auto h3 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}, {h3}}); - - createLb(); - auto worker_lb = createWorkerLb(); - - publishHealthyWeights({0.2, 0.3, 0.5}); - - // The hash draw (chooseHostSet) consumes the first random value even with a single priority, - // so the second draw must drive locality selection. - forceRandomValues({0, std::numeric_limits::max()}); - EXPECT_EQ(h3, worker_lb->chooseHost(nullptr).host); + std::vector{{h_p0}}); + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(h_p0, worker_lb->chooseHost(nullptr).host); + } } // Advisory weights are keyed by Locality identity, so adding a locality that shifts another's @@ -1467,70 +1044,24 @@ TEST_F(LoadAwareLocalityLbTest, DedicatedStatsAllOverloadedTotalIncOncePerTickWh createLb(); EXPECT_EQ(0, counterValue(cluster_info_, "load_aware_locality.all_overloaded_total")); - // Make all localities fully saturated — headroom = 0. setUtilizationForHosts(h_a, 1.0); setUtilizationForHosts(h_b, 1.0); recomputeWeights(); EXPECT_EQ(1, counterValue(cluster_info_, "load_aware_locality.all_overloaded_total")); - // Next tick, still saturated — counter goes to 2 (one per tick, not per source). expectCounterIncrementsPerTick("load_aware_locality.all_overloaded_total", 1); } TEST_F(LoadAwareLocalityLbTest, DedicatedStatsLocalPreferredTotalIncOncePerTickOnVarianceSnap) { - auto h_local = makeWeightTrackingMockHost(); - auto h_remote = makeWeightTrackingMockHost(); - setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); - - // High threshold so local is NOT preferred when both are saturated (utils=1.0). - createLb(/*variance_threshold=*/0.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0); - // Cold start: no data yet, both utils=0 → utils[0]=0 <= 0+0.0=0 → snap fires on the initial - // tick, deterministically. - const uint64_t after_create = - counterValue(cluster_info_, "load_aware_locality.local_preferred_total"); - EXPECT_EQ(1u, after_create); - - // Force high local utilization so snap does NOT fire. - setHostUtilization(h_local, 0.9); - setHostUtilization(h_remote, 0.1); - recomputeWeights(); - // 0.9 > 0.1 + 0.0 → no snap. - const uint64_t after_no_snap = - counterValue(cluster_info_, "load_aware_locality.local_preferred_total"); - EXPECT_EQ(after_create, after_no_snap); - - // Now snap fires (local <= remote + threshold). - setHostUtilization(h_local, 0.1); - setHostUtilization(h_remote, 0.9); - recomputeWeights(); - EXPECT_EQ(after_no_snap + 1, - counterValue(cluster_info_, "load_aware_locality.local_preferred_total")); - - // Another snap tick. - expectCounterIncrementsPerTick("load_aware_locality.local_preferred_total", 1); + expectTwoHostSnapCounter("load_aware_locality.local_preferred_total", + /*variance_threshold=*/0.0, /*remote_probe_fraction=*/0.0, + std::make_pair(0.9, 0.1), std::make_pair(0.1, 0.9)); } TEST_F(LoadAwareLocalityLbTest, DedicatedStatsProbeActiveTotalIncOncePerTickWhenProbeKicksIn) { - auto h_local = makeWeightTrackingMockHost(); - auto h_remote = makeWeightTrackingMockHost(); - setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); - - // remote_probe_fraction=0.05 → probe redistribution fires when remote < 5% of total. - // variance_threshold=1.0 so local is always preferred. - createLb(/*variance_threshold=*/1.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.05); - // Cold start tick: utils=0, local preferred, probe fires on the initial tick. - const uint64_t after_create = - counterValue(cluster_info_, "load_aware_locality.probe_active_total"); - EXPECT_EQ(1u, after_create); - - setHostUtilization(h_local, 0.3); - setHostUtilization(h_remote, 1.0); - recomputeWeights(); - // all_local fires (local preferred), then probe redistribution shifts weight to remote. - EXPECT_EQ(after_create + 1, - counterValue(cluster_info_, "load_aware_locality.probe_active_total")); - - expectCounterIncrementsPerTick("load_aware_locality.probe_active_total", 1); + expectTwoHostSnapCounter("load_aware_locality.probe_active_total", + /*variance_threshold=*/1.0, /*remote_probe_fraction=*/0.05, std::nullopt, + std::make_pair(0.3, 1.0)); } TEST_F(LoadAwareLocalityLbTest, DedicatedStatsStaleLocalityTotalCountsOncePerStalePerTick) { @@ -1542,26 +1073,20 @@ TEST_F(LoadAwareLocalityLbTest, DedicatedStatsStaleLocalityTotalCountsOncePerSta createLb(/*variance_threshold=*/0.1, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.0, std::chrono::milliseconds(5000)); - // Never-reported localities are cold, not stale: nothing counted before first data arrives. EXPECT_EQ(0, counterValue(cluster_info_, "load_aware_locality.stale_locality_total")); expectCounterIncrementsPerTick("load_aware_locality.stale_locality_total", 0); - // Give all localities fresh data and re-initialize so we have a clean baseline. setUtilizationForHosts(locality_a, 0.5); setUtilizationForHosts(locality_b, 0.4); setUtilizationForHosts(locality_c, 0.6); recomputeWeights(); - // Record counter after the tick where all data is fresh (should add 0). const uint64_t baseline = counterValue(cluster_info_, "load_aware_locality.stale_locality_total"); - // Age locality_b and locality_c past expiration; locality_a stays fresh. context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); setUtilizationForHosts(locality_a, 0.5); recomputeWeights(); - // 2 stale localities (b and c) — counter adds 2 per tick, NOT 6 (not 3× for 3 sources). EXPECT_EQ(baseline + 2, counterValue(cluster_info_, "load_aware_locality.stale_locality_total")); - // Next tick, same state → adds another 2. expectCounterIncrementsPerTick("load_aware_locality.stale_locality_total", 2); } @@ -1625,19 +1150,17 @@ TEST_F(LoadAwareLocalityLbTest, PeekAndChooseReplayTheSameRandomSequence) { auto h2 = makeWeightTrackingMockHost(); setupLocalities({{h1}, {h2}}); createLb(); - - setHostUtilization(h1, 0.9); - setHostUtilization(h2, 0.1); - recomputeWeights(); auto worker_lb = createWorkerLb(); - // Peek draws (priority hash + locality draw) are stashed: 0 → locality 0 (weight 0.1 of 1.0). + publishHealthyWeights({0.0, 1.0}); + auto counts = countPicks(*worker_lb, 100); + EXPECT_GE(counts[h2.get()], 98); + + publishHealthyWeights({0.1, 0.9}); forceRandomValues({0, 0}); auto peeked = worker_lb->peekAnotherHost(nullptr); ASSERT_NE(nullptr, peeked); - // If chooseHost drew fresh randoms these values would steer it to locality 1; replaying the - // stashed peek draws must land on the same locality (and host) the peek returned. forceRandomValues( {std::numeric_limits::max() - 1, std::numeric_limits::max() - 1}); auto chosen = worker_lb->chooseHost(nullptr).host; @@ -1672,41 +1195,7 @@ TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverHealthyRoutingPaths) { EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_cross_zone_.value(), 0); } -TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverDegradedAndPanicSources) { - auto h_local = makeWeightTrackingMockHost(); - auto h_remote = makeWeightTrackingMockHost(); - // No healthy hosts, all degraded (but enough to avoid panic): live priority targets the - // Degraded source. - setupLocalities({{h_local}, {h_remote}}, - /*has_local_locality=*/true, std::vector{{}, {}}, - std::vector{{h_local}, {h_remote}}); - - createLb(); - recomputeWeights(); - - auto worker_lb = createWorkerLb(); - auto degraded_counts = countPicks(*worker_lb, 400); - EXPECT_GT(degraded_counts[h_local.get()] + degraded_counts[h_remote.get()], 0); - - auto h_remote_2 = makeWeightTrackingMockHost(); - setupLocalities({{h_local}, {h_remote, h_remote_2}}, - /*has_local_locality=*/true, std::vector{{h_local}, {}}, - std::vector{{}, {}}); - - createLb(); - recomputeWeights(); - worker_lb = createWorkerLb(); - auto panic_counts = countPicks(*worker_lb, 400); - EXPECT_GT(panic_counts[h_local.get()] + panic_counts[h_remote.get()] + - panic_counts[h_remote_2.get()], - 0); - EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_sampled_.value() + - cluster_info_.lbStats().lb_zone_routing_cross_zone_.value() + - cluster_info_.lbStats().lb_zone_routing_all_directly_.value(), - 0u); -} - -TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsStayZeroWithoutLocalLocality) { +TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverNoLocalDegradedAndPanicSources) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost(); setupLocalities({{h1}, {h2}}, /*has_local_locality=*/false); @@ -1720,36 +1209,34 @@ TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsStayZeroWithoutLocalLocality) { for (int i = 0; i < 100; ++i) { worker_lb->chooseHost(nullptr); } + EXPECT_EQ(0u, zoneRoutingStatTotal()); - EXPECT_EQ(0, cluster_info_.lbStats().lb_zone_routing_all_directly_.value()); - EXPECT_EQ(0, cluster_info_.lbStats().lb_zone_routing_cross_zone_.value()); - EXPECT_EQ(0, cluster_info_.lbStats().lb_zone_routing_sampled_.value()); -} - -TEST_F(LoadAwareLocalityLbTest, EmptyLocalityHostsDoNotCrash) { - auto* host_set = priority_set_.getMockHostSet(0); - host_set->hosts_.clear(); - host_set->healthy_hosts_.clear(); - host_set->hosts_per_locality_ = - Upstream::makeHostsPerLocality({}, /*force_no_local_locality=*/true); - host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, + /*has_local_locality=*/true, std::vector{{}, {}}, + std::vector{{h_local}, {h_remote}}); createLb(); - auto worker_lb = createWorkerLb(); - EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); -} + recomputeWeights(); -TEST_F(LoadAwareLocalityLbTest, WorkerFallsBackWithoutPublishedRoutingSnapshot) { - auto h1 = makeWeightTrackingMockHost(); - auto h2 = makeWeightTrackingMockHost(); - setupLocalities({{h1}, {h2}}); + worker_lb = createWorkerLb(); + auto degraded_counts = countPicks(*worker_lb, 400); + EXPECT_GT(degraded_counts[h_local.get()] + degraded_counts[h_remote.get()], 0); + EXPECT_GT(zoneRoutingStatTotal(), 0u); - auto factory = createRoundRobinWorkerFactory(); - auto worker_lb = factory->create({priority_set_, nullptr}); - ASSERT_NE(nullptr, worker_lb); + const uint64_t before_panic_zone_stats = zoneRoutingStatTotal(); + const auto panic_hosts = setupPanicHosts(/*has_local_locality=*/true); - expectOnlyHost(*worker_lb, h1, 20); - EXPECT_EQ(h1, worker_lb->peekAnotherHost(nullptr)); + createLb(); + recomputeWeights(); + worker_lb = createWorkerLb(); + auto panic_counts = countPicks(*worker_lb, 400); + EXPECT_GT(panic_counts[panic_hosts.healthy.get()] + panic_counts[panic_hosts.unhealthy1.get()] + + panic_counts[panic_hosts.unhealthy2.get()], + 0); + EXPECT_GT(cluster_info_.lbStats().lb_healthy_panic_.value(), 0u); + EXPECT_EQ(before_panic_zone_stats, zoneRoutingStatTotal()); } TEST_F(LoadAwareLocalityLbTest, InitializePropagatesChildFactoryError) { @@ -1802,87 +1289,45 @@ TEST_F(LoadAwareLocalityLbTest, MetricNamesDriveUtilization) { // (when the default orca_weight_manager_use_named_metrics_first flag is enabled). Event::GlobalTimeSystem ts; - { - // Named metric present → use its value. - LocalityLbHostData slot(ts, makeMetricNames({"named_metrics.foo"})); - xds::data::orca::v3::OrcaLoadReport report; - (*report.mutable_named_metrics())["foo"] = 0.6; - ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); - EXPECT_DOUBLE_EQ(0.6, slot.utilization()); - } - - { - // Named metric absent and no application_utilization → cpu_utilization fallback. - LocalityLbHostData slot(ts, makeMetricNames({"named_metrics.foo"})); - xds::data::orca::v3::OrcaLoadReport report; - report.set_cpu_utilization(0.3); - ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); - EXPECT_DOUBLE_EQ(0.3, slot.utilization()); - } + struct Case { + std::string name; + std::vector metric_names; + std::optional named_metric; + std::optional application_utilization; + std::optional cpu_utilization; + double expected_utilization; + }; - { - // Empty metric_names → application_utilization used when set. - LocalityLbHostData slot(ts, makeMetricNames({})); - xds::data::orca::v3::OrcaLoadReport report; - report.set_application_utilization(0.7); - ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); - EXPECT_DOUBLE_EQ(0.7, slot.utilization()); - } + const std::vector cases = { + {"NamedMetric", {"named_metrics.foo"}, 0.6, std::nullopt, std::nullopt, 0.6}, + {"NamedMetricAbsentFallsBackToCpu", + {"named_metrics.foo"}, + std::nullopt, + std::nullopt, + 0.3, + 0.3}, + {"ApplicationUtilization", {}, std::nullopt, 0.7, std::nullopt, 0.7}, + {"CpuFallback", {}, std::nullopt, std::nullopt, 0.4, 0.4}, + }; - { - // Empty metric_names and no application_utilization → cpu_utilization fallback. - LocalityLbHostData slot(ts, makeMetricNames({})); + for (const Case& c : cases) { + SCOPED_TRACE(c.name); + LocalityLbHostData slot(ts, makeMetricNames(c.metric_names)); xds::data::orca::v3::OrcaLoadReport report; - report.set_cpu_utilization(0.4); + if (c.named_metric.has_value()) { + (*report.mutable_named_metrics())["foo"] = *c.named_metric; + } + if (c.application_utilization.has_value()) { + report.set_application_utilization(*c.application_utilization); + } + if (c.cpu_utilization.has_value()) { + report.set_cpu_utilization(*c.cpu_utilization); + } ASSERT_TRUE(slot.onOrcaLoadReport(report, stream_info_).ok()); - EXPECT_DOUBLE_EQ(0.4, slot.utilization()); + EXPECT_DOUBLE_EQ(c.expected_utilization, slot.utilization()); } } -// Tests 1+2: chooseHost and peekAnotherHost empty-priority-set guard. -// The guards at the top of each function return early when hostSetsPerPriority() is empty. -// We construct the LB against a real priority set (so LoadBalancerBase initialises safely), then -// override the mock to present an empty set for subsequent calls. -TEST_F(LoadAwareLocalityLbTest, EmptyPrioritySetGuardInChooseHostAndPeekAnotherHost) { - auto h1 = makeWeightTrackingMockHost(); - setupLocalities({{h1}}); - - createLb(); - auto worker_lb = createWorkerLb(); - - // Override so that hostSetsPerPriority() returns empty from this point on. - static const std::vector kEmptyHostSets; - ON_CALL(testing::Const(priority_set_), hostSetsPerPriority()) - .WillByDefault(testing::ReturnRef(kEmptyHostSets)); - - // Both guards fire: `if (priority_set_.hostSetsPerPriority().empty()) return {nullptr/nullptr}`. - EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); - EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); -} - -// Test 3: empty-delta (no-rebuild) priority-count-mismatch early-return. -// When the priority count in the live set grows but the worker hasn't rebuilt yet, an empty-delta -// callback runs syncPriority with allow_rebuild=false. The size mismatch guard returns early -// without crashing. -TEST_F(LoadAwareLocalityLbTest, InPlaceHostUpdateSizeMismatchReturnsEarly) { - auto h1 = makeWeightTrackingMockHost(); - setupLocalities({{h1}}); - - createLb(); - auto worker_lb = createWorkerLb(); - expectOnlyHost(*worker_lb, h1, 10); - - // Grow the backing priority set to 6 priorities without a membership-delta rebuild first. - priority_set_.getMockHostSet(5); - // Empty-delta update on priority 5: syncPriority(5, allow_rebuild=false). - // host_sets.size()=6, per_priority_locality_.size()=1 → size-mismatch guard returns early. - priority_set_.runUpdateCallbacks(5, {}, {}); - - // The LB is still functional for priority 0 (early-return did not corrupt state). - expectOnlyHost(*worker_lb, h1, 10); -} - -// Test 4: storeUtilization non-finite guard. // When getUtilizationFromOrcaReport returns a non-finite value (Inf/NaN), storeUtilization // discards it: utilization() stays 0.0 and lastUpdateTimeMs() stays 0. TEST_F(LoadAwareLocalityLbTest, StoreUtilizationRejectsNonFiniteValue) { From 358f46a00e44d2ac6baf6f8ab10932bcf0fb220f Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Sun, 5 Jul 2026 20:34:15 -0600 Subject: [PATCH 12/16] spellcheck and coverage Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_aware_locality_lb.h | 2 +- .../load_aware_locality/integration_test.cc | 18 +- .../load_aware_locality_lb_test.cc | 189 +++++++++++++++--- 3 files changed, 170 insertions(+), 39 deletions(-) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h index 53f714bd4b886..9c253f5113b01 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -142,7 +142,7 @@ class LoadAwareLocalityLbConfig : public Upstream::LoadBalancerConfig { }; // Key locality state by identity rather than HostsPerLocality position so add/remove index shifts -// cannot mis-map main-thread snapshots to worker-local membership. +// cannot map main-thread snapshots to the wrong worker-local membership. using LocalityWeightsMap = absl::flat_hash_map; diff --git a/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc index 4237b781a5aae..f0543ea560adc 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc @@ -137,17 +137,13 @@ class LoadAwareLocalityIntegrationTest : public testing::TestWithParamwaitForCounter("cluster.cluster_0.lb_recalculate_zone_structures", - testing::Ge(expected_counter)); - } + void advanceWeightTick() { timeSystem().advanceTimeWait(std::chrono::seconds(11)); } - void seedWithTwoCycles(const std::vector& utilizations, uint64_t starting_count = 1) { + void seedWithTwoCycles(const std::vector& utilizations) { sendRequestsAndTrack(30, utilizations); - advanceWeightTick(starting_count + 1); + advanceWeightTick(); sendRequestsAndTrack(30, utilizations); - advanceWeightTick(starting_count + 2); + advanceWeightTick(); } uint64_t zoneTraffic(const std::vector& usage, size_t zone_index) const { @@ -179,7 +175,7 @@ TEST_P(LoadAwareLocalityIntegrationTest, StrictLocalWithRoundRobinWithinLocality const std::vector utilizations = {0.05, 0.05, 0.05, 0.05}; sendRequestsAndTrack(30, utilizations); - advanceWeightTick(2); + advanceWeightTick(); const auto usage = sendRequestsAndTrack(100, utilizations); @@ -225,7 +221,7 @@ TEST_P(LoadAwareLocalityIntegrationTest, AdaptiveSpillAndRecovery) { const std::vector rebalanced = {0.3, 0.3, 0.3, 0.3}; sendRequestsAndTrack(60, rebalanced); - advanceWeightTick(4); + advanceWeightTick(); const auto phase2_usage = sendRequestsAndTrack(100, rebalanced); const uint64_t phase2_local = localTraffic(phase2_usage); @@ -243,7 +239,7 @@ TEST_P(LoadAwareLocalityIntegrationTest, EwmaDampensSpike) { const std::vector spiked_local = {0.9, 0.9, 0.2, 0.2}; sendRequestsAndTrack(30, spiked_local); - advanceWeightTick(4); + advanceWeightTick(); const auto usage = sendRequestsAndTrack(100, spiked_local); const uint64_t remote = remoteTraffic(usage); diff --git a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc index a088ba49d2195..acbda6c7f6399 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -37,9 +38,10 @@ namespace { using testing::NiceMock; using testing::Return; +using testing::ReturnRef; -// Distinct Locality identity for locality-group index `i`. Region strings sort lexicographically by -// index ("r00", "r01", ...) so the default setup order matches a local-first/lexicographic +// Distinct Locality identity for locality-group index `i`. Region strings sort in lexicographic +// order by index ("r00", "r01", ...) so the default setup order matches a local-first/lexicographic // ordering. envoy::config::core::v3::Locality localityForIndex(size_t i) { return Upstream::Locality(absl::StrCat("r", absl::Dec(i, absl::kZeroPad2)), "z", "sz"); @@ -67,12 +69,34 @@ Upstream::HostVector makeHosts(uint32_t count, uint32_t weight = 1) { class RecordingWorkerChildFactory : public Upstream::LoadBalancerFactory { public: - explicit RecordingWorkerChildFactory(bool recreate_on_host_change) - : recreate_on_host_change_(recreate_on_host_change) {} + using Inspector = std::function; - Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams) override { + explicit RecordingWorkerChildFactory(bool recreate_on_host_change, Inspector inspector = nullptr) + : recreate_on_host_change_(recreate_on_host_change), inspector_(std::move(inspector)) {} + + Upstream::LoadBalancerPtr create(Upstream::LoadBalancerParams params) override { ++create_count_; - return std::make_unique>(); + auto lb = std::make_unique>(); + Upstream::HostConstSharedPtr host; + const auto& host_sets = params.priority_set.hostSetsPerPriority(); + if (!host_sets.empty() && !host_sets[0]->hosts().empty()) { + host = host_sets[0]->hosts()[0]; + } + ON_CALL(*lb, chooseHost(testing::_)) + .WillByDefault([host, inspector = inspector_](Upstream::LoadBalancerContext* context) { + if (context != nullptr && inspector) { + inspector(*context); + } + return Upstream::HostSelectionResponse{host}; + }); + ON_CALL(*lb, peekAnotherHost(testing::_)) + .WillByDefault([host, inspector = inspector_](Upstream::LoadBalancerContext* context) { + if (context != nullptr && inspector) { + inspector(*context); + } + return host; + }); + return lb; } bool recreateOnHostChangeDeprecated() const override { return recreate_on_host_change_; } @@ -81,10 +105,11 @@ class RecordingWorkerChildFactory : public Upstream::LoadBalancerFactory { private: const bool recreate_on_host_change_; + const Inspector inspector_; uint32_t create_count_{0}; }; -// A child LB that always selects asynchronously: chooseHost returns no host plus a cancelable +// A child LB that always selects asynchronously: chooseHost returns no host plus a cancellation // handle and (in production) would retain the context to call back later, mimicking dynamic modules // or dynamic forward proxy. Flips `cancelled` when its handle is cancelled. class AsyncSelectingChildLb : public Upstream::LoadBalancer { @@ -766,6 +791,28 @@ TEST_F(LoadAwareLocalityLbTest, HealthyWeightsCanDivergeFromAllHostWeights) { expectWeights(PriorityRoutingWeights::SelectionSource::AllHosts, {4.0, 2.0}); } +TEST_F(LoadAwareLocalityLbTest, WeightComputationSkipsEmptyLocalitiesAndHostsWithoutPolicyData) { + auto h1 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {}}); + createLb(); + + setHostUtilization(h1, 1.0); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 0.0}); + + auto h_without_data = makeWeightTrackingMockHost(); + setHostLocality(h_without_data, 2); + auto* host_set = priority_set_.getMockHostSet(0); + host_set->hosts_ = {h1, h_without_data}; + host_set->healthy_hosts_ = {h1, h_without_data}; + host_set->hosts_per_locality_ = Upstream::makeHostsPerLocality({{h1}, {}, {h_without_data}}, + /*force_no_local_locality=*/true); + host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + + timer_->invokeCallback(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.0, 0.0, 1.0}); +} + TEST_F(LoadAwareLocalityLbTest, EwmaLifecycleDampensSpikesAndClearsExpiredState) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost(); @@ -885,7 +932,7 @@ TEST_F(LoadAwareLocalityLbTest, PanicUsesAllHosts) { const auto hosts = setupPanicHosts(); createLb(); - // Only 1 of 3 hosts healthy → live LoadBalancerBase panic on priority 0; all hosts selectable. + // Only 1 of 3 hosts healthy → live LoadBalancerBase panic on priority 0; all hosts eligible. expectWeights(PriorityRoutingWeights::SelectionSource::AllHosts, {1.0, 2.0}); auto worker_lb = createWorkerLb(); @@ -988,7 +1035,7 @@ TEST_F(LoadAwareLocalityLbTest, FailoverIsCallbackFreshNotTimerGated) { } // Advisory weights are keyed by Locality identity, so adding a locality that shifts another's -// lexicographic index between the snapshot tick and the worker's live membership must NOT mis-map +// lexicographic index between the snapshot tick and the worker's live membership must NOT misplace // the weight. Set up local A + remote C with distinct weights, publish a snapshot, then add remote // B (sorts between A and C, shifting C from index 1 to index 2) WITHOUT recomputing weights, and // assert C's traffic still follows C's weight (and the new B, absent from the snapshot, gets none). @@ -1103,31 +1150,59 @@ TEST_F(LoadAwareLocalityLbTest, WeightUpdateTimerRearmsOnEachTick) { EXPECT_TRUE(timer_->enabled_); } -TEST_F(LoadAwareLocalityLbTest, ChildContextShieldsRetryPriorityWithoutComputingHash) { - auto h_a = makeWeightTrackingMockHost(); - auto h_b = makeWeightTrackingMockHost(); - setupLocalities({{h_a}, {h_b}}); - createLb(); - auto worker_lb = createWorkerLb(); +TEST_F(LoadAwareLocalityLbTest, ChildContextForwardsMethodsExceptRetryPriorityAndAsyncCallback) { + auto h = makeWeightTrackingMockHost(); + setupLocalities({{h}}); + + int inspected = 0; + Upstream::PrioritySetImpl child_priority_set; + auto child_factory = std::make_shared( + /*recreate_on_host_change=*/false, + [&child_priority_set, h, &inspected](Upstream::LoadBalancerContext& context) { + ++inspected; + std::ignore = context.computeHashKey(); + const Upstream::HealthyAndDegradedLoad default_priority_load{Upstream::HealthyLoad({100}), + Upstream::DegradedLoad({0})}; + EXPECT_EQ(&default_priority_load, + &context.determinePriorityLoad(child_priority_set, default_priority_load, + Upstream::RetryPriority::defaultPriorityMapping)); + std::ignore = context.metadataMatchCriteria(); + std::ignore = context.downstreamConnection(); + std::ignore = context.requestStreamInfo(); + std::ignore = context.downstreamHeaders(); + std::ignore = context.shouldSelectAnotherHost(*h); + std::ignore = context.hostSelectionRetryCount(); + std::ignore = context.upstreamSocketOptions(); + std::ignore = context.upstreamTransportSocketOptions(); + std::ignore = context.overrideHostToSelect(); + Upstream::HostConstSharedPtr async_host; + context.onAsyncHostSelection(std::move(async_host), "ignored"); + context.setHeadersModifier(nullptr); + }); + auto factory = makeWorkerFactoryWithChild(child_factory); + auto worker_lb = factory->create({priority_set_, nullptr}); + ASSERT_NE(nullptr, worker_lb); NiceMock ctx; - // The (possibly stateful) retry-priority hook must be consulted exactly once per pick, by the - // parent against the cluster priority set, never re-invoked by the per-locality child LB whose - // single-priority view would corrupt retry-priority state. Route-level hash policies should not - // be evaluated by non-hash children. + // The parent consults retry-priority once per selection. The child wrapper returns its default + // priority load locally and never forwards the child's single-priority view to the parent. EXPECT_CALL(ctx, determinePriorityLoad(testing::Ref(priority_set_), testing::_, testing::_)) - .WillOnce([](const Upstream::PrioritySet&, - const Upstream::HealthyAndDegradedLoad& default_priority_load, - const Upstream::RetryPriority::PriorityMappingFunc&) - -> const Upstream::HealthyAndDegradedLoad& { return default_priority_load; }); - EXPECT_CALL(ctx, computeHashKey()).Times(0); - - EXPECT_NE(nullptr, worker_lb->chooseHost(&ctx).host); + .Times(2) + .WillRepeatedly( + [](const Upstream::PrioritySet&, + const Upstream::HealthyAndDegradedLoad& default_priority_load, + const Upstream::RetryPriority::PriorityMappingFunc&) + -> const Upstream::HealthyAndDegradedLoad& { return default_priority_load; }); + EXPECT_CALL(ctx, onAsyncHostSelection(testing::_, testing::_)).Times(0); + + EXPECT_EQ(h, worker_lb->chooseHost(&ctx).host); + EXPECT_EQ(h, worker_lb->peekAnotherHost(&ctx)); + EXPECT_EQ(2, inspected); } // An asynchronous child LB would retain the stack-scoped child context and call back after // chooseHost returns (use-after-free). The policy must cancel the async selection and fail -// synchronously, never propagating the cancelable handle to its caller. +// synchronously, never propagating the cancellation handle to its caller. TEST_F(LoadAwareLocalityLbTest, AsyncChildSelectionIsCancelledAndFailsSynchronously) { auto h = makeWeightTrackingMockHost(); setupLocalities({{h}}); @@ -1167,6 +1242,66 @@ TEST_F(LoadAwareLocalityLbTest, PeekAndChooseReplayTheSameRandomSequence) { EXPECT_EQ(peeked.get(), chosen.get()); } +TEST_F(LoadAwareLocalityLbTest, WorkerFallbackEdgeCases) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + EXPECT_EQ(h1, worker_lb->chooseHost(nullptr).host); + + publishSnapshot(makeSnapshot({PriorityRoutingWeights{}})); + EXPECT_EQ(h1, worker_lb->chooseHost(nullptr).host); + + publishHealthyWeights({0.0, 0.0}); + EXPECT_EQ(h1, worker_lb->chooseHost(nullptr).host); + + publishHealthyWeights({1.0, 1.0}); + forceRandomValues({0, std::numeric_limits::max()}); + EXPECT_EQ(h2, worker_lb->chooseHost(nullptr).host); + + worker_lb = createWorkerLb(); + Upstream::HealthyAndDegradedLoad degraded_load{Upstream::HealthyLoad({0}), + Upstream::DegradedLoad({100})}; + NiceMock ctx; + EXPECT_CALL(ctx, determinePriorityLoad(testing::Ref(priority_set_), testing::_, testing::_)) + .WillOnce(ReturnRef(degraded_load)); + EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(&ctx)); + + publishHealthyWeights({0.0, 1.0}); + EXPECT_EQ(h2, worker_lb->chooseHost(nullptr).host); + + reshapeLocalities(0, {{h1}, {}}, /*added=*/{}, /*removed=*/{h2}); + EXPECT_EQ(h1, worker_lb->chooseHost(nullptr).host); + + reshapeLocalities(0, {{}, {}}, /*added=*/{}, /*removed=*/{h1}); + EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); + + auto h_p1 = makeWeightTrackingMockHost(); + setupPriorityLocalities(1, {{h_p1}}); + priority_set_.runUpdateCallbacks(1, {}, {}); + Upstream::HealthyAndDegradedLoad force_priority_1{Upstream::HealthyLoad({0, 100}), + Upstream::DegradedLoad({0, 0})}; + NiceMock priority_ctx; + EXPECT_CALL(priority_ctx, + determinePriorityLoad(testing::Ref(priority_set_), testing::_, testing::_)) + .WillOnce(ReturnRef(force_priority_1)); + EXPECT_EQ(nullptr, worker_lb->chooseHost(&priority_ctx).host); +} + +TEST_F(LoadAwareLocalityLbTest, PeekAnotherHostCapsPreconnectLookahead) { + auto h = makeWeightTrackingMockHost(); + setupLocalities({{h}}); + + createLb(); + auto worker_lb = createWorkerLb(); + + EXPECT_EQ(h, worker_lb->peekAnotherHost(nullptr)); + EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); +} + TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverHealthyRoutingPaths) { auto h_local = makeWeightTrackingMockHost(); auto h_remote = makeWeightTrackingMockHost(); From c857553b017a488336568eec432b4591d2bf096d Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:25:01 -0600 Subject: [PATCH 13/16] priority-growth sync fixes, weight-cache refresh on raw compare, spill_active stat, review cleanups Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_balancing/load_aware_locality.rst | 37 +++-- .../load_aware_locality/BUILD | 1 - .../load_aware_locality/config.cc | 14 +- .../load_aware_locality_lb.cc | 155 ++++++++++++------ .../load_aware_locality_lb.h | 40 +++-- .../load_aware_locality/config_test.cc | 30 ++-- .../load_aware_locality/integration_test.cc | 56 +++++-- .../load_aware_locality_lb_test.cc | 140 +++++++++++----- 8 files changed, 332 insertions(+), 141 deletions(-) diff --git a/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst b/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst index 92e6c9dc6853e..2f99f30a66658 100644 --- a/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst +++ b/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst @@ -207,17 +207,23 @@ above: # Per-tick smoothing factor (consistent settling regardless of tick rate) alpha = 1 - exp(-weight_update_period / smoothing_time_constant) - # Per-host sample validity filter (excludes hosts whose last ORCA report - # is older than weight_expiration_period; if disabled, all hosts qualify) - valid(h) = (now - last_report_time(h)) <= weight_expiration_period + # Per-host sample validity filter (excludes hosts that have never reported, + # plus hosts whose last ORCA report is older than weight_expiration_period; + # if expiration is disabled, all reporting hosts qualify) + valid(h) = has_reported(h) and + (now - last_report_time(h)) <= weight_expiration_period valid_hosts(L) = { h in hosts(L) : valid(h) } # Per-locality utilization (EWMA smoothed; first sample applied raw so # the policy reacts within one tick instead of waiting ~5 time constants # to converge from the cold-start prior of 0) if valid_hosts(L) is empty: - smoothed_util(L) = prev_smoothed_util(L) # carry prior value - stale(L) = true + if prev_smoothed_util(L) exists: + smoothed_util(L) = prev_smoothed_util(L) # carry prior value + stale(L) = true + else: # never sampled (cold start) + smoothed_util(L) = 0 + stale(L) = false else: raw_util(L) = avg over h in valid_hosts(L) of util(h) if no prior smoothed_util(L): # first sample for L @@ -235,12 +241,13 @@ above: base_weight(L) = host_count(L) * headroom(L) total_base_weight = sum(base_weight(L_i) for all L_i) + total_host_count = sum(host_count(L_i) for all L_i) remote_host_count = sum(host_count(R_i) for all remote R_i) - # All-overloaded fallback - if total_base_weight == 0: + # All-overloaded fallback (skipped when the subset has no hosts at all) + if total_base_weight == 0 and total_host_count > 0: adjusted_weight(L_i) = host_count(L_i) - else: + elif total_base_weight > 0: adjusted_weight(L_i) = base_weight(L_i) if local exists and remote_host_count > 0: @@ -363,10 +370,8 @@ priority load calculation; locality selection then applies within the chosen priority. Unlike zone-aware routing (priority 0 only), this policy applies at all priority levels. -The policy honors ``common_lb_config.healthy_panic_threshold`` and -``common_lb_config.zone_aware_lb_config.fail_traffic_on_panic``. When -in panic mode the policy routes to all hosts; when ``fail_traffic_on_panic`` -is set it returns no host instead. +The policy honors ``common_lb_config.healthy_panic_threshold``. When +in panic mode the policy routes to all hosts. Three independent weight sets are maintained per priority: @@ -478,6 +483,9 @@ The policy emits stats under ``cluster..load_aware_locality.*``: 100% local. * - ``probe_active_total`` - Per tick where ``remote_probe_fraction`` redistribution kicked in. + * - ``spill_active_total`` + - Per tick where local utilization exceeded the remote average by more + than the variance threshold, spilling traffic to remote localities. * - ``stale_locality_total`` - Incremented once per stale locality per priority per tick: a locality that previously had fresh reports but whose hosts' samples have all @@ -487,6 +495,11 @@ The policy emits stats under ``cluster..load_aware_locality.*``: single-priority 5-locality cluster with 2 stale localities adds 2 each tick. +The four condition counters increment at most once per tick, ORed across +every priority and each of the three host subsets (healthy, degraded, +all-hosts) the policy weighs independently; ``stale_locality_total`` is +derived from the all-hosts subset. + Migrating from zone-aware routing? The per-request zone routing counters (``lb_zone_routing_all_directly``, ``lb_zone_routing_sampled``, ``lb_zone_routing_cross_zone``) are still emitted with equivalent semantics, diff --git a/source/extensions/load_balancing_policies/load_aware_locality/BUILD b/source/extensions/load_balancing_policies/load_aware_locality/BUILD index e114f1ffad8f5..95f316eb2c5af 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/BUILD +++ b/source/extensions/load_balancing_policies/load_aware_locality/BUILD @@ -40,7 +40,6 @@ envoy_cc_library( "//source/extensions/load_balancing_policies/common:orca_weight_manager_lib", "@abseil-cpp//absl/container:flat_hash_map", "@abseil-cpp//absl/container:flat_hash_set", - "@abseil-cpp//absl/container:inlined_vector", "@abseil-cpp//absl/strings", ], ) diff --git a/source/extensions/load_balancing_policies/load_aware_locality/config.cc b/source/extensions/load_balancing_policies/load_aware_locality/config.cc index 448c3923f4f16..b9e3253e0593f 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/config.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/config.cc @@ -22,7 +22,8 @@ Factory::create(OptRef lb_config, absl::StatusOr Factory::loadConfig(Server::Configuration::ServerFactoryContext& context, const Protobuf::Message& config) { - const auto& lb_config = dynamic_cast(config); + const auto& lb_config = + Envoy::Protobuf::DynamicCastMessage(config); // Validate policy-level knobs before resolving the child policy so a config error is reported // for what it is rather than masked by an unresolvable child. @@ -30,6 +31,10 @@ Factory::loadConfig(Server::Configuration::ServerFactoryContext& context, return absl::InvalidArgumentError( "load_aware_locality: enable_oob_load_report is not yet supported"); } + if (lb_config.has_oob_reporting_period()) { + return absl::InvalidArgumentError("load_aware_locality: oob_reporting_period requires " + "enable_oob_load_report, which is not yet supported"); + } const std::chrono::milliseconds weight_update_period = std::chrono::milliseconds(PROTOBUF_GET_MS_OR_DEFAULT(lb_config, weight_update_period, 1000)); // PGV constraints are not enforced on unpacked LB configs, so the documented floor is owned @@ -72,6 +77,11 @@ Factory::loadConfig(Server::Configuration::ServerFactoryContext& context, /*is_optional=*/true); if (endpoint_picking_policy_factory != nullptr) { + // Self-nesting double-registers stats and drops the outer instance's metric names. + if (endpoint_picking_policy_factory->name() == name()) { + return absl::InvalidArgumentError( + "load_aware_locality cannot be its own endpoint_picking_policy"); + } // Load and validate the child policy configuration. auto sub_lb_proto_message = endpoint_picking_policy_factory->createEmptyConfigProto(); RETURN_IF_NOT_OK(Config::Utility::translateOpaqueConfig( @@ -83,7 +93,7 @@ Factory::loadConfig(Server::Configuration::ServerFactoryContext& context, RETURN_IF_NOT_OK(lb_config_or_error.status()); return std::make_unique( - *endpoint_picking_policy_factory, endpoint_picking_policy_factory->name(), + *endpoint_picking_policy_factory, LoadBalancerConfigSharedPtr(std::move(lb_config_or_error.value())), weight_update_period, utilization_variance_threshold, ewma_alpha, remote_probe_fraction, weight_expiration_period, std::move(metric_names), context.mainThreadDispatcher(), diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc index c9875365edde5..26603dcd2ebe3 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -14,7 +14,6 @@ #include "source/extensions/load_balancing_policies/common/orca_weight_manager.h" #include "absl/container/flat_hash_set.h" -#include "absl/container/inlined_vector.h" #include "absl/strings/str_cat.h" namespace Envoy { @@ -22,6 +21,16 @@ namespace Extensions { namespace LoadBalancingPolicies { namespace LoadAwareLocality { +namespace { + +int64_t monotonicMs(TimeSource& time_source) { + return std::chrono::duration_cast( + time_source.monotonicTime().time_since_epoch()) + .count(); +} + +} // namespace + absl::Status LocalityLbHostData::onOrcaLoadReport(const Upstream::OrcaLoadReport& report, const StreamInfo::StreamInfo&) { const double util = @@ -30,9 +39,7 @@ absl::Status LocalityLbHostData::onOrcaLoadReport(const Upstream::OrcaLoadReport if (util <= 0.0) { return absl::OkStatus(); } - const int64_t now_ms = std::chrono::duration_cast( - time_source_.monotonicTime().time_since_epoch()) - .count(); + const int64_t now_ms = monotonicMs(time_source_); storeUtilization(util, now_ms); return absl::OkStatus(); } @@ -115,13 +122,12 @@ void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { per_source_state.resize(host_sets.size()); } - const int64_t now_ms = std::chrono::duration_cast( - time_source_.monotonicTime().time_since_epoch()) - .count(); + const int64_t now_ms = monotonicMs(time_source_); bool tick_all_overloaded = false; bool tick_local_preferred = false; bool tick_probe_active = false; + bool tick_spill_active = false; uint32_t tick_stale_localities = 0; for (size_t priority = 0; priority < host_sets.size(); ++priority) { @@ -158,6 +164,8 @@ void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { all_hosts_result.local_preferred; tick_probe_active |= healthy_result.probe_active || degraded_result.probe_active || all_hosts_result.probe_active; + tick_spill_active |= healthy_result.spill_active || degraded_result.spill_active || + all_hosts_result.spill_active; tick_stale_localities += all_hosts_result.stale_localities; } @@ -170,6 +178,9 @@ void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { if (tick_probe_active) { lb_stats_.probe_active_total_.inc(); } + if (tick_spill_active) { + lb_stats_.spill_active_total_.inc(); + } if (tick_stale_localities > 0) { lb_stats_.stale_locality_total_.add(tick_stale_localities); } @@ -287,6 +298,10 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( } } + if (remote_hosts > 0 && !all_local) { + result.spill_active = true; + } + if (remote_hosts > 0 && remote_probe_fraction_ > 0.0) { const double total = std::accumulate(weights.begin(), weights.end(), 0.0); const double remote_target = total * remote_probe_fraction_; @@ -324,9 +339,7 @@ WorkerLocalLbFactory::WorkerLocalLbFactory( : child_worker_factory_(std::move(child_worker_factory)), child_config_(std::move(child_config)), cluster_info_(cluster_info), random_(random), runtime_(runtime), healthy_panic_threshold_(PROTOBUF_PERCENT_TO_ROUNDED_INTEGER_OR_DEFAULT( - cluster_info.lbConfig(), healthy_panic_threshold, 100, 50)), - fail_traffic_on_panic_( - cluster_info.lbConfig().zone_aware_lb_config().fail_traffic_on_panic()) { + cluster_info.lbConfig(), healthy_panic_threshold, 100, 50)) { tls_ = ThreadLocal::TypedSlot::makeUnique(tls_slot_allocator); tls_->set([](Event::Dispatcher&) { return std::make_shared(); }); } @@ -457,10 +470,11 @@ void WorkerLocalLb::updateLocalityHosts(PerSourceLocalityState& state, void WorkerLocalLb::buildPerPriorityLocalities() { const auto& host_sets = priority_set_.hostSetsPerPriority(); - per_priority_locality_.clear(); + // Resize preserves existing priorities' child LBs; build only the newly added ones. + const size_t old_size = per_priority_locality_.size(); per_priority_locality_.resize(host_sets.size()); - for (size_t priority = 0; priority < host_sets.size(); ++priority) { + for (size_t priority = old_size; priority < host_sets.size(); ++priority) { buildPerLocality(priority, *host_sets[priority]); } } @@ -502,18 +516,23 @@ void WorkerLocalLb::syncLocalityState(PerLocalityState& state, const Upstream::H } const auto& old_hosts = source_state.priority_set->hostSetsPerPriority()[0]->hosts(); - absl::flat_hash_set old_set(old_hosts.begin(), old_hosts.end()); + absl::flat_hash_set old_set; + old_set.reserve(old_hosts.size()); + for (const auto& host : old_hosts) { + old_set.insert(host.get()); + } Upstream::HostVector hosts_added; Upstream::HostVector hosts_removed; for (const auto& host : new_hosts) { - if (!old_set.erase(host)) { + if (!old_set.erase(host.get())) { hosts_added.push_back(host); } } - hosts_removed.reserve(old_set.size()); - for (const auto& host : old_set) { - hosts_removed.push_back(std::const_pointer_cast(host)); + for (const auto& host : old_hosts) { + if (old_set.contains(host.get())) { + hosts_removed.push_back(host); + } } const bool membership_changed = !hosts_added.empty() || !hosts_removed.empty(); @@ -539,6 +558,8 @@ void WorkerLocalLb::buildPerLocality(uint32_t priority, const Upstream::HostSet& // Locality routing structures for this priority are (re)generated on this worker, the same // event the zone-aware counter tracks. stats_.lb_recalculate_zone_structures_.inc(); + // Topology changed: force refreshLocalityWeights to rebuild the index-aligned weights next pick. + built_snapshot_.reset(); const auto& locality_hosts = host_set.hostsPerLocality().get(); auto& per_locality = per_priority_locality_[priority].localities; per_locality.clear(); @@ -555,12 +576,16 @@ void WorkerLocalLb::syncPriority(uint32_t priority, bool allow_rebuild) { return; } - // Priority count changed; full rebuild when permitted, else wait for a membership-delta update. + // Priority count changed; build the newly added priorities when permitted. if (host_sets.size() != per_priority_locality_.size()) { + const size_t built = per_priority_locality_.size(); if (allow_rebuild) { buildPerPriorityLocalities(); } - return; + if (priority >= built) { + return; // Newly built or still-pending priority; the build read current membership. + } + // An existing priority triggered the callback; fall through to apply its update. } const auto& host_set = host_sets[priority]; @@ -579,6 +604,11 @@ void WorkerLocalLb::syncPriority(uint32_t priority, bool allow_rebuild) { for (size_t i = 0; i < locality_hosts.size(); ++i) { syncLocalityState(per_locality[i], *host_set, i, recreate_child); } + // A membership delta can flip a locality's identity (e.g. draining to empty), so rebuild the + // index-aligned weights on the next pick. + if (allow_rebuild) { + built_snapshot_.reset(); + } } Upstream::LoadBalancer* @@ -617,8 +647,7 @@ WorkerLocalLb::resolvePrioritySource(Upstream::LoadBalancerContext* context, boo const uint32_t priority = host_set_and_availability.first.priority(); if (isInPanic(priority)) { - return {priority, SelectionSource::AllHosts, /*in_panic=*/true, - /*fail=*/factory_.failTrafficOnPanic()}; + return {priority, SelectionSource::AllHosts, /*in_panic=*/true, /*fail=*/false}; } const SelectionSource source = host_set_and_availability.second == Upstream::LoadBalancerBase::HostAvailability::Healthy @@ -634,36 +663,33 @@ WorkerLocalLb::LocalityLbSelection WorkerLocalLb::selectLocalityLb(const Priorit return selection; } - selection.priority = pick.priority; - selection.source = pick.source; - // A custom retry-priority plugin can route to a just-created priority that // per_priority_locality_ doesn't cover yet (empty-delta update, rebuild pending). - if (selection.priority >= per_priority_locality_.size()) { + if (pick.priority >= per_priority_locality_.size()) { return selection; } - auto& per_locality = per_priority_locality_[selection.priority].localities; + auto& per_locality = per_priority_locality_[pick.priority].localities; if (per_locality.empty()) { return selection; } - selection.snapshot = factory_.routingWeights(); - const size_t preferred_idx = - chooseLocality(selection.snapshot, selection.priority, selection.source, peeking); + const RoutingWeightsSnapshot* snapshot = factory_.routingWeights(); + selection.snapshot = snapshot; + refreshLocalityWeights(snapshot); + const size_t preferred_idx = chooseLocality(pick.priority, pick.source, peeking); selection.locality_idx = preferred_idx; - selection.lb = - pickLocalityLb(per_locality, selection.source, preferred_idx, selection.locality_idx); + selection.lb = pickLocalityLb(per_locality, pick.source, preferred_idx, selection.locality_idx); return selection; } void WorkerLocalLb::recordZoneRoutingStats(const PrioritySourcePick& pick, const LocalityLbSelection& selection) { if (pick.in_panic || selection.snapshot == nullptr || - selection.priority >= selection.snapshot->priority_weights.size()) { + pick.priority >= selection.snapshot->priority_weights.size()) { return; } - const auto& priority_snapshot = selection.snapshot->priority_weights[selection.priority]; + const auto& priority_snapshot = selection.snapshot->priority_weights[pick.priority]; if (!priority_snapshot.has_local_locality) { return; } @@ -671,7 +697,7 @@ void WorkerLocalLb::recordZoneRoutingStats(const PrioritySourcePick& pick, stats_.lb_zone_routing_cross_zone_.inc(); return; } - if (priority_snapshot.allLocalFor(selection.source)) { + if (priority_snapshot.allLocalFor(pick.source)) { stats_.lb_zone_routing_all_directly_.inc(); } else { stats_.lb_zone_routing_sampled_.inc(); @@ -699,29 +725,52 @@ Upstream::HostSelectionResponse WorkerLocalLb::chooseHost(Upstream::LoadBalancer return failOnAsyncSelection(selection.lb->chooseHost(&child_context)); } -size_t WorkerLocalLb::chooseLocality(const RoutingWeightsSnapshot* snapshot, uint32_t priority, +void WorkerLocalLb::refreshLocalityWeights(const RoutingWeightsSnapshot* snapshot) { + if (snapshot == built_snapshot_.get()) { + return; // Already built for this snapshot; a topology change reset built_snapshot_. + } + // Cold path: take shared ownership so built_snapshot_ pins this snapshot's address and the raw + // compare above stays ABA-free. + built_snapshot_ = factory_.routingWeightsShared(); + ASSERT(built_snapshot_.get() == snapshot); + + for (size_t priority = 0; priority < per_priority_locality_.size(); ++priority) { + auto& pstate = per_priority_locality_[priority]; + const size_t locality_count = pstate.localities.size(); + const bool have_priority = snapshot != nullptr && priority < snapshot->priority_weights.size(); + + for (size_t s = 0; s < pstate.source_weights.size(); ++s) { + auto& index_weights = pstate.source_weights[s]; + index_weights.assign(locality_count, 0.0); + double total = 0.0; + if (have_priority) { + // Once-per-publish replacement for chooseLocality's old per-pick weights.find(locality). + // Snapshot-missing live localities get weight 0 until the next recompute. + const auto& weights = snapshot->priority_weights[priority].weightsFor( + static_cast(s)); + for (size_t i = 0; i < locality_count; ++i) { + const auto weight = weights.find(pstate.localities[i].locality); + index_weights[i] = weight != weights.end() ? weight->second : 0.0; + total += index_weights[i]; + } + } + pstate.source_total[s] = total; + } + } +} + +size_t WorkerLocalLb::chooseLocality(uint32_t priority, PriorityRoutingWeights::SelectionSource source, bool peeking) { - const auto& per_locality = per_priority_locality_[priority].localities; + const auto& pstate = per_priority_locality_[priority]; + const auto& per_locality = pstate.localities; if (per_locality.size() <= 1) { return 0; } - if (snapshot == nullptr || priority >= snapshot->priority_weights.size()) { - return 0; - } - const auto& weights = snapshot->priority_weights[priority].weightsFor(source); - if (weights.empty()) { - return 0; - } - - // Snapshot-missing live localities get weight 0 until the next recompute. - absl::InlinedVector locality_weights(per_locality.size()); - double effective_total = 0.0; - for (size_t i = 0; i < per_locality.size(); ++i) { - const auto weight = weights.find(per_locality[i].locality); - locality_weights[i] = weight != weights.end() ? weight->second : 0.0; - effective_total += locality_weights[i]; - } - if (effective_total <= 0.0) { + const size_t s = static_cast(source); + const auto& locality_weights = pstate.source_weights[s]; + const double effective_total = pstate.source_total[s]; + // Defensive: only reachable before the first snapshot publish, where source_total is 0. + if (locality_weights.size() != per_locality.size() || effective_total <= 0.0) { return 0; } diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h index 9c253f5113b01..9945ba4ea4ca1 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -34,6 +34,7 @@ namespace LoadAwareLocality { COUNTER(all_overloaded_total) \ COUNTER(local_preferred_total) \ COUNTER(probe_active_total) \ + COUNTER(spill_active_total) \ COUNTER(stale_locality_total) struct LoadAwareLocalityStats { @@ -85,7 +86,6 @@ using LoadBalancerConfigSharedPtr = std::shared_ptrrouting_weights.get() : nullptr; } + // Shared-ownership variant: a worker holds the snapshot it derived its index-aligned weights + // from, so the identity comparison in refreshLocalityWeights is free of pointer-reuse (ABA). + RoutingWeightsSnapshotConstSharedPtr routingWeightsShared() const { + auto shim = tls_->get(); + return shim.has_value() ? shim->routing_weights : nullptr; + } + Upstream::LoadBalancerPtr createWorkerChildLb(Upstream::PrioritySetImpl& per_locality_priority_set); @@ -219,7 +225,6 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { Envoy::Random::RandomGenerator& random() const { return random_; } Runtime::Loader& runtime() const { return runtime_; } uint32_t healthyPanicThreshold() const { return healthy_panic_threshold_; } - bool failTrafficOnPanic() const { return fail_traffic_on_panic_; } Upstream::ClusterLbStats& lbStats() const { return cluster_info_.lbStats(); } private: @@ -227,12 +232,13 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { // ownership keeps it valid on workers after the child LB is destroyed, per the // ThreadAwareLoadBalancer factory contract. Upstream::LoadBalancerFactorySharedPtr child_worker_factory_; + // Co-owns the child endpoint-picking config so it outlives the child LB on worker threads, even + // when the LoadAwareLocalityLbConfig that also owns it is destroyed first. LoadBalancerConfigSharedPtr child_config_; const Upstream::ClusterInfo& cluster_info_; Envoy::Random::RandomGenerator& random_; Runtime::Loader& runtime_; uint32_t healthy_panic_threshold_; - const bool fail_traffic_on_panic_; std::unique_ptr> tls_; }; @@ -257,6 +263,11 @@ struct PerLocalityState { struct PerPriorityLocalityState { std::vector localities; + // Index-aligned routing weights derived from the snapshot once per publish, so the per-pick path + // reads by locality index instead of hashing locality identity. Rebuilt on snapshot or topology + // change; source_total[s] is the sum of source_weights[s]. + std::array, 3> source_weights; + std::array source_total{{0.0, 0.0, 0.0}}; }; // Worker-local LB: picks live priority/source, then delegates to a per-locality child LB. @@ -294,9 +305,6 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { struct LocalityLbSelection { Upstream::LoadBalancer* lb{nullptr}; const RoutingWeightsSnapshot* snapshot{nullptr}; - uint32_t priority{0}; - PriorityRoutingWeights::SelectionSource source{ - PriorityRoutingWeights::SelectionSource::Healthy}; size_t locality_idx{0}; }; @@ -307,8 +315,12 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { void recordZoneRoutingStats(const PrioritySourcePick& pick, const LocalityLbSelection& selection); // Returns 0 for single locality, no/stale snapshot, or zero effective total. - size_t chooseLocality(const RoutingWeightsSnapshot* snapshot, uint32_t priority, - PriorityRoutingWeights::SelectionSource source, bool peeking); + size_t chooseLocality(uint32_t priority, PriorityRoutingWeights::SelectionSource source, + bool peeking); + + // Rebuilds per-priority index-aligned weights when the published snapshot changes (or a topology + // change reset built_snapshot_), so chooseLocality reads by index instead of hashing per pick. + void refreshLocalityWeights(const RoutingWeightsSnapshot* snapshot); // Falls back to any usable child LB when the routing snapshot lags membership. Upstream::LoadBalancer* pickLocalityLb(const std::vector& per_locality, @@ -317,6 +329,9 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { WorkerLocalLbFactory& factory_; std::vector per_priority_locality_; + // Snapshot the index-aligned weights in per_priority_locality_ were built from; reset to force a + // rebuild after a topology change. + RoutingWeightsSnapshotConstSharedPtr built_snapshot_; // Destroyed explicitly in the destructor before other members so the callback doesn't fire // during destruction and access freed per-locality state. Envoy::Common::CallbackHandlePtr priority_sync_cb_; @@ -345,6 +360,7 @@ class LoadAwareLocalityLoadBalancer : public Upstream::ThreadAwareLoadBalancer, bool all_overloaded{false}; bool local_preferred{false}; bool probe_active{false}; + bool spill_active{false}; uint32_t stale_localities{0}; }; diff --git a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc index b6826cd03f9af..50f62fb881e45 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc @@ -58,8 +58,6 @@ void expectFactoryCreateSucceeds(const Upstream::LoadBalancerConfigPtr& config, Server::Configuration::MockServerFactoryContext& context) { NiceMock cluster_info; NiceMock priority_set; - NiceMock dispatcher; - ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); Factory factory; auto lb = factory.create(*config, cluster_info, priority_set, context.runtime_loader_, @@ -93,8 +91,6 @@ class RejectingEndpointValidationConfig : public Upstream::LoadBalancerConfig { TEST(LoadAwareLocalityConfigTest, Defaults) { NiceMock context; - NiceMock dispatcher; - ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); Factory factory; auto result = factory.loadConfig(context, loadAwareConfig({roundRobinEndpointPolicy()})); @@ -113,8 +109,6 @@ TEST(LoadAwareLocalityConfigTest, Defaults) { TEST(LoadAwareLocalityConfigTest, Overrides) { NiceMock context; - NiceMock dispatcher; - ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); config_proto.mutable_weight_update_period()->set_seconds(2); @@ -148,10 +142,21 @@ TEST(LoadAwareLocalityConfigTest, MissingChildPolicy) { EXPECT_EQ(result.status(), absl::InvalidArgumentError("No supported endpoint picking policy.")); } +TEST(LoadAwareLocalityConfigTest, SelfNestedChildPolicyRejected) { + NiceMock context; + + Factory factory; + auto result = factory.loadConfig( + context, + loadAwareConfig({typedExtensionConfig("envoy.load_balancing_policies.load_aware_locality", + loadAwareConfig({roundRobinEndpointPolicy()}))})); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(result.status().message(), testing::HasSubstr("its own endpoint_picking_policy")); +} + TEST(LoadAwareLocalityConfigTest, FirstSupportedChildWins) { NiceMock context; - NiceMock dispatcher; - ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); envoy::config::core::v3::TypedExtensionConfig unknown_policy; unknown_policy.set_name("envoy.load_balancing_policies.does_not_exist"); @@ -216,6 +221,11 @@ TEST(LoadAwareLocalityConfigTest, InvalidPolicyKnobsRejected) { config.mutable_enable_oob_load_report()->set_value(true); }, "enable_oob_load_report"}, + {"OobReportingPeriodWithoutOobEnabled", + [](LoadAwareLocalityLbProto& config) { + config.mutable_oob_reporting_period()->set_seconds(10); + }, + "oob_reporting_period"}, {"NegativeWeightExpiration", [](LoadAwareLocalityLbProto& config) { config.mutable_weight_expiration_period()->set_seconds(-1); @@ -297,7 +307,7 @@ TEST(LoadAwareLocalityConfigTest, EndpointValidationDelegatesToChildConfig) { NiceMock dispatcher; LoadAwareLocalityLbConfig config( - child_factory, "mock_child", std::make_shared(), + child_factory, std::make_shared(), std::chrono::milliseconds(1000), 0.1, 0.3, 0.03, std::chrono::milliseconds(180000), std::vector{}, dispatcher, context.thread_local_); @@ -309,8 +319,6 @@ TEST(LoadAwareLocalityConfigTest, EndpointValidationDelegatesToChildConfig) { TEST(LoadAwareLocalityConfigTest, MetricNamesReachConfig) { NiceMock context; - NiceMock dispatcher; - ON_CALL(context, mainThreadDispatcher()).WillByDefault(ReturnRef(dispatcher)); auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); config_proto.add_metric_names_for_computing_utilization("named_metrics.foo"); diff --git a/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc index f0543ea560adc..e6cd4a0302884 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc @@ -22,7 +22,6 @@ class LoadAwareLocalityIntegrationTest : public testing::TestWithParamcounter(counter); + before = c != nullptr ? c->value() : 0; + } + timeSystem().advanceTimeWait(std::chrono::seconds(11)); + if (!counter.empty()) { + test_server_->waitForCounter(counter, testing::Ge(before + 1)); + } + } + + // Sends requests until every upstream has served one, so every host has reported utilization + // and the next tick's weight-computation branch is deterministic instead of sampling-dependent. + void seedAllUpstreams(const std::vector& utilizations) { + std::vector seen(num_upstreams_, false); + uint32_t covered = 0; + for (int i = 0; i < 500 && covered < num_upstreams_; ++i) { + const uint64_t idx = sendRequestWithOrcaResponse(utilizations); + if (!seen[idx]) { + seen[idx] = true; + ++covered; + } + } + ASSERT_EQ(num_upstreams_, covered); + } - void seedWithTwoCycles(const std::vector& utilizations) { - sendRequestsAndTrack(30, utilizations); - advanceWeightTick(); + void seedWithTwoCycles(const std::vector& utilizations, + const std::string& counter_name = "") { + seedAllUpstreams(utilizations); + advanceWeightTick(counter_name); sendRequestsAndTrack(30, utilizations); - advanceWeightTick(); + advanceWeightTick(counter_name); } uint64_t zoneTraffic(const std::vector& usage, size_t zone_index) const { @@ -175,7 +204,7 @@ TEST_P(LoadAwareLocalityIntegrationTest, StrictLocalWithRoundRobinWithinLocality const std::vector utilizations = {0.05, 0.05, 0.05, 0.05}; sendRequestsAndTrack(30, utilizations); - advanceWeightTick(); + advanceWeightTick("local_preferred_total"); const auto usage = sendRequestsAndTrack(100, utilizations); @@ -194,7 +223,7 @@ TEST_P(LoadAwareLocalityIntegrationTest, ProbeFloorKeepsRemoteSampledInAllLocalM /*weight_update_period_seconds=*/10); const std::vector balanced = {0.4, 0.4, 0.4, 0.4}; - seedWithTwoCycles(balanced); + seedWithTwoCycles(balanced, "probe_active_total"); const auto usage = sendRequestsAndTrack(400, balanced); const uint64_t local = localTraffic(usage); @@ -211,7 +240,7 @@ TEST_P(LoadAwareLocalityIntegrationTest, AdaptiveSpillAndRecovery) { /*weight_update_period_seconds=*/10); const std::vector overloaded_local = {0.9, 0.9, 0.2, 0.2}; - seedWithTwoCycles(overloaded_local); + seedWithTwoCycles(overloaded_local, "spill_active_total"); const auto phase1_usage = sendRequestsAndTrack(100, overloaded_local); const uint64_t phase1_local = localTraffic(phase1_usage); @@ -221,7 +250,8 @@ TEST_P(LoadAwareLocalityIntegrationTest, AdaptiveSpillAndRecovery) { const std::vector rebalanced = {0.3, 0.3, 0.3, 0.3}; sendRequestsAndTrack(60, rebalanced); - advanceWeightTick(); + // Rebalanced utilization snaps local, incrementing local_preferred_total. + advanceWeightTick("local_preferred_total"); const auto phase2_usage = sendRequestsAndTrack(100, rebalanced); const uint64_t phase2_local = localTraffic(phase2_usage); @@ -235,11 +265,11 @@ TEST_P(LoadAwareLocalityIntegrationTest, EwmaDampensSpike) { /*weight_update_period_seconds=*/10, /*smoothing_time_constant_seconds=*/14); const std::vector baseline = {0.2, 0.2, 0.2, 0.2}; - seedWithTwoCycles(baseline); + seedWithTwoCycles(baseline, "local_preferred_total"); const std::vector spiked_local = {0.9, 0.9, 0.2, 0.2}; sendRequestsAndTrack(30, spiked_local); - advanceWeightTick(); + advanceWeightTick("spill_active_total"); const auto usage = sendRequestsAndTrack(100, spiked_local); const uint64_t remote = remoteTraffic(usage); @@ -253,7 +283,7 @@ TEST_P(LoadAwareLocalityIntegrationTest, ThreeLocalityDistribution) { /*remote_zones=*/{"zone-b", "zone-c"}); const std::vector utilizations = {0.8, 0.8, 0.3, 0.3, 0.5, 0.5}; - seedWithTwoCycles(utilizations); + seedWithTwoCycles(utilizations, "spill_active_total"); const auto usage = sendRequestsAndTrack(400, utilizations); const uint64_t zone_a = zoneTraffic(usage, 0); diff --git a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc index acbda6c7f6399..a856b68bcbf60 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -205,9 +205,9 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes round_robin_factory.loadConfig(context_, *round_robin_proto).value(); auto lb_config = std::make_unique( - round_robin_factory, round_robin_factory.name(), std::move(round_robin_lb_config), - weight_update_period, variance_threshold, ewma_alpha, remote_probe_fraction, - weight_expiration_period, std::vector{}, dispatcher_, context_.thread_local_); + round_robin_factory, std::move(round_robin_lb_config), weight_update_period, + variance_threshold, ewma_alpha, remote_probe_fraction, weight_expiration_period, + std::vector{}, dispatcher_, context_.thread_local_); timer_ = new NiceMock(&dispatcher_); @@ -327,16 +327,15 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes void recomputeWeights() { ASSERT_TRUE(thread_aware_lb_->initialize().ok()); } // Reshapes a priority's host set, then fires the membership update callback. Two modes: - // - full membership reshape: sets hosts_/healthy_hosts_ and both per-locality vectors from - // `localities` (healthy mirrors all hosts). + // - full membership reshape: delegates to setupPriorityLocalities (healthy mirrors all hosts). // - health-only flip (when `healthy_override` is set): leaves hosts_/hosts_per_locality_ intact // and only sets healthy_hosts_/healthy_hosts_per_locality_ from `healthy_override`. void reshapeLocalities( uint32_t priority, std::vector localities, Upstream::HostVector added = {}, Upstream::HostVector removed = {}, bool has_local = false, std::optional> healthy_override = std::nullopt) { - auto* host_set = priority_set_.getMockHostSet(priority); if (healthy_override.has_value()) { + auto* host_set = priority_set_.getMockHostSet(priority); Upstream::HostVector healthy_hosts; for (const auto& locality : *healthy_override) { healthy_hosts.insert(healthy_hosts.end(), locality.begin(), locality.end()); @@ -345,20 +344,7 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes host_set->healthy_hosts_per_locality_ = Upstream::makeHostsPerLocality(std::move(*healthy_override), !has_local); } else { - for (size_t i = 0; i < localities.size(); ++i) { - for (const auto& host : localities[i]) { - setHostLocality(host, i); - } - } - Upstream::HostVector all_hosts; - for (const auto& locality : localities) { - all_hosts.insert(all_hosts.end(), locality.begin(), locality.end()); - } - host_set->hosts_ = all_hosts; - host_set->healthy_hosts_ = all_hosts; - host_set->hosts_per_locality_ = - Upstream::makeHostsPerLocality(std::move(localities), !has_local); - host_set->healthy_hosts_per_locality_ = host_set->hosts_per_locality_; + setupPriorityLocalities(priority, std::move(localities), has_local); } priority_set_.runUpdateCallbacks(priority, added, removed); } @@ -505,14 +491,15 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes void expectTwoHostSnapCounter(const std::string& counter, double variance_threshold, double remote_probe_fraction, std::optional> quiet_utilizations, - std::pair active_utilizations) { + std::pair active_utilizations, + uint64_t startup_count = 1) { auto h_local = makeWeightTrackingMockHost(); auto h_remote = makeWeightTrackingMockHost(); setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); createLb(variance_threshold, /*ewma_alpha=*/1.0, remote_probe_fraction); const uint64_t after_create = counterValue(cluster_info_, counter); - EXPECT_EQ(1u, after_create); + EXPECT_EQ(startup_count, after_create); if (quiet_utilizations.has_value()) { setHostUtilization(h_local, quiet_utilizations->first); @@ -628,6 +615,15 @@ INSTANTIATE_TEST_SUITE_P(WeightComputation, WeightComputationTest, true, {0.94, 0.02, 0.04}, 0.01}, + {"ProbeHostCountProportionalUnequalUtils", + {1, 2, 4}, + true, + 1.0, + 0.06, + {0.5, 0.8, 0.2}, + true, + {0.94, 0.02, 0.04}, + 0.01}, {"ProbeRedistributionClampsLocalAtZero", {1, 1}, true, @@ -702,6 +698,12 @@ TEST_F(LoadAwareLocalityLbTest, WeightExpirationIgnoresStaleData) { setHostUtilization(h2, 0.1); recomputeWeights(); expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.9}); + + // Advance time past the 5s expiration window. The stale data is ignored, and weights + // fall back to host counts (1.0, 1.0). + context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {1.0, 1.0}); } TEST_F(LoadAwareLocalityLbTest, WeightExpirationDisabledUsesStaleData) { @@ -942,19 +944,6 @@ TEST_F(LoadAwareLocalityLbTest, PanicUsesAllHosts) { EXPECT_GT(cluster_info_.lbStats().lb_healthy_panic_.value(), 0u); } -TEST_F(LoadAwareLocalityLbTest, FailTrafficOnPanicReturnsNoHost) { - cluster_info_.lb_config_.mutable_zone_aware_lb_config()->set_fail_traffic_on_panic(true); - - setupPanicHosts(); - createLb(); - - auto worker_lb = createWorkerLb(); - for (int i = 0; i < 20; ++i) { - EXPECT_EQ(nullptr, worker_lb->chooseHost(nullptr).host); - } - EXPECT_GT(cluster_info_.lbStats().lb_healthy_panic_.value(), 0u); -} - TEST_F(LoadAwareLocalityLbTest, MembershipAndTopologyUpdatesShareTheSameWorkerLb) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost(); @@ -1111,6 +1100,15 @@ TEST_F(LoadAwareLocalityLbTest, DedicatedStatsProbeActiveTotalIncOncePerTickWhen std::make_pair(0.3, 1.0)); } +TEST_F(LoadAwareLocalityLbTest, DedicatedStatsSpillActiveTotalIncOncePerTickWhenSpilling) { + // Startup no-data tick snaps local, so the spill counter starts at 0; the balanced quiet phase + // also snaps; only the overloaded-local phase spills. + expectTwoHostSnapCounter("load_aware_locality.spill_active_total", + /*variance_threshold=*/0.1, /*remote_probe_fraction=*/0.0, + std::make_pair(0.2, 0.2), std::make_pair(0.9, 0.2), + /*startup_count=*/0); +} + TEST_F(LoadAwareLocalityLbTest, DedicatedStatsStaleLocalityTotalCountsOncePerStalePerTick) { auto locality_a = makeHosts(2); auto locality_b = makeHosts(1); @@ -1382,8 +1380,8 @@ TEST_F(LoadAwareLocalityLbTest, InitializePropagatesChildFactoryError) { .WillByDefault(testing::Return(testing::ByMove(nullptr))); auto lb_config = std::make_unique( - null_child_factory, "mock_null_factory", nullptr, std::chrono::milliseconds(1000), 0.1, 0.3, - 0.03, std::chrono::milliseconds(180000), std::vector{}, dispatcher_, + null_child_factory, nullptr, std::chrono::milliseconds(1000), 0.1, 0.3, 0.03, + std::chrono::milliseconds(180000), std::vector{}, dispatcher_, context_.thread_local_); LoadAwareLocalityLoadBalancer lb(*lb_config, cluster_info_, priority_set_, @@ -1412,6 +1410,74 @@ TEST_F(LoadAwareLocalityLbTest, ChildRecreatedWhenUnderlyingFactoryRequiresIt) { EXPECT_GT(child_factory->createCount(), initial_create_count); } +TEST_F(LoadAwareLocalityLbTest, AddedPriorityBuildsOnlyItsOwnChildLbs) { + auto local = makeWeightTrackingMockHost(); + auto remote = makeWeightTrackingMockHost(); + setupLocalities({{local}, {remote}}); + + auto child_factory = std::make_shared(false); + auto factory = makeWorkerFactoryWithChild(child_factory); + auto worker_lb = factory->create({priority_set_, nullptr}); + ASSERT_NE(nullptr, worker_lb); + // P0: two localities, each with healthy and all-hosts children. + const uint32_t p0_create_count = child_factory->createCount(); + EXPECT_EQ(p0_create_count, 4); + + auto p1_host = makeWeightTrackingMockHost(); + setupPriorityLocalities(1, {{p1_host}}); + priority_set_.runUpdateCallbacks(1, {p1_host}, {}); + + // Only the new priority's children are built; P0's live child LBs are preserved. + EXPECT_EQ(child_factory->createCount(), p0_create_count + 2); + auto picked = worker_lb->chooseHost(nullptr).host; + EXPECT_TRUE(picked == local || picked == remote); +} + +TEST_F(LoadAwareLocalityLbTest, ExistingPriorityDeltaSyncsWhenPriorityCountGrows) { + auto h1 = makeWeightTrackingMockHost(); + auto h2a = makeWeightTrackingMockHost(); + auto h2b = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2a, h2b}}); + + auto child_factory = std::make_shared(true); + auto factory = makeWorkerFactoryWithChild(child_factory); + auto worker_lb = factory->create({priority_set_, nullptr}); + ASSERT_NE(nullptr, worker_lb); + const uint32_t initial_create_count = child_factory->createCount(); + + // Grow the priority set without a callback (as intermediate host-set creation does), then + // deliver a P0 membership delta while the priority counts still disagree. + setupPriorityLocalities(1, {{makeWeightTrackingMockHost()}}); + reshapeLocalities(0, {{h1}, {h2a}}, /*added=*/{}, /*removed=*/{h2b}); + + // P1 is built (+2) and P0's own delta is not dropped: locality 1's healthy and all-hosts + // children are recreated (+2) under the recreate-on-host-change contract. + EXPECT_EQ(child_factory->createCount(), initial_create_count + 4); +} + +TEST_F(LoadAwareLocalityLbTest, HealthOnlyUpdateSyncsExistingPriorityWhileGrowthPending) { + auto h1 = makeWeightTrackingMockHost(); + auto h2a = makeWeightTrackingMockHost(); + auto h2b = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2a, h2b}}); + + auto child_factory = std::make_shared(true); + auto factory = makeWorkerFactoryWithChild(child_factory); + auto worker_lb = factory->create({priority_set_, nullptr}); + ASSERT_NE(nullptr, worker_lb); + const uint32_t initial_create_count = child_factory->createCount(); + + // Grow the priority set without a callback, then deliver a P0 health-only (empty-delta) update + // while the priority counts still disagree. + setupPriorityLocalities(1, {{makeWeightTrackingMockHost()}}); + reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, + std::vector{{h1}, {h2a}}); + + // P1 stays pending (an empty delta cannot rebuild), but P0's healthy flip is applied: locality + // 2's healthy child is recreated under the recreate-on-host-change contract. + EXPECT_EQ(child_factory->createCount(), initial_create_count + 1); +} + // Wraps metric names for LocalityLbHostData's shared_ptr constructor. std::shared_ptr> makeMetricNames(std::vector names) { return std::make_shared>(std::move(names)); From 2650f89d243dae1416024e387fd87bbf307d3888 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:45:18 -0600 Subject: [PATCH 14/16] Flip back to WIP Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_balancing__load-aware-locality-lb-policy.rst | 3 ++- .../upstream/load_balancing/load_aware_locality.rst | 5 +++++ source/extensions/extensions_metadata.yaml | 2 +- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst b/changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst index 74680f49ad817..3cb3d633b795a 100644 --- a/changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst +++ b/changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst @@ -2,4 +2,5 @@ load_balancing: implemented the :ref:`envoy.load_balancing_policies.load_aware_locality ` locality-picking load balancer. It weights localities by ORCA-derived utilization headroom, -consumes in-band ORCA reporting, and applies at all priority levels. +consumes in-band ORCA reporting, and applies at all priority levels. The extension is +work-in-progress and not intended for production use. diff --git a/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst b/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst index 2f99f30a66658..41c14c800a50b 100644 --- a/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst +++ b/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst @@ -3,6 +3,11 @@ Load-aware locality load balancing ----------------------------------- +.. attention:: + + This extension is **work-in-progress**. Functionality is incomplete and it + is not intended for production use. + The load-aware locality LB policy (:ref:`envoy.load_balancing_policies.load_aware_locality `) diff --git a/source/extensions/extensions_metadata.yaml b/source/extensions/extensions_metadata.yaml index dd77435f1a098..3c71e27bee679 100644 --- a/source/extensions/extensions_metadata.yaml +++ b/source/extensions/extensions_metadata.yaml @@ -2131,7 +2131,7 @@ envoy.load_balancing_policies.load_aware_locality: categories: - envoy.load_balancing_policies security_posture: unknown - status: alpha + status: wip type_urls: - envoy.extensions.load_balancing_policies.load_aware_locality.v3.LoadAwareLocality envoy.load_balancing_policies.random: From 21fbca1d85f5cfab4b24f04723c4cd507c925fc6 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:48:59 -0600 Subject: [PATCH 15/16] Include round_robin in test extension_names Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_balancing_policies/load_aware_locality/BUILD | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/extensions/load_balancing_policies/load_aware_locality/BUILD b/test/extensions/load_balancing_policies/load_aware_locality/BUILD index 10604f86bb1d8..467c0ca9020f2 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/BUILD +++ b/test/extensions/load_balancing_policies/load_aware_locality/BUILD @@ -58,7 +58,10 @@ envoy_extension_cc_test( name = "integration_test", size = "large", srcs = ["integration_test.cc"], - extension_names = ["envoy.load_balancing_policies.load_aware_locality"], + extension_names = [ + "envoy.load_balancing_policies.load_aware_locality", + "envoy.load_balancing_policies.round_robin", + ], rbe_pool = "6gig", deps = [ "//source/common/protobuf", From fb02adbfb6b8fb7429c4876c436d89d47a30a9b2 Mon Sep 17 00:00:00 2001 From: jukie <10012479+jukie@users.noreply.github.com> Date: Thu, 9 Jul 2026 12:25:53 -0400 Subject: [PATCH 16/16] more review fixes Signed-off-by: jukie <10012479+jukie@users.noreply.github.com> --- .../load_balancing/load_aware_locality.rst | 3 + .../load_aware_locality/BUILD | 1 + .../load_aware_locality/config.cc | 4 - .../load_aware_locality_lb.cc | 104 +++++++++--------- .../load_aware_locality_lb.h | 55 +++++---- .../load_aware_locality/config_test.cc | 17 ++- .../load_aware_locality_lb_test.cc | 72 ++++++++++-- 7 files changed, 165 insertions(+), 91 deletions(-) diff --git a/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst b/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst index 41c14c800a50b..d9eb91a0475c1 100644 --- a/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst +++ b/docs/root/intro/arch_overview/upstream/load_balancing/load_aware_locality.rst @@ -480,6 +480,9 @@ The policy emits stats under ``cluster..load_aware_locality.*``: * - Counter - Increments when + * - ``recompute_total`` + - Every main-thread weight-recompute tick. A liveness signal: the + expected rate is one per ``weight_update_period``. * - ``all_overloaded_total`` - Per tick where every locality's headroom was 0 (fallback to host-count weighting). diff --git a/source/extensions/load_balancing_policies/load_aware_locality/BUILD b/source/extensions/load_balancing_policies/load_aware_locality/BUILD index 95f316eb2c5af..764397bbe7f9e 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/BUILD +++ b/source/extensions/load_balancing_policies/load_aware_locality/BUILD @@ -35,6 +35,7 @@ envoy_cc_library( "//envoy/upstream:upstream_interface", "//source/common/common:minimal_logger_lib", "//source/common/protobuf:utility_lib", + "//source/common/upstream:load_balancer_context_base_lib", "//source/common/upstream:upstream_includes", "//source/extensions/load_balancing_policies/common:load_balancer_lib", "//source/extensions/load_balancing_policies/common:orca_weight_manager_lib", diff --git a/source/extensions/load_balancing_policies/load_aware_locality/config.cc b/source/extensions/load_balancing_policies/load_aware_locality/config.cc index b9e3253e0593f..1e12f7fa65d4e 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/config.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/config.cc @@ -31,10 +31,6 @@ Factory::loadConfig(Server::Configuration::ServerFactoryContext& context, return absl::InvalidArgumentError( "load_aware_locality: enable_oob_load_report is not yet supported"); } - if (lb_config.has_oob_reporting_period()) { - return absl::InvalidArgumentError("load_aware_locality: oob_reporting_period requires " - "enable_oob_load_report, which is not yet supported"); - } const std::chrono::milliseconds weight_update_period = std::chrono::milliseconds(PROTOBUF_GET_MS_OR_DEFAULT(lb_config, weight_update_period, 1000)); // PGV constraints are not enforced on unpacked LB configs, so the documented floor is owned diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc index 26603dcd2ebe3..02469a9f89e5d 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -10,6 +10,7 @@ #include "envoy/stats/stats_macros.h" #include "source/common/protobuf/utility.h" +#include "source/common/upstream/load_balancer_context_base.h" #include "source/extensions/load_balancing_policies/common/load_balancer_impl.h" #include "source/extensions/load_balancing_policies/common/orca_weight_manager.h" @@ -184,6 +185,7 @@ void LoadAwareLocalityLoadBalancer::computeLocalityRoutingWeights() { if (tick_stale_localities > 0) { lb_stats_.stale_locality_total_.add(tick_stale_localities); } + lb_stats_.recompute_total_.inc(); ENVOY_LOG(trace, "computeLocalityRoutingWeights: {} priorities", snapshot->priority_weights.size()); @@ -194,7 +196,7 @@ LoadAwareLocalityLoadBalancer::SourceComputeResult LoadAwareLocalityLoadBalancer::computeSourceWeights( const Upstream::HostsPerLocality& all_hosts_per_locality, const std::vector& eligible_hosts_per_locality, int64_t now_ms, - LocalityWeightsMap& weights_map, bool& all_local, LocalityEwmaMap& ewma_state) { + LocalityRoutingWeightsMap& weights_map, bool& all_local, LocalityEwmaMap& ewma_state) { SourceComputeResult result; const auto& locality_hosts = all_hosts_per_locality.get(); const size_t locality_count = locality_hosts.size(); @@ -313,6 +315,8 @@ LoadAwareLocalityLoadBalancer::computeSourceWeights( weights[i] += take_from_local * static_cast(host_counts[i]) / remote_hosts; } result.probe_active = true; + // Weights are no longer 100% local, so local picks count as sampled, not all-directly. + all_local = false; } } } @@ -364,21 +368,16 @@ Upstream::LoadBalancerPtr WorkerLocalLbFactory::create(Upstream::LoadBalancerPar namespace { -// Context for per-locality child LBs. Priority selection stays fixed to the child's default because -// the parent already ran retry-priority against cluster-wide priorities; rerunning it against the -// single-priority child set can corrupt plugin state. Other methods pass through; children must not -// retain this stack object. -class ChildLoadBalancerContext : public Upstream::LoadBalancerContext { +// Context for per-locality child LBs. The base's non-delegating determinePriorityLoad keeps +// priority selection fixed to the child's default: the parent already ran retry-priority against +// cluster-wide priorities, and rerunning it against the single-priority child set can corrupt +// plugin state. The base's no-op onAsyncHostSelection matches chooseHost cancelling async child +// selection. Other methods pass through; children must not retain this stack object. +class ChildLoadBalancerContext : public Upstream::LoadBalancerContextBase { public: explicit ChildLoadBalancerContext(Upstream::LoadBalancerContext& wrapped) : wrapped_(wrapped) {} std::optional computeHashKey() override { return wrapped_.computeHashKey(); } - const Upstream::HealthyAndDegradedLoad& - determinePriorityLoad(const Upstream::PrioritySet&, - const Upstream::HealthyAndDegradedLoad& default_priority_load, - const Upstream::RetryPriority::PriorityMappingFunc&) override { - return default_priority_load; - } const Router::MetadataMatchCriteria* metadataMatchCriteria() override { return wrapped_.metadataMatchCriteria(); } @@ -404,9 +403,6 @@ class ChildLoadBalancerContext : public Upstream::LoadBalancerContext { OptRef overrideHostToSelect() const override { return wrapped_.overrideHostToSelect(); } - // Async child selection is unsupported (the async handle is cancelled in chooseHost), so this is - // never invoked; a no-op matches the subset LB and avoids implying otherwise. - void onAsyncHostSelection(Upstream::HostConstSharedPtr&&, std::string&&) override {} void setHeadersModifier(std::function modifier) override { wrapped_.setHeadersModifier(std::move(modifier)); } @@ -415,6 +411,15 @@ class ChildLoadBalancerContext : public Upstream::LoadBalancerContext { Upstream::LoadBalancerContext& wrapped_; }; +// splitMix64 (golden-ratio increment + finishing mix) makes the locality target independent of +// choosePriority's hash % 100 so both decisions can share the pick's single random draw. +uint64_t splitMix64(uint64_t x) { + x += 0x9e3779b97f4a7c15; + x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9; + x = (x ^ (x >> 27)) * 0x94d049bb133111eb; + return x ^ (x >> 31); +} + // Async child selection would retain ChildLoadBalancerContext after chooseHost returns. Cancel any // async handle and fail synchronously, matching the subset LB. Upstream::HostSelectionResponse failOnAsyncSelection(Upstream::HostSelectionResponse response) { @@ -637,7 +642,7 @@ WorkerLocalLb::PrioritySourcePick WorkerLocalLb::resolvePrioritySource(Upstream::LoadBalancerContext* context, bool peeking) { using SelectionSource = PriorityRoutingWeights::SelectionSource; if (priority_set_.hostSetsPerPriority().empty()) { - return {0, SelectionSource::Healthy, /*in_panic=*/false, /*fail=*/true}; + return {0, SelectionSource::Healthy, 0, /*in_panic=*/false, /*fail=*/true}; } const uint64_t hash = random(peeking); @@ -647,17 +652,16 @@ WorkerLocalLb::resolvePrioritySource(Upstream::LoadBalancerContext* context, boo const uint32_t priority = host_set_and_availability.first.priority(); if (isInPanic(priority)) { - return {priority, SelectionSource::AllHosts, /*in_panic=*/true, /*fail=*/false}; + return {priority, SelectionSource::AllHosts, hash, /*in_panic=*/true, /*fail=*/false}; } const SelectionSource source = host_set_and_availability.second == Upstream::LoadBalancerBase::HostAvailability::Healthy ? SelectionSource::Healthy : SelectionSource::Degraded; - return {priority, source, /*in_panic=*/false, /*fail=*/false}; + return {priority, source, hash, /*in_panic=*/false, /*fail=*/false}; } -WorkerLocalLb::LocalityLbSelection WorkerLocalLb::selectLocalityLb(const PrioritySourcePick& pick, - bool peeking) { +WorkerLocalLb::LocalityLbSelection WorkerLocalLb::selectLocalityLb(const PrioritySourcePick& pick) { LocalityLbSelection selection; if (pick.fail) { return selection; @@ -673,10 +677,14 @@ WorkerLocalLb::LocalityLbSelection WorkerLocalLb::selectLocalityLb(const Priorit return selection; } - const RoutingWeightsSnapshot* snapshot = factory_.routingWeights(); + if (shim_ == nullptr) { + shim_ = factory_.tlsShim(); + } + const RoutingWeightsSnapshot* snapshot = + shim_ != nullptr ? shim_->routing_weights.get() : nullptr; selection.snapshot = snapshot; refreshLocalityWeights(snapshot); - const size_t preferred_idx = chooseLocality(pick.priority, pick.source, peeking); + const size_t preferred_idx = chooseLocality(pick.priority, pick.source, pick.hash); selection.locality_idx = preferred_idx; selection.lb = pickLocalityLb(per_locality, pick.source, preferred_idx, selection.locality_idx); return selection; @@ -711,7 +719,7 @@ Upstream::HostSelectionResponse WorkerLocalLb::chooseHost(Upstream::LoadBalancer stats_.lb_healthy_panic_.inc(); } - const auto selection = selectLocalityLb(pick, /*peeking=*/false); + const auto selection = selectLocalityLb(pick); if (selection.lb == nullptr) { return {nullptr}; } @@ -744,63 +752,59 @@ void WorkerLocalLb::refreshLocalityWeights(const RoutingWeightsSnapshot* snapsho index_weights.assign(locality_count, 0.0); double total = 0.0; if (have_priority) { - // Once-per-publish replacement for chooseLocality's old per-pick weights.find(locality). - // Snapshot-missing live localities get weight 0 until the next recompute. + // Once-per-publish replacement for chooseLocality's old per-pick weights.find(locality), + // stored as prefix sums so the pick is a binary search. Snapshot-missing live localities + // get weight 0 until the next recompute. const auto& weights = snapshot->priority_weights[priority].weightsFor( static_cast(s)); for (size_t i = 0; i < locality_count; ++i) { const auto weight = weights.find(pstate.localities[i].locality); - index_weights[i] = weight != weights.end() ? weight->second : 0.0; - total += index_weights[i]; + total += weight != weights.end() ? weight->second : 0.0; + index_weights[i] = total; } } - pstate.source_total[s] = total; } } } size_t WorkerLocalLb::chooseLocality(uint32_t priority, - PriorityRoutingWeights::SelectionSource source, bool peeking) { + PriorityRoutingWeights::SelectionSource source, + uint64_t hash) { const auto& pstate = per_priority_locality_[priority]; const auto& per_locality = pstate.localities; if (per_locality.size() <= 1) { return 0; } const size_t s = static_cast(source); - const auto& locality_weights = pstate.source_weights[s]; - const double effective_total = pstate.source_total[s]; - // Defensive: only reachable before the first snapshot publish, where source_total is 0. - if (locality_weights.size() != per_locality.size() || effective_total <= 0.0) { + const auto& cumulative_weights = pstate.source_weights[s]; + if (cumulative_weights.size() != per_locality.size()) { + return 0; + } + // Zero total is only reachable before the first snapshot publish, where the weights are all + // zero. + const double effective_total = cumulative_weights.back(); + if (effective_total <= 0.0) { return 0; } - // This draw is skipped by the early returns above, so a snapshot published between a paired peek - // and choose can offset the stashed-draw replay by one (harmless; membership changes clear it). - const double target = (static_cast(random(peeking)) / + const double target = (static_cast(splitMix64(hash)) / static_cast(std::numeric_limits::max())) * effective_total; - double cumulative = 0.0; - for (size_t i = 0; i < per_locality.size(); ++i) { - cumulative += locality_weights[i]; - if (target < cumulative) { - return i; - } - } - - // Floating-point rounding guard: return the last locality, not an arbitrary one. - return per_locality.size() - 1; + const auto it = std::upper_bound(cumulative_weights.begin(), cumulative_weights.end(), target); + // Floating-point rounding guard: clamp to the last locality, not an arbitrary one. + return std::min(static_cast(it - cumulative_weights.begin()), + per_locality.size() - 1); } Upstream::HostConstSharedPtr WorkerLocalLb::peekAnotherHost(Upstream::LoadBalancerContext* context) { - // Cap preconnect look-ahead as the sibling LBs do, so stashed_random_ cannot grow unbounded and - // the peek/choose replay stream stays bounded. + // Cap preconnect look-ahead as the sibling LBs do; each peek stashes exactly one draw, so the + // stash size is the number of outstanding peeked picks. if (Upstream::tooManyPreconnects(stashed_random_.size(), total_healthy_hosts_)) { return nullptr; } - const auto selection = selectLocalityLb(resolvePrioritySource(context, /*peeking=*/true), - /*peeking=*/true); + const auto selection = selectLocalityLb(resolvePrioritySource(context, /*peeking=*/true)); if (selection.lb == nullptr) { return nullptr; } diff --git a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h index 9945ba4ea4ca1..633197c9d9053 100644 --- a/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -34,6 +34,7 @@ namespace LoadAwareLocality { COUNTER(all_overloaded_total) \ COUNTER(local_preferred_total) \ COUNTER(probe_active_total) \ + COUNTER(recompute_total) \ COUNTER(spill_active_total) \ COUNTER(stale_locality_total) @@ -140,28 +141,30 @@ class LoadAwareLocalityLbConfig : public Upstream::LoadBalancerConfig { }; // Key locality state by identity rather than HostsPerLocality position so add/remove index shifts -// cannot map main-thread snapshots to the wrong worker-local membership. -using LocalityWeightsMap = absl::flat_hash_map; +// cannot map main-thread snapshots to the wrong worker-local membership. Named to avoid shadowing +// Upstream::LocalityWeightsMap, which is a different type. +using LocalityRoutingWeightsMap = + absl::flat_hash_map; // EWMA state uses the same identity key: present-with-no-valid-hosts is stale; absent is cold. -using LocalityEwmaMap = absl::flat_hash_map; +using LocalityEwmaMap = LocalityRoutingWeightsMap; struct PriorityRoutingWeights { enum class SelectionSource : uint8_t { Healthy = 0, Degraded = 1, AllHosts = 2 }; + static constexpr size_t kSourceCount = 3; struct SourceWeights { - LocalityWeightsMap weights; + LocalityRoutingWeightsMap weights; // Does not affect the routing decision (fully encoded in weights); read on the hot path only to // pick the zone-routing stat counter. bool all_local{false}; }; - std::array by_source; + std::array by_source; bool has_local_locality{false}; - const LocalityWeightsMap& weightsFor(SelectionSource s) const { + const LocalityRoutingWeightsMap& weightsFor(SelectionSource s) const { return by_source[static_cast(s)].weights; } bool allLocalFor(SelectionSource s) const { return by_source[static_cast(s)].all_local; } @@ -194,7 +197,7 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { bool recreateOnHostChangeDeprecated() const override { return false; } void updateRoutingWeights(RoutingWeightsSnapshotConstSharedPtr snapshot) { - tls_->runOnAllThreads([snapshot](OptRef shim) { + tls_->runOnAllThreads([snapshot = std::move(snapshot)](OptRef shim) { if (shim.has_value()) { shim->routing_weights = snapshot; } @@ -202,11 +205,12 @@ class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { } // SAFETY: Must only be called on a thread that owns a TLS slot instance (worker or main - // thread). The returned pointer is valid only for the duration of the current task; do not - // store it across yield points or after the TLS slot could be updated. - const RoutingWeightsSnapshot* routingWeights() const { + // thread). The shim object is created once per thread at factory construction and never + // replaced (updateRoutingWeights mutates its field), and the factory outlives any worker LB it + // created, so callers on their own thread may cache the returned pointer for their lifetime. + const ThreadLocalShim* tlsShim() const { auto shim = tls_->get(); - return shim.has_value() ? shim->routing_weights.get() : nullptr; + return shim.ptr(); } // Shared-ownership variant: a worker holds the snapshot it derived its index-aligned weights @@ -249,7 +253,7 @@ struct PerSourceLocalityState { }; struct PerLocalityState { - std::array by_source; + std::array by_source; // Empty for a locality with no hosts. envoy::config::core::v3::Locality locality; @@ -263,11 +267,10 @@ struct PerLocalityState { struct PerPriorityLocalityState { std::vector localities; - // Index-aligned routing weights derived from the snapshot once per publish, so the per-pick path - // reads by locality index instead of hashing locality identity. Rebuilt on snapshot or topology - // change; source_total[s] is the sum of source_weights[s]. - std::array, 3> source_weights; - std::array source_total{{0.0, 0.0, 0.0}}; + // Index-aligned cumulative routing weights (prefix sums) derived from the snapshot once per + // publish, so the per-pick path binary-searches by locality index instead of hashing locality + // identity. Rebuilt on snapshot or topology change. + std::array, PriorityRoutingWeights::kSourceCount> source_weights; }; // Worker-local LB: picks live priority/source, then delegates to a per-locality child LB. @@ -298,6 +301,9 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { struct PrioritySourcePick { uint32_t priority; PriorityRoutingWeights::SelectionSource source; + // The pick's single random draw; chooseLocality derives its target from this so each peek + // stashes exactly one entry, keeping the preconnect gate and peek/choose replay aligned. + uint64_t hash; bool in_panic; bool fail; }; @@ -310,13 +316,13 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { PrioritySourcePick resolvePrioritySource(Upstream::LoadBalancerContext* context, bool peeking); - LocalityLbSelection selectLocalityLb(const PrioritySourcePick& pick, bool peeking); + LocalityLbSelection selectLocalityLb(const PrioritySourcePick& pick); void recordZoneRoutingStats(const PrioritySourcePick& pick, const LocalityLbSelection& selection); // Returns 0 for single locality, no/stale snapshot, or zero effective total. size_t chooseLocality(uint32_t priority, PriorityRoutingWeights::SelectionSource source, - bool peeking); + uint64_t hash); // Rebuilds per-priority index-aligned weights when the published snapshot changes (or a topology // change reset built_snapshot_), so chooseLocality reads by index instead of hashing per pick. @@ -328,6 +334,9 @@ class WorkerLocalLb : public Upstream::LoadBalancerBase { size_t preferred_idx, size_t& actual_idx) const; WorkerLocalLbFactory& factory_; + // Lazily cached per-thread shim so the per-pick snapshot read is a plain pointer load instead + // of a virtual TLS lookup plus shared_ptr copy; see tlsShim() for the lifetime contract. + const ThreadLocalShim* shim_{nullptr}; std::vector per_priority_locality_; // Snapshot the index-aligned weights in per_priority_locality_ were built from; reset to force a // rebuild after a topology change. @@ -367,7 +376,7 @@ class LoadAwareLocalityLoadBalancer : public Upstream::ThreadAwareLoadBalancer, SourceComputeResult computeSourceWeights(const Upstream::HostsPerLocality& all_hosts_per_locality, const std::vector& eligible_hosts_per_locality, - int64_t now_ms, LocalityWeightsMap& weights_map, bool& all_local, + int64_t now_ms, LocalityRoutingWeightsMap& weights_map, bool& all_local, LocalityEwmaMap& ewma_state); const Upstream::PrioritySet& priority_set_; @@ -379,7 +388,7 @@ class LoadAwareLocalityLoadBalancer : public Upstream::ThreadAwareLoadBalancer, std::chrono::milliseconds weight_expiration_period_; std::shared_ptr> metric_names_; // Per-source, per-priority EWMA state keyed by locality identity (main thread only). - std::array, 3> ewma_state_; + std::array, PriorityRoutingWeights::kSourceCount> ewma_state_; Event::TimerPtr weight_update_timer_; std::chrono::milliseconds weight_update_period_; std::string child_factory_name_; diff --git a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc index 50f62fb881e45..641d79896f02b 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/config_test.cc @@ -221,11 +221,6 @@ TEST(LoadAwareLocalityConfigTest, InvalidPolicyKnobsRejected) { config.mutable_enable_oob_load_report()->set_value(true); }, "enable_oob_load_report"}, - {"OobReportingPeriodWithoutOobEnabled", - [](LoadAwareLocalityLbProto& config) { - config.mutable_oob_reporting_period()->set_seconds(10); - }, - "oob_reporting_period"}, {"NegativeWeightExpiration", [](LoadAwareLocalityLbProto& config) { config.mutable_weight_expiration_period()->set_seconds(-1); @@ -253,6 +248,18 @@ TEST(LoadAwareLocalityConfigTest, InvalidPolicyKnobsRejected) { } } +TEST(LoadAwareLocalityConfigTest, OobReportingPeriodAloneIsInert) { + NiceMock context; + + // Per the API contract the field is used only when enable_oob_load_report is true, so setting + // it alone must not reject the config. + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_oob_reporting_period()->set_seconds(10); + + Factory factory; + EXPECT_TRUE(factory.loadConfig(context, config_proto).ok()); +} + TEST(LoadAwareLocalityConfigTest, PolicyValidationPrecedesChildResolution) { NiceMock context; diff --git a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc index a856b68bcbf60..6b59665b426a7 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -316,8 +316,8 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes ON_CALL(*mock, locality()).WillByDefault(testing::ReturnRef(localities_.back())); } - static LocalityWeightsMap makeWeightsMap(const std::vector& weights) { - LocalityWeightsMap map; + static LocalityRoutingWeightsMap makeWeightsMap(const std::vector& weights) { + LocalityRoutingWeightsMap map; for (size_t i = 0; i < weights.size(); ++i) { map[localityForIndex(i)] = weights[i]; } @@ -377,7 +377,8 @@ class LoadAwareLocalityLbTest : public Event::TestUsingSimulatedTime, public tes const RoutingWeightsSnapshot* routingSnapshot() const { auto* typed_factory = dynamic_cast(factory_.get()); EXPECT_NE(nullptr, typed_factory); - return typed_factory != nullptr ? typed_factory->routingWeights() : nullptr; + const auto* shim = typed_factory != nullptr ? typed_factory->tlsShim() : nullptr; + return shim != nullptr ? shim->routing_weights.get() : nullptr; } void publishSnapshot(std::shared_ptr snapshot) { @@ -612,7 +613,7 @@ INSTANTIATE_TEST_SUITE_P(WeightComputation, WeightComputationTest, 1.0, 0.06, {0.5, 0.5, 0.5}, - true, + false, {0.94, 0.02, 0.04}, 0.01}, {"ProbeHostCountProportionalUnequalUtils", @@ -621,7 +622,7 @@ INSTANTIATE_TEST_SUITE_P(WeightComputation, WeightComputationTest, 1.0, 0.06, {0.5, 0.8, 0.2}, - true, + false, {0.94, 0.02, 0.04}, 0.01}, {"ProbeRedistributionClampsLocalAtZero", @@ -1230,12 +1231,12 @@ TEST_F(LoadAwareLocalityLbTest, PeekAndChooseReplayTheSameRandomSequence) { EXPECT_GE(counts[h2.get()], 98); publishHealthyWeights({0.1, 0.9}); - forceRandomValues({0, 0}); + forceRandomValues({0}); auto peeked = worker_lb->peekAnotherHost(nullptr); ASSERT_NE(nullptr, peeked); - forceRandomValues( - {std::numeric_limits::max() - 1, std::numeric_limits::max() - 1}); + // The choose must replay the peek's stashed draw, leaving this decoy value unconsumed. + forceRandomValues({std::numeric_limits::max() - 1}); auto chosen = worker_lb->chooseHost(nullptr).host; EXPECT_EQ(peeked.get(), chosen.get()); } @@ -1248,6 +1249,8 @@ TEST_F(LoadAwareLocalityLbTest, WorkerFallbackEdgeCases) { createLb(); auto worker_lb = createWorkerLb(); + // splitMix64(3)/max ~= 0.11 -> locality 0 under the equal host-count fallback weights. + forceRandomValues({3}); EXPECT_EQ(h1, worker_lb->chooseHost(nullptr).host); publishSnapshot(makeSnapshot({PriorityRoutingWeights{}})); @@ -1257,7 +1260,8 @@ TEST_F(LoadAwareLocalityLbTest, WorkerFallbackEdgeCases) { EXPECT_EQ(h1, worker_lb->chooseHost(nullptr).host); publishHealthyWeights({1.0, 1.0}); - forceRandomValues({0, std::numeric_limits::max()}); + // splitMix64(6)/max ~= 0.74, in the upper half of the equal-weight distribution -> locality 1. + forceRandomValues({6}); EXPECT_EQ(h2, worker_lb->chooseHost(nullptr).host); worker_lb = createWorkerLb(); @@ -1300,6 +1304,34 @@ TEST_F(LoadAwareLocalityLbTest, PeekAnotherHostCapsPreconnectLookahead) { EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); } +TEST_F(LoadAwareLocalityLbTest, PeekLookaheadCountsPeeksNotDraws) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + createLb(); + auto worker_lb = createWorkerLb(); + publishHealthyWeights({1.0, 1.0}); + + // Multi-locality picks make both a priority and a locality decision from one stashed draw, so + // two healthy hosts allow exactly two peeks. Forced hashes route the peeks to distinct + // localities (splitMix64 fractions 0.11 and 0.74) so the per-locality child caps don't gate. + forceRandomValues({3, 6}); + EXPECT_EQ(h1, worker_lb->peekAnotherHost(nullptr)); + EXPECT_EQ(h2, worker_lb->peekAnotherHost(nullptr)); + EXPECT_EQ(nullptr, worker_lb->peekAnotherHost(nullptr)); +} + +TEST_F(LoadAwareLocalityLbTest, RecomputeTotalIncrementsEveryTick) { + auto h = makeWeightTrackingMockHost(); + setupLocalities({{h}}); + + createLb(); + EXPECT_EQ(1, counterValue(cluster_info_, "load_aware_locality.recompute_total")); + expectCounterIncrementsPerTick("load_aware_locality.recompute_total", 1); + expectCounterIncrementsPerTick("load_aware_locality.recompute_total", 1); +} + TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverHealthyRoutingPaths) { auto h_local = makeWeightTrackingMockHost(); auto h_remote = makeWeightTrackingMockHost(); @@ -1328,6 +1360,28 @@ TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverHealthyRoutingPaths) { EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_cross_zone_.value(), 0); } +TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsProbeCountsLocalPicksAsSampled) { + auto h_local = makeWeightTrackingMockHost(); + auto h_remote = makeWeightTrackingMockHost(); + setupLocalities({{h_local}, {h_remote}}, /*has_local_locality=*/true); + + // Local preference snaps to all-local, then the probe redistribution diverts traffic: local + // picks are sampled from a <100% local distribution, never all-directly. + createLb(/*variance_threshold=*/1.0, /*ewma_alpha=*/1.0, /*remote_probe_fraction=*/0.03); + setHostUtilization(h_local, 0.3); + setHostUtilization(h_remote, 0.3); + recomputeWeights(); + expectAllLocal(false); + + auto worker_lb = createWorkerLb(); + for (int i = 0; i < 200; ++i) { + worker_lb->chooseHost(nullptr); + } + EXPECT_EQ(0, static_cast(cluster_info_.lbStats().lb_zone_routing_all_directly_.value())); + EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_sampled_.value(), 0); + EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_cross_zone_.value(), 0); +} + TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverNoLocalDegradedAndPanicSources) { auto h1 = makeWeightTrackingMockHost(); auto h2 = makeWeightTrackingMockHost();