diff --git a/changelogs/current/new_features/dynamic_modules__native-load-balancing-clusters.rst b/changelogs/current/new_features/dynamic_modules__native-load-balancing-clusters.rst new file mode 100644 index 0000000000000..5c34aa6def9e3 --- /dev/null +++ b/changelogs/current/new_features/dynamic_modules__native-load-balancing-clusters.rst @@ -0,0 +1,5 @@ +Dynamic module clusters (``envoy.clusters.dynamic_modules``) can now use +Envoy's built-in load balancers. In addition to ``CLUSTER_PROVIDED``, +``lb_policy`` may be set to ``LEAST_REQUEST``, ``ROUND_ROBIN``, ``RANDOM``, +``RING_HASH``, or ``MAGLEV``; the module then supplies only host discovery and +Envoy performs host selection. diff --git a/source/extensions/clusters/dynamic_modules/cluster.cc b/source/extensions/clusters/dynamic_modules/cluster.cc index 768c5b87848b1..a0aeb1ac0f899 100644 --- a/source/extensions/clusters/dynamic_modules/cluster.cc +++ b/source/extensions/clusters/dynamic_modules/cluster.cc @@ -903,12 +903,22 @@ DynamicModuleClusterFactory::createClusterWithConfig( const envoy::extensions::clusters::dynamic_modules::v3::ClusterConfig& proto_config, Upstream::ClusterFactoryContext& context) { - // Validate that CLUSTER_PROVIDED LB policy is used. - if (cluster.lb_policy() != envoy::config::cluster::v3::Cluster::CLUSTER_PROVIDED) { + // CLUSTER_PROVIDED uses the module's load balancer; the native policies use Envoy's factory + // load balancer, with the module supplying only host discovery. + const auto policy = cluster.lb_policy(); + const bool is_module_lb = (policy == envoy::config::cluster::v3::Cluster::CLUSTER_PROVIDED); + const bool is_native_lb = (policy == envoy::config::cluster::v3::Cluster::LEAST_REQUEST || + policy == envoy::config::cluster::v3::Cluster::ROUND_ROBIN || + policy == envoy::config::cluster::v3::Cluster::RANDOM || + policy == envoy::config::cluster::v3::Cluster::RING_HASH || + policy == envoy::config::cluster::v3::Cluster::MAGLEV); + + if (!is_module_lb && !is_native_lb) { return absl::InvalidArgumentError( fmt::format("cluster: LB policy {} is not valid for cluster type " - "'envoy.clusters.dynamic_modules'. Only 'CLUSTER_PROVIDED' is allowed.", - envoy::config::cluster::v3::Cluster::LbPolicy_Name(cluster.lb_policy()))); + "'envoy.clusters.dynamic_modules'. Supported policies are CLUSTER_PROVIDED, " + "LEAST_REQUEST, ROUND_ROBIN, RANDOM, RING_HASH, and MAGLEV.", + envoy::config::cluster::v3::Cluster::LbPolicy_Name(policy))); } Server::Configuration::ServerFactoryContext& server_context = context.serverFactoryContext(); @@ -952,9 +962,13 @@ DynamicModuleClusterFactory::createClusterWithConfig( cluster, std::move(config_or_error.value()), context, creation_status)); RETURN_IF_NOT_OK(creation_status); - // Create the thread-aware load balancer. - auto handle = std::make_shared(new_cluster); - auto lb = std::make_unique(handle); + // Create the thread-aware load balancer only if the module provides LB (CLUSTER_PROVIDED). + // For native LB policies, return nullptr so the cluster manager builds the native factory LB. + Upstream::ThreadAwareLoadBalancerPtr lb; + if (cluster.lb_policy() == envoy::config::cluster::v3::Cluster::CLUSTER_PROVIDED) { + auto handle = std::make_shared(new_cluster); + lb = std::make_unique(handle); + } return std::make_pair(std::move(new_cluster), std::move(lb)); } diff --git a/source/extensions/dynamic_modules/sdk/rust/src/cluster.rs b/source/extensions/dynamic_modules/sdk/rust/src/cluster.rs index c6235a55ae1ae..21a4f30488615 100644 --- a/source/extensions/dynamic_modules/sdk/rust/src/cluster.rs +++ b/source/extensions/dynamic_modules/sdk/rust/src/cluster.rs @@ -48,7 +48,16 @@ pub trait Cluster: Send + Sync { /// /// Each worker thread gets its own load balancer instance. The `envoy_lb` /// provides thread-local access to the cluster's host set. - fn new_load_balancer(&self, envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box; + /// + /// Return `Some(lb)` if this cluster implements host selection, or `None` if the cluster + /// only provides host discovery and Envoy should use its native load balancer. + /// When returning `None`, Envoy will use the standard load balancer factory based on + /// `lb_policy` + `common_lb_config` (e.g., zone-aware or locality-weighted routing). + /// The module's `choose_host` hook will never be called if this returns `None`. + fn new_load_balancer( + &self, + envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option>; /// Called on the main thread when a new event is scheduled via /// [`EnvoyClusterScheduler::commit`] for this [`Cluster`]. @@ -2560,6 +2569,9 @@ pub unsafe extern "C" fn envoy_dynamic_module_on_cluster_destroy( /// Wrapper that pairs a module-side load balancer with the Envoy-side LB pointer. /// The `lb_envoy_ptr` is needed by [`ClusterLbContextRef::should_select_another_host`] to /// resolve host pointers from the priority set. +/// +/// If the module returns None from `new_load_balancer()`, this wrapper is not created and +/// a null pointer is returned, signaling to Envoy to use the native factory load balancer. struct ClusterLbWrapper { lb: Box, lb_envoy_ptr: abi::envoy_dynamic_module_type_cluster_lb_envoy_ptr, @@ -2578,9 +2590,16 @@ pub unsafe extern "C" fn envoy_dynamic_module_on_cluster_lb_new( let cluster = cluster_module_ptr as *const *const dyn Cluster; let cluster = &**cluster; let envoy_lb = EnvoyClusterLoadBalancerImpl::new(lb_envoy_ptr); - let lb = cluster.new_load_balancer(&envoy_lb); - let wrapper = Box::new(ClusterLbWrapper { lb, lb_envoy_ptr }); - Box::into_raw(wrapper) as abi::envoy_dynamic_module_type_cluster_lb_module_ptr + match cluster.new_load_balancer(&envoy_lb) { + Some(lb) => { + let wrapper = Box::new(ClusterLbWrapper { lb, lb_envoy_ptr }); + Box::into_raw(wrapper) as abi::envoy_dynamic_module_type_cluster_lb_module_ptr + }, + None => { + // Module does not provide a load balancer; return null so Envoy uses native LB. + std::ptr::null() + }, + } })) .unwrap_or_else(|panic| { crate::log_ffi_panic("envoy_dynamic_module_on_cluster_lb_new", panic); diff --git a/source/extensions/dynamic_modules/sdk/rust/src/lib.rs b/source/extensions/dynamic_modules/sdk/rust/src/lib.rs index 19a0e7b42bc47..39feb2f6a1823 100644 --- a/source/extensions/dynamic_modules/sdk/rust/src/lib.rs +++ b/source/extensions/dynamic_modules/sdk/rust/src/lib.rs @@ -1420,8 +1420,8 @@ pub static NEW_CLUSTER_CONFIG_FUNCTION: OnceLock = Onc /// envoy_cluster.pre_init_complete(); /// } /// -/// fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { -/// Box::new(MyClusterLb {}) +/// fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Option> { +/// Some(Box::new(MyClusterLb {})) /// } /// } /// diff --git a/source/extensions/dynamic_modules/sdk/rust/src/lib_test.rs b/source/extensions/dynamic_modules/sdk/rust/src/lib_test.rs index 1d70a4bfe7153..81faea08aa041 100644 --- a/source/extensions/dynamic_modules/sdk/rust/src/lib_test.rs +++ b/source/extensions/dynamic_modules/sdk/rust/src/lib_test.rs @@ -6973,7 +6973,10 @@ fn test_cluster_callout_done_with_null_buffers_yields_none() { struct TestCluster; impl Cluster for TestCluster { fn on_init(&mut self, _envoy_cluster: &dyn EnvoyCluster) {} - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { unimplemented!("not exercised by this test") } diff --git a/test/extensions/clusters/dynamic_modules/BUILD b/test/extensions/clusters/dynamic_modules/BUILD index 9f454206966ab..a11cd83590a14 100644 --- a/test/extensions/clusters/dynamic_modules/BUILD +++ b/test/extensions/clusters/dynamic_modules/BUILD @@ -27,6 +27,11 @@ envoy_cc_test( "//source/extensions/clusters/dynamic_modules:cluster_lib", "//source/extensions/dynamic_modules:abi_impl", "//source/extensions/load_balancing_policies/cluster_provided:config", + "//source/extensions/load_balancing_policies/least_request:config", + "//source/extensions/load_balancing_policies/maglev:config", + "//source/extensions/load_balancing_policies/random:config", + "//source/extensions/load_balancing_policies/ring_hash:config", + "//source/extensions/load_balancing_policies/round_robin:config", "//source/extensions/transport_sockets/raw_buffer:config", "//test/common/upstream:utility_lib", "//test/extensions/dynamic_modules:util", diff --git a/test/extensions/clusters/dynamic_modules/cluster_test.cc b/test/extensions/clusters/dynamic_modules/cluster_test.cc index fd179273a48ed..0ee6914c2b65d 100644 --- a/test/extensions/clusters/dynamic_modules/cluster_test.cc +++ b/test/extensions/clusters/dynamic_modules/cluster_test.cc @@ -207,8 +207,52 @@ TEST_F(DynamicModuleClusterTest, CreationWithClusterConfig) { EXPECT_NE(nullptr, result->second); } -// Test that a non-CLUSTER_PROVIDED lb_policy is rejected. -TEST_F(DynamicModuleClusterTest, InvalidLbPolicy) { +// Test that CLUSTER_PROVIDED lb_policy returns a non-null thread-aware LB. +TEST_F(DynamicModuleClusterTest, ClusterProvidedLbPolicy) { + const std::string yaml = R"EOF( +name: test_cluster +connect_timeout: 0.25s +lb_policy: CLUSTER_PROVIDED +cluster_type: + name: envoy.clusters.dynamic_modules + typed_config: + "@type": type.googleapis.com/envoy.extensions.clusters.dynamic_modules.v3.ClusterConfig + dynamic_module_config: + name: cluster_no_op + cluster_name: test +)EOF"; + + auto result = createCluster(yaml); + ASSERT_TRUE(result.ok()) << result.status().message(); + EXPECT_NE(nullptr, result->first); + // CLUSTER_PROVIDED should return a non-null thread-aware LB (module LB). + EXPECT_NE(nullptr, result->second); +} + +// Test that LEAST_REQUEST lb_policy is accepted and returns nullptr thread-aware LB. +TEST_F(DynamicModuleClusterTest, LeastRequestLbPolicy) { + const std::string yaml = R"EOF( +name: test_cluster +connect_timeout: 0.25s +lb_policy: LEAST_REQUEST +cluster_type: + name: envoy.clusters.dynamic_modules + typed_config: + "@type": type.googleapis.com/envoy.extensions.clusters.dynamic_modules.v3.ClusterConfig + dynamic_module_config: + name: cluster_no_op + cluster_name: test +)EOF"; + + auto result = createCluster(yaml); + ASSERT_TRUE(result.ok()) << result.status().message(); + EXPECT_NE(nullptr, result->first); + // Native LB policies should return nullptr thread-aware LB. + EXPECT_EQ(nullptr, result->second); +} + +// Test that ROUND_ROBIN lb_policy is accepted and returns nullptr thread-aware LB. +TEST_F(DynamicModuleClusterTest, RoundRobinLbPolicy) { const std::string yaml = R"EOF( name: test_cluster connect_timeout: 0.25s @@ -223,7 +267,76 @@ lb_policy: ROUND_ROBIN )EOF"; auto result = createCluster(yaml); - ASSERT_THAT(result, HasStatusMessage(testing::HasSubstr("CLUSTER_PROVIDED"))); + ASSERT_TRUE(result.ok()) << result.status().message(); + EXPECT_NE(nullptr, result->first); + // Native LB policies should return nullptr thread-aware LB. + EXPECT_EQ(nullptr, result->second); +} + +// Test that RANDOM lb_policy is accepted and returns nullptr thread-aware LB. +TEST_F(DynamicModuleClusterTest, RandomLbPolicy) { + const std::string yaml = R"EOF( +name: test_cluster +connect_timeout: 0.25s +lb_policy: RANDOM +cluster_type: + name: envoy.clusters.dynamic_modules + typed_config: + "@type": type.googleapis.com/envoy.extensions.clusters.dynamic_modules.v3.ClusterConfig + dynamic_module_config: + name: cluster_no_op + cluster_name: test +)EOF"; + + auto result = createCluster(yaml); + ASSERT_TRUE(result.ok()) << result.status().message(); + EXPECT_NE(nullptr, result->first); + // Native LB policies should return nullptr thread-aware LB. + EXPECT_EQ(nullptr, result->second); +} + +// Test that RING_HASH lb_policy is accepted and returns nullptr thread-aware LB. +TEST_F(DynamicModuleClusterTest, RingHashLbPolicy) { + const std::string yaml = R"EOF( +name: test_cluster +connect_timeout: 0.25s +lb_policy: RING_HASH +cluster_type: + name: envoy.clusters.dynamic_modules + typed_config: + "@type": type.googleapis.com/envoy.extensions.clusters.dynamic_modules.v3.ClusterConfig + dynamic_module_config: + name: cluster_no_op + cluster_name: test +)EOF"; + + auto result = createCluster(yaml); + ASSERT_TRUE(result.ok()) << result.status().message(); + EXPECT_NE(nullptr, result->first); + // Native LB policies should return nullptr thread-aware LB. + EXPECT_EQ(nullptr, result->second); +} + +// Test that MAGLEV lb_policy is accepted and returns nullptr thread-aware LB. +TEST_F(DynamicModuleClusterTest, MaglevLbPolicy) { + const std::string yaml = R"EOF( +name: test_cluster +connect_timeout: 0.25s +lb_policy: MAGLEV +cluster_type: + name: envoy.clusters.dynamic_modules + typed_config: + "@type": type.googleapis.com/envoy.extensions.clusters.dynamic_modules.v3.ClusterConfig + dynamic_module_config: + name: cluster_no_op + cluster_name: test +)EOF"; + + auto result = createCluster(yaml); + ASSERT_TRUE(result.ok()) << result.status().message(); + EXPECT_NE(nullptr, result->first); + // Native LB policies should return nullptr thread-aware LB. + EXPECT_EQ(nullptr, result->second); } // Test that a missing module fails gracefully. diff --git a/test/extensions/clusters/dynamic_modules/integration_test.cc b/test/extensions/clusters/dynamic_modules/integration_test.cc index e3692ef021c9f..bfe80e4162907 100644 --- a/test/extensions/clusters/dynamic_modules/integration_test.cc +++ b/test/extensions/clusters/dynamic_modules/integration_test.cc @@ -510,6 +510,71 @@ TEST_P(DynamicModuleClusterDynamicMetadataIntegrationTest, SetsDynamicMetadataDu EXPECT_EQ("1234 test_value", log); } +// ============================================================================= +// Native LB test: cluster with LEAST_REQUEST policy (native LB, not module LB). +// ============================================================================= + +class DynamicModuleClusterNativeLbIntegrationTest + : public testing::TestWithParam, + public HttpIntegrationTest { +public: + DynamicModuleClusterNativeLbIntegrationTest() + : HttpIntegrationTest(Http::CodecType::HTTP1, GetParam()) {} + + void initializeWithNativeLb() { + TestEnvironment::setEnvVar( + "ENVOY_DYNAMIC_MODULES_SEARCH_PATH", + TestEnvironment::runfilesPath("test/extensions/dynamic_modules/test_data/rust"), 1); + + config_helper_.addConfigModifier([this](envoy::config::bootstrap::v3::Bootstrap& bootstrap) { + auto* cluster = bootstrap.mutable_static_resources()->mutable_clusters(0); + const std::string upstream_address = fake_upstreams_[0]->localAddress()->asString(); + + cluster->set_name("cluster_0"); + // Use LEAST_REQUEST policy instead of CLUSTER_PROVIDED to test native LB. + cluster->set_lb_policy(envoy::config::cluster::v3::Cluster::LEAST_REQUEST); + cluster->clear_load_assignment(); + + envoy::extensions::clusters::dynamic_modules::v3::ClusterConfig dec_config; + dec_config.mutable_dynamic_module_config()->set_name("cluster_integration_test"); + dec_config.set_cluster_name("native_lb_test"); + + Protobuf::StringValue config_proto; + config_proto.set_value(upstream_address); + std::ignore = dec_config.mutable_cluster_config()->PackFrom(config_proto); + + cluster->mutable_cluster_type()->set_name("envoy.clusters.dynamic_modules"); + std::ignore = cluster->mutable_cluster_type()->mutable_typed_config()->PackFrom(dec_config); + }); + + HttpIntegrationTest::initialize(); + } +}; + +INSTANTIATE_TEST_SUITE_P(IpVersions, DynamicModuleClusterNativeLbIntegrationTest, + testing::ValuesIn(TestEnvironment::getIpVersionsForTest()), + TestUtility::ipTestParamsToString); + +// A native lb_policy makes Envoy build its factory load balancer instead of the module's. The +// request routes to the upstream, and the module's choose_host must not run: the module provides +// a load balancer, but native_lb_requests stays 0 because Envoy never calls into it. +TEST_P(DynamicModuleClusterNativeLbIntegrationTest, NativeLbWithLeastRequestPolicy) { + initializeWithNativeLb(); + codec_client_ = makeHttpConnection(makeClientConnection(lookupPort("http"))); + + auto response = + sendRequestAndWaitForResponse(default_request_headers_, 0, default_response_headers_, 0); + + EXPECT_TRUE(upstream_request_->complete()); + EXPECT_TRUE(response->complete()); + EXPECT_EQ("200", response->headers().getStatusValue()); + + // Native LB selected the host; the module's choose_host was never called, so its + // native_lb_requests counter is unincremented (absent or 0). + auto choose_host = test_server_->counter("dynamicmodulescustom.native_lb_requests"); + EXPECT_TRUE(choose_host == nullptr || choose_host->value() == 0); +} + } // namespace DynamicModules } // namespace Clusters } // namespace Extensions diff --git a/test/extensions/dynamic_modules/test_data/rust/cluster_dynamic_metadata_test.rs b/test/extensions/dynamic_modules/test_data/rust/cluster_dynamic_metadata_test.rs index 381991dc3f3c9..04253b519b8e8 100644 --- a/test/extensions/dynamic_modules/test_data/rust/cluster_dynamic_metadata_test.rs +++ b/test/extensions/dynamic_modules/test_data/rust/cluster_dynamic_metadata_test.rs @@ -69,10 +69,13 @@ impl Cluster for DynamicMetadataWriterCluster { envoy_cluster.pre_init_complete(); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(DynamicMetadataWriterLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(DynamicMetadataWriterLb { hosts: self.hosts.clone(), - }) + })) } } diff --git a/test/extensions/dynamic_modules/test_data/rust/cluster_filter_state_test.rs b/test/extensions/dynamic_modules/test_data/rust/cluster_filter_state_test.rs index 1fb007a254660..7ba5daa7633c7 100644 --- a/test/extensions/dynamic_modules/test_data/rust/cluster_filter_state_test.rs +++ b/test/extensions/dynamic_modules/test_data/rust/cluster_filter_state_test.rs @@ -128,10 +128,13 @@ impl Cluster for FilterStateReaderCluster { envoy_cluster.pre_init_complete(); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(FilterStateReaderLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(FilterStateReaderLb { hosts: self.hosts.clone(), - }) + })) } } @@ -200,10 +203,13 @@ impl Cluster for FilterStateWriterCluster { envoy_cluster.pre_init_complete(); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(FilterStateWriterLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(FilterStateWriterLb { hosts: self.hosts.clone(), - }) + })) } } diff --git a/test/extensions/dynamic_modules/test_data/rust/cluster_integration_test.rs b/test/extensions/dynamic_modules/test_data/rust/cluster_integration_test.rs index ec3b87e7a6497..479dc20ada9da 100644 --- a/test/extensions/dynamic_modules/test_data/rust/cluster_integration_test.rs +++ b/test/extensions/dynamic_modules/test_data/rust/cluster_integration_test.rs @@ -89,6 +89,10 @@ fn new_cluster_config( metrics: envoy_cluster_metrics, })) }, + "native_lb_test" => Some(Box::new(NativeLbTestClusterConfig { + upstream_address: config_str.to_string(), + metrics: envoy_cluster_metrics, + })), _ => None, } } @@ -140,13 +144,16 @@ impl Cluster for SyncHostSelectionCluster { envoy_cluster.pre_init_complete(); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(SyncHostSelectionLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(SyncHostSelectionLb { hosts: self.hosts.clone(), index: AtomicUsize::new(0), counter_id: self.counter_id, metrics: self.metrics.clone(), - }) + })) } } @@ -207,10 +214,13 @@ impl Cluster for AsyncHostSelectionCluster { envoy_cluster.pre_init_complete(); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(AsyncHostSelectionLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(AsyncHostSelectionLb { hosts: self.hosts.clone(), - }) + })) } } @@ -297,11 +307,14 @@ impl Cluster for SchedulerHostUpdateCluster { scheduler.commit(ADD_HOST_EVENT_ID); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(SchedulerHostUpdateLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(SchedulerHostUpdateLb { hosts: self.hosts.clone(), membership_update_count: AtomicUsize::new(0), - }) + })) } fn on_scheduled(&self, envoy_cluster: &dyn EnvoyCluster, event_id: u64) { @@ -376,10 +389,13 @@ impl Cluster for LifecycleCallbacksCluster { envoy_cluster.pre_init_complete(); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(LifecycleCallbacksLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(LifecycleCallbacksLb { hosts: self.hosts.clone(), - }) + })) } fn on_server_initialized(&mut self, _envoy_cluster: &dyn EnvoyCluster) { @@ -472,10 +488,13 @@ impl Cluster for RunOnAllWorkersCluster { scheduler.commit(RUN_ON_ALL_WORKERS_TRIGGER_EVENT_ID); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(RunOnAllWorkersLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(RunOnAllWorkersLb { hosts: self.hosts.clone(), - }) + })) } fn on_scheduled(&self, envoy_cluster: &dyn EnvoyCluster, event_id: u64) { @@ -557,13 +576,16 @@ impl Cluster for WorkerLocalRebuildCluster { scheduler.commit(WORKER_LOCAL_REBUILD_ADD_EVENT_ID); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(WorkerLocalRebuildLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(WorkerLocalRebuildLb { hosts: Vec::new(), index: 0, counter_id: self.counter_id, metrics: self.metrics.clone(), - }) + })) } fn on_scheduled(&self, envoy_cluster: &dyn EnvoyCluster, event_id: u64) { @@ -671,13 +693,16 @@ impl Cluster for MemberUpdatePackedAddressCluster { scheduler.commit(PACKED_ADDRESS_ADD_EVENT_ID); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(MemberUpdatePackedAddressLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(MemberUpdatePackedAddressLb { hosts: Vec::new(), index: 0, counter_id: self.counter_id, metrics: self.metrics.clone(), - }) + })) } fn on_scheduled(&self, envoy_cluster: &dyn EnvoyCluster, event_id: u64) { @@ -806,14 +831,17 @@ impl Cluster for WorkerTimerCluster { envoy_cluster.pre_init_complete(); } - fn new_load_balancer(&self, _envoy_lb: &dyn EnvoyClusterLoadBalancer) -> Box { - Box::new(WorkerTimerLb { + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(WorkerTimerLb { hosts: self.hosts.clone(), timer: None, armed_id: self.armed_id, fired_id: self.fired_id, metrics: self.metrics.clone(), - }) + })) } } @@ -868,3 +896,83 @@ impl ClusterLb for WorkerTimerLb { timer.enable(std::time::Duration::from_millis(WORKER_TIMER_INTERVAL_MS)); } } + +// ============================================================================= +// Native LB test: module provides a load balancer that must be bypassed. +// ============================================================================= +// +// choose_host increments native_lb_requests. Tests configure a native lb_policy +// (LEAST_REQUEST, ROUND_ROBIN, RANDOM) so Envoy uses its factory LB; the module's +// choose_host must never run, so native_lb_requests stays 0. + +struct NativeLbTestClusterConfig { + upstream_address: String, + metrics: Arc, +} + +impl ClusterConfig for NativeLbTestClusterConfig { + fn new_cluster(&self, _envoy_cluster: &dyn EnvoyCluster) -> Box { + let counter_id = self.metrics.define_counter("native_lb_requests").ok(); + Box::new(NativeLbTestCluster { + upstream_address: self.upstream_address.clone(), + hosts: Arc::new(Mutex::new(HostList(Vec::new()))), + counter_id, + metrics: self.metrics.clone(), + }) + } +} + +struct NativeLbTestCluster { + upstream_address: String, + hosts: SharedHostList, + counter_id: Option, + metrics: Arc, +} + +impl Cluster for NativeLbTestCluster { + fn on_init(&mut self, envoy_cluster: &dyn EnvoyCluster) { + let addresses = vec![self.upstream_address.clone()]; + let weights = vec![1u32]; + if let Some(host_ptrs) = envoy_cluster.add_hosts(&addresses, &weights) { + self.hosts.lock().unwrap().0 = host_ptrs; + } + envoy_cluster.pre_init_complete(); + } + + fn new_load_balancer( + &self, + _envoy_lb: &dyn EnvoyClusterLoadBalancer, + ) -> Option> { + Some(Box::new(NativeLbTestLb { + hosts: self.hosts.clone(), + index: AtomicUsize::new(0), + counter_id: self.counter_id, + metrics: self.metrics.clone(), + })) + } +} + +struct NativeLbTestLb { + hosts: SharedHostList, + index: AtomicUsize, + counter_id: Option, + metrics: Arc, +} + +impl ClusterLb for NativeLbTestLb { + fn choose_host( + &mut self, + _context: Option<&dyn ClusterLbContext>, + _async_completion: Box, + ) -> HostSelectionResult { + let hosts = self.hosts.lock().unwrap(); + if hosts.0.is_empty() { + return HostSelectionResult::NoHost; + } + let idx = self.index.fetch_add(1, Ordering::Relaxed) % hosts.0.len(); + if let Some(counter_id) = self.counter_id { + let _ = self.metrics.increment_counter(counter_id, 1); + } + HostSelectionResult::Selected(hosts.0[idx]) + } +}