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..3cb3d633b795a --- /dev/null +++ b/changelogs/current/new_features/load_balancing__load-aware-locality-lb-policy.rst @@ -0,0 +1,6 @@ +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. 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 5cb3acb8dadb7..2e120e177bacd 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,5 +1,3 @@ -:orphan: - .. _arch_overview_load_balancing_load_aware_locality: Load-aware locality load balancing @@ -7,7 +5,8 @@ Load-aware locality load balancing .. attention:: - This extension is **work-in-progress** and is not yet implemented. + 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 @@ -42,7 +41,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 +90,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 ` @@ -133,15 +128,23 @@ per-endpoint capacity. Each consumer attaches independent Utilization is derived from each host's ORCA report using the same extraction as CSWRR, which takes the first source whose value is greater -than 0 (precedence may be flipped by the +than 0. 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. Named metrics via ``metric_names_for_computing_utilization`` -- max of present values. 2. ``application_utilization``. 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 +163,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 @@ -209,17 +213,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 @@ -237,19 +247,22 @@ 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: - # Local preference (one-sided: local must not be too far ABOVE remote) + # Local preference (one-sided: local must not be too far ABOVE remote; + # skipped when local has no eligible hosts) remote_weighted_avg = sum(smoothed_util(R_i) * host_count(R_i)) / remote_host_count - if smoothed_util(local) <= remote_weighted_avg + utilization_variance_threshold: + if host_count(local) > 0 and + smoothed_util(local) <= remote_weighted_avg + utilization_variance_threshold: adjusted_weight(local) = total_base_weight adjusted_weight(R_i) = 0 @@ -307,9 +320,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 @@ -342,9 +355,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 +367,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 +378,9 @@ 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``. When +in panic mode the policy routes to all hosts. + Three independent weight sets are maintained per priority: - **Healthy** -- common case, healthy hosts only. @@ -395,17 +397,37 @@ 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 +- **Out-of-band ORCA reporting is not yet supported.** Only in-band + (per-response) ORCA reports are consumed. ``enable_oob_load_report`` and + ``oob_reporting_period`` are accepted but currently have no effect. +- **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 +456,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 @@ -462,7 +487,8 @@ The policy emits stats under ``cluster..load_aware_locality.*``: * - Counter - Increments when * - ``recompute_total`` - - Per main-thread tick that recomputes weights. + - 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). @@ -471,17 +497,30 @@ 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 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. + +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, +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 diff --git a/source/extensions/load_balancing_policies/load_aware_locality/BUILD b/source/extensions/load_balancing_policies/load_aware_locality/BUILD index 890af1964ef23..764397bbe7f9e 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,33 @@ 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/config:utility_lib", + "//source/extensions/load_balancing_policies/common:factory_base", "@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: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", + "@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..008d742715572 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,103 @@ #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 = 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. + 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 + // 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)); + 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) { + // 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( + 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, + 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, 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..8d197d9790e17 --- /dev/null +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.cc @@ -0,0 +1,805 @@ +#include "source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h" + +#include +#include +#include +#include + +#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" + +#include "absl/container/flat_hash_set.h" +#include "absl/status/status.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_); + // Signal-less reports do not refresh freshness; they age out instead of pinning idle. + if (util <= 0.0) { + return absl::OkStatus(); + } + storeUtilization(util, time_source_.monotonicTime()); + 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()); + + 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( + 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(); }); +} + +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() { + 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()); + } + + 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() { + 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()); + } + + const MonotonicTime now = time_source_.monotonicTime(); + + 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) { + 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; + } + + 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, 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()); + 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; + tick_spill_active |= healthy_result.spill_active || degraded_result.spill_active || + all_hosts_result.spill_active; + 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_spill_active) { + lb_stats_.spill_active_total_.inc(); + } + 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()); + factory_->updateRoutingWeights(std::move(snapshot)); +} + +LoadAwareLocalityLoadBalancer::SourceComputeResult +LoadAwareLocalityLoadBalancer::computeSourceWeights( + const Upstream::HostsPerLocality& all_hosts_per_locality, + const std::vector& eligible_hosts_per_locality, MonotonicTime now, + 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(); + 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 MonotonicTime last_update = host_data->lastUpdateTime(); + if (last_update == LocalityLbHostData::kNeverReported) { + continue; + } + + if (weight_expiration_period_.count() > 0 && + (now - last_update) > weight_expiration_period_) { + 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; + } + + 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()) { + new_state[locality_hosts[i][0]->locality()] = prev->second; + utilizations[i] = prev->second; + stale[i] = true; + ++result.stale_localities; + } + } + 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); + weights[0] = 1.0; + }; + + 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 if (total_base_weight > 0.0) { + 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]; + } + + // Comparison needs eligible local hosts: an empty local slice spills for health, not load, + // and does not count as spill. + 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 { + 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_; + const double remote_sum = total - weights[0]; + 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) { + 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; + } + } + } + + 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( + 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) + : 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)) { + tls_ = ThreadLocal::TypedSlot::makeUnique(tls_slot_allocator); + tls_->set([](Event::Dispatcher&) { return std::make_shared(); }); +} + +Upstream::LoadBalancerPtr +WorkerLocalLbFactory::createWorkerChildLb(Upstream::PrioritySetImpl& per_locality_priority_set) { + // 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_worker_factory_->create(child_params); +} + +bool WorkerLocalLbFactory::recreateChildOnHostChange() const { + ASSERT(child_worker_factory_ != nullptr); + return child_worker_factory_->recreateOnHostChangeDeprecated(); +} + +Upstream::LoadBalancerPtr WorkerLocalLbFactory::create(Upstream::LoadBalancerParams params) { + return std::make_unique(*this, params.priority_set); +} + +namespace { + +// 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 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(); + } + void setHeadersModifier(std::function modifier) override { + wrapped_.setHeadersModifier(std::move(modifier)); + } + +private: + 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) { + 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) { + // Empty-delta updates refresh child host attributes without rebuilding topology. + syncPriority(priority, /*allow_rebuild=*/!hosts_added.empty() || !hosts_removed.empty()); + }); +} + +WorkerLocalLb::~WorkerLocalLb() { 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); + // 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, + 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, + std::nullopt, std::nullopt); +} + +void WorkerLocalLb::buildPerPriorityLocalities() { + const auto& host_sets = priority_set_.hostSetsPerPriority(); + // 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 = old_size; 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 has none. + 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_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.get())) { + hosts_added.push_back(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(); + 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(); + // Topology changed: force refreshLocalityWeights to rebuild the index-aligned weights next pick. + built_snapshot_.reset(); + const auto& locality_hosts = host_set.hostsPerLocality().get(); + per_priority_locality_[priority].has_local_locality = + host_set.hostsPerLocality().hasLocalLocality(); + 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; 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(); + } + 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]; + const auto& locality_hosts = host_set->hostsPerLocality().get(); + auto& per_locality = per_priority_locality_[priority].localities; + + // Topology changed; rebuild this priority when permitted, else wait. + if (locality_hosts.size() != per_locality.size()) { + if (allow_rebuild) { + buildPerLocality(priority, *host_set); + } + return; + } + + per_priority_locality_[priority].has_local_locality = + host_set->hostsPerLocality().hasLocalLocality(); + 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); + } + // 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* +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()) { + // per_priority_locality_ is empty too, so selectLocalityLb's bounds guard yields no selection. + return {0, SelectionSource::Healthy, 0, /*in_panic=*/false}; + } + + 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, hash, /*in_panic=*/true}; + } + const SelectionSource source = + host_set_and_availability.second == Upstream::LoadBalancerBase::HostAvailability::Healthy + ? SelectionSource::Healthy + : SelectionSource::Degraded; + return {priority, source, hash, /*in_panic=*/false}; +} + +WorkerLocalLb::LocalityLbSelection WorkerLocalLb::selectLocalityLb(const PrioritySourcePick& pick) { + LocalityLbSelection selection; + // Covers the empty priority set and a custom retry-priority plugin routing to a just-created + // priority that per_priority_locality_ doesn't cover yet (empty-delta update, rebuild pending). + if (pick.priority >= per_priority_locality_.size()) { + return selection; + } + auto& per_locality = per_priority_locality_[pick.priority].localities; + if (per_locality.empty()) { + return selection; + } + + 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, pick.hash); + 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 || + pick.priority >= selection.snapshot->priority_weights.size()) { + return; + } + + // "Is index 0 local" is worker-live state (the snapshot may lag topology); the snapshot, + // whose coverage is gated above, only supplies the all-local flag for the split below. + if (!per_priority_locality_[pick.priority].has_local_locality) { + return; + } + if (selection.locality_idx != 0) { + stats_.lb_zone_routing_cross_zone_.inc(); + return; + } + if (selection.snapshot->priority_weights[pick.priority].allLocalFor(pick.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); + if (selection.lb == nullptr) { + return {nullptr}; + } + + recordZoneRoutingStats(pick, selection); + + if (context == nullptr) { + return failOnAsyncSelection(selection.lb->chooseHost(nullptr)); + } + ChildLoadBalancerContext child_context(*context); + return failOnAsyncSelection(selection.lb->chooseHost(&child_context)); +} + +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_ = shim_->routing_weights; + 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), + // 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); + total += weight != weights.end() ? weight->second : 0.0; + index_weights[i] = total; + } + } + } + } +} + +size_t WorkerLocalLb::chooseLocality(uint32_t priority, + 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& cumulative_weights = pstate.source_weights[s]; + if (cumulative_weights.size() != per_locality.size()) { + return 0; + } + // Zero total: no snapshot yet, the snapshot lags this worker's membership (every identity + // lookup missed), or the published weights for this priority are all zero. + const double effective_total = cumulative_weights.back(); + if (effective_total <= 0.0) { + return 0; + } + + const double target = (static_cast(splitMix64(hash)) / + static_cast(std::numeric_limits::max())) * + effective_total; + 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; 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)); + if (selection.lb == nullptr) { + return nullptr; + } + if (context == nullptr) { + return selection.lb->peekAnotherHost(nullptr); + } + ChildLoadBalancerContext child_context(*context); + return selection.lb->peekAnotherHost(&child_context); +} + +} // 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..e85764813d455 --- /dev/null +++ b/source/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb.h @@ -0,0 +1,402 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "envoy/common/random_generator.h" +#include "envoy/common/time.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" +#include "absl/status/status.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(recompute_total) \ + COUNTER(spill_active_total) \ + COUNTER(stale_locality_total) + +struct LoadAwareLocalityStats { + ALL_LOAD_AWARE_LOCALITY_STATS(GENERATE_COUNTER_STRUCT) +}; + +// 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: + // Out-of-band sentinel so monotonic time 0 remains a valid report timestamp. + static constexpr MonotonicTime kNeverReported = MonotonicTime::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); } + MonotonicTime lastUpdateTime() const { return last_update_time_.load(std::memory_order_acquire); } + +private: + void storeUtilization(double util, MonotonicTime now) { + if (!std::isfinite(util)) { + return; + } + utilization_.store(std::clamp(util, 0.0, 1.0), std::memory_order_relaxed); + last_update_time_.store(now, std::memory_order_release); + } + + static_assert(std::atomic::is_always_lock_free, + "std::atomic must be lock-free for safe cross-thread utilization updates"); + static_assert(std::atomic::is_always_lock_free, + "std::atomic must be lock-free for safe cross-thread freshness " + "updates"); + std::atomic utilization_{0.0}; + std::atomic last_update_time_{kNeverReported}; + TimeSource& time_source_; + const std::shared_ptr> metric_names_; +}; + +// Shared between config and worker factory; must outlive the child LB. +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, + 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_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_; + } + std::string endpointPickingPolicyName() const { return endpoint_picking_policy_factory_.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 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_; +}; + +// 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. 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 = LocalityRoutingWeightsMap; + +struct PriorityRoutingWeights { + enum class SelectionSource : uint8_t { Healthy = 0, Degraded = 1, AllHosts = 2 }; + static constexpr size_t kSourceCount = 3; + + struct SourceWeights { + 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; + + 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; } +}; + +// Advisory per-priority locality weights. Priority/health/panic selection stays live on workers. +struct RoutingWeightsSnapshot { + std::vector priority_weights; +}; + +using RoutingWeightsSnapshotConstSharedPtr = std::shared_ptr; + +struct ThreadLocalShim : public ThreadLocal::ThreadLocalObject { + RoutingWeightsSnapshotConstSharedPtr routing_weights; +}; + +class WorkerLocalLb; + +// Factory shared across workers; publishes routing-weight snapshots via TLS. +class WorkerLocalLbFactory : public Upstream::LoadBalancerFactory { +public: + 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; } + + void updateRoutingWeights(RoutingWeightsSnapshotConstSharedPtr snapshot) { + tls_->runOnAllThreads([snapshot = std::move(snapshot)](OptRef shim) { + if (shim.has_value()) { + shim->routing_weights = snapshot; + } + }); + } + + // SAFETY: Must only be called on a thread that owns a TLS slot instance (worker or main + // 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.ptr(); + } + + 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_; } + Upstream::ClusterLbStats& lbStats() const { return cluster_info_.lbStats(); } + +private: + // 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_; + // 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_; + + std::unique_ptr> tls_; +}; + +struct PerSourceLocalityState { + std::unique_ptr priority_set; + Upstream::LoadBalancerPtr lb; +}; + +struct PerLocalityState { + std::array by_source; + // Empty for a locality with no hosts. + 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; + // Whether index 0 is the local locality on this worker. + bool has_local_locality{false}; + // 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. +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: + void buildPerPriorityLocalities(); + + void buildPerLocality(uint32_t priority, const Upstream::HostSet& host_set); + + // allow_rebuild gates topology rebuilds for empty-delta in-place host updates. + void syncPriority(uint32_t priority, bool allow_rebuild); + + void updateLocalityHosts(PerSourceLocalityState& state, const Upstream::HostVector& hosts, + bool is_local, const Upstream::HostVector& hosts_added, + const Upstream::HostVector& hosts_removed); + + void syncLocalityState(PerLocalityState& state, const Upstream::HostSet& host_set, + size_t locality_index, bool recreate_child); + + 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; + }; + + struct LocalityLbSelection { + Upstream::LoadBalancer* lb{nullptr}; + const RoutingWeightsSnapshot* snapshot{nullptr}; + size_t locality_idx{0}; + }; + + PrioritySourcePick resolvePrioritySource(Upstream::LoadBalancerContext* context, 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, + 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. + 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, + PriorityRoutingWeights::SelectionSource source, + 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. + 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_; +}; + +// Main-thread LB: computes locality weights and publishes snapshots to workers. +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: + void computeLocalityRoutingWeights(); + + void addLbPolicyDataToHosts(const Upstream::HostVector& hosts); + + struct SourceComputeResult { + bool all_overloaded{false}; + bool local_preferred{false}; + bool probe_active{false}; + bool spill_active{false}; + uint32_t stale_localities{0}; + }; + + SourceComputeResult + computeSourceWeights(const Upstream::HostsPerLocality& all_hosts_per_locality, + const std::vector& eligible_hosts_per_locality, + MonotonicTime now, LocalityRoutingWeightsMap& 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, PriorityRoutingWeights::kSourceCount> 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_; +}; + +} // namespace LoadAwareLocality +} // namespace LoadBalancingPolicies +} // namespace Extensions +} // namespace Envoy diff --git a/test/extensions/load_balancing_policies/load_aware_locality/BUILD b/test/extensions/load_balancing_policies/load_aware_locality/BUILD index 39023a8a56f60..22d1d6e3cbf1f 100644 --- a/test/extensions/load_balancing_policies/load_aware_locality/BUILD +++ b/test/extensions/load_balancing_policies/load_aware_locality/BUILD @@ -18,10 +18,57 @@ envoy_extension_cc_test( rbe_pool = "6gig", deps = [ "//source/extensions/load_balancing_policies/load_aware_locality:config", - "//test/mocks/server:factory_context_mocks", + "//source/extensions/load_balancing_policies/round_robin:config", + "//test/mocks/server:server_factory_context_mocks", "//test/mocks/upstream:cluster_info_mocks", "//test/mocks/upstream:priority_set_mocks", - "//test/test_common:status_utility_lib", + "//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: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", + "envoy.load_balancing_policies.round_robin", + ], + 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 a201879097f3c..c73bbdf2b1ed3 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,11 +1,23 @@ +#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" +#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/server/server_factory_context.h" #include "test/mocks/upstream/cluster_info.h" #include "test/mocks/upstream/priority_set.h" -#include "test/test_common/status_utility.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 { @@ -13,30 +25,318 @@ namespace LoadBalancingPolicies { namespace LoadAwareLocality { namespace { -using ::Envoy::StatusHelpers::HasStatusCode; +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; +} -TEST(LoadAwareLocalityConfigTest, CreateFactory) { - NiceMock context; +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; - 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_THAT(lb_config, HasStatusCode(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; + + 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; + + 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, 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; + + 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, 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"}, + {"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, OobFieldsAreAcceptedAsNoOp) { + NiceMock context; + + // Out-of-band reporting is not implemented yet; the fields must be accepted so configs can be + // rolled out ahead of the implementation. + auto config_proto = loadAwareConfig({roundRobinEndpointPolicy()}); + config_proto.mutable_enable_oob_load_report()->set_value(true); + config_proto.mutable_oob_reporting_period()->set_seconds(10); + + Factory factory; + EXPECT_TRUE(factory.loadConfig(context, config_proto).ok()); +} + +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_weight_update_period()->set_nanos(50000000); + + 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, 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, 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; + + 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); } } // 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..e6cd4a0302884 --- /dev/null +++ b/test/extensions/load_balancing_policies/load_aware_locality/integration_test.cc @@ -0,0 +1,304 @@ +#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()) { + 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(const std::string& counter_name = "") { + // Capture before advancing: initialize() already ran a no-data tick, so an absolute + // threshold would be satisfied one tick early and provide no ordering. + uint64_t before = 0; + std::string counter; + if (!counter_name.empty()) { + counter = fmt::format("cluster.cluster_0.load_aware_locality.{}", counter_name); + const auto c = test_server_->counter(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, + const std::string& counter_name = "") { + seedAllUpstreams(utilizations); + advanceWeightTick(counter_name); + sendRequestsAndTrack(30, utilizations); + advanceWeightTick(counter_name); + } + + 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("local_preferred_total"); + + 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, "probe_active_total"); + + 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, "spill_active_total"); + + 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); + // 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); + 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, "local_preferred_total"); + + const std::vector spiked_local = {0.9, 0.9, 0.2, 0.2}; + sendRequestsAndTrack(30, spiked_local); + advanceWeightTick("spill_active_total"); + + 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, "spill_active_total"); + + 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..86fbee90a18fb --- /dev/null +++ b/test/extensions/load_balancing_policies/load_aware_locality/load_aware_locality_lb_test.cc @@ -0,0 +1,1668 @@ +#include +#include +#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/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; +using testing::ReturnRef; + +// 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"); +} + +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: + using Inspector = std::function; + + 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_; + 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_; } + + uint32_t createCount() const { return create_count_; } + +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 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 { +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 {}; + } + std::optional + selectExistingConnection(Upstream::LoadBalancerContext*, const Upstream::Host&, + std::vector&) override { + return std::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_; +}; + +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: + 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. + 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 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, 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 + makeWorkerFactoryWithChild(Upstream::LoadBalancerFactorySharedPtr child) { + return std::make_shared( + std::move(child), + std::make_shared(), + cluster_info_, context_.runtime_loader_, random_, context_.thread_local_); + } + + void setupPriorityLocalities( + uint32_t priority, std::vector localities, + bool has_local_locality = false, + 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) { + 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, + 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)); + } + + 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)); + } + + 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())); + } + + static LocalityRoutingWeightsMap makeWeightsMap(const std::vector& weights) { + LocalityRoutingWeightsMap map; + for (size_t i = 0; i < weights.size(); ++i) { + map[localityForIndex(i)] = weights[i]; + } + return map; + } + + 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: 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) { + 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()); + } + host_set->healthy_hosts_ = healthy_hosts; + host_set->healthy_hosts_per_locality_ = + Upstream::makeHostsPerLocality(std::move(*healthy_override), !has_local); + } else { + setupPriorityLocalities(priority, std::move(localities), has_local); + } + 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); + } + } + + 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); + const auto* shim = typed_factory != nullptr ? typed_factory->tlsShim() : nullptr; + return shim != nullptr ? shim->routing_weights.get() : 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; + } + + 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)); + } + + 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, + 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(startup_count, 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_; + 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 references returned by mock hosts. + std::deque localities_; + + Upstream::ThreadAwareLoadBalancerPtr thread_aware_lb_; + Upstream::LoadBalancerFactorySharedPtr factory_; +}; + +// 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 + 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}, + {"AllLocalitiesSaturatedFallBacksToHostCount", + {3, 1}, + false, + 0.1, + 0.0, + {1.0, 1.0}, + false, + {3.0, 1.0}, + 0.01}, + // All-overloaded fallback is checked before the local-preference + // snap: a saturated local locality spreads by host count, not + // all-local. + {"AllSaturatedWithLocalSpreadsByHostCountNotSnap", + {3, 1}, + true, + 0.1, + 0.0, + {1.0, 1.0}, + false, + {3.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}, + false, + {0.94, 0.02, 0.04}, + 0.01}, + {"ProbeHostCountProportionalUnequalUtils", + {1, 2, 4}, + true, + 1.0, + 0.06, + {0.5, 0.8, 0.2}, + false, + {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); + + auto h1 = makeWeightTrackingMockHost(); + setupLocalities({{h1}}); + recomputeWeights(); + + auto populated_lb = createWorkerLb(); + ASSERT_NE(nullptr, populated_lb); + EXPECT_EQ(h1, populated_lb->chooseHost(nullptr).host); + 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, 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}); + + // Zero-utilization reports carry no positive signal and do not refresh host data. + setHostUtilization(h1, 0.0); + setHostUtilization(h2, 0.9); + recomputeWeights(); + expectWeights(PriorityRoutingWeights::SelectionSource::Healthy, {0.1, 0.1}); +} + +TEST_F(LoadAwareLocalityLbTest, WeightExpirationIgnoresStaleData) { + auto [h1, h2] = setupTwoLocalityLb(std::chrono::milliseconds(5000)); + + recomputeWeights(); + + 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}); + + // 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) { + // Expiration disabled (0ms) — data is always used regardless of age. + 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, h2] = setupTwoLocalityLb(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, h_remote] = + setupTwoLocalityLb(std::chrono::milliseconds(5000), /*has_local=*/true); + + setHostUtilization(h_local, 0.5); + setHostUtilization(h_remote, 0.6); + recomputeWeights(); + + expectAllLocal(true); + + // 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(); + + 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(); + + // 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, 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, 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(); + 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, 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) { + const auto hosts = setupPanicHosts(); + createLb(); + + // 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(); + auto counts = countPicks(*worker_lb, 400); + 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); + + // peekAnotherHost resolves panic through the same path but must not double-count the stat. + const uint64_t panic_count = cluster_info_.lbStats().lb_healthy_panic_.value(); + worker_lb->peekAnotherHost(nullptr); + EXPECT_EQ(panic_count, cluster_info_.lbStats().lb_healthy_panic_.value()); +} + +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, 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, 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); + } + + reshapeLocalities(0, /*localities=*/{}, /*added=*/{}, /*removed=*/{}, /*has_local=*/false, + 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 +// 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). +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); + // With a local locality present, the fallback must still win over the local-preference snap. + setupLocalities({h_a, h_b}, /*has_local_locality=*/true); + + createLb(); + EXPECT_EQ(0, counterValue(cluster_info_, "load_aware_locality.all_overloaded_total")); + + setUtilizationForHosts(h_a, 1.0); + setUtilizationForHosts(h_b, 1.0); + recomputeWeights(); + EXPECT_EQ(1, counterValue(cluster_info_, "load_aware_locality.all_overloaded_total")); + + expectCounterIncrementsPerTick("load_aware_locality.all_overloaded_total", 1); +} + +TEST_F(LoadAwareLocalityLbTest, DedicatedStatsLocalPreferredTotalIncOncePerTickOnVarianceSnap) { + 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) { + 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, 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, DedicatedStatsSpillActiveNotCountedWithoutEligibleLocalHosts) { + auto h_local = makeHosts(1); + auto h_remote = makeHosts(1); + // Local hosts are ineligible in the healthy view: spill requires the utilization comparison to + // have run and failed; an empty local slice spills for health, not load. + setupLocalities({h_local, h_remote}, /*has_local_locality=*/true, + std::vector{{}, h_remote}); + + createLb(); + setUtilizationForHosts(h_remote, 0.2); + recomputeWeights(); + EXPECT_EQ(0, counterValue(cluster_info_, "load_aware_locality.spill_active_total")); + expectCounterIncrementsPerTick("load_aware_locality.spill_active_total", 0); +} + +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)); + + EXPECT_EQ(0, counterValue(cluster_info_, "load_aware_locality.stale_locality_total")); + expectCounterIncrementsPerTick("load_aware_locality.stale_locality_total", 0); + + setUtilizationForHosts(locality_a, 0.5); + setUtilizationForHosts(locality_b, 0.4); + setUtilizationForHosts(locality_c, 0.6); + recomputeWeights(); + const uint64_t baseline = counterValue(cluster_info_, "load_aware_locality.stale_locality_total"); + + context_.time_system_.advanceTimeWait(std::chrono::milliseconds(6000)); + setUtilizationForHosts(locality_a, 0.5); + recomputeWeights(); + EXPECT_EQ(baseline + 2, counterValue(cluster_info_, "load_aware_locality.stale_locality_total")); + + 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, 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 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::_)) + .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 cancellation 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); + 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(); + auto worker_lb = createWorkerLb(); + + publishHealthyWeights({0.0, 1.0}); + auto counts = countPicks(*worker_lb, 100); + EXPECT_GE(counts[h2.get()], 98); + + publishHealthyWeights({0.1, 0.9}); + forceRandomValues({0}); + auto peeked = worker_lb->peekAnotherHost(nullptr); + ASSERT_NE(nullptr, peeked); + + // 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()); +} + +TEST_F(LoadAwareLocalityLbTest, WorkerFallbackEdgeCases) { + auto h1 = makeWeightTrackingMockHost(); + auto h2 = makeWeightTrackingMockHost(); + setupLocalities({{h1}, {h2}}); + + 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{}})); + 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}); + // 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(); + 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, 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(); + 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, 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, ZoneRoutingStatsFollowWorkerTopologyOnLocalityFlip) { + 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(0u, zoneRoutingStatTotal()); + + // The local locality appears on this worker before the next snapshot: attribution follows the + // worker's live topology immediately rather than the stale snapshot. + reshapeLocalities(0, {{h1}, {h2}}, /*added=*/{}, /*removed=*/{}, /*has_local=*/true); + for (int i = 0; i < 100; ++i) { + worker_lb->chooseHost(nullptr); + } + EXPECT_GT(zoneRoutingStatTotal(), 0u); +} + +TEST_F(LoadAwareLocalityLbTest, ZoneRoutingStatsCoverNoLocalDegradedAndPanicSources) { + 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(0u, zoneRoutingStatTotal()); + + 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(); + // Degraded-source spill: local degraded utilization above the remote average, so local picks + // count as sampled and remote picks as cross-zone. + setHostUtilization(h_local, 0.9); + setHostUtilization(h_remote, 0.2); + recomputeWeights(); + + worker_lb = createWorkerLb(); + auto degraded_counts = countPicks(*worker_lb, 400); + EXPECT_GT(degraded_counts[h_local.get()], 0); + EXPECT_GT(degraded_counts[h_remote.get()], 0); + EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_sampled_.value(), 0); + EXPECT_GT(cluster_info_.lbStats().lb_zone_routing_cross_zone_.value(), 0); + + const uint64_t before_panic_zone_stats = zoneRoutingStatTotal(); + const auto panic_hosts = setupPanicHosts(/*has_local_locality=*/true); + + 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) { + 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, 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); + + 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); +} + +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)); +} + +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). + struct Case { + std::string name; + std::vector metric_names; + std::optional named_metric; + std::optional application_utilization; + std::optional cpu_utilization; + double expected_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}, + }; + + for (const Case& c : cases) { + SCOPED_TRACE(c.name); + LocalityLbHostData slot(simTime(), makeMetricNames(c.metric_names)); + xds::data::orca::v3::OrcaLoadReport report; + 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(c.expected_utilization, slot.utilization()); + } +} + +// When getUtilizationFromOrcaReport returns a non-finite value (Inf/NaN), storeUtilization +// discards it: utilization() stays 0.0 and lastUpdateTime() stays kNeverReported. +TEST_F(LoadAwareLocalityLbTest, StoreUtilizationRejectsNonFiniteValue) { + LocalityLbHostData slot(simTime(), 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.lastUpdateTime()); +} + +} // namespace +} // namespace LoadAwareLocality +} // namespace LoadBalancingPolicies +} // namespace Extensions +} // namespace Envoy