v0.25.1

src/balance/pool_selection.rs

@@ -0,0 +1,132 @@
+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;
+        }
+    }
+}