mirror of
https://github.com/telemt/telemt.git
synced 2026-09-13 22:14:08 +03:00
WEB: websocket + websocket-lanes as Carrier
This commit is contained in:
Alexey
@@ -0,0 +1,313 @@
|
||||
use std::collections::HashMap;
|
||||
use std::net::IpAddr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::OwnedSemaphorePermit;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::{ManagerError, ProfileKey, WebProcessRuntime, WebSocketBudgetLease};
|
||||
|
||||
/// One process-owned WebSocket carrier class used for eviction priority.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum WebSocketKind {
|
||||
/// One connection multiplexes every logical stream in a session.
|
||||
Multiplex,
|
||||
/// One connection owns exactly one logical stream lane.
|
||||
Lane(u32),
|
||||
}
|
||||
|
||||
pub(super) struct WebSocketEntry {
|
||||
id: u64,
|
||||
owner: ProfileKey,
|
||||
session_id: u64,
|
||||
client_ip: IpAddr,
|
||||
kind: WebSocketKind,
|
||||
liveness_interval_ms: u64,
|
||||
created_tick: u64,
|
||||
last_peer_tick: AtomicU64,
|
||||
last_progress_tick: AtomicU64,
|
||||
opened: AtomicBool,
|
||||
cancel: CancellationToken,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct WebSocketRegistry {
|
||||
entries: HashMap<u64, Arc<WebSocketEntry>>,
|
||||
closed: bool,
|
||||
}
|
||||
|
||||
/// Exact process-owned admission retained through the upgraded socket lifetime.
|
||||
pub(crate) struct WebSocketConnection {
|
||||
runtime: std::sync::Weak<WebProcessRuntime>,
|
||||
entry: Arc<WebSocketEntry>,
|
||||
slot: Option<OwnedSemaphorePermit>,
|
||||
base_budget: Option<WebSocketBudgetLease>,
|
||||
}
|
||||
|
||||
impl WebSocketConnection {
|
||||
/// Returns the cancellation signal used by shutdown and pressure eviction.
|
||||
pub(crate) fn cancellation(&self) -> CancellationToken {
|
||||
self.entry.cancel.clone()
|
||||
}
|
||||
|
||||
/// Returns the process-unique connection identifier used only for debugging.
|
||||
pub(crate) fn id(&self) -> u64 {
|
||||
self.entry.id
|
||||
}
|
||||
|
||||
/// Returns the creation-time transport liveness interval.
|
||||
pub(crate) fn liveness_interval(&self) -> Duration {
|
||||
Duration::from_millis(self.entry.liveness_interval_ms)
|
||||
}
|
||||
|
||||
/// Marks successful ownership transfer from HTTP to the WebSocket codec.
|
||||
pub(crate) fn mark_opened(&self) {
|
||||
self.entry.opened.store(true, Ordering::Release);
|
||||
self.mark_progress();
|
||||
}
|
||||
|
||||
/// Refreshes the peer-liveness deadline after any received WebSocket message.
|
||||
pub(crate) fn mark_peer_activity(&self) {
|
||||
if let Some(runtime) = self.runtime.upgrade() {
|
||||
let now = runtime.websocket_tick();
|
||||
self.entry.last_peer_tick.store(now, Ordering::Release);
|
||||
self.entry.last_progress_tick.store(now, Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
||||
/// Refreshes least-recently-progressed ordering after a committed write.
|
||||
pub(crate) fn mark_progress(&self) {
|
||||
if let Some(runtime) = self.runtime.upgrade() {
|
||||
self.entry
|
||||
.last_progress_tick
|
||||
.store(runtime.websocket_tick(), Ordering::Release);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WebSocketConnection {
|
||||
fn drop(&mut self) {
|
||||
if let Some(runtime) = self.runtime.upgrade() {
|
||||
runtime.websockets.lock().entries.remove(&self.entry.id);
|
||||
drop(self.base_budget.take());
|
||||
drop(self.slot.take());
|
||||
runtime.websocket_notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn admit(
|
||||
runtime: &Arc<WebProcessRuntime>,
|
||||
owner: ProfileKey,
|
||||
session_id: u64,
|
||||
client_ip: IpAddr,
|
||||
kind: WebSocketKind,
|
||||
base_bytes: usize,
|
||||
liveness_interval: Duration,
|
||||
eviction_timeout: Duration,
|
||||
) -> Result<WebSocketConnection, ManagerError> {
|
||||
let liveness_interval_ms = liveness_interval.as_millis().min(u128::from(u64::MAX)) as u64;
|
||||
if let Some(connection) = try_admit(
|
||||
runtime,
|
||||
owner,
|
||||
session_id,
|
||||
client_ip,
|
||||
kind,
|
||||
base_bytes,
|
||||
liveness_interval_ms,
|
||||
) {
|
||||
return Ok(connection);
|
||||
}
|
||||
let Some(victim) = select_victim(runtime, owner, session_id, client_ip, None) else {
|
||||
runtime.record_limit_hit();
|
||||
return Err(ManagerError::Limit);
|
||||
};
|
||||
let released = runtime.websocket_notify.notified();
|
||||
victim.cancel.cancel();
|
||||
let _ = tokio::time::timeout(eviction_timeout, released).await;
|
||||
try_admit(
|
||||
runtime,
|
||||
owner,
|
||||
session_id,
|
||||
client_ip,
|
||||
kind,
|
||||
base_bytes,
|
||||
liveness_interval_ms,
|
||||
)
|
||||
.ok_or_else(|| {
|
||||
runtime.record_limit_hit();
|
||||
ManagerError::Limit
|
||||
})
|
||||
}
|
||||
|
||||
fn try_admit(
|
||||
runtime: &Arc<WebProcessRuntime>,
|
||||
owner: ProfileKey,
|
||||
session_id: u64,
|
||||
client_ip: IpAddr,
|
||||
kind: WebSocketKind,
|
||||
base_bytes: usize,
|
||||
liveness_interval_ms: u64,
|
||||
) -> Option<WebSocketConnection> {
|
||||
let slot = Arc::clone(&runtime.websocket_connections)
|
||||
.try_acquire_owned()
|
||||
.ok()?;
|
||||
let base_budget = runtime.try_websocket_base_budget(owner, base_bytes)?;
|
||||
let id = runtime.websocket_next_id.fetch_add(1, Ordering::Relaxed);
|
||||
let now = runtime.websocket_tick();
|
||||
let entry = Arc::new(WebSocketEntry {
|
||||
id,
|
||||
owner,
|
||||
session_id,
|
||||
client_ip,
|
||||
kind,
|
||||
liveness_interval_ms,
|
||||
created_tick: now,
|
||||
last_peer_tick: AtomicU64::new(now),
|
||||
last_progress_tick: AtomicU64::new(now),
|
||||
opened: AtomicBool::new(false),
|
||||
cancel: CancellationToken::new(),
|
||||
});
|
||||
let mut registry = runtime.websockets.lock();
|
||||
if registry.closed {
|
||||
return None;
|
||||
}
|
||||
registry.entries.insert(id, Arc::clone(&entry));
|
||||
drop(registry);
|
||||
Some(WebSocketConnection {
|
||||
runtime: Arc::downgrade(runtime),
|
||||
entry,
|
||||
slot: Some(slot),
|
||||
base_budget: Some(base_budget),
|
||||
})
|
||||
}
|
||||
|
||||
impl WebProcessRuntime {
|
||||
pub(super) fn websocket_tick(&self) -> u64 {
|
||||
self.websocket_clock.elapsed().as_millis() as u64
|
||||
}
|
||||
|
||||
pub(super) fn cleanup_websockets(&self) {
|
||||
let now = self.websocket_tick();
|
||||
let mut victims = self
|
||||
.websockets
|
||||
.lock()
|
||||
.entries
|
||||
.values()
|
||||
.filter(|entry| {
|
||||
now.saturating_sub(entry.last_peer_tick.load(Ordering::Acquire))
|
||||
>= dead_after(entry)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
if victims.is_empty()
|
||||
&& self.data_budget.take_pressure()
|
||||
&& let Some(victim) = select_pressure_victim(self, now)
|
||||
{
|
||||
victims.push(victim);
|
||||
}
|
||||
for victim in victims {
|
||||
victim.cancel.cancel();
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn close_websockets(&self) {
|
||||
let victims = {
|
||||
let mut registry = self.websockets.lock();
|
||||
registry.closed = true;
|
||||
registry.entries.values().cloned().collect::<Vec<_>>()
|
||||
};
|
||||
for victim in victims {
|
||||
victim.cancel.cancel();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn select_victim(
|
||||
runtime: &WebProcessRuntime,
|
||||
owner: ProfileKey,
|
||||
session_id: u64,
|
||||
client_ip: IpAddr,
|
||||
excluded_id: Option<u64>,
|
||||
) -> Option<Arc<WebSocketEntry>> {
|
||||
let fair_share = runtime.data_budget.fair_share(Some(owner));
|
||||
let requester_usage = runtime.data_budget.owner_usage(owner);
|
||||
let now = runtime.websocket_tick();
|
||||
runtime
|
||||
.websockets
|
||||
.lock()
|
||||
.entries
|
||||
.values()
|
||||
.filter(|entry| Some(entry.id) != excluded_id)
|
||||
.filter_map(|entry| {
|
||||
let owner_rank = if entry.session_id == session_id {
|
||||
0
|
||||
} else if entry.owner == owner {
|
||||
1
|
||||
} else if entry.client_ip == client_ip {
|
||||
2
|
||||
} else {
|
||||
if requester_usage >= fair_share
|
||||
|| runtime.data_budget.owner_usage(entry.owner) <= fair_share
|
||||
{
|
||||
return None;
|
||||
}
|
||||
3
|
||||
};
|
||||
let priority = entry_priority(entry, now);
|
||||
Some((
|
||||
(
|
||||
owner_rank,
|
||||
priority,
|
||||
entry.last_progress_tick.load(Ordering::Acquire),
|
||||
entry.created_tick,
|
||||
entry.id,
|
||||
),
|
||||
Arc::clone(entry),
|
||||
))
|
||||
})
|
||||
.min_by_key(|(key, _)| *key)
|
||||
.map(|(_, entry)| entry)
|
||||
}
|
||||
|
||||
fn select_pressure_victim(runtime: &WebProcessRuntime, now: u64) -> Option<Arc<WebSocketEntry>> {
|
||||
runtime
|
||||
.websockets
|
||||
.lock()
|
||||
.entries
|
||||
.values()
|
||||
.map(|entry| {
|
||||
(
|
||||
(
|
||||
entry_priority(entry, now),
|
||||
entry.last_progress_tick.load(Ordering::Acquire),
|
||||
entry.created_tick,
|
||||
entry.id,
|
||||
),
|
||||
Arc::clone(entry),
|
||||
)
|
||||
})
|
||||
.min_by_key(|(key, _)| *key)
|
||||
.map(|(_, entry)| entry)
|
||||
}
|
||||
|
||||
fn entry_priority(entry: &WebSocketEntry, now: u64) -> u8 {
|
||||
if !entry.opened.load(Ordering::Acquire)
|
||||
|| now.saturating_sub(entry.last_peer_tick.load(Ordering::Acquire)) >= dead_after(entry)
|
||||
{
|
||||
0
|
||||
} else if matches!(entry.kind, WebSocketKind::Lane(_)) {
|
||||
1
|
||||
} else {
|
||||
2
|
||||
}
|
||||
}
|
||||
|
||||
fn dead_after(entry: &WebSocketEntry) -> u64 {
|
||||
entry.liveness_interval_ms.saturating_mul(2)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
Reference in New Issue
Block a user