ME Pool Runtime API

src/transport/middle_proxy/pool_runtime_api.rs

@@ -0,0 +1,128 @@
+use std::collections::HashMap;
+use std::time::Instant;
+
+use super::pool::{MePool, RefillDcKey};
+use crate::network::IpFamily;
+
+#[derive(Clone, Debug)]
+pub(crate) struct MeApiRefillDcSnapshot {
+    pub dc: i16,
+    pub family: &'static str,
+    pub inflight: usize,
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct MeApiRefillSnapshot {
+    pub inflight_endpoints_total: usize,
+    pub inflight_dc_total: usize,
+    pub by_dc: Vec<MeApiRefillDcSnapshot>,
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct MeApiNatReflectionSnapshot {
+    pub addr: std::net::SocketAddr,
+    pub age_secs: u64,
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct MeApiNatStunSnapshot {
+    pub nat_probe_enabled: bool,
+    pub nat_probe_disabled_runtime: bool,
+    pub nat_probe_attempts: u8,
+    pub configured_servers: Vec<String>,
+    pub live_servers: Vec<String>,
+    pub reflection_v4: Option<MeApiNatReflectionSnapshot>,
+    pub reflection_v6: Option<MeApiNatReflectionSnapshot>,
+    pub stun_backoff_remaining_ms: Option<u64>,
+}
+
+impl MePool {
+    pub(crate) async fn api_refill_snapshot(&self) -> MeApiRefillSnapshot {
+        let inflight_endpoints_total = self.refill_inflight.lock().await.len();
+        let inflight_dc_keys = self
+            .refill_inflight_dc
+            .lock()
+            .await
+            .iter()
+            .copied()
+            .collect::<Vec<RefillDcKey>>();
+
+        let mut by_dc_map = HashMap::<(i16, &'static str), usize>::new();
+        for key in inflight_dc_keys {
+            let family = match key.family {
+                IpFamily::V4 => "v4",
+                IpFamily::V6 => "v6",
+            };
+            let dc = key.dc as i16;
+            *by_dc_map.entry((dc, family)).or_insert(0) += 1;
+        }
+
+        let mut by_dc = by_dc_map
+            .into_iter()
+            .map(|((dc, family), inflight)| MeApiRefillDcSnapshot {
+                dc,
+                family,
+                inflight,
+            })
+            .collect::<Vec<_>>();
+        by_dc.sort_by_key(|entry| (entry.dc, entry.family));
+
+        MeApiRefillSnapshot {
+            inflight_endpoints_total,
+            inflight_dc_total: by_dc.len(),
+            by_dc,
+        }
+    }
+
+    pub(crate) async fn api_nat_stun_snapshot(&self) -> MeApiNatStunSnapshot {
+        let now = Instant::now();
+        let mut configured_servers = if !self.nat_stun_servers.is_empty() {
+            self.nat_stun_servers.clone()
+        } else if let Some(stun) = &self.nat_stun {
+            if stun.trim().is_empty() {
+                Vec::new()
+            } else {
+                vec![stun.clone()]
+            }
+        } else {
+            Vec::new()
+        };
+        configured_servers.sort();
+        configured_servers.dedup();
+
+        let mut live_servers = self.nat_stun_live_servers.read().await.clone();
+        live_servers.sort();
+        live_servers.dedup();
+
+        let reflection = self.nat_reflection_cache.lock().await;
+        let reflection_v4 = reflection.v4.map(|(ts, addr)| MeApiNatReflectionSnapshot {
+            addr,
+            age_secs: now.saturating_duration_since(ts).as_secs(),
+        });
+        let reflection_v6 = reflection.v6.map(|(ts, addr)| MeApiNatReflectionSnapshot {
+            addr,
+            age_secs: now.saturating_duration_since(ts).as_secs(),
+        });
+        drop(reflection);
+
+        let backoff_until = *self.stun_backoff_until.read().await;
+        let stun_backoff_remaining_ms = backoff_until.and_then(|until| {
+            (until > now).then_some(until.duration_since(now).as_millis() as u64)
+        });
+
+        MeApiNatStunSnapshot {
+            nat_probe_enabled: self.nat_probe,
+            nat_probe_disabled_runtime: self
+                .nat_probe_disabled
+                .load(std::sync::atomic::Ordering::Relaxed),
+            nat_probe_attempts: self
+                .nat_probe_attempts
+                .load(std::sync::atomic::Ordering::Relaxed),
+            configured_servers,
+            live_servers,
+            reflection_v4,
+            reflection_v6,
+            stun_backoff_remaining_ms,
+        }
+    }
+}