use dashmap::DashMap;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex as StdMutex};
use tokio::sync::{broadcast, watch, Mutex, RwLock};
use crate::ilink::{QrLoginUiEvent, UpstreamSink};
use crate::store::Store;
use super::*;
pub const MAX_CONCURRENT_POLLS_PER_VTOKEN: usize = 3;
pub const MAX_HUB_POLLS_DEFAULT: usize = 8192;
pub const HISTOGRAM_BUCKETS_MS: &[u64] = &[1, 5, 25, 100, 500, 2_500, 10_000];
pub const HISTOGRAM_BUCKET_COUNT: usize = HISTOGRAM_BUCKETS_MS.len() + 1;
#[derive(Debug)]
pub struct LatencyHistogram {
pub count: AtomicU64,
pub sum_us: AtomicU64,
pub buckets: [AtomicU64; HISTOGRAM_BUCKET_COUNT],
}
impl Default for LatencyHistogram {
fn default() -> Self {
Self::new()
}
}
impl LatencyHistogram {
pub fn new() -> Self {
Self {
count: AtomicU64::new(0),
sum_us: AtomicU64::new(0),
buckets: std::array::from_fn(|_| AtomicU64::new(0)),
}
}
pub fn observe(&self, elapsed: std::time::Duration) {
let us = u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX);
let ms = elapsed.as_millis() as u64;
self.count.fetch_add(1, Ordering::Relaxed);
self.sum_us.fetch_add(us, Ordering::Relaxed);
for (i, boundary) in HISTOGRAM_BUCKETS_MS.iter().enumerate() {
if ms <= *boundary {
self.buckets[i].fetch_add(1, Ordering::Relaxed);
return;
}
}
self.buckets[HISTOGRAM_BUCKETS_MS.len()].fetch_add(1, Ordering::Relaxed);
}
}
pub struct Metrics {
pub messages_dispatched: AtomicU64,
pub messages_dropped: AtomicU64,
pub upstream_user_messages: AtomicU64,
pub sendmessage_total: AtomicU64,
pub sendmessage_errors: AtomicU64,
pub relogin_attempts: AtomicU64,
pub messages_persist_dropped: AtomicU64,
pub getupdates_latency_ms: LatencyHistogram,
pub sendmessage_upstream_latency_ms: LatencyHistogram,
pub dispatch_latency_ms: LatencyHistogram,
pub process_start_unix_secs: f64,
pub dispatcher_lagged: AtomicU64,
pub persist_fire_and_forget_failures_forward: AtomicU64,
pub persist_fire_and_forget_failures_broadcast: AtomicU64,
}
impl Metrics {
pub fn new() -> Self {
let process_start_unix_secs = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0);
Self {
messages_dispatched: AtomicU64::new(0),
messages_dropped: AtomicU64::new(0),
upstream_user_messages: AtomicU64::new(0),
sendmessage_total: AtomicU64::new(0),
sendmessage_errors: AtomicU64::new(0),
relogin_attempts: AtomicU64::new(0),
messages_persist_dropped: AtomicU64::new(0),
getupdates_latency_ms: LatencyHistogram::new(),
sendmessage_upstream_latency_ms: LatencyHistogram::new(),
dispatch_latency_ms: LatencyHistogram::new(),
process_start_unix_secs,
dispatcher_lagged: AtomicU64::new(0),
persist_fire_and_forget_failures_forward: AtomicU64::new(0),
persist_fire_and_forget_failures_broadcast: AtomicU64::new(0),
}
}
}
impl Default for Metrics {
fn default() -> Self {
Self::new()
}
}
pub struct LatencyGuard<'a> {
start: std::time::Instant,
histogram: &'a LatencyHistogram,
}
impl<'a> LatencyGuard<'a> {
pub fn new(histogram: &'a LatencyHistogram) -> Self {
Self {
start: std::time::Instant::now(),
histogram,
}
}
}
impl Drop for LatencyGuard<'_> {
fn drop(&mut self) {
self.histogram.observe(self.start.elapsed());
}
}
#[derive(Debug, Default)]
pub struct PollTracker {
pub counts: StdMutex<HashMap<String, usize>>,
total: AtomicUsize,
hub_cap: AtomicUsize,
}
impl PollTracker {
pub fn set_hub_cap(&self, cap: usize) {
self.hub_cap.store(cap, Ordering::Relaxed);
}
pub fn total_polls(&self) -> usize {
self.total.load(Ordering::Relaxed)
}
pub fn enter(self: &Arc<Self>, vtoken: &str) -> EnterOutcome {
let cap = self.hub_cap.load(Ordering::Relaxed);
let prev_total = self.total.fetch_add(1, Ordering::AcqRel);
if prev_total >= cap {
self.total.fetch_sub(1, Ordering::AcqRel);
return EnterOutcome::HubLimitReached {
total: prev_total,
cap,
};
}
let count = {
let Ok(mut counts) = self.counts.lock() else {
return EnterOutcome::Poisoned {
guard: PollGuard {
tracker: Arc::clone(self),
vtoken: vtoken.to_string(),
},
};
};
let c = counts.entry(vtoken.to_string()).or_insert(0);
*c += 1;
*c
};
EnterOutcome::Ok {
per_vtoken: count,
guard: PollGuard {
tracker: Arc::clone(self),
vtoken: vtoken.to_string(),
},
}
}
}
#[derive(Debug)]
pub enum EnterOutcome {
Ok { per_vtoken: usize, guard: PollGuard },
HubLimitReached { total: usize, cap: usize },
Poisoned { guard: PollGuard },
}
impl EnterOutcome {
#[allow(dead_code)]
pub fn guard(self) -> Option<PollGuard> {
match self {
EnterOutcome::Ok { guard, .. } | EnterOutcome::Poisoned { guard } => Some(guard),
EnterOutcome::HubLimitReached { .. } => None,
}
}
}
#[derive(Debug)]
pub struct PollGuard {
tracker: Arc<PollTracker>,
vtoken: String,
}
impl Drop for PollGuard {
fn drop(&mut self) {
self.tracker.total.fetch_sub(1, Ordering::AcqRel);
let Ok(mut counts) = self.tracker.counts.lock() else {
return;
};
if let Some(c) = counts.get_mut(&self.vtoken) {
*c = c.saturating_sub(1);
if *c == 0 {
counts.remove(&self.vtoken);
}
}
}
}
pub struct IlinkConnState {
pub upstream: Arc<dyn UpstreamSink>,
pub shutdown: watch::Receiver<bool>,
pub ilink_status: Arc<AtomicU8>,
pub qr_tx: broadcast::Sender<QrLoginUiEvent>,
pub qr_last_ready: Arc<std::sync::Mutex<Option<QrLoginUiEvent>>>,
pub relogin_tx: broadcast::Sender<()>,
pub qr_ticket: crate::server::sse_ticket::SseTicketStore,
}
impl IlinkConnState {
pub(crate) fn new(upstream: Arc<dyn UpstreamSink>, shutdown: watch::Receiver<bool>) -> Self {
let (qr_tx, _) = broadcast::channel(16);
let (relogin_tx, _) = broadcast::channel(4);
Self {
upstream,
shutdown,
ilink_status: Arc::new(AtomicU8::new(ilink_status::UNKNOWN)),
qr_tx,
qr_last_ready: Arc::new(std::sync::Mutex::new(None)),
relogin_tx,
qr_ticket: crate::server::sse_ticket::SseTicketStore::new(),
}
}
}
pub struct RoutingState {
pub router: Mutex<Router>,
pub quote_index: Mutex<QuoteRouteIndex>,
}
impl RoutingState {
pub(crate) fn new() -> Self {
Self {
router: Mutex::new(Router::new(None)),
quote_index: Mutex::new(QuoteRouteIndex::default()),
}
}
}
pub struct ClientState {
pub registry: RwLock<ClientRegistry>,
pub pairing: RwLock<PairingRegistry>,
pub pairing_notify: Arc<tokio::sync::Notify>,
pub queue: Arc<dyn MessageQueue>,
pub poll_tracker: Arc<PollTracker>,
pub last_seen: Arc<DashMap<String, AtomicU64>>,
}
impl ClientState {
pub(crate) fn new(queue: Arc<dyn MessageQueue>) -> Self {
let poll_tracker = Arc::new(PollTracker::default());
poll_tracker.set_hub_cap(MAX_HUB_POLLS_DEFAULT);
Self {
registry: RwLock::new(ClientRegistry::new()),
pairing: RwLock::new(PairingRegistry::new()),
pairing_notify: Arc::new(tokio::sync::Notify::new()),
queue,
poll_tracker,
last_seen: Arc::new(DashMap::new()),
}
}
}
const MAX_CONCURRENT_PERSIST_TASKS: usize = 32;
#[derive(Debug, Clone)]
pub struct AdminConfig {
pub token: Option<String>,
pub insecure_no_auth: bool,
pub outbound_origin_label: Option<String>,
}
impl AdminConfig {
pub fn from_env() -> Self {
let token = std::env::var("ILINK_ADMIN_TOKEN")
.ok()
.filter(|s| !s.is_empty());
let insecure_no_auth = token.is_none()
&& std::env::var("ILINK_ADMIN_INSECURE_NO_AUTH")
.ok()
.map(|v| matches!(v.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes"))
.unwrap_or(false);
let outbound_origin_label = std::env::var("ILINKHUB_OUTBOUND_ORIGIN_LABEL").ok();
Self {
token,
insecure_no_auth,
outbound_origin_label,
}
}
}
pub struct HubState {
pub ilink: IlinkConnState,
pub routing: RoutingState,
pub clients: ClientState,
pub store: Arc<Store>,
pub metrics: Arc<Metrics>,
pub persist_sem: Arc<tokio::sync::Semaphore>,
pub relay_secret: String,
pub admin: AdminConfig,
pub quote_index_warmed: Arc<AtomicBool>,
}
impl HubState {
pub fn new(
upstream: Arc<dyn UpstreamSink>,
store: Arc<Store>,
queue: Arc<dyn MessageQueue>,
shutdown: watch::Receiver<bool>,
relay_secret: String,
admin: AdminConfig,
) -> Arc<Self> {
Arc::new(Self {
ilink: IlinkConnState::new(upstream, shutdown),
routing: RoutingState::new(),
clients: ClientState::new(queue),
store,
metrics: Arc::new(Metrics::new()),
persist_sem: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_PERSIST_TASKS)),
relay_secret,
admin,
quote_index_warmed: Arc::new(AtomicBool::new(false)),
})
}
pub async fn with_router_and_registry<R>(
self: &Arc<Self>,
f: impl FnOnce(&mut Router, &ClientRegistry) -> R,
) -> Option<R> {
let mut router_guard = self.routing.router.lock().await;
let registry_guard = self.clients.registry.read().await;
let result = f(&mut router_guard, ®istry_guard);
drop(registry_guard);
drop(router_guard);
Some(result)
}
}