diff --git a/grpc/src/client/channel.rs b/grpc/src/client/channel.rs index b44baeff1..5176d54e7 100644 --- a/grpc/src/client/channel.rs +++ b/grpc/src/client/channel.rs @@ -55,6 +55,7 @@ use crate::client::load_balancing::QueuingPicker; use crate::client::load_balancing::WorkData; use crate::client::load_balancing::WorkScheduler; use crate::client::load_balancing::graceful_switch::GracefulSwitchPolicy; +use crate::client::load_balancing::least_request; use crate::client::load_balancing::pick_first; use crate::client::load_balancing::round_robin; use crate::client::load_balancing::subchannel::Subchannel; @@ -171,6 +172,7 @@ impl Channel { ) -> Self { pick_first::reg(); round_robin::reg(); + least_request::reg(); dns::reg(); #[cfg(unix)] name_resolution::unix::reg(); @@ -380,11 +382,20 @@ impl Invoke for Arc { "channel has been closed", )); }; - let result = &state.picker.pick(&headers); + let result = state.picker.pick(&headers); match result { - PickResult::Pick(pr) => { + PickResult::Pick(mut pr) => { if let Some(sc) = pr.subchannel.downcast_ref::() { - return sc.dyn_invoke(headers, options.clone()).await; + let (tx, rx) = sc.dyn_invoke(headers, options.clone()).await; + let rx = if let Some(on_complete) = pr.on_complete.take() { + Box::new(CallbackRecvStream { + delegate: rx, + on_complete: Some(on_complete), + }) as Box + } else { + rx + }; + return (tx, rx); } else { panic!( "picked subchannel is not an implementation provided by the channel" @@ -395,7 +406,7 @@ impl Invoke for Arc { // Continue and retry the RPC with the next picker. } PickResult::Fail(status) => { - return FailingRecvStream::new_stream_pair(status.clone()); + return FailingRecvStream::new_stream_pair(status); } PickResult::Drop(status) => { todo!("dropped pick: {:?}", status); @@ -411,6 +422,39 @@ impl Drop for ActiveChannel { } } +struct CallbackRecvStream { + delegate: Box, + on_complete: Option>, +} + +#[tonic::async_trait] +impl DynRecvStream for CallbackRecvStream { + async fn dyn_recv( + &mut self, + msg: &mut dyn crate::core::RecvMessage, + ) -> crate::client::ResponseStreamItem { + let item = self.delegate.dyn_recv(msg).await; + if matches!( + item, + crate::client::ResponseStreamItem::Trailers(_) + | crate::client::ResponseStreamItem::StreamClosed + ) { + if let Some(cb) = self.on_complete.take() { + cb(); + } + } + item + } +} + +impl Drop for CallbackRecvStream { + fn drop(&mut self) { + if let Some(cb) = self.on_complete.take() { + cb(); + } + } +} + struct ResolverWorkScheduler { wqtx: WorkQueueTx, } @@ -462,13 +506,22 @@ impl ResolverChannelController { impl name_resolution::ChannelController for ResolverChannelController { fn update(&mut self, update: ResolverUpdate) -> Result<(), String> { - let json_config = if let Ok(Some(service_config)) = update.service_config.as_ref() - && service_config + let json_config = if let Ok(Some(service_config)) = update.service_config.as_ref() { + if service_config .load_balancing_policy .as_ref() .is_some_and(|p| *p == LbPolicyType::RoundRobin) - { - json!([{round_robin::POLICY_NAME: {}}]) + { + json!([{round_robin::POLICY_NAME: {}}]) + } else if service_config + .load_balancing_policy + .as_ref() + .is_some_and(|p| *p == LbPolicyType::LeastRequest) + { + json!([{least_request::POLICY_NAME: {}}]) + } else { + json!([{pick_first::POLICY_NAME: {"shuffleAddressList": true, "unknown_field": false}}]) + } } else { json!([{pick_first::POLICY_NAME: {"shuffleAddressList": true, "unknown_field": false}}]) }; diff --git a/grpc/src/client/load_balancing/child_manager.rs b/grpc/src/client/load_balancing/child_manager.rs index 341cd9446..070fed7df 100644 --- a/grpc/src/client/load_balancing/child_manager.rs +++ b/grpc/src/client/load_balancing/child_manager.rs @@ -65,10 +65,17 @@ pub(crate) struct Child { pub identifier: T, pub builder: Arc, pub state: LbState, + pub subchannels: Vec, policy: Box, work_scheduler: Arc, } +impl Child { + pub fn subchannels(&self) -> impl Iterator> + '_ { + self.subchannels.iter().filter_map(|weak| weak.upgrade()) + } +} + /// A collection of data sent to a child of the ChildManager. pub(crate) struct ChildUpdate<'a, T> { /// The identifier the ChildManager should use for this child. @@ -158,6 +165,7 @@ where for csc in channel_controller.created_subchannels { self.subchannel_to_child_idx .insert((&csc).into(), child_idx); + self.children[child_idx].subchannels.push((&csc).into()); } // Update the tracked state if the child produced an update. if let Some(state) = channel_controller.picker_update { @@ -220,6 +228,7 @@ where policy: e.policy, builder: e.builder, state: e.state, + subchannels: e.subchannels, work_scheduler: e.work_scheduler, }, ) @@ -237,9 +246,11 @@ where if let Some(old_child) = old_children.remove(&k) { let old_idx = old_child.identifier; let new_child_idx = self.children.len(); + let mut subchannels = Vec::new(); for subchannel in mem::take(&mut old_child_subchannels[old_idx]) { self.subchannel_to_child_idx - .insert(subchannel, new_child_idx); + .insert(subchannel.clone(), new_child_idx); + subchannels.push(subchannel); } self.handle_to_child_idx .insert(old_child.work_scheduler.handle.clone(), new_child_idx); @@ -247,6 +258,7 @@ where builder, identifier: k.1, state: old_child.state, + subchannels, policy: old_child.policy, work_scheduler: old_child.work_scheduler, }); @@ -267,6 +279,7 @@ where builder, identifier: k.1, state: LbState::initial(), + subchannels: Vec::new(), policy, work_scheduler, }); diff --git a/grpc/src/client/load_balancing/least_request.rs b/grpc/src/client/load_balancing/least_request.rs new file mode 100644 index 000000000..a3728e590 --- /dev/null +++ b/grpc/src/client/load_balancing/least_request.rs @@ -0,0 +1,1195 @@ +/* + * + * Copyright 2026 gRPC authors. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + * + */ + +use std::collections::HashMap; +use std::collections::HashSet; +use std::fmt::Debug; +use std::sync::Arc; +use std::sync::Once; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use crate::attributes::Attributes; +use crate::client::ConnectivityState; +use crate::client::load_balancing::ChannelController; +use crate::client::load_balancing::DynLbPolicyBuilder; +use crate::client::load_balancing::FailingPicker; +use crate::client::load_balancing::GLOBAL_LB_REGISTRY; +use crate::client::load_balancing::LbPolicy; +use crate::client::load_balancing::LbPolicyBuilder; +use crate::client::load_balancing::LbPolicyOptions; +use crate::client::load_balancing::LbState; +use crate::client::load_balancing::ParsedJsonLbConfig; +use crate::client::load_balancing::Pick; +use crate::client::load_balancing::PickResult; +use crate::client::load_balancing::Picker; +use crate::client::load_balancing::QueuingPicker; +use crate::client::load_balancing::Subchannel; +use crate::client::load_balancing::SubchannelState; +use crate::client::load_balancing::WorkData; +use crate::client::load_balancing::child_manager::ChildManager; +use crate::client::load_balancing::child_manager::ChildUpdate; +use crate::client::load_balancing::pick_first; +use crate::client::load_balancing::round_robin::RoundRobinPicker; +use crate::client::load_balancing::subchannel::WeakSubchannel; +use crate::client::name_resolution::Endpoint; +use crate::client::name_resolution::ResolverUpdate; +use crate::core::RequestHeaders; +use crate::metadata::MetadataMap; + +pub(crate) static POLICY_NAME: &str = "least_request_experimental"; +static START: Once = Once::new(); + +#[derive(serde::Deserialize, Debug, Clone, PartialEq)] +#[serde(rename_all = "camelCase")] +pub(crate) struct LeastRequestLoadBalancingConfig { + #[serde(default = "default_choice_count")] + pub choice_count: u32, +} + +fn default_choice_count() -> u32 { + 2 +} + +#[derive(Debug)] +pub(crate) struct LeastRequestBuilder {} + +impl LbPolicyBuilder for LeastRequestBuilder { + type LbPolicy = LeastRequestPolicy; + + fn build(&self, options: LbPolicyOptions) -> Self::LbPolicy { + let child_manager = ChildManager::new(options.runtime, options.work_scheduler); + LeastRequestPolicy::new( + child_manager, + GLOBAL_LB_REGISTRY + .get_policy(pick_first::POLICY_NAME) + .unwrap(), + ) + } + + fn name(&self) -> &'static str { + POLICY_NAME + } + + fn parse_config( + &self, + config: &ParsedJsonLbConfig, + ) -> Result, String> { + let mut parsed: LeastRequestLoadBalancingConfig = config + .convert_to() + .map_err(|e| format!("failed to parse least_request config: {e}"))?; + + if parsed.choice_count < 2 { + return Err("choice_count must be at least 2".to_string()); + } + + parsed.choice_count = (parsed.choice_count).min(10); + Ok(Some(parsed)) + } +} + +#[derive(Debug)] +pub(crate) struct LeastRequestPolicy { + child_manager: ChildManager, + pick_first_builder: Arc, + config: Option, + subchannel_counters: HashMap>, +} + +impl LeastRequestPolicy { + fn new( + child_manager: ChildManager, + pick_first_builder: Arc, + ) -> Self { + Self { + child_manager, + pick_first_builder, + config: None, + subchannel_counters: HashMap::new(), + } + } + + // Sets the policy's state to TRANSIENT_FAILURE with a picker returning the + // error string provided, then requests re-resolution from the channel. + fn move_to_transient_failure( + &mut self, + error: String, + channel_controller: &mut dyn ChannelController, + ) { + channel_controller.update_picker(LbState { + connectivity_state: ConnectivityState::TransientFailure, + picker: Arc::new(FailingPicker { error }), + }); + channel_controller.request_resolution(); + } + + // Sends an aggregate picker based on states of children. + fn update_picker(&mut self, channel_controller: &mut dyn ChannelController, force: bool) { + if !force && !self.child_manager.child_updated() { + return; + } + let aggregate_state = self.child_manager.aggregate_states(); + + let picker: Arc = match aggregate_state { + ConnectivityState::Ready => { + let mut ready_subchannels = Vec::new(); + let mut child_weaks = HashSet::new(); + + for child in self.child_manager.children() { + for subchannel in child.subchannels() { + let weak = WeakSubchannel::new(&subchannel); + child_weaks.insert(weak.clone()); + + if child.state.connectivity_state == ConnectivityState::Ready { + let counter = self + .subchannel_counters + .entry(weak.clone()) + .or_insert_with(|| Arc::new(AtomicUsize::new(0))) + .clone(); + ready_subchannels.push(SubchannelWithCounter { + subchannel, + active_requests: counter, + }); + } + } + } + + // Clean up stale counters - retain all subchannels owned by the child manager + self.subchannel_counters + .retain(|weak, _| child_weaks.contains(weak)); + + let choice_count = self + .config + .as_ref() + .map(|cfg| cfg.choice_count as usize) + .unwrap_or_else(|| default_choice_count() as usize); + + Arc::new(LeastRequestPicker { + subchannels: ready_subchannels, + choice_count, + }) + } + ConnectivityState::Connecting + | ConnectivityState::Idle + | ConnectivityState::TransientFailure => { + let pickers: Vec> = self + .child_manager + .children() + .filter(|cs| cs.state.connectivity_state == aggregate_state) + .map(|cs| cs.state.picker.clone()) + .collect(); + + if pickers.is_empty() { + match aggregate_state { + ConnectivityState::Connecting => Arc::new(QueuingPicker {}), + ConnectivityState::Idle => Arc::new(QueuingPicker {}), + ConnectivityState::TransientFailure => Arc::new(FailingPicker { + error: "No children in TransientFailure state".to_string(), + }), + _ => unreachable!(), + } + } else if pickers.len() == 1 { + pickers[0].clone() + } else { + Arc::new(RoundRobinPicker::new(pickers)) + } + } + }; + + channel_controller.update_picker(LbState { + connectivity_state: aggregate_state, + picker, + }); + } + + // Responds to an incoming ResolverUpdate containing an Err in endpoints by + // forwarding it to all children unconditionally. Updates the picker as + // needed. + fn handle_resolver_error( + &mut self, + resolver_update: ResolverUpdate, + channel_controller: &mut dyn ChannelController, + ) -> Result<(), String> { + let err = format!( + "Received error from name resolver: {}", + resolver_update.endpoints.as_ref().unwrap_err() + ); + if self.child_manager.children().next().is_none() { + // We had no children so we must produce an erroring picker. + self.move_to_transient_failure(err.clone(), channel_controller); + return Err(err); + } + // Forward the error to each child, ignoring their responses. + let _ = self + .child_manager + .resolver_update(resolver_update, None, channel_controller); + self.update_picker(channel_controller, false); + Err(err) + } +} + +impl LbPolicy for LeastRequestPolicy { + type LbConfig = LeastRequestLoadBalancingConfig; + + fn resolver_update( + &mut self, + update: ResolverUpdate, + config: Option<&Self::LbConfig>, + channel_controller: &mut dyn ChannelController, + ) -> Result<(), String> { + let mut config_changed = false; + if let Some(cfg) = config.filter(|&cfg| self.config.as_ref() != Some(cfg)) { + self.config = Some(cfg.clone()); + config_changed = true; + } + + if update.endpoints.is_err() { + return self.handle_resolver_error(update, channel_controller); + } + + // De-duplicate endpoints using a HashSet to maintain uniqueness + let mut unique_endpoints = Vec::new(); + let mut seen = HashSet::new(); + for e in update.endpoints.unwrap() { + if !seen.contains(&e) { + seen.insert(e.clone()); + unique_endpoints.push(e); + } + } + + // Shard the update by endpoint. + let service_config = update.service_config; + let updates = unique_endpoints.into_iter().map(|e| { + let update = ResolverUpdate { + attributes: Attributes::default(), + endpoints: Ok(vec![e.clone()]), + service_config: service_config.clone(), + resolution_note: None, + }; + ChildUpdate { + child_identifier: e, + child_policy_builder: self.pick_first_builder.clone(), + child_update: Some((update, None)), + } + }); + self.child_manager + .update(updates, channel_controller) + .unwrap(); + + if self.child_manager.children().next().is_none() { + // There are no children remaining, so report this error and produce + // an erroring picker. + let err = "Received empty address list from the name resolver"; + self.move_to_transient_failure(err.into(), channel_controller); + return Err(err.into()); + } + + self.update_picker(channel_controller, config_changed); + Ok(()) + } + + fn subchannel_update( + &mut self, + subchannel: Arc, + state: &SubchannelState, + channel_controller: &mut dyn ChannelController, + ) { + self.child_manager + .subchannel_update(subchannel, state, channel_controller); + self.update_picker(channel_controller, false); + } + + fn work(&mut self, data: Option, channel_controller: &mut dyn ChannelController) { + self.child_manager.work(data, channel_controller); + self.update_picker(channel_controller, false); + } + + fn exit_idle(&mut self, channel_controller: &mut dyn ChannelController) { + self.child_manager.exit_idle(channel_controller); + self.update_picker(channel_controller, false); + } +} + +/// Register least request as a LbPolicy. +pub(crate) fn reg() { + START.call_once(|| { + GLOBAL_LB_REGISTRY.add_builder(LeastRequestBuilder {}); + }); +} + +#[derive(Clone, Debug)] +struct SubchannelWithCounter { + subchannel: Arc, + active_requests: Arc, +} + +#[derive(Debug)] +struct LeastRequestPicker { + subchannels: Vec, + choice_count: usize, +} + +impl Picker for LeastRequestPicker { + fn pick(&self, _request_headers: &RequestHeaders) -> PickResult { + let len = self.subchannels.len(); + if len == 0 { + return PickResult::Queue; + } + + let mut best_idx: Option = None; + let mut best_active_requests = usize::MAX; + + for _ in 0..self.choice_count { + let idx = if len == 1 { + 0 + } else { + rand::random_range(0..len) + }; + let active_reqs = self.subchannels[idx] + .active_requests + .load(Ordering::Relaxed); + if best_idx.is_none() || active_reqs < best_active_requests { + best_idx = Some(idx); + best_active_requests = active_reqs; + } + } + + let selected_idx = best_idx.unwrap(); + let selected = &self.subchannels[selected_idx]; + + selected.active_requests.fetch_add(1, Ordering::Relaxed); + + struct CompletionState { + counter: Arc, + active: AtomicBool, + } + + impl Drop for CompletionState { + fn drop(&mut self) { + if self.active.load(Ordering::Relaxed) { + self.counter.fetch_sub(1, Ordering::Relaxed); + } + } + } + + let state = CompletionState { + counter: selected.active_requests.clone(), + active: AtomicBool::new(true), + }; + + let on_complete = Box::new(move || { + if state.active.swap(false, Ordering::Relaxed) { + state.counter.fetch_sub(1, Ordering::Relaxed); + } + }); + + PickResult::Pick(Pick { + subchannel: selected.subchannel.clone(), + metadata: MetadataMap::new(), + on_complete: Some(on_complete), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::load_balancing::test_utils::{TestChannelController, TestWorkScheduler}; + use crate::client::name_resolution::Address; + use crate::rt::default_runtime; + use std::sync::atomic::Ordering; + use std::sync::mpsc; + + use crate::client::load_balancing::FailingPicker; + use crate::client::load_balancing::GLOBAL_LB_REGISTRY; + use crate::client::load_balancing::LbPolicy; + use crate::client::load_balancing::LbState; + use crate::client::load_balancing::Pick; + use crate::client::load_balancing::PickResult; + use crate::client::load_balancing::Picker; + use crate::client::load_balancing::QueuingPicker; + use crate::client::load_balancing::pick_first; + use crate::client::load_balancing::test_utils; + use crate::client::load_balancing::test_utils::{StubPolicyData, StubPolicyFuncs, TestEvent}; + use crate::metadata::MetadataMap; + + fn setup( + test_name: &'static str, + ) -> ( + mpsc::Receiver, + LeastRequestPolicy, + Box, + ) { + pick_first::reg(); + super::reg(); + test_utils::reg_stub_policy(test_name, create_funcs_for_leastrequest_tests()); + + let (tx_events, rx_events) = mpsc::channel(); + let work_scheduler = Arc::new(TestWorkScheduler { + tx_events: tx_events.clone(), + }); + let child_manager = ChildManager::new(default_runtime(), work_scheduler); + let tcc = Box::new(TestChannelController { tx_events }); + let child_policy_builder = GLOBAL_LB_REGISTRY.get_policy(test_name).unwrap(); + let lb_policy = LeastRequestPolicy::new(child_manager, child_policy_builder); + (rx_events, lb_policy, tcc) + } + + struct TestSubchannelList { + subchannels: Vec>, + } + + impl TestSubchannelList { + fn new(addresses: &[Address], channel_controller: &mut dyn ChannelController) -> Self { + TestSubchannelList { + subchannels: addresses + .iter() + .map(|a| channel_controller.new_subchannel(a).0) + .collect(), + } + } + + fn contains(&self, sc: &Arc) -> bool { + self.subchannels.contains(sc) + } + } + + fn create_endpoints(num_endpoints: usize, num_addresses: usize) -> Vec { + let mut endpoints = Vec::with_capacity(num_endpoints); + for i in 0..num_endpoints { + let mut addresses: Vec
= Vec::with_capacity(num_addresses); + for j in 0..num_addresses { + addresses.push(Address { + address: format!("{}.{}.{}.{}:{}", i + 1, i + 1, i + 1, i + 1, j).into(), + ..Default::default() + }); + } + endpoints.push(Endpoint { + addresses, + ..Default::default() + }) + } + endpoints + } + + fn verify_subchannel_creation( + rx_events: &mut mpsc::Receiver, + number_of_subchannels: usize, + ) -> Vec> { + let mut subchannels = Vec::new(); + for _ in 0..number_of_subchannels { + match rx_events.recv().unwrap() { + TestEvent::NewSubchannel(sc) => { + subchannels.push(sc); + } + other => panic!("unexpected event {:?}", other), + }; + } + subchannels + } + + fn verify_connecting_picker(rx_events: &mut mpsc::Receiver) -> Arc { + match rx_events.recv().unwrap() { + TestEvent::UpdatePicker(update) => { + assert_eq!(update.connectivity_state, ConnectivityState::Connecting); + let req = test_utils::new_request_headers(); + assert_eq!(update.picker.pick(&req), PickResult::Queue); + update.picker + } + other => panic!("unexpected event {:?}", other), + } + } + + fn verify_ready_picker(rx_events: &mut mpsc::Receiver) -> Arc { + match rx_events.recv().unwrap() { + TestEvent::UpdatePicker(update) => { + assert_eq!(update.connectivity_state, ConnectivityState::Ready); + update.picker.clone() + } + other => panic!("unexpected event {:?}", other), + } + } + + fn verify_transient_failure_picker( + rx_events: &mut mpsc::Receiver, + ) -> Arc { + match rx_events.recv().unwrap() { + TestEvent::UpdatePicker(update) => { + assert_eq!( + update.connectivity_state, + ConnectivityState::TransientFailure + ); + update.picker.clone() + } + other => panic!("unexpected event {:?}", other), + } + } + + struct PickFirstState { + subchannel_list: Option, + addresses: Vec
, + connectivity_state: ConnectivityState, + } + + fn create_funcs_for_leastrequest_tests() -> StubPolicyFuncs { + StubPolicyFuncs { + resolver_update: Some(Arc::new( + |data: &mut StubPolicyData, update: ResolverUpdate, _, channel_controller| { + let state = data + .test_data + .get_or_insert_with(|| { + Box::new(PickFirstState { + subchannel_list: None, + addresses: vec![], + connectivity_state: ConnectivityState::Connecting, + }) + }) + .downcast_mut::() + .unwrap(); + if let Err(error) = update.endpoints { + channel_controller.update_picker(LbState { + connectivity_state: ConnectivityState::TransientFailure, + picker: Arc::new(FailingPicker { + error: error.to_string(), + }), + }); + state.connectivity_state = ConnectivityState::TransientFailure; + channel_controller.request_resolution(); + return Ok(()); + }; + let endpoints = update.endpoints.unwrap(); + let mut addresses = Vec::new(); + for ep in endpoints { + addresses.extend(ep.addresses.clone()); + } + if addresses.is_empty() { + channel_controller.update_picker(LbState { + connectivity_state: ConnectivityState::TransientFailure, + picker: Arc::new(FailingPicker { + error: "Received empty address list from the name resolver" + .to_string(), + }), + }); + state.connectivity_state = ConnectivityState::TransientFailure; + channel_controller.request_resolution(); + return Err("Received empty address list from the name resolver".into()); + } + + if state.connectivity_state != ConnectivityState::Idle { + state.subchannel_list = + Some(TestSubchannelList::new(&addresses, channel_controller)); + } + state.addresses = addresses; + Ok(()) + }, + )), + subchannel_update: Some(Arc::new( + |data: &mut StubPolicyData, subchannel, state, channel_controller| { + let test_data = data.test_data.as_mut().unwrap(); + let test_state = test_data.downcast_mut::().unwrap(); + let scl = &mut test_state.subchannel_list.as_ref().unwrap(); + assert!( + scl.contains(&subchannel), + "subchannel_update received an update for a subchannel it does not own." + ); + test_state.connectivity_state = state.connectivity_state; + match state.connectivity_state { + ConnectivityState::Ready => { + channel_controller.update_picker(LbState { + connectivity_state: state.connectivity_state, + picker: Arc::new(OneSubchannelPicker { sc: subchannel }), + }); + } + ConnectivityState::Idle => {} + ConnectivityState::Connecting => { + channel_controller.update_picker(LbState { + connectivity_state: state.connectivity_state, + picker: Arc::new(QueuingPicker {}), + }); + } + ConnectivityState::TransientFailure => { + channel_controller.update_picker(LbState { + connectivity_state: state.connectivity_state, + picker: Arc::new(FailingPicker { + error: state + .last_connection_error + .as_ref() + .unwrap() + .to_string(), + }), + }); + } + } + }, + )), + ..Default::default() + } + } + + #[derive(Debug)] + struct OneSubchannelPicker { + sc: Arc, + } + + impl Picker for OneSubchannelPicker { + fn pick(&self, _: &RequestHeaders) -> PickResult { + PickResult::Pick(Pick { + subchannel: self.sc.clone(), + on_complete: None, + metadata: MetadataMap::new(), + }) + } + } + + #[derive(Debug, Clone)] + struct MockSubchannel { + address: Address, + } + + impl crate::client::load_balancing::subchannel::private::Sealed for MockSubchannel {} + impl crate::client::load_balancing::subchannel::DynHash for MockSubchannel { + fn dyn_hash(&self, state: &mut Box<&mut dyn std::hash::Hasher>) { + use std::hash::Hash; + self.address.hash(state); + } + } + impl crate::client::load_balancing::subchannel::DynPartialEq for MockSubchannel { + fn dyn_eq(&self, other: &&dyn std::any::Any) -> bool { + if let Some(other) = other.downcast_ref::() { + self.address == other.address + } else { + false + } + } + } + impl Subchannel for MockSubchannel { + fn address(&self) -> Address { + self.address.clone() + } + fn get_attribute_dyn(&self, _id: std::any::TypeId) -> Option<&dyn std::any::Any> { + None + } + fn connect(&self) {} + } + + #[test] + fn test_config_parsing() { + let builder = LeastRequestBuilder {}; + + // Default choice count + let default_config = ParsedJsonLbConfig::new("{}").unwrap(); + let parsed = builder.parse_config(&default_config).unwrap().unwrap(); + assert_eq!(parsed.choice_count, 2); + + // Explicit valid choice count + let valid_config = ParsedJsonLbConfig::new("{\"choiceCount\": 5}").unwrap(); + let parsed = builder.parse_config(&valid_config).unwrap().unwrap(); + assert_eq!(parsed.choice_count, 5); + + // Clamped choice count + let high_config = ParsedJsonLbConfig::new("{\"choiceCount\": 15}").unwrap(); + let parsed = builder.parse_config(&high_config).unwrap().unwrap(); + assert_eq!(parsed.choice_count, 10); + + // Rejected choice count + let low_config = ParsedJsonLbConfig::new("{\"choiceCount\": 1}").unwrap(); + assert!(builder.parse_config(&low_config).is_err()); + } + + #[test] + fn test_picker_least_request_selection() { + let sc1 = Arc::new(MockSubchannel { + address: Address { + address: "127.0.0.1:80".to_string().into(), + ..Default::default() + }, + }) as Arc; + + let sc2 = Arc::new(MockSubchannel { + address: Address { + address: "127.0.0.1:81".to_string().into(), + ..Default::default() + }, + }) as Arc; + + let count1 = Arc::new(AtomicUsize::new(5)); + let count2 = Arc::new(AtomicUsize::new(2)); + + let picker = LeastRequestPicker { + subchannels: vec![ + SubchannelWithCounter { + subchannel: sc1.clone(), + active_requests: count1.clone(), + }, + SubchannelWithCounter { + subchannel: sc2.clone(), + active_requests: count2.clone(), + }, + ], + choice_count: 2, + }; + + let mut picked_sc2 = false; + for _ in 0..20 { + let res = picker.pick(&RequestHeaders::default()); + let pick = res.unwrap_pick(); + if pick.subchannel.address().address == "127.0.0.1:81".to_string().into() { + picked_sc2 = true; + assert_eq!(count2.load(Ordering::Relaxed), 3); + let on_complete = pick.on_complete.unwrap(); + on_complete(); + assert_eq!(count2.load(Ordering::Relaxed), 2); + break; + } + } + assert!( + picked_sc2, + "sc2 (with fewer requests) was never picked in 20 attempts" + ); + } + + #[test] + fn test_picker_tie_breaking() { + let sc1 = Arc::new(MockSubchannel { + address: Address { + address: "127.0.0.1:80".to_string().into(), + ..Default::default() + }, + }) as Arc; + + let sc2 = Arc::new(MockSubchannel { + address: Address { + address: "127.0.0.1:81".to_string().into(), + ..Default::default() + }, + }) as Arc; + + let count1 = Arc::new(AtomicUsize::new(2)); + let count2 = Arc::new(AtomicUsize::new(2)); + + let picker = LeastRequestPicker { + subchannels: vec![ + SubchannelWithCounter { + subchannel: sc1.clone(), + active_requests: count1.clone(), + }, + SubchannelWithCounter { + subchannel: sc2.clone(), + active_requests: count2.clone(), + }, + ], + choice_count: 2, + }; + + let res = picker.pick(&RequestHeaders::default()); + let pick = res.unwrap_pick(); + let chosen_addr = pick.subchannel.address().address.to_string(); + assert!(chosen_addr == "127.0.0.1:80" || chosen_addr == "127.0.0.1:81"); + } + + #[test] + fn test_picker_fewer_subchannels_than_choice_count() { + let sc1 = Arc::new(MockSubchannel { + address: Address { + address: "127.0.0.1:80".to_string().into(), + ..Default::default() + }, + }) as Arc; + + let sc2 = Arc::new(MockSubchannel { + address: Address { + address: "127.0.0.1:81".to_string().into(), + ..Default::default() + }, + }) as Arc; + + let count1 = Arc::new(AtomicUsize::new(5)); + let count2 = Arc::new(AtomicUsize::new(2)); + + let picker = LeastRequestPicker { + subchannels: vec![ + SubchannelWithCounter { + subchannel: sc1.clone(), + active_requests: count1.clone(), + }, + SubchannelWithCounter { + subchannel: sc2.clone(), + active_requests: count2.clone(), + }, + ], + choice_count: 3, + }; + + let mut picked_sc2 = false; + for _ in 0..20 { + let res = picker.pick(&RequestHeaders::default()); + let pick = res.unwrap_pick(); + if pick.subchannel.address().address == "127.0.0.1:81".to_string().into() { + picked_sc2 = true; + assert_eq!(count2.load(Ordering::Relaxed), 3); + let on_complete = pick.on_complete.unwrap(); + on_complete(); + assert_eq!(count2.load(Ordering::Relaxed), 2); + break; + } + } + assert!( + picked_sc2, + "sc2 (with fewer requests) was never picked in 20 attempts" + ); + } + + #[test] + fn test_picker_cancellation_drop_guard() { + let sc = Arc::new(MockSubchannel { + address: Address { + address: "127.0.0.1:80".to_string().into(), + ..Default::default() + }, + }) as Arc; + + let count = Arc::new(AtomicUsize::new(5)); + + let picker = LeastRequestPicker { + subchannels: vec![SubchannelWithCounter { + subchannel: sc, + active_requests: count.clone(), + }], + choice_count: 1, + }; + + let res = picker.pick(&RequestHeaders::default()); + assert_eq!(count.load(Ordering::Relaxed), 6); + drop(res); + assert_eq!(count.load(Ordering::Relaxed), 5); + } + + #[test] + fn test_policy_empty_resolver_update() { + let (tx_events, _rx_events) = mpsc::channel(); + let work_scheduler = Arc::new(TestWorkScheduler { + tx_events: tx_events.clone(), + }); + let child_manager = ChildManager::new(default_runtime(), work_scheduler); + pick_first::reg(); + let pick_first_builder = GLOBAL_LB_REGISTRY + .get_policy(pick_first::POLICY_NAME) + .unwrap(); + + let mut policy = LeastRequestPolicy::new(child_manager, pick_first_builder); + let mut tcc = TestChannelController { tx_events }; + + let update = ResolverUpdate { + endpoints: Ok(vec![]), + ..Default::default() + }; + + let res = policy.resolver_update(update, None, &mut tcc); + assert!(res.is_err()); + } + + #[test] + fn test_endpoint_deduplication() { + let (mut rx_events, mut lb_policy, mut tcc) = setup("stub-test_endpoint_deduplication"); + let tcc = tcc.as_mut(); + + let endpoint = create_endpoints(1, 1)[0].clone(); + let endpoints = vec![endpoint.clone(), endpoint]; + + let update = ResolverUpdate { + endpoints: Ok(endpoints), + ..Default::default() + }; + let _ = lb_policy.resolver_update(update, None, tcc); + + let subchannels = verify_subchannel_creation(&mut rx_events, 1); + assert_eq!(subchannels.len(), 1); + } + + #[test] + fn test_connectivity_state_aggregation_all_cases() { + let (mut rx_events, mut lb_policy, mut tcc) = + setup("stub-test_connectivity_state_aggregation"); + let tcc = tcc.as_mut(); + + let endpoints = create_endpoints(2, 1); + let update = ResolverUpdate { + endpoints: Ok(endpoints), + ..Default::default() + }; + let _ = lb_policy.resolver_update(update, None, tcc); + + let subchannels = verify_subchannel_creation(&mut rx_events, 2); + + lb_policy.subchannel_update(subchannels[0].clone(), &SubchannelState::connecting(), tcc); + lb_policy.subchannel_update(subchannels[1].clone(), &SubchannelState::connecting(), tcc); + + for _ in 0..2 { + verify_connecting_picker(&mut rx_events); + } + + lb_policy.subchannel_update(subchannels[0].clone(), &SubchannelState::ready(), tcc); + verify_ready_picker(&mut rx_events); + } + + #[test] + fn test_picker_roundrobin_non_ready_states() { + let (mut rx_events, mut lb_policy, mut tcc) = + setup("stub-test_picker_roundrobin_non_ready"); + let tcc = tcc.as_mut(); + + let endpoints = create_endpoints(2, 1); + let update = ResolverUpdate { + endpoints: Ok(endpoints), + ..Default::default() + }; + let _ = lb_policy.resolver_update(update, None, tcc); + + let subchannels = verify_subchannel_creation(&mut rx_events, 2); + + lb_policy.subchannel_update( + subchannels[0].clone(), + &SubchannelState { + connectivity_state: ConnectivityState::TransientFailure, + last_connection_error: Some("connection error 0".into()), + }, + tcc, + ); + verify_connecting_picker(&mut rx_events); + + lb_policy.subchannel_update( + subchannels[1].clone(), + &SubchannelState { + connectivity_state: ConnectivityState::TransientFailure, + last_connection_error: Some("connection error 1".into()), + }, + tcc, + ); + + let picker = verify_transient_failure_picker(&mut rx_events); + let headers = RequestHeaders::default(); + + let res1 = picker.pick(&headers); + let res2 = picker.pick(&headers); + let res3 = picker.pick(&headers); + + let msg1 = match res1 { + PickResult::Fail(status) => status.message().to_string(), + other => panic!("expected Fail, got {:?}", other), + }; + let msg2 = match res2 { + PickResult::Fail(status) => status.message().to_string(), + other => panic!("expected Fail, got {:?}", other), + }; + let msg3 = match res3 { + PickResult::Fail(status) => status.message().to_string(), + other => panic!("expected Fail, got {:?}", other), + }; + + assert_ne!(msg1, msg2); + assert_eq!(msg1, msg3); + assert!(msg1.contains("connection error 0") || msg1.contains("connection error 1")); + assert!(msg2.contains("connection error 0") || msg2.contains("connection error 1")); + } + + #[test] + fn stress_test_p2c_selection_uniformity() { + // Create 5 subchannels + let mut subchannels = Vec::new(); + let mut counters = Vec::new(); + for i in 0..5 { + let sc = Arc::new(MockSubchannel { + address: Address { + address: format!("127.0.0.1:{}", 80 + i).into(), + ..Default::default() + }, + }) as Arc; + let count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + counters.push(count.clone()); + subchannels.push(SubchannelWithCounter { + subchannel: sc, + active_requests: count, + }); + } + + let picker = LeastRequestPicker { + subchannels, + choice_count: 2, + }; + + let mut selection_counts = vec![0; 5]; + for _ in 0..10000 { + // Reset counters to ensure they are perfectly equal + for count in &counters { + count.store(0, Ordering::Relaxed); + } + + let res = picker.pick(&RequestHeaders::default()); + let pick = res.unwrap_pick(); + let port = pick + .subchannel + .address() + .address + .split(':') + .next_back() + .unwrap() + .parse::() + .unwrap(); + let idx = port - 80; + selection_counts[idx] += 1; + } + + println!( + "Selection counts for 5 subchannels over 10000 runs: {:?}", + selection_counts + ); + + // Expected count = 2000. Standard deviation = 40. + // Assert that all counts are within 4 std dev of expected (1840 to 2160). + for (i, count) in selection_counts.iter().enumerate() { + assert!( + *count >= 1840 && *count <= 2160, + "Subchannel {} selection count {} was outside the highly-probable range [1840, 2160]", + i, + count + ); + } + } + + #[test] + fn stress_test_completion_state_thread_safety() { + use std::sync::Barrier; + use std::thread; + + let sc = Arc::new(MockSubchannel { + address: Address { + address: "127.0.0.1:80".to_string().into(), + ..Default::default() + }, + }) as Arc; + + let counter = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let picker = Arc::new(LeastRequestPicker { + subchannels: vec![SubchannelWithCounter { + subchannel: sc, + active_requests: counter.clone(), + }], + choice_count: 1, + }); + + const NUM_THREADS: usize = 20; + const ITERATIONS: usize = 500; + let barrier = Arc::new(Barrier::new(NUM_THREADS)); + let mut handles = Vec::new(); + + for _ in 0..NUM_THREADS { + let picker = picker.clone(); + let barrier = barrier.clone(); + handles.push(thread::spawn(move || { + barrier.wait(); + for i in 0..ITERATIONS { + let res = picker.pick(&RequestHeaders::default()); + let pick = res.unwrap_pick(); + + if (i + rand::random_range(0..100)) % 2 == 0 { + if let Some(on_complete) = pick.on_complete { + on_complete(); + } + } else { + // Just let it drop naturally + } + } + })); + } + + for handle in handles { + handle.join().unwrap(); + } + + // After all threads completed, the counter MUST be exactly 0! + assert_eq!( + counter.load(Ordering::SeqCst), + 0, + "Counter leaked or underflowed!" + ); + } + + #[test] + fn stress_test_dynamic_config_invalidation() { + let (mut rx_events, mut lb_policy, mut tcc) = + setup("stub-test_dynamic_config_invalidation"); + let tcc = tcc.as_mut(); + + let endpoints = create_endpoints(3, 1); + let update = ResolverUpdate { + endpoints: Ok(endpoints), + ..Default::default() + }; + + let config2 = LeastRequestLoadBalancingConfig { choice_count: 2 }; + let _ = lb_policy.resolver_update(update.clone(), Some(&config2), tcc); + + let subchannels = verify_subchannel_creation(&mut rx_events, 3); + verify_connecting_picker(&mut rx_events); // Consume the initial Connecting picker from resolver_update + + lb_policy.subchannel_update(subchannels[0].clone(), &SubchannelState::ready(), tcc); + let _ = verify_ready_picker(&mut rx_events); // Consume the first Ready picker + + lb_policy.subchannel_update(subchannels[1].clone(), &SubchannelState::ready(), tcc); + let _ = verify_ready_picker(&mut rx_events); // Consume the second Ready picker + + lb_policy.subchannel_update(subchannels[2].clone(), &SubchannelState::ready(), tcc); + let picker = verify_ready_picker(&mut rx_events); // Consume the third Ready picker + + let dbg_format = format!("{:?}", picker); + assert!( + dbg_format.contains("choice_count: 2"), + "Expected choice_count to be 2, got {:?}", + dbg_format + ); + + let config3 = LeastRequestLoadBalancingConfig { choice_count: 3 }; + let _ = lb_policy.resolver_update(update.clone(), Some(&config3), tcc); + + let _ = verify_subchannel_creation(&mut rx_events, 3); // Consume the new subchannel events from stub policy + let picker = verify_ready_picker(&mut rx_events); + let dbg_format = format!("{:?}", picker); + assert!( + dbg_format.contains("choice_count: 3"), + "Expected choice_count to be 3, got {:?}", + dbg_format + ); + + let _ = lb_policy.resolver_update(update, Some(&config2), tcc); + let _ = verify_subchannel_creation(&mut rx_events, 3); // Consume the new subchannel events from stub policy + let picker = verify_ready_picker(&mut rx_events); + let dbg_format = format!("{:?}", picker); + assert!( + dbg_format.contains("choice_count: 2"), + "Expected choice_count to be 2, got {:?}", + dbg_format + ); + } +} diff --git a/grpc/src/client/load_balancing/mod.rs b/grpc/src/client/load_balancing/mod.rs index 1d2dba2f9..c6cddbd09 100644 --- a/grpc/src/client/load_balancing/mod.rs +++ b/grpc/src/client/load_balancing/mod.rs @@ -43,6 +43,7 @@ use crate::rt::GrpcRuntime; pub(crate) mod child_manager; pub(crate) mod graceful_switch; pub(crate) mod lazy; +pub(crate) mod least_request; pub(crate) mod pick_first; pub(crate) mod round_robin; pub(crate) mod subchannel; diff --git a/grpc/src/client/load_balancing/round_robin.rs b/grpc/src/client/load_balancing/round_robin.rs index 505fb9a87..f4f87e669 100644 --- a/grpc/src/client/load_balancing/round_robin.rs +++ b/grpc/src/client/load_balancing/round_robin.rs @@ -228,13 +228,13 @@ pub(crate) fn reg() { } #[derive(Debug)] -struct RoundRobinPicker { +pub(crate) struct RoundRobinPicker { pickers: Vec>, next: AtomicUsize, } impl RoundRobinPicker { - fn new(pickers: Vec>) -> Self { + pub(crate) fn new(pickers: Vec>) -> Self { let random_index: usize = rand::random_range(..pickers.len()); Self { pickers, diff --git a/grpc/src/client/load_balancing/subchannel.rs b/grpc/src/client/load_balancing/subchannel.rs index dab785474..b78ac4e0d 100644 --- a/grpc/src/client/load_balancing/subchannel.rs +++ b/grpc/src/client/load_balancing/subchannel.rs @@ -194,6 +194,12 @@ impl Display for dyn Subchannel { #[derive(Debug)] pub(crate) struct WeakSubchannel(Weak); +impl Clone for WeakSubchannel { + fn clone(&self) -> Self { + WeakSubchannel(self.0.clone()) + } +} + impl From<&Arc> for WeakSubchannel { fn from(subchannel: &Arc) -> Self { WeakSubchannel(Arc::downgrade(subchannel)) diff --git a/grpc/src/client/service_config.rs b/grpc/src/client/service_config.rs index dc7464920..2ecbe0abc 100644 --- a/grpc/src/client/service_config.rs +++ b/grpc/src/client/service_config.rs @@ -35,4 +35,6 @@ pub enum LbPolicyType { #[default] PickFirst, RoundRobin, + // TODO: remove LeastRequest. It's here because it won't build without it. + LeastRequest, } diff --git a/grpc/src/inmemory/mod.rs b/grpc/src/inmemory/mod.rs index 4bba7357f..15e9f7f4a 100644 --- a/grpc/src/inmemory/mod.rs +++ b/grpc/src/inmemory/mod.rs @@ -390,9 +390,23 @@ pub struct InMemoryResolverBuilder {} impl ResolverBuilder for InMemoryResolverBuilder { fn build(&self, target: &Target, options: ResolverOptions) -> Box { let path = target.path().strip_prefix('/').unwrap_or(target.path()); - let ids: Vec = path.split(',').map(|s| s.to_string()).collect(); + // TODO: The LB policy *should* get set by specifying a default service config, + // not through the in memory resolver. We don't have that yet, though, so this is + // fine for now since this module is unexported and unused except in tests. + let (lb_policy, rest) = if let Some(stripped) = path.strip_prefix("leastrequest/") { + ( + crate::client::service_config::LbPolicyType::LeastRequest, + stripped, + ) + } else { + ( + crate::client::service_config::LbPolicyType::RoundRobin, + path, + ) + }; + let ids: Vec = rest.split(',').map(|s| s.to_string()).collect(); options.work_scheduler.schedule_work(); - Box::new(InMemoryResolver { ids }) + Box::new(InMemoryResolver { ids, lb_policy }) } fn scheme(&self) -> &str { @@ -406,6 +420,7 @@ impl ResolverBuilder for InMemoryResolverBuilder { struct InMemoryResolver { ids: Vec, + lb_policy: crate::client::service_config::LbPolicyType, } impl Resolver for InMemoryResolver { @@ -428,9 +443,7 @@ impl Resolver for InMemoryResolver { let _ = channel_controller.update(ResolverUpdate { endpoints: Ok(endpoints), service_config: Ok(Some(ServiceConfig { - load_balancing_policy: Some( - crate::client::service_config::LbPolicyType::RoundRobin, - ), + load_balancing_policy: Some(self.lb_policy.clone()), })), ..Default::default() }); @@ -491,4 +504,96 @@ mod tests { _ => panic!("expected trailers with error, got {:?}", item), } } + + #[tokio::test] + async fn test_in_memory_least_request_load_balancing() { + reg(); // Register transport and resolver + crate::client::load_balancing::least_request::reg(); // Register least request policy + + let backend1 = InMemoryListener::new(); + let backend2 = InMemoryListener::new(); + + let b1_id = backend1.id(); + let b2_id = backend2.id(); + + let (tx1, mut rx1) = tokio::sync::mpsc::channel(10); + let (tx2, mut rx2) = tokio::sync::mpsc::channel(10); + + let b1 = backend1.clone(); + let handle1 = tokio::spawn(async move { + if let Some(call) = b1.accept().await { + tx1.send(()).await.unwrap(); + // Keep the connection/stream open until the test is done + tokio::time::sleep(std::time::Duration::from_secs(5)).await; + drop(call); + } + }); + + let b2 = backend2.clone(); + let handle2 = tokio::spawn(async move { + if let Some(call) = b2.accept().await { + tx2.send(()).await.unwrap(); + // Keep the connection/stream open until the test is done + tokio::time::sleep(std::time::Duration::from_secs(5)).await; + drop(call); + } + }); + + // Construct target URI using the newly supported prefix + let target = format!("inmemory:///leastrequest/{},{}", b1_id, b2_id); + + let channel = crate::client::Channel::new( + &target, + crate::credentials::LocalChannelCredentials::new_arc(), + crate::client::ChannelOptions::default(), + ); + + // Make first invoke to establish an active request on whichever backend is ready first + let (_send1, _recv1) = channel + .invoke( + crate::core::RequestHeaders::new().with_method_name("/test/method"), + crate::client::CallOptions::default(), + ) + .await; + + // Loop and retry the second invoke until the other backend is ready and picked. + // Since the first backend is kept active (1 active request), the Least Request policy + // will immediately pick Backend 2 once it becomes ready (0 active requests). + let mut b1_called = false; + let mut b2_called = false; + + for _ in 0..50 { + // Check which backends have been called so far + if rx1.try_recv().is_ok() { + b1_called = true; + } + if rx2.try_recv().is_ok() { + b2_called = true; + } + + if b1_called && b2_called { + break; + } + + // Make a short-lived invoke. If it goes to the same backend, it will be dropped immediately, + // returning its active request count back to 1. + let invoke_future = channel.invoke( + crate::core::RequestHeaders::new().with_method_name("/test/method"), + crate::client::CallOptions::default(), + ); + if let Ok((_send_tmp, _recv_tmp)) = + tokio::time::timeout(std::time::Duration::from_millis(100), invoke_future).await + { + // Successfully made call + } + + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + + assert!(b1_called, "Backend 1 was not called"); + assert!(b2_called, "Backend 2 was not called"); + + handle1.abort(); + handle2.abort(); + } }