133 lines
4.6 KiB
Rust
133 lines
4.6 KiB
Rust
impl Pool {
|
|
pub fn select(&self, peer: IpAddr, headers: &HeaderMap, tls: bool) -> Option<Selection> {
|
|
self.select_excluding(peer, headers, tls, &[])
|
|
}
|
|
pub fn select_retry(
|
|
&self,
|
|
peer: IpAddr,
|
|
headers: &HeaderMap,
|
|
tls: bool,
|
|
excluded: &[String],
|
|
) -> Option<Selection> {
|
|
self.select_excluding(peer, headers, tls, excluded)
|
|
}
|
|
fn eligible_healthy_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()
|
|
&& !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<u64> {
|
|
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<Selection, SelectionFailure> {
|
|
self.select_retry_wait(peer, headers, tls, &[]).await
|
|
}
|
|
pub async fn select_member_wait(&self, id: &str) -> Result<Selection, SelectionFailure> {
|
|
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<Selection, SelectionFailure> {
|
|
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;
|
|
}
|
|
}
|
|
}
|