diff --git a/grpc/src/client/load_balancing/pick_first.rs b/grpc/src/client/load_balancing/pick_first.rs index 383d971bb..468a3084f 100644 --- a/grpc/src/client/load_balancing/pick_first.rs +++ b/grpc/src/client/load_balancing/pick_first.rs @@ -22,7 +22,6 @@ * */ -use std::collections::HashMap; use std::collections::HashSet; use std::fmt::Debug; use std::sync::Arc; @@ -75,15 +74,10 @@ impl LbPolicyBuilder for PickFirstBuilder { PickFirstPolicy { work_scheduler: options.work_scheduler, runtime: options.runtime, - connectivity_state: ConnectivityState::Idle, - subchannels: Vec::default(), - subchannel_states: HashMap::default(), - selected: None, - frontier_index: 0, - last_connection_error: None, shuffler: build_shuffler(), - timer: None, - steady_state: None, + state: PickFirstState::Idle(IdleState { + addresses: Vec::new(), + }), } } @@ -98,233 +92,341 @@ impl LbPolicyBuilder for PickFirstBuilder { } pub(crate) fn reg() { - super::GLOBAL_LB_REGISTRY.add_builder(PickFirstBuilder {}) + super::GLOBAL_LB_REGISTRY.add_builder(PickFirstBuilder {}); } +/// A load balancing policy that receives endpoints from the name resolver and +/// connects to the first available backend using the [Happy Eyeballs](https://datatracker.ietf.org/doc/html/rfc8305) +/// connection algorithm. pub struct PickFirstPolicy { work_scheduler: Arc, runtime: GrpcRuntime, - connectivity_state: ConnectivityState, - - // Subchannel information. - subchannels: Vec>, - subchannel_states: HashMap, // Cached states for all subchannels by address. - selected: Option>, - frontier_index: usize, - - // Detailed error tracking. - last_connection_error: Option, - - // Injectable shuffler for deterministic testing. shuffler: Arc, - - // Timer state tracks when the last connect attempt was started. - timer: Option, - - // Steady state tracking for continuous retries after pass exhaustion. - // TODO: should steady_state be a "mode" selector enum, e.g.: - // - FirstPass - holds the timer above - // - SteadyState - existing info - // - Ready - holds selected subchannel - // - Error - for zero addresses; holds no data - steady_state: Option, + state: PickFirstState, } impl Debug for PickFirstPolicy { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PickFirstPolicy") - .field("subchannels", &self.subchannels) - .field("selected", &self.selected) - .field("frontier_index", &self.frontier_index) - .field("connectivity_state", &self.connectivity_state) - .field("last_connection_error", &self.last_connection_error) - .finish() + .field("state", &self.state) + .finish_non_exhaustive() } } impl PickFirstPolicy { - fn rebuild_subchannels( - &mut self, - new_addresses: Vec
, - channel_controller: &mut dyn ChannelController, - ) -> Option> { - // Map existing subchannels by address. - let mut existing_subchannels: HashMap> = self - .subchannels - .drain(..) - .map(|sc| (sc.address(), sc)) - .collect(); - - let mut new_subchannels = Vec::with_capacity(new_addresses.len()); - let mut new_states = HashMap::with_capacity(new_addresses.len()); - let mut ready_subchannel = None; - - for addr in new_addresses { - let (sc, state) = if let Some(sc) = existing_subchannels.remove(&addr) { - let state = self.subchannel_states.get(&addr).unwrap().clone(); - (sc, state) - } else { - // Get a new subchannel handle from the controller if we don't - // have an existing one. - channel_controller.new_subchannel(&addr) - }; + fn compile_address( + &self, + mut endpoints: Vec, + config: Option<&PickFirstConfig>, + ) -> Result, String> { + if config.is_some_and(|c| c.shuffle_address_list) { + (self.shuffler)(&mut endpoints); + } - // Track the best candidate for immediate activation: - // 1. Priority: The currently selected subchannel if still READY. - // 2. Fallback: The first generic READY subchannel encountered. - if state.connectivity_state == ConnectivityState::Ready { - if self.subchannel_is_selected(&sc) { - // Sticky channel wins immediately and overrides any - // fallback candidates. - ready_subchannel = Some(sc.clone()); - } else if ready_subchannel.is_none() { - // Capture fallback candidate, but does not overwrite if a - //sticky channel was already found. - ready_subchannel = Some(sc.clone()); + let mut seen = HashSet::new(); + let mut ipv6 = Vec::new(); + let mut ipv4 = Vec::new(); + let mut unknown = Vec::new(); + + for ep in endpoints { + for addr in ep.addresses { + if seen.insert(addr.clone()) { + if addr.network_type == crate::client::name_resolution::TCP_IP_NETWORK_TYPE { + if addr.address.contains(':') { + ipv6.push(addr); + } else { + ipv4.push(addr); + } + } else { + unknown.push(addr); + } } } + } - new_subchannels.push(sc); - new_states.insert(addr, state); + let mut interleaved = Vec::with_capacity(ipv6.len() + ipv4.len() + unknown.len()); + let mut v6_iter = ipv6.into_iter(); + let mut v4_iter = ipv4.into_iter(); + let mut unknown_iter = unknown.into_iter(); + + loop { + let mut more = false; + + if let Some(v6) = v6_iter.next() { + interleaved.push(v6); + more = true; + } + if let Some(v4) = v4_iter.next() { + interleaved.push(v4); + more = true; + } + if let Some(unknown) = unknown_iter.next() { + interleaved.push(unknown); + more = true; + } + + if !more { + break; + } } - self.subchannels = new_subchannels; - self.subchannel_states = new_states; - ready_subchannel - } + if interleaved.is_empty() { + return Err("empty address list".to_string()); + } - /// Call when the selected subchannel loses connection. - // This causes the LB to go IDLE. - fn subchannel_drop(&mut self, channel_controller: &mut dyn ChannelController) { - self.selected = None; - self.update_picker( - ConnectivityState::Idle, - Arc::new(IdlePicker::new(self.work_scheduler.clone())), - channel_controller, - ); + Ok(interleaved) } +} + +// The `PickFirstPolicy` is structured as a discrete finite state machine +// (`PickFirstState`): +// - `Idle`: Initial state or post-disconnect state waiting for traffic to trigger +// resolution. +// - `FirstPass`: Happy Eyeballs connection pass staggering connection attempts +// across resolved addresses. +// - `SteadyState`: All addresses failed; holds sticky TRANSIENT_FAILURE and +// retries connections as backoffs expire. +// - `Ready`: Successfully connected to a subchannel, routing all picks to it. +impl LbPolicy for PickFirstPolicy { + type LbConfig = PickFirstConfig; - fn update_picker( + fn resolver_update( &mut self, - connectivity_state: ConnectivityState, - picker: Arc, - channel_controller: &mut dyn ChannelController, - ) { - if self.connectivity_state == connectivity_state - && connectivity_state == ConnectivityState::Connecting - { - // Prevent redundant connecting updates. - // TODO: prevent redundant IDLE updates? - return; + update: ResolverUpdate, + config: Option<&Self::LbConfig>, + controller: &mut dyn ChannelController, + ) -> Result<(), String> { + let mut ctx = PickFirstContext { + runtime: &self.runtime, + work_scheduler: &self.work_scheduler, + controller, + }; + + match update.endpoints { + Ok(endpoints) => { + let addresses = match self.compile_address(endpoints, config) { + Ok(addrs) => addrs, + Err(e) => { + self.state = PickFirstState::Idle(IdleState { + addresses: Vec::new(), + }); + ctx.set_failing_picker(&e); + ctx.controller.request_resolution(); + return Err(e); + } + }; + let current = std::mem::take(&mut self.state); + self.state = current.resolver_update(&mut ctx, addresses); + Ok(()) + } + Err(error) => { + let has_addresses = match &self.state { + PickFirstState::Idle(s) => !s.addresses.is_empty(), + PickFirstState::FirstPass(s) => !s.addresses.is_empty(), + PickFirstState::SteadyState(s) => !s.addresses.is_empty(), + PickFirstState::Ready(s) => !s.addresses.is_empty(), + }; + let is_tf = matches!(self.state, PickFirstState::SteadyState(_)); + + if !has_addresses || is_tf { + ctx.set_failing_picker(&error); + ctx.controller.request_resolution(); + return Err(error); + } + Ok(()) + } } - self.connectivity_state = connectivity_state; - channel_controller.update_picker(LbState { - connectivity_state, - picker, - }); } - fn subchannel_activate( + fn subchannel_update( &mut self, subchannel: Arc, - channel_controller: &mut dyn ChannelController, + state: &SubchannelState, + controller: &mut dyn ChannelController, ) { - if self.subchannel_is_selected(&subchannel) { - // Already selected; skip activation. - return; - } - self.selected = Some(subchannel.clone()); - self.subchannels = vec![subchannel.clone()]; // Keep only the winner. - self.steady_state = None; // Reset mode to First Pass. - self.timer = None; // Stop the happy eyeballs timer. - - self.update_picker( - ConnectivityState::Ready, - Arc::new(OneSubchannelPicker { sc: subchannel }), - channel_controller, - ); + let mut ctx = PickFirstContext { + runtime: &self.runtime, + work_scheduler: &self.work_scheduler, + controller, + }; + let current = std::mem::take(&mut self.state); + self.state = current.subchannel_update(&mut ctx, subchannel, state); } - /// Starts a connection pass through the address list. - // This clears the selected subchannel. - fn start_connection_pass(&mut self, channel_controller: &mut dyn ChannelController) { - self.selected = None; + fn work(&mut self, data: Option, controller: &mut dyn ChannelController) { + debug_assert!(data.is_none(), "expected no data but got {data:?}"); + let mut ctx = PickFirstContext { + runtime: &self.runtime, + work_scheduler: &self.work_scheduler, + controller, + }; + let current = std::mem::take(&mut self.state); + self.state = current.work(&mut ctx); + } - // If there is a viable subchannel at the frontier, connect to it and - // update picker to CONNECTING. - if let Some(sc) = self.advance_frontier(true) { - self.trigger_subchannel_connection(sc, channel_controller); + fn exit_idle(&mut self, controller: &mut dyn ChannelController) { + let mut ctx = PickFirstContext { + runtime: &self.runtime, + work_scheduler: &self.work_scheduler, + controller, + }; + let current = std::mem::take(&mut self.state); + self.state = current.exit_idle(&mut ctx); + } +} - // TODO: avoid this update if we are in TF (i.e. sticky TF)? - self.update_picker( - ConnectivityState::Connecting, - Arc::new(QueuingPicker {}), - channel_controller, - ); - } else { - // Otherwise all addresses are in transient failure: update picker. - // TODO: set the last connection error? Is it correct to do so, - // e.g. if the address it came from was removed from the address - // list and a name resolver update triggered this connection pass? - _ = self.set_transient_failure(channel_controller, None); - } +struct PickFirstContext<'a> { + runtime: &'a GrpcRuntime, + work_scheduler: &'a Arc, + controller: &'a mut dyn ChannelController, +} + +impl PickFirstContext<'_> { + fn set_failing_picker(&mut self, error: &str) { + self.controller.update_picker(LbState { + connectivity_state: ConnectivityState::TransientFailure, + picker: Arc::new(FailingPicker { + error: error.to_string(), + }), + }); } +} - // Book-keeping for tracking progress on the first pass through the address - // list. Assumes the subchannel is in a non-READY state. - // If the failure is from the subchannel at the frontier, advances the - // frontier and triggers a connection on the next subchannel. - fn update_first_pass( - &mut self, - subchannel: Arc, - state: &SubchannelState, - channel_controller: &mut dyn ChannelController, - ) { - // Advance frontier if this failure is from the active frontier subchannel. - if let Some(attempting) = self.subchannels.get(self.frontier_index) - && attempting.address() == subchannel.address() - && state.connectivity_state == ConnectivityState::TransientFailure - && let Some(next_sc) = self.advance_frontier(false) - { - self.trigger_subchannel_connection(next_sc, channel_controller); +// State node: Idle. +#[derive(Debug, Default)] +struct IdleState { + addresses: Vec
, +} + +impl IdleState { + fn enter(ctx: &mut PickFirstContext<'_>, addresses: Vec
) -> PickFirstState { + ctx.controller.update_picker(LbState { + connectivity_state: ConnectivityState::Idle, + picker: Arc::new(IdlePicker::new(ctx.work_scheduler.clone())), + }); + PickFirstState::Idle(IdleState { addresses }) + } + + #[allow(clippy::unused_self)] + fn resolver_update( + self, + ctx: &mut PickFirstContext<'_>, + addresses: Vec
, + ) -> PickFirstState { + FirstPassState::fresh_enter(ctx, addresses) + } + + #[allow(clippy::unused_self)] + fn subchannel_update( + self, + _ctx: &mut PickFirstContext<'_>, + _subchannel: Arc, + _state: &SubchannelState, + ) -> PickFirstState { + PickFirstState::Idle(self) + } + + fn work(self, ctx: &mut PickFirstContext<'_>) -> PickFirstState { + self.exit_idle(ctx) + } + + fn exit_idle(self, ctx: &mut PickFirstContext<'_>) -> PickFirstState { + if self.addresses.is_empty() { + ctx.controller.request_resolution(); + return PickFirstState::Idle(self); } + FirstPassState::fresh_enter(ctx, self.addresses) + } +} - // Check if First Pass is fully exhausted (frontier exhausted AND zero connecting). - if self.frontier_index >= self.subchannels.len() { - let any_connecting = self.subchannels.iter().any(|sc| { - self.subchannel_states - .get(&sc.address()) - .is_some_and(|s| s.connectivity_state == ConnectivityState::Connecting) - }); +// An entry associating an address, its created subchannel handle, and its cached +// connectivity state. +#[derive(Debug, Clone)] +struct SubchannelEntry { + address: Address, + subchannel: Arc, + state: SubchannelState, +} - if !any_connecting { - // Nothing currently connecting; first pass complete. Enter - // steady state. - self.timer = None; - let error = self.last_connection_error.clone(); - _ = self.set_transient_failure(channel_controller, error); - self.steady_state = Some(SteadyState::new(self.subchannels.len())); - - // Trigger connection attempts on any subchannels that - // transitioned to IDLE during the first pass, ensuring they - // don't get stuck. - for sc in &self.subchannels { - let is_idle = self - .subchannel_states - .get(&sc.address()) - .is_some_and(|s| s.connectivity_state == ConnectivityState::Idle); - if is_idle { - sc.connect(); - } +// State node: FirstPass (Happy Eyeballs). +struct FirstPassState { + addresses: Vec
, + subchannels: Vec, + frontier_index: usize, + timer: Timer, + last_connection_error: Option, +} + +impl FirstPassState { + fn fresh_enter(ctx: &mut PickFirstContext<'_>, addresses: Vec
) -> PickFirstState { + Self::enter(ctx, addresses, Vec::new()) + } + + fn enter( + ctx: &mut PickFirstContext<'_>, + addresses: Vec
, + mut existing: Vec, + ) -> PickFirstState { + let mut new_subchannels = Vec::with_capacity(addresses.len()); + + for addr in &addresses { + let entry = if let Some(pos) = existing.iter().position(|e| &e.address == addr) { + existing.swap_remove(pos) + } else { + let (subchannel, state) = ctx.controller.new_subchannel(addr); + SubchannelEntry { + address: addr.clone(), + subchannel, + state, } + }; + + if entry.state.connectivity_state == ConnectivityState::Ready { + return ReadyState::enter(ctx, addresses, entry.subchannel); } + + new_subchannels.push(entry); + } + + let mut first_pass = FirstPassState { + addresses, + subchannels: new_subchannels, + frontier_index: 0, + timer: Timer::start(ctx.runtime.clone(), ctx.work_scheduler.clone()), + last_connection_error: None, + }; + + if let Some(sc) = first_pass.advance_frontier(true) { + first_pass.trigger_connection(ctx, &sc); + + ctx.controller.update_picker(LbState { + connectivity_state: ConnectivityState::Connecting, + picker: Arc::new(QueuingPicker {}), + }); + + PickFirstState::FirstPass(first_pass) + } else { + let err = first_pass + .last_connection_error + .clone() + .unwrap_or_else(|| "all addresses in transient failure".to_string()); + + SteadyState::enter(ctx, first_pass.addresses, first_pass.subchannels, err) } } - /// Advances the frontier to the next non-TransientFailure subchannel and - /// returns it. If `reset` is true, starts the scan from index 0. - // The frontier is the latest index in which connectivity has been attempted. - // Returns a cloned Subchannel to avoid borrowing from self. + fn trigger_connection(&mut self, ctx: &mut PickFirstContext<'_>, sc: &Arc) { + let addr = sc.address(); + if let Some(entry) = self.subchannels.iter_mut().find(|e| e.address == addr) { + entry.state = SubchannelState { + connectivity_state: ConnectivityState::Connecting, + last_connection_error: None, + }; + } + sc.connect(); + self.timer = Timer::start(ctx.runtime.clone(), ctx.work_scheduler.clone()); + } + fn advance_frontier(&mut self, reset: bool) -> Option> { if reset { self.frontier_index = 0; @@ -333,287 +435,351 @@ impl PickFirstPolicy { } while self.frontier_index < self.subchannels.len() { - let sc = &self.subchannels[self.frontier_index]; - let addr = sc.address(); - let state = self - .subchannel_states - .get(&addr) - .map(|s| s.connectivity_state) - .expect("Expected non-None subchannel state"); - - match state { - // Push the frontier if sc is in TransientFailure + let entry = &self.subchannels[self.frontier_index]; + match entry.state.connectivity_state { ConnectivityState::TransientFailure => self.frontier_index += 1, - // Otherwise return the subchannel. - _ => return Some(sc.clone()), + _ => return Some(entry.subchannel.clone()), } } None } - /// Returns true if the given subchannel matches the currently selected - /// active subchannel. - fn subchannel_is_selected(&self, subchannel: &Arc) -> bool { - self.selected - .as_ref() - .is_some_and(|sel| sel.address() == subchannel.address()) + fn resolver_update( + self, + ctx: &mut PickFirstContext<'_>, + addresses: Vec
, + ) -> PickFirstState { + FirstPassState::enter(ctx, addresses, self.subchannels) } - /// Returns true if the subchannel's address is present in the most recently - /// received address list. - // This compares against the current list of subchannels the LB is - // attempting to connect to. To see if the LB already connected to the - // channel, use 'subchannel_is_selected'. - fn subchannel_is_current(&self, subchannel: &Arc) -> bool { - self.subchannels - .iter() - .any(|sc| sc.address() == subchannel.address()) - } + fn subchannel_update( + mut self, + ctx: &mut PickFirstContext<'_>, + subchannel: Arc, + state: &SubchannelState, + ) -> PickFirstState { + let addr = subchannel.address(); - /// Triggers a connection on the subchannel, and starts the 250ms timer. - /// If no connection succeeds before the timer expires, the frontier will - /// advance to the next subchannel. - fn trigger_subchannel_connection( - &mut self, - sc: Arc, - channel_controller: &mut dyn ChannelController, - ) { - self.subchannel_states.insert( - sc.address(), - SubchannelState { - connectivity_state: ConnectivityState::Connecting, - last_connection_error: None, - }, - ); - sc.connect(); + let Some(entry) = self.subchannels.iter_mut().find(|e| e.address == addr) else { + return PickFirstState::FirstPass(self); + }; - // Start happy eyeballs timer; replacing any pre-existing timer. - self.timer = Some(Timer::start( - self.runtime.clone(), - self.work_scheduler.clone(), - )); - } + entry.state = state.clone(); - // Converts the update endpoints to an address list. - // Shuffles endpoints (if enabled) before flattening and de-duplication. - fn compile_address( - &mut self, - mut endpoints: Vec, - config: Option<&PickFirstConfig>, - channel_controller: &mut dyn ChannelController, - ) -> Vec
{ - // Shuffle endpoints if enabled. - if config.is_some_and(|c| c.shuffle_address_list) { - (self.shuffler)(&mut endpoints); + if state.connectivity_state == ConnectivityState::Ready { + return ReadyState::enter(ctx, self.addresses, subchannel); } - // Flatten and de-duplicate unique addresses in order. - let mut seen = HashSet::new(); - let unique_addresses: Vec
= endpoints - .into_iter() - .flat_map(|ep| ep.addresses) - .filter(|addr| seen.insert(addr.clone())) - .collect(); - - // Partition out all 'unknown' non-TCP addresses. - // This is to remain consistent with similar behavior in C++ and Java. - let (tcp_addresses, unknown): (Vec
, Vec
) = - unique_addresses.into_iter().partition(|addr| { - addr.network_type == crate::client::name_resolution::TCP_IP_NETWORK_TYPE - }); - - // Partition by family (Basic IPv6 detection via colon). - let (ipv6, ipv4): (Vec
, Vec
) = tcp_addresses - .into_iter() - .partition(|addr| addr.address.contains(':')); + if state.connectivity_state == ConnectivityState::TransientFailure { + if let Some(err) = &state.last_connection_error { + self.last_connection_error = Some(err.clone()); + } - // Interleave the two lists so ipv6 and ipv4 addresses are alternated. - let mut interleaved = Vec::with_capacity(ipv6.len() + ipv4.len() + unknown.len()); - let mut v6_iter = ipv6.into_iter(); - let mut v4_iter = ipv4.into_iter(); - let mut unknown_iter = unknown.into_iter(); + if let Some(attempting) = self.subchannels.get(self.frontier_index) + && attempting.address == addr + && let Some(next_sc) = self.advance_frontier(false) + { + self.trigger_connection(ctx, &next_sc); + } + } - loop { - let mut more = false; + if self.frontier_index >= self.subchannels.len() { + let any_connecting = self + .subchannels + .iter() + .any(|e| e.state.connectivity_state == ConnectivityState::Connecting); - if let Some(v6) = v6_iter.next() { - interleaved.push(v6); - more = true; - } - if let Some(v4) = v4_iter.next() { - interleaved.push(v4); - more = true; - } - if let Some(unknown) = unknown_iter.next() { - interleaved.push(unknown); - more = true; - } + if !any_connecting { + let err = self + .last_connection_error + .clone() + .unwrap_or_else(|| "all addresses in transient failure".to_string()); - if !more { - break; + return SteadyState::enter(ctx, self.addresses, self.subchannels, err); } } - interleaved + + PickFirstState::FirstPass(self) } - // Sets LB state to TRANSIENT_FAILURE and updates picker with error. - // Triggers a re-resolution request. - // TODO: make error mandatory. - fn set_transient_failure( - &mut self, - channel_controller: &mut dyn ChannelController, - error: Option, - ) -> Result<(), String> { - // Replace the last connection error if we have a new one. - if let Some(e) = error { - self.last_connection_error = Some(e); + fn work(mut self, ctx: &mut PickFirstContext<'_>) -> PickFirstState { + if self.timer.expired() + && let Some(next_sc) = self.advance_frontier(false) + { + self.trigger_connection(ctx, &next_sc); } - let err = self - .last_connection_error - .clone() - .expect("no last connection error set"); - self.update_picker( - ConnectivityState::TransientFailure, - Arc::new(FailingPicker { error: err.clone() }), - channel_controller, - ); - channel_controller.request_resolution(); - Err(err.clone()) + PickFirstState::FirstPass(self) } } -impl LbPolicy for PickFirstPolicy { - type LbConfig = PickFirstConfig; +// State node: SteadyState (Sticky TRANSIENT_FAILURE). +#[derive(Debug)] +struct SteadyState { + addresses: Vec
, + subchannels: Vec, + failure_threshold: usize, + failure_count: usize, + last_connection_error: String, +} - fn resolver_update( - &mut self, - update: ResolverUpdate, - config: Option<&Self::LbConfig>, - channel_controller: &mut dyn ChannelController, - ) -> Result<(), String> { - self.timer = None; +impl SteadyState { + fn enter( + ctx: &mut PickFirstContext<'_>, + addresses: Vec
, + subchannels: Vec, + last_error: String, + ) -> PickFirstState { + ctx.set_failing_picker(&last_error); + ctx.controller.request_resolution(); + + for entry in &subchannels { + if entry.state.connectivity_state == ConnectivityState::Idle { + entry.subchannel.connect(); + } + } - // Reset steady state on new update - self.steady_state = None; + PickFirstState::SteadyState(SteadyState { + failure_threshold: subchannels.len(), + failure_count: 0, + last_connection_error: last_error, + addresses, + subchannels, + }) + } - match update.endpoints { - Ok(endpoints) => { - let new_addresses = self.compile_address(endpoints, config, channel_controller); - // If we have no addresses, clear subchannels and set TRANSIENT_FAILURE. - if new_addresses.is_empty() { - self.subchannels.clear(); - self.selected = None; - self.set_transient_failure( - channel_controller, - Some("empty address list".to_string()), - )?; - } + fn resolver_update( + self, + ctx: &mut PickFirstContext<'_>, + addresses: Vec
, + ) -> PickFirstState { + FirstPassState::enter(ctx, addresses, self.subchannels) + } - if let Some(ready_subchannel) = - self.rebuild_subchannels(new_addresses, channel_controller) - { - self.subchannel_activate(ready_subchannel, channel_controller); - } else { - self.start_connection_pass(channel_controller); - } + fn subchannel_update( + mut self, + ctx: &mut PickFirstContext<'_>, + subchannel: Arc, + state: &SubchannelState, + ) -> PickFirstState { + let addr = subchannel.address(); + let Some(entry) = self.subchannels.iter_mut().find(|e| e.address == addr) else { + return PickFirstState::SteadyState(self); + }; + entry.state = state.clone(); + + match state.connectivity_state { + ConnectivityState::Ready => ReadyState::enter(ctx, self.addresses, subchannel), + ConnectivityState::Idle => { + subchannel.connect(); + PickFirstState::SteadyState(self) } - Err(e) => { - let error = e.to_string(); - if self.subchannels.is_empty() - || self.connectivity_state == ConnectivityState::TransientFailure - { - self.set_transient_failure(channel_controller, Some(error))?; + ConnectivityState::TransientFailure => { + if let Some(err) = &state.last_connection_error { + self.last_connection_error.clone_from(err); + ctx.set_failing_picker(&self.last_connection_error); } + self.failure_count += 1; + if self.failure_count >= self.failure_threshold { + self.failure_count = 0; + ctx.controller.request_resolution(); + } + PickFirstState::SteadyState(self) } + ConnectivityState::Connecting => PickFirstState::SteadyState(self), } + } +} + +// State node: Ready. +struct ReadyState { + addresses: Vec
, + selected: Arc, +} - Ok(()) +impl ReadyState { + fn enter( + ctx: &mut PickFirstContext<'_>, + addresses: Vec
, + selected: Arc, + ) -> PickFirstState { + ctx.controller.update_picker(LbState { + connectivity_state: ConnectivityState::Ready, + picker: Arc::new(OneSubchannelPicker { + sc: selected.clone(), + }), + }); + PickFirstState::Ready(ReadyState { + addresses, + selected, + }) } + fn resolver_update( + mut self, + ctx: &mut PickFirstContext<'_>, + addresses: Vec
, + ) -> PickFirstState { + let selected_addr = self.selected.address(); + if addresses.contains(&selected_addr) { + self.addresses = addresses; + return PickFirstState::Ready(self); + } + + FirstPassState::fresh_enter(ctx, addresses) + } + + #[allow(clippy::needless_pass_by_value)] fn subchannel_update( - &mut self, + self, + ctx: &mut PickFirstContext<'_>, subchannel: Arc, state: &SubchannelState, - channel_controller: &mut dyn ChannelController, - ) { - if !self - .subchannels - .iter() - .any(|sc| sc.address() == subchannel.address()) - { - // This update is from an outdated subchannel that is no longer in - // the address list. Ignore it. - return; + ) -> PickFirstState { + let addr = subchannel.address(); + if self.selected.address() == addr && state.connectivity_state != ConnectivityState::Ready { + IdleState::enter(ctx, self.addresses) + } else { + PickFirstState::Ready(self) } + } +} - // Update the cache for all updates. - self.subchannel_states - .insert(subchannel.address(), state.clone()); +// The Pick First State Machine enum. +enum PickFirstState { + Idle(IdleState), + FirstPass(FirstPassState), + SteadyState(SteadyState), + Ready(ReadyState), +} - match ( - // Does the load balancer have an active subchannel already? - self.subchannel_is_selected(&subchannel), - // What is the updating subchannel's state? - state.connectivity_state, - ) { - (true, ConnectivityState::Ready) => { - // The selected subchannel is still ready; do nothing w/update. - } - (true, _) => { - // The selected subchannel has failed (is no longer READY); - // drop the connection. - self.subchannel_drop(channel_controller); - } - (false, ConnectivityState::Ready) => { - // The updating subchannel is READY; activate it. - self.subchannel_activate(subchannel, channel_controller); - } - (false, _) => { - // Always capture freshest unselected error. - if state.connectivity_state == ConnectivityState::TransientFailure - && let Some(err) = &state.last_connection_error - { - self.last_connection_error = Some(err.clone()); - } +impl Default for PickFirstState { + fn default() -> Self { + Self::Idle(IdleState::default()) + } +} - // Track progress based on whether we are in a connection pass. - if let Some(steady) = self.steady_state.as_mut() { - steady.subchannel_nonready(channel_controller, subchannel, state); - } else { - self.update_first_pass(subchannel, state, channel_controller); - } - } +impl PickFirstState { + fn resolver_update( + self, + ctx: &mut PickFirstContext<'_>, + addresses: Vec
, + ) -> PickFirstState { + match self { + Self::Idle(s) => s.resolver_update(ctx, addresses), + Self::FirstPass(s) => s.resolver_update(ctx, addresses), + Self::SteadyState(s) => s.resolver_update(ctx, addresses), + Self::Ready(s) => s.resolver_update(ctx, addresses), } } - fn work(&mut self, data: Option, channel_controller: &mut dyn ChannelController) { - debug_assert!(data.is_none(), "expected no data but got {data:?}"); - if self.connectivity_state == ConnectivityState::Idle { - // TODO: is it safe to assume any call to work() while idle means we - // should connect? - self.exit_idle(channel_controller); - } else if self.timer.as_ref().is_some_and(|t| t.expired()) { - // Advance frontier and trigger next connection. - if let Some(next_sc) = self.advance_frontier(false) { - self.trigger_subchannel_connection(next_sc, channel_controller); - } + fn subchannel_update( + self, + ctx: &mut PickFirstContext<'_>, + subchannel: Arc, + state: &SubchannelState, + ) -> PickFirstState { + match self { + Self::Idle(s) => s.subchannel_update(ctx, subchannel, state), + Self::FirstPass(s) => s.subchannel_update(ctx, subchannel, state), + Self::SteadyState(s) => s.subchannel_update(ctx, subchannel, state), + Self::Ready(s) => s.subchannel_update(ctx, subchannel, state), + } + } + + fn work(self, ctx: &mut PickFirstContext<'_>) -> PickFirstState { + match self { + Self::Idle(s) => s.work(ctx), + Self::FirstPass(s) => s.work(ctx), + _ => self, } } - fn exit_idle(&mut self, channel_controller: &mut dyn ChannelController) { - self.start_connection_pass(channel_controller); + fn exit_idle(self, ctx: &mut PickFirstContext<'_>) -> PickFirstState { + match self { + Self::Idle(s) => s.exit_idle(ctx), + _ => self, + } } } -/// Implements the happy eyeballs timer task. `expired` becomes set when it -/// fires. When dropped, the timer is cancelled. +impl Debug for PickFirstState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Idle(s) => f + .debug_struct("Idle") + .field("addresses", &s.addresses) + .finish(), + Self::FirstPass(fp) => f + .debug_struct("FirstPass") + .field("frontier_index", &fp.frontier_index) + .field("addresses", &fp.addresses) + .finish(), + Self::SteadyState(ss) => f.debug_tuple("SteadyState").field(ss).finish(), + Self::Ready(r) => f + .debug_struct("Ready") + .field("selected", &r.selected.address()) + .finish(), + } + } +} + +#[cfg(test)] +impl PickFirstPolicy { + fn selected(&self) -> Option<&Arc> { + match &self.state { + PickFirstState::Ready(r) => Some(&r.selected), + _ => None, + } + } + + fn timer(&self) -> Option<&Timer> { + match &self.state { + PickFirstState::FirstPass(fp) => Some(&fp.timer), + _ => None, + } + } + + fn steady_state(&self) -> Option<&SteadyState> { + match &self.state { + PickFirstState::SteadyState(ss) => Some(ss), + _ => None, + } + } + + fn subchannels(&self) -> Vec> { + match &self.state { + PickFirstState::FirstPass(fp) => fp + .subchannels + .iter() + .map(|e| e.subchannel.clone()) + .collect(), + PickFirstState::SteadyState(ss) => ss + .subchannels + .iter() + .map(|e| e.subchannel.clone()) + .collect(), + _ => Vec::new(), + } + } + + fn last_connection_error(&self) -> Option<&str> { + match &self.state { + PickFirstState::FirstPass(fp) => fp.last_connection_error.as_deref(), + PickFirstState::SteadyState(ss) => Some(&ss.last_connection_error), + _ => None, + } + } +} + +// Implements the happy eyeballs timer task. `expired` becomes set when it +// fires. When dropped, the timer is cancelled. struct Timer { expired: Arc, handle: BoxedTaskHandle, } impl Timer { - /// Starts a new timer, returning it. + // Starts a new timer, returning it. fn start(runtime: GrpcRuntime, work_scheduler: Arc) -> Self { let expired = Arc::new(AtomicBool::new(false)); let expired_clone = expired.clone(); @@ -625,7 +791,7 @@ impl Timer { Self { expired, handle } } - /// Returns whether the timer has expired yet. + // Returns whether the timer has expired yet. fn expired(&self) -> bool { self.expired.load(Ordering::SeqCst) } @@ -683,70 +849,22 @@ fn build_shuffler() -> Arc { }) } -/// Tracks a the 'steady state' pass of subchannels when looking for a ready -/// connection. If the number of reported subchannel failures reaches the -/// failure threshold, this will ask the Name Resolver to re-resolve. -#[derive(Debug)] -struct SteadyState { - /// The number of failures before triggering a re-resolution of addresses. - /// This is a rough heuristic to approximate if all subchannels have failed - /// since we entered steady state, and can be tuned as needed. - failure_threshold: usize, - /// The number of failures connecting, used to roughly approximate if a - /// re-resolution needs to happen. - failure_count: usize, -} - -impl SteadyState { - fn new(threshold: usize) -> Self { - Self { - failure_threshold: threshold, - failure_count: 0, - } - } - - /// Handles non-ready subchannel updates when the LB is in 'steady state' - /// connection mode. - fn subchannel_nonready( - &mut self, - channel_controller: &mut dyn ChannelController, - subchannel: Arc, - state: &SubchannelState, - ) { - match state.connectivity_state { - ConnectivityState::Idle => { - // Subchannel backoff expired: trigger reconnection attempt. - subchannel.connect(); - } - ConnectivityState::TransientFailure => { - // Track failures. If all known subchannels have failed, - // request new addresses. - self.failure_count += 1; - if self.failure_count >= self.failure_threshold { - self.failure_count = 0; - channel_controller.request_resolution(); - } - } - _ => {} - } - } -} - #[cfg(test)] mod test { use std::sync::mpsc; use std::time::Duration; use super::*; - use crate::client::load_balancing::test_utils::{ - TestChannelController, TestEvent, TestWorkScheduler, - }; + use crate::client::load_balancing::test_utils::TestChannelController; + use crate::client::load_balancing::test_utils::TestEvent; + use crate::client::load_balancing::test_utils::TestWorkScheduler; const DEFAULT_TEST_DURATION: Duration = Duration::from_secs(10); // Helper to create endpoints from a list of address strings. // If attrs are provided, they will be added to each endpoint; otherwise, // default attributes will be used. + #[allow(clippy::needless_pass_by_value)] fn create_endpoints( addrs: Vec<&str>, attrs: Option, @@ -795,40 +913,40 @@ mod test { fn expect_new_subchannel(rx: &mpsc::Receiver) -> Arc { match rx.try_recv() { Ok(TestEvent::NewSubchannel(sc)) => sc, - Ok(other) => panic!("expected NewSubchannel, got {:?}", other), - Err(e) => panic!("expected NewSubchannel, got error: {:?}", e), + Ok(other) => panic!("expected NewSubchannel, got {other:?}"), + Err(e) => panic!("expected NewSubchannel, got error: {e:?}"), } } fn expect_connect(rx: &mpsc::Receiver) -> Address { match rx.try_recv() { Ok(TestEvent::Connect(addr)) => addr, - Ok(other) => panic!("expected Connect, got {:?}", other), - Err(e) => panic!("expected Connect, got error: {:?}", e), + Ok(other) => panic!("expected Connect, got {other:?}"), + Err(e) => panic!("expected Connect, got error: {e:?}"), } } fn expect_picker_update(rx: &mpsc::Receiver) -> LbState { match rx.try_recv() { Ok(TestEvent::UpdatePicker(state)) => state, - Ok(other) => panic!("expected UpdatePicker, got {:?}", other), - Err(e) => panic!("expected UpdatePicker, got error: {:?}", e), + Ok(other) => panic!("expected UpdatePicker, got {other:?}"), + Err(e) => panic!("expected UpdatePicker, got error: {e:?}"), } } fn expect_request_resolution(rx: &mpsc::Receiver) { match rx.try_recv() { Ok(TestEvent::RequestResolution) => {} - Ok(other) => panic!("expected RequestResolution, got {:?}", other), - Err(e) => panic!("expected RequestResolution, got error: {:?}", e), + Ok(other) => panic!("expected RequestResolution, got {other:?}"), + Err(e) => panic!("expected RequestResolution, got error: {e:?}"), } } fn expect_schedule_work(rx: &mpsc::Receiver) { match rx.try_recv() { Ok(TestEvent::ScheduleWork(_)) => {} - Ok(other) => panic!("expected ScheduleWork, got {:?}", other), - Err(e) => panic!("expected ScheduleWork, got error: {:?}", e), + Ok(other) => panic!("expected ScheduleWork, got {other:?}"), + Err(e) => panic!("expected ScheduleWork, got error: {e:?}"), } } @@ -880,7 +998,7 @@ mod test { let (rx, mut policy, mut controller) = simulate_connection(addrs, attrs); // Simulating READY for addr1. - let sc1 = policy.subchannels[0].clone(); + let sc1 = policy.subchannels()[0].clone(); policy.subchannel_update( sc1.clone(), &SubchannelState { @@ -903,7 +1021,7 @@ mod test { let (rx, mut policy, mut controller) = simulate_connection(addrs, attrs); // Simulating TransientFailure for addr1. - let sc1 = policy.subchannels[0].clone(); + let sc1 = policy.subchannels()[0].clone(); policy.subchannel_update( sc1.clone(), &SubchannelState { @@ -928,9 +1046,9 @@ mod test { let res = state.picker.pick(&RequestHeaders::default()); match res { PickResult::Pick(pick) => { - assert_eq!(pick.subchannel.address().address.to_string(), "addr1") + assert_eq!(pick.subchannel.address().address.to_string(), "addr1"); } - other => panic!("unexpected pick result {:?}", other), + other => panic!("unexpected pick result {other:?}"), } } @@ -945,7 +1063,7 @@ mod test { assert_eq!(addr.address.to_string(), "addr2"); // Simulate addr2 succeeding. - let sc2 = policy.subchannels[1].clone(); + let sc2 = policy.subchannels()[1].clone(); policy.subchannel_update( sc2, &SubchannelState { @@ -984,20 +1102,13 @@ mod test { ) .unwrap(); - // Should create new subchannel for addr2 (was cleared by cleanup). - let sc2 = expect_new_subchannel(&rx); - assert_eq!(sc2.address().address.to_string(), "addr2"); - // Should create new subchannel for addr3 (was not in previous list). - let sc3 = expect_new_subchannel(&rx); - assert_eq!(sc3.address().address.to_string(), "addr3"); - - // Should NOT have any more events (no Connect, no UpdatePicker), - // because it stuck to the original selected subchannel. + // Should NOT have any new subchannels or connect events because it stuck to the + // original selected subchannel. assert!(rx.try_recv().is_err(), "unexpected event"); assert_eq!( policy - .selected + .selected() .as_ref() .unwrap() .address() @@ -1028,6 +1139,7 @@ mod test { // based on the provided shuffler and config. #[tokio::test] async fn test_pick_first_shuffling_and_interleaving_deterministic() { + const NUM_ADDRS: usize = 4; let (rx, mut policy, mut controller) = setup(); // Enable shuffling in config. @@ -1084,7 +1196,6 @@ mod test { ) .unwrap(); - const NUM_ADDRS: usize = 4; let mut resulting_addrs = Vec::with_capacity(NUM_ADDRS); for _ in 0..NUM_ADDRS { let sc = expect_new_subchannel(&rx); @@ -1168,7 +1279,7 @@ mod test { } } - assert_eq!(policy.subchannels.len(), 2, "De-duplication failed"); + assert_eq!(policy.subchannels().len(), 2, "De-duplication failed"); } // If the resolver update contains no addresses, the LB should clear @@ -1215,7 +1326,7 @@ mod test { // Simulate timer expiration by setting the flag directly. policy - .timer + .timer() .as_ref() .unwrap() .expired @@ -1234,7 +1345,7 @@ mod test { #[tokio::test] async fn test_pick_first_steady_state_retries() { let (rx, mut policy, mut controller) = simulate_failed_connection(vec!["addr1"], None); - let sc1 = policy.subchannels[0].clone(); + let sc1 = policy.subchannels()[0].clone(); // Expect UpdatePicker(TransientFailure) and RequestResolution. let state = expect_picker_update(&rx); @@ -1245,7 +1356,7 @@ mod test { expect_request_resolution(&rx); // Ensure steady state was entered. - assert!(policy.steady_state.is_some()); + assert!(policy.steady_state().is_some()); // Simulate addr1 transitioning to IDLE (backoff over). policy.subchannel_update( @@ -1270,7 +1381,7 @@ mod test { async fn test_pick_first_steady_state_multi_backend() { let (rx, mut policy, mut controller) = simulate_failed_connection(vec!["addr1", "addr2"], None); - let sc1 = policy.subchannels[0].clone(); + let sc1 = policy.subchannels()[0].clone(); // Should failover to addr2: expect Connect(addr2). let addr = expect_connect(&rx); @@ -1287,11 +1398,10 @@ mod test { ); // We should NOT reconnect to addr1 during the first pass. - // Wait a bit to ensure no event is sent. assert!(rx.try_recv().is_err(), "unexpected event"); // Now fail addr2 to complete first pass. - let sc2 = policy.subchannels[1].clone(); + let sc2 = policy.subchannels()[1].clone(); policy.subchannel_update( sc2.clone(), &SubchannelState { @@ -1301,7 +1411,8 @@ mod test { controller.as_mut(), ); - // Expect UpdatePicker(TransientFailure), RequestResolution, and Connect(addr1) from first pass exhaustion. + // Expect UpdatePicker(TransientFailure), RequestResolution, and Connect(addr1) + // from first pass exhaustion. let state = expect_picker_update(&rx); assert_eq!( state.connectivity_state, @@ -1312,7 +1423,7 @@ mod test { assert_eq!(addr.address.to_string(), "addr1"); // Confirm LB is in steady state. - assert!(policy.steady_state.is_some()); + assert!(policy.steady_state().is_some()); // Simulate addr1 going IDLE again. policy.subchannel_update( @@ -1346,7 +1457,7 @@ mod test { PickResult::Pick(pick) => { assert_eq!(pick.subchannel.address().address.to_string(), "addr1"); } - other => panic!("unexpected pick result {:?}", other), + other => panic!("unexpected pick result {other:?}"), } } @@ -1359,7 +1470,7 @@ mod test { async fn test_pick_first_steady_state_stuck_idle_prevention() { let (rx, mut policy, mut controller) = simulate_failed_connection(vec!["addr1", "addr2"], None); - let sc1 = policy.subchannels[0].clone(); + let sc1 = policy.subchannels()[0].clone(); // Expect Connect(addr2). let addr = expect_connect(&rx); @@ -1380,7 +1491,7 @@ mod test { assert!(rx.try_recv().is_err(), "unexpected event during first pass"); // Fail addr2 to exhaust the first pass. - let sc2 = policy.subchannels[1].clone(); + let sc2 = policy.subchannels()[1].clone(); policy.subchannel_update( sc2, &SubchannelState { @@ -1390,8 +1501,7 @@ mod test { controller.as_mut(), ); - // Expect UpdatePicker(TransientFailure) and RequestResolution from - // exhaustion. + // Expect UpdatePicker(TransientFailure) and RequestResolution from exhaustion. let state = expect_picker_update(&rx); assert_eq!( state.connectivity_state, @@ -1430,7 +1540,6 @@ mod test { .unwrap(); // This should be a different subchannel due to different attributes. - // Therefore, expect a new TestEvent::NewSubchannel event to be emitted. let mut found_new_subchannel = false; while let Ok(event) = rx.try_recv() { if let TestEvent::NewSubchannel(_) = event { @@ -1448,8 +1557,7 @@ mod test { // If a resolver error is received while the LB is in the process of // connecting to addresses, it should not abort the connection attempt or // clear the existing addresses, as long as there are still valid addresses - // in the LB. This ensures that transient resolver errors do not cause - // unnecessary disruption to active connection attempts. + // in the LB. #[tokio::test] async fn test_pick_first_resolver_error_during_connecting() { let (rx, mut policy, mut controller) = simulate_connection(vec!["addr1"], None); @@ -1467,33 +1575,24 @@ mod test { ) .unwrap(); - // Resolver errors received during active connection attempts should NOT - // abort the attempt or force TransientFailure immediately if the load - // balancer still has valid addresses. - // Expect NO events to be emitted (no UpdatePicker/RequestResolution). assert!( rx.try_recv().is_err(), "Unexpected event after resolver error" ); - // Verify internal state did not clear endpoints. assert!( - !policy.subchannels.is_empty(), + !policy.subchannels().is_empty(), "Subchannels erroneously cleared by resolver error." ); } - // Out-of-Order Failure Detection - // Ensures the policy waits for all parallel connection attempts to drop - // before failing the channel. + // Out-of-Order Failure Detection. #[tokio::test] async fn test_pick_first_happy_eyeballs_out_of_order_failure() { let (rx, mut policy, mut controller) = simulate_connection(vec!["addr1", "addr2"], None); - // 1. Simulate Happy Eyeballs timer firing to launch parallel connection - // to addr2. policy - .timer + .timer() .as_ref() .unwrap() .expired @@ -1503,8 +1602,7 @@ mod test { let addr = expect_connect(&rx); assert_eq!(addr.address.to_string(), "addr2"); - // 2. Simulate addr2 failing first while addr1 is still in flight. - let sc2 = policy.subchannels[1].clone(); + let sc2 = policy.subchannels()[1].clone(); policy.subchannel_update( sc2, &SubchannelState { @@ -1514,11 +1612,9 @@ mod test { controller.as_mut(), ); - // Verify policy does NOT enter TransientFailure yet. assert!(rx.try_recv().is_err(), "unexpected premature event"); - // 3. Simulate addr1 failing. Pass is now fully exhausted. - let sc1 = policy.subchannels[0].clone(); + let sc1 = policy.subchannels()[0].clone(); policy.subchannel_update( sc1, &SubchannelState { @@ -1535,24 +1631,20 @@ mod test { ); } - // Freshest Error Caching (Steady State) - // Ensures background failures during Steady State continuously overwrite - // stale connection errors. + // Freshest Error Caching (Steady State). #[tokio::test] async fn test_pick_first_steady_state_freshest_error() { let (rx, mut policy, mut controller) = simulate_failed_connection(vec!["addr1"], None); - // Consume exhaustion events. let state = expect_picker_update(&rx); assert_eq!( state.connectivity_state, ConnectivityState::TransientFailure ); expect_request_resolution(&rx); - assert!(policy.steady_state.is_some()); + assert!(policy.steady_state().is_some()); - // Simulate background failure during Steady State with net-new error telemetry. - let sc1 = policy.subchannels[0].clone(); + let sc1 = policy.subchannels()[0].clone(); policy.subchannel_update( sc1, &SubchannelState { @@ -1562,22 +1654,18 @@ mod test { controller.as_mut(), ); - // Verify policy caches the freshest unselected error. assert_eq!( - policy.last_connection_error.as_deref(), + policy.last_connection_error(), Some("steady state network drop") ); } // Tests that when a selected subchannel disconnects (transitions to Idle), // the policy reports an Idle state and uses an IdlePicker. - // When an RPC occurs, the IdlePicker schedules work, and the policy - // reconnects when the work scheduler runs. #[tokio::test] async fn test_pick_first_disconnect_to_idle_and_reconnect() { let (rx, mut policy, mut controller) = simulate_successful_connection(vec!["addr1"], None); - // 1. Consume the initial Ready picker update. let state = expect_picker_update(&rx); assert_eq!(state.connectivity_state, ConnectivityState::Ready); let res = state.picker.pick(&RequestHeaders::default()); @@ -1586,10 +1674,9 @@ mod test { assert_eq!(pick.subchannel.address().address.to_string(), "addr1"); pick.subchannel } - other => panic!("unexpected pick result {:?}", other), + other => panic!("unexpected pick result {other:?}"), }; - // 2. Simulate the subchannel disconnecting (transitioning to Idle). policy.subchannel_update( sc1.clone(), &SubchannelState { @@ -1599,29 +1686,157 @@ mod test { controller.as_mut(), ); - // 3. Verify the policy updates the picker to Idle state. let state = expect_picker_update(&rx); assert_eq!(state.connectivity_state, ConnectivityState::Idle); let idle_picker = state.picker; - // At this point, there should be no more events, as we are waiting for an RPC. assert!(rx.try_recv().is_err(), "unexpected event"); - // 4. Simulate an RPC (pick) happening. let pick_result = idle_picker.pick(&RequestHeaders::default()); assert!(matches!(pick_result, PickResult::Queue)); - // 5. The picker should schedule work. expect_schedule_work(&rx); - // 6. Call work to execute the scheduled connection attempt. policy.work(None, controller.as_mut()); - // 7. Verify that the policy initiates a reconnection to addr1. + expect_new_subchannel(&rx); let addr = expect_connect(&rx); assert_eq!(addr.address.to_string(), "addr1"); - // And the picker goes to Connecting. + let state = expect_picker_update(&rx); + assert_eq!(state.connectivity_state, ConnectivityState::Connecting); + } + + // Tests that when connected to multi-address endpoints, if the winner drops + // to Idle and reconnects on RPC, the policy attempts to connect to all + // original addresses (falling over to addr2 if addr1 fails). + #[tokio::test] + async fn test_pick_first_disconnect_to_idle_multi_address_reconnect() { + let (rx, mut policy, mut controller) = + simulate_successful_connection(vec!["addr1", "addr2"], None); + + let state = expect_picker_update(&rx); + assert_eq!(state.connectivity_state, ConnectivityState::Ready); + let res = state.picker.pick(&RequestHeaders::default()); + let sc1 = match res { + PickResult::Pick(pick) => { + assert_eq!(pick.subchannel.address().address.to_string(), "addr1"); + pick.subchannel + } + other => panic!("unexpected pick result {other:?}"), + }; + + // 2. Simulate addr1 disconnecting (transitioning to Idle). + policy.subchannel_update( + sc1.clone(), + &SubchannelState { + connectivity_state: ConnectivityState::Idle, + last_connection_error: None, + }, + controller.as_mut(), + ); + + // 3. Verify the policy updates the picker to Idle state. + let state = expect_picker_update(&rx); + assert_eq!(state.connectivity_state, ConnectivityState::Idle); + let idle_picker = state.picker; + + assert!(rx.try_recv().is_err(), "unexpected event"); + + // 4. Simulate an RPC happening. + let pick_result = idle_picker.pick(&RequestHeaders::default()); + assert!(matches!(pick_result, PickResult::Queue)); + expect_schedule_work(&rx); + + // 5. Call work to execute the scheduled connection attempt. + policy.work(None, controller.as_mut()); + + // Reconnect creates subchannels for both addr1 and addr2. + expect_new_subchannel(&rx); + expect_new_subchannel(&rx); + + // 6. Verify that the policy initiates reconnection to addr1. + let addr = expect_connect(&rx); + assert_eq!(addr.address.to_string(), "addr1"); + + let state = expect_picker_update(&rx); + assert_eq!(state.connectivity_state, ConnectivityState::Connecting); + + // 7. Simulate addr1 failing on reconnect. + let sc1_reconnect = policy.subchannels()[0].clone(); + policy.subchannel_update( + sc1_reconnect, + &SubchannelState { + connectivity_state: ConnectivityState::TransientFailure, + last_connection_error: Some("connection refused".to_string()), + }, + controller.as_mut(), + ); + + // 8. Policy should failover to addr2. + let addr = expect_connect(&rx); + assert_eq!(addr.address.to_string(), "addr2"); + } + + // Tests that when a resolver update arrives while in FirstPass, any subchannels + // that have already failed (TransientFailure) preserve their state so the policy + // does not redundantly re-attempt them before trying remaining addresses. + #[tokio::test] + async fn test_pick_first_resolver_update_preserves_subchannel_failure_state() { + let (rx, mut policy, mut controller) = setup(); + + // 1. Initial endpoints [addr1, addr2]. + let endpoints = create_endpoints(vec!["addr1", "addr2"], None); + policy + .resolver_update( + ResolverUpdate { + endpoints: Ok(endpoints), + ..Default::default() + }, + None, + controller.as_mut(), + ) + .unwrap(); + + let sc1 = expect_new_subchannel(&rx); + let _sc2 = expect_new_subchannel(&rx); + let addr = expect_connect(&rx); + assert_eq!(addr.address.to_string(), "addr1"); + let state = expect_picker_update(&rx); + assert_eq!(state.connectivity_state, ConnectivityState::Connecting); + + // 2. addr1 fails with TransientFailure. Frontier advances to addr2. + policy.subchannel_update( + sc1, + &SubchannelState { + connectivity_state: ConnectivityState::TransientFailure, + last_connection_error: Some("connection refused".to_string()), + }, + controller.as_mut(), + ); + let addr = expect_connect(&rx); + assert_eq!(addr.address.to_string(), "addr2"); + + // 3. New resolver update [addr1, addr3]. (Retains addr1 which already failed). + let endpoints_new = create_endpoints(vec!["addr1", "addr3"], None); + policy + .resolver_update( + ResolverUpdate { + endpoints: Ok(endpoints_new), + ..Default::default() + }, + None, + controller.as_mut(), + ) + .unwrap(); + + // 4. addr3 is created as a new subchannel. + let _sc3 = expect_new_subchannel(&rx); + + // 5. Because addr1 is already in TransientFailure, the policy should skip addr1 + // and immediately connect addr3. + let addr = expect_connect(&rx); + assert_eq!(addr.address.to_string(), "addr3"); let state = expect_picker_update(&rx); assert_eq!(state.connectivity_state, ConnectivityState::Connecting); }