impl Pool { pub fn new( id: String, settings: Settings, definitions: &[PoolMember], backends: Vec>, ) -> Result { ensure!( definitions.len() == backends.len() && !backends.is_empty(), "invalid compiled pool" ); let fingerprint = Sha256::digest(serde_json::to_vec(&(definitions, &settings))?).to_vec(); let cookie_name = settings .sticky .cookie_name .clone() .unwrap_or_else(|| format!("RP_AFFINITY_{id}")); let members = definitions .iter() .cloned() .zip(backends) .map(|(definition, backend)| { let health_settings = Arc::new(settings.health_for(&definition)); let passive_settings = Arc::new(settings.passive_for(&definition)); Arc::new(Member { state: Arc::new(BackendState::new(health_settings.mode)), health_probe: backend.health_probe.clone(), definition, backend, health_settings, passive_settings, token: rand::rng() .sample_iter(Alphanumeric) .take(32) .map(char::from) .collect(), }) }) .collect(); Ok(Self { id, settings, members, fingerprint, cookie_name, cursor: AtomicU64::new(0), scheduling: Mutex::new(std::collections::HashMap::new()), }) } pub fn reuse_observations(&mut self, previous: &Pool) { for member in &mut self.members { let Some(old) = previous .members .iter() .find(|old| old.definition.id == member.definition.id) else { continue; }; let same_origin = serde_json::to_vec(&member.definition.backend()).ok() == serde_json::to_vec(&old.definition.backend()).ok(); let same_health = serde_json::to_vec(&( self.settings.health_for(&member.definition), self.settings.passive_for(&member.definition), )) .ok() == serde_json::to_vec(&( previous.settings.health_for(&old.definition), previous.settings.passive_for(&old.definition), )) .ok(); let same_clients = member.backend.client.reusable_with(&old.backend.client) && member .backend .websocket_client .reusable_with(&old.backend.websocket_client); if let Some(member) = Arc::get_mut(member) { if same_clients { member.backend = old.backend.clone(); } if same_origin && same_health && (old.definition.enabled || !member.definition.enabled) { member.state = old.state.clone(); member.token = old.token.clone(); } } } } fn candidates_excluding(&self, excluded: &[String]) -> Vec<&Arc> { let mut eligible: Vec<_> = self .members .iter() .filter(|m| { m.definition.pool_eligible && m.definition.enabled && !m.definition.drain && m.state.healthy() && m.state.has_capacity(m.definition.max_concurrent_requests) && !excluded.iter().any(|id| id == &m.definition.id) }) .collect(); if eligible.iter().any(|m| !m.definition.backup) { eligible.retain(|m| !m.definition.backup); } eligible } pub fn ready(&self) -> bool { !self.eligible_healthy_excluding(&[]).is_empty() } pub fn cookie_affinity(&self) -> bool { self.settings.sticky.mode == StickyMode::Cookie } fn affinity<'a>(&self, headers: &'a HeaderMap) -> Option<&'a str> { let mut values = headers .get_all(::http::header::COOKIE) .iter() .filter_map(|v| v.to_str().ok()) .flat_map(|v| v.split(';')) .filter_map(|pair| pair.trim().split_once('=')) .filter(|(name, _)| *name == self.cookie_name) .map(|(_, value)| value); let first = values.next()?; if values.next().is_some() { None } else { Some(first) } } pub fn strip_affinity_response(&self, headers: &mut HeaderMap) { if !self.cookie_affinity() { return; } let retained: Vec<_> = headers .get_all(::http::header::SET_COOKIE) .iter() .filter(|value| { value .to_str() .ok() .and_then(|v| v.split_once('=')) .is_none_or(|(name, _)| name.trim() != self.cookie_name) }) .cloned() .collect(); headers.remove(::http::header::SET_COOKIE); for value in retained { headers.append(::http::header::SET_COOKIE, value); } } pub fn strip_affinity(&self, headers: &mut HeaderMap) { if !self.cookie_affinity() { return; } let values: Vec = headers .get_all(::http::header::COOKIE) .iter() .filter_map(|value| { let Ok(text) = value.to_str() else { return Some(value.clone()); }; let clean = text .split(';') .filter(|part| { part.trim() .split_once('=') .is_none_or(|(n, _)| n != self.cookie_name) }) .map(str::trim) .collect::>() .join("; "); if clean.is_empty() { None } else { clean.parse().ok() } }) .collect(); headers.remove(::http::header::COOKIE); for value in values { headers.append(::http::header::COOKIE, value); } } pub fn select_member(&self, id: &str) -> Option { let member = self.members.iter().find(|m| { m.definition.id == id && m.definition.enabled && !m.definition.drain && m.state.healthy() })?; if !member .state .try_acquire_active(member.definition.max_concurrent_requests) { return None; } Some(Selection { lease: Arc::new(Lease { member: member.clone(), health: member.health_settings.clone(), passive_health: member.passive_settings.clone(), error_recorded: AtomicBool::new(false), }), cookie: None, }) } fn select_excluding( &self, peer: IpAddr, headers: &HeaderMap, tls: bool, excluded: &[String], ) -> Option { if self.members.len() == 1 && !self.cookie_affinity() { let member = &self.members[0]; if !member.definition.pool_eligible || !member.definition.enabled || member.definition.drain || !member.state.healthy() || excluded.iter().any(|id| id == &member.definition.id) || !member .state .try_acquire_active(member.definition.max_concurrent_requests) { return None; } return Some(Selection { lease: Arc::new(Lease { member: member.clone(), health: member.health_settings.clone(), passive_health: member.passive_settings.clone(), error_recorded: AtomicBool::new(false), }), cookie: None, }); } let members = self.candidates_excluding(excluded); let existing = if self.cookie_affinity() { self.affinity(headers) } else { None }; let matched = existing.and_then(|token| { self.members.iter().find(|m| { m.token == token && m.definition.pool_eligible && m.definition.enabled && m.state.healthy() && m.state.has_capacity(m.definition.max_concurrent_requests) && !excluded.iter().any(|id| id == &m.definition.id) }) }); if members.is_empty() && matched.is_none() { return None; } let cursor = self.cursor.fetch_add(1, Ordering::Relaxed); let selected = if let Some(member) = matched { member } else { if self.settings.algorithm == Algorithm::RoundRobin && self.settings.sticky.mode != StickyMode::IpHash { let mut scheduling = lock(&self.scheduling); let mut chosen = members[0]; let mut best = f64::NEG_INFINITY; let mut total = 0.0; for member in &members { let weight = member .state .weight(member.definition.weight, self.settings.slow_start_seconds); let current = if let Some(current) = scheduling.get_mut(member.definition.id.as_str()) { current } else { scheduling .entry(member.definition.id.clone()) .or_insert(0.0) }; *current += weight; total += weight; if *current > best { best = *current; chosen = *member; } } if let Some(current) = scheduling.get_mut(&chosen.definition.id) { *current -= total; } chosen } else { let mut chosen = members[(cursor as usize) % members.len()]; let mut best = f64::INFINITY; let algorithm = if self.settings.sticky.mode == StickyMode::IpHash { Algorithm::IpHash } else { self.settings.algorithm }; let hash_ip = (algorithm == Algorithm::IpHash) .then(|| crate::vhost::canonical_ip(peer).to_string()); for offset in 0..members.len() { let index = ((cursor as usize) % members.len() + offset) % members.len(); let member = members[index]; let active = member.state.active.load(Ordering::Relaxed) as f64; let score = match algorithm { Algorithm::IpHash => { let mut hash = Sha256::new(); hash.update(self.id.as_bytes()); hash.update(b"|"); hash.update(hash_ip.as_deref().expect("hash address").as_bytes()); hash.update(b"|"); hash.update(member.definition.id.as_bytes()); let digest = hash.finalize(); let n = u64::from_be_bytes([ digest[0], digest[1], digest[2], digest[3], digest[4], digest[5], digest[6], digest[7], ]); let unit = (n as f64 + 1.0) / (u64::MAX as f64 + 2.0); -unit.max(f64::MIN_POSITIVE).ln() / member.definition.weight as f64 } _ => { let weight = member .state .weight(member.definition.weight, self.settings.slow_start_seconds); match algorithm { Algorithm::LeastTraffic => { (member.state.traffic_rate() + (active + 1.0) * 1024.0) / weight } Algorithm::Adaptive => { ((active + 1.0) * member.state.latency() + member.state.traffic_rate() / 16_384.0) / weight } _ => (active + 1.0) / weight, } } }; if score < best { best = score; chosen = member; } } chosen } }; if !selected .state .try_acquire_active(selected.definition.max_concurrent_requests) { return None; } let cookie = if self.cookie_affinity() && existing != Some(selected.token.as_str()) { format!( "{}={}; Path=/; Max-Age={}; HttpOnly; SameSite=Lax{}", self.cookie_name, selected.token, self.settings.sticky.ttl_seconds, if tls { "; Secure" } else { "" } ) .parse() .ok() } else { None }; Some(Selection { lease: Arc::new(Lease { member: selected.clone(), health: selected.health_settings.clone(), passive_health: selected.passive_settings.clone(), error_recorded: AtomicBool::new(false), }), cookie, }) } }