375 lines
14 KiB
Rust
375 lines
14 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)| {
|
|
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<Member>> {
|
|
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<HeaderValue> = 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::<Vec<_>>()
|
|
.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<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: 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<Selection> {
|
|
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,
|
|
})
|
|
}
|
|
}
|