v0.25.1
This commit is contained in:
Mateusz Gruszczyński committed
1 parent
87e571c170
commit
fa7abc9e53
308 files changed
+49746
-50237
No files matched your search
@@ -0,0 +1,334 @@
|
||||
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,
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user