333 lines
12 KiB
Rust
333 lines
12 KiB
Rust
impl Pool {
|
|
pub fn new(
|
|
id: String,
|
|
settings: Settings,
|
|
definitions: &[PoolMember],
|
|
backends: Vec<Arc<CompiledBackend>>,
|
|
) -> Result<Self> {
|
|
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)| {
|
|
Arc::new(Member {
|
|
state: Arc::new(BackendState::new(settings.health_check.mode)),
|
|
definition,
|
|
backend,
|
|
token: rand::thread_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_check, &self.settings.passive_health))
|
|
.ok()
|
|
== serde_json::to_vec(&(
|
|
&previous.settings.health_check,
|
|
&previous.settings.passive_health,
|
|
))
|
|
.ok();
|
|
if same_origin && same_health && (old.definition.enabled || !member.definition.enabled)
|
|
{
|
|
// This pool has not been published or cloned into routes yet.
|
|
if let Some(member) = Arc::get_mut(member) {
|
|
member.state = old.state.clone();
|
|
member.token = old.token.clone();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
fn candidates_excluding(&self, excluded: &[String]) -> Vec<&Arc<Member>> {
|
|
let all: 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();
|
|
let primary = all.iter().any(|m| !m.definition.backup);
|
|
all.into_iter()
|
|
.filter(|m| !primary || !m.definition.backup)
|
|
.collect()
|
|
}
|
|
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("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("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("set-cookie");
|
|
for value in retained {
|
|
headers.append("set-cookie", value);
|
|
}
|
|
}
|
|
pub fn strip_affinity(&self, headers: &mut HeaderMap) {
|
|
if !self.cookie_affinity() {
|
|
return;
|
|
}
|
|
let values: Vec<HeaderValue> = headers
|
|
.get_all("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::<Vec<_>>()
|
|
.join("; ");
|
|
if clean.is_empty() {
|
|
None
|
|
} else {
|
|
clean.parse().ok()
|
|
}
|
|
})
|
|
.collect();
|
|
headers.remove("cookie");
|
|
for value in values {
|
|
headers.append("cookie", value);
|
|
}
|
|
}
|
|
pub fn select_member(&self, id: &str) -> Option<Selection> {
|
|
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: self.settings.health_check.clone(),
|
|
passive_health: self.settings.passive_health.clone(),
|
|
error_recorded: AtomicBool::new(false),
|
|
}),
|
|
cookie: None,
|
|
})
|
|
}
|
|
fn select_excluding(
|
|
&self,
|
|
peer: IpAddr,
|
|
headers: &HeaderMap,
|
|
tls: bool,
|
|
excluded: &[String],
|
|
) -> Option<Selection> {
|
|
let mut scheduling = lock(&self.scheduling);
|
|
let members = self.candidates_excluding(excluded);
|
|
let existing = if self.cookie_affinity() {
|
|
self.affinity(headers)
|
|
} else {
|
|
None
|
|
};
|
|
// Drain stops new assignments while existing cookie-affine clients may
|
|
// finish their session. Disabled/unhealthy members never receive traffic.
|
|
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 {
|
|
let weights: Vec<f64> = members
|
|
.iter()
|
|
.map(|m| {
|
|
m.state
|
|
.weight(m.definition.weight, self.settings.slow_start_seconds)
|
|
})
|
|
.collect();
|
|
if self.settings.algorithm == Algorithm::RoundRobin
|
|
&& self.settings.sticky.mode != StickyMode::IpHash
|
|
{
|
|
// Smooth weighted round-robin: distribute weighted requests,
|
|
// rather than sending long consecutive bursts to one member.
|
|
let mut chosen = members[0];
|
|
let mut best = f64::NEG_INFINITY;
|
|
let mut total = 0.0;
|
|
for (member, weight) in members.iter().zip(&weights) {
|
|
let current = 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;
|
|
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 algorithm = if self.settings.sticky.mode == StickyMode::IpHash {
|
|
Algorithm::IpHash
|
|
} else {
|
|
self.settings.algorithm
|
|
};
|
|
let score = match algorithm {
|
|
Algorithm::IpHash => {
|
|
let digest = Sha256::digest(
|
|
format!(
|
|
"{}|{}|{}",
|
|
self.id,
|
|
crate::vhost::canonical_ip(peer),
|
|
member.definition.id
|
|
)
|
|
.as_bytes(),
|
|
);
|
|
let n = u64::from_be_bytes([
|
|
digest[0], digest[1], digest[2], digest[3], digest[4], digest[5],
|
|
digest[6], digest[7],
|
|
]);
|
|
// Weighted rendezvous: stable across requests, no client-controlled forwarding headers.
|
|
let unit = (n as f64 + 1.0) / (u64::MAX as f64 + 2.0);
|
|
-unit.max(f64::MIN_POSITIVE).ln() / member.definition.weight as f64
|
|
}
|
|
Algorithm::LeastTraffic => {
|
|
(member.state.traffic_rate() + (active + 1.0) * 1024.0) / weights[index]
|
|
}
|
|
Algorithm::Adaptive => {
|
|
((active + 1.0) * member.state.latency()
|
|
+ member.state.traffic_rate() / 16_384.0)
|
|
/ weights[index]
|
|
}
|
|
_ => (active + 1.0) / weights[index],
|
|
};
|
|
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: self.settings.health_check.clone(),
|
|
passive_health: self.settings.passive_health.clone(),
|
|
error_recorded: AtomicBool::new(false),
|
|
}),
|
|
cookie,
|
|
})
|
|
}
|
|
}
|