impl Pool { pub fn select(&self, peer: IpAddr, headers: &HeaderMap, tls: bool) -> Option { self.select_excluding(peer, headers, tls, &[]) } pub fn select_retry( &self, peer: IpAddr, headers: &HeaderMap, tls: bool, excluded: &[String], ) -> Option { self.select_excluding(peer, headers, tls, excluded) } fn eligible_healthy_excluding(&self, excluded: &[String]) -> Vec<&Arc> { let all: Vec<_> = self .members .iter() .filter(|m| { m.definition.pool_eligible && m.definition.enabled && !m.definition.drain && m.state.healthy() && !excluded.iter().any(|id| id == &m.definition.id) }) .collect(); let primary = all.iter().any(|m| !m.definition.backup); all.into_iter() .filter(|m| !primary || !m.definition.backup) .collect() } fn capacity_wait_ms(&self, excluded: &[String]) -> Option { let members = self.eligible_healthy_excluding(excluded); if members.is_empty() { return None; } Some( members .iter() .map(|m| m.definition.queue_timeout_ms) .max() .unwrap_or(0), ) } fn note_capacity_rejection(&self, excluded: &[String], timed_out: bool) { if let Some(member) = self .eligible_healthy_excluding(excluded) .into_iter() .min_by_key(|m| m.state.active.load(Ordering::Relaxed)) { member .state .capacity_rejections_total .fetch_add(1, Ordering::Relaxed); if timed_out { member .state .queue_timeouts_total .fetch_add(1, Ordering::Relaxed); } } } pub async fn select_wait( &self, peer: IpAddr, headers: &HeaderMap, tls: bool, ) -> Result { self.select_retry_wait(peer, headers, tls, &[]).await } pub async fn select_member_wait(&self, id: &str) -> Result { let Some(member) = self.members.iter().find(|m| { m.definition.id == id && m.definition.enabled && !m.definition.drain && m.state.healthy() }) else { return Err(SelectionFailure::Unavailable); }; let wait_ms = member.definition.queue_timeout_ms; let started = Instant::now(); loop { if !member.definition.enabled || member.definition.drain || !member.state.healthy() { return Err(SelectionFailure::Unavailable); } if let Some(selection) = self.select_member(id) { return Ok(selection); } if wait_ms == 0 || started.elapsed() >= Duration::from_millis(wait_ms) { member .state .capacity_rejections_total .fetch_add(1, Ordering::Relaxed); if wait_ms > 0 { member .state .queue_timeouts_total .fetch_add(1, Ordering::Relaxed); } return Err(SelectionFailure::Capacity); } tokio::time::sleep(Duration::from_millis(5)).await; } } pub async fn select_retry_wait( &self, peer: IpAddr, headers: &HeaderMap, tls: bool, excluded: &[String], ) -> Result { if let Some(selection) = self.select_retry(peer, headers, tls, excluded) { return Ok(selection); } let Some(wait_ms) = self.capacity_wait_ms(excluded) else { return Err(SelectionFailure::Unavailable); }; let started = Instant::now(); loop { if self.eligible_healthy_excluding(excluded).is_empty() { return Err(SelectionFailure::Unavailable); } if let Some(selection) = self.select_retry(peer, headers, tls, excluded) { return Ok(selection); } if wait_ms == 0 || started.elapsed() >= Duration::from_millis(wait_ms) { self.note_capacity_rejection(excluded, wait_ms > 0); return Err(SelectionFailure::Capacity); } tokio::time::sleep(Duration::from_millis(5)).await; } } }