mirror of
https://github.com/telemt/telemt.git
synced 2026-09-11 13:04:08 +03:00
WEB: websocket + websocket-lanes as Carrier
This commit is contained in:
@@ -84,7 +84,7 @@ mod tests {
|
||||
async fn global_downlink_budget_preserves_one_maximum_uplink_batch() {
|
||||
let generation = test_runtime_generation(1, ProxyConfig::default());
|
||||
let runtime = WebProcessRuntime::start(Arc::new(ArcSwap::from(generation)));
|
||||
let control_items = super::super::state::control_item_reserve(&runtime.limits);
|
||||
let control_items = super::super::budget::control_item_reserve(&runtime.limits);
|
||||
let data_bytes = runtime
|
||||
.limits
|
||||
.pending_bytes_global
|
||||
@@ -97,20 +97,30 @@ mod tests {
|
||||
.limits
|
||||
.max_body_bytes
|
||||
.saturating_add(runtime.limits.max_frames_per_body * QUEUE_ITEM_COST);
|
||||
let downlink_bytes = data_bytes - uplink_bytes;
|
||||
let websocket_bytes = runtime.limits.carrier_batch_bytes;
|
||||
let downlink_bytes = data_bytes - uplink_bytes - websocket_bytes;
|
||||
let downlink_items = data_items - runtime.limits.max_frames_per_body;
|
||||
|
||||
assert!(runtime.try_reserve_pending(downlink_bytes, downlink_items, false, true,));
|
||||
assert!(runtime.try_reserve_pending([0; 32], downlink_bytes, downlink_items, false, true,));
|
||||
assert!(runtime.try_reserve_pending(
|
||||
[0; 32],
|
||||
uplink_bytes,
|
||||
runtime.limits.max_frames_per_body,
|
||||
false,
|
||||
false,
|
||||
));
|
||||
assert!(!runtime.try_reserve_pending(1, 1, false, true));
|
||||
let websocket = runtime.try_websocket_data_budget([0; 32], websocket_bytes);
|
||||
assert!(websocket.is_some());
|
||||
assert!(!runtime.try_reserve_pending([0; 32], 1, 1, false, true));
|
||||
|
||||
runtime.release_pending(downlink_bytes, downlink_items, false);
|
||||
runtime.release_pending(uplink_bytes, runtime.limits.max_frames_per_body, false);
|
||||
drop(websocket);
|
||||
runtime.release_pending([0; 32], downlink_bytes, downlink_items, false);
|
||||
runtime.release_pending(
|
||||
[0; 32],
|
||||
uplink_bytes,
|
||||
runtime.limits.max_frames_per_body,
|
||||
false,
|
||||
);
|
||||
runtime.shutdown().await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,366 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use super::ProfileKey;
|
||||
use crate::config::WebLimitsConfig;
|
||||
use crate::web::session::QUEUE_ITEM_COST;
|
||||
|
||||
/// WebSocket allocation class with a distinct pressure watermark.
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) enum WebSocketBudgetClass {
|
||||
/// Long-lived codec and driver memory acquired before an upgrade commits.
|
||||
Base,
|
||||
/// One bounded inbound message or outbound write staging allocation.
|
||||
Data,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct BudgetState {
|
||||
queue_bytes: usize,
|
||||
queue_items: usize,
|
||||
queue_control_bytes: usize,
|
||||
queue_control_items: usize,
|
||||
websocket_bytes: usize,
|
||||
owner_bytes: HashMap<ProfileKey, usize>,
|
||||
high_water_bytes: usize,
|
||||
closed: bool,
|
||||
}
|
||||
|
||||
/// Process-owned byte and item governor shared by queues and WebSocket I/O.
|
||||
pub(super) struct WebDataBudget {
|
||||
limits: WebLimitsConfig,
|
||||
state: Mutex<BudgetState>,
|
||||
notify: Arc<Notify>,
|
||||
pressured: AtomicBool,
|
||||
}
|
||||
|
||||
/// One exact WebSocket allocation released on every cancellation path.
|
||||
pub(crate) struct WebSocketBudgetLease {
|
||||
budget: Arc<WebDataBudget>,
|
||||
owner: ProfileKey,
|
||||
bytes: usize,
|
||||
}
|
||||
|
||||
/// Lock-free diagnostic snapshot of one short locked budget state.
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub(crate) struct WebDataBudgetSnapshot {
|
||||
/// Total queue bytes currently retained.
|
||||
pub(crate) queue_bytes: usize,
|
||||
/// Total queue items currently retained.
|
||||
pub(crate) queue_items: usize,
|
||||
/// Total WebSocket bytes currently retained.
|
||||
pub(crate) websocket_bytes: usize,
|
||||
/// Largest combined byte usage observed since process start.
|
||||
pub(crate) high_water_bytes: usize,
|
||||
}
|
||||
|
||||
impl WebDataBudget {
|
||||
pub(super) fn new(limits: WebLimitsConfig) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
limits,
|
||||
state: Mutex::new(BudgetState::default()),
|
||||
notify: Arc::new(Notify::new()),
|
||||
pressured: AtomicBool::new(false),
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn try_reserve_queue(
|
||||
&self,
|
||||
owner: ProfileKey,
|
||||
bytes: usize,
|
||||
items: usize,
|
||||
control: bool,
|
||||
downlink: bool,
|
||||
) -> bool {
|
||||
let mut state = self.state.lock();
|
||||
if state.closed {
|
||||
return false;
|
||||
}
|
||||
let control_item_reserve = control_item_reserve(&self.limits);
|
||||
let data_byte_limit = self
|
||||
.limits
|
||||
.pending_bytes_global
|
||||
.saturating_sub(self.limits.control_bytes_global);
|
||||
let data_item_limit = self
|
||||
.limits
|
||||
.pending_items_global
|
||||
.saturating_sub(control_item_reserve);
|
||||
let (fits, websocket_byte_pressure) = if control {
|
||||
let byte_pressure = state.websocket_bytes != 0
|
||||
&& state.queue_bytes.saturating_add(state.websocket_bytes)
|
||||
> self.limits.pending_bytes_global.saturating_sub(bytes);
|
||||
let fits = bytes <= self.limits.control_bytes_global
|
||||
&& items <= control_item_reserve
|
||||
&& state.queue_bytes.saturating_add(state.websocket_bytes)
|
||||
<= self.limits.pending_bytes_global.saturating_sub(bytes)
|
||||
&& state.queue_items <= self.limits.pending_items_global.saturating_sub(items)
|
||||
&& state.queue_control_bytes
|
||||
<= self.limits.control_bytes_global.saturating_sub(bytes)
|
||||
&& state.queue_control_items <= control_item_reserve.saturating_sub(items);
|
||||
(fits, byte_pressure)
|
||||
} else {
|
||||
let queue_data_bytes = state.queue_bytes.saturating_sub(state.queue_control_bytes);
|
||||
let queue_data_items = state.queue_items.saturating_sub(state.queue_control_items);
|
||||
let (byte_limit, item_limit) = if downlink {
|
||||
let uplink_bytes = self.limits.max_body_bytes.saturating_add(
|
||||
self.limits
|
||||
.max_frames_per_body
|
||||
.saturating_mul(QUEUE_ITEM_COST),
|
||||
);
|
||||
(
|
||||
data_byte_limit
|
||||
.saturating_sub(uplink_bytes)
|
||||
.saturating_sub(self.limits.carrier_batch_bytes),
|
||||
data_item_limit.saturating_sub(self.limits.max_frames_per_body),
|
||||
)
|
||||
} else {
|
||||
(data_byte_limit, data_item_limit)
|
||||
};
|
||||
let byte_pressure = state.websocket_bytes != 0
|
||||
&& queue_data_bytes.saturating_add(state.websocket_bytes)
|
||||
> byte_limit.saturating_sub(bytes);
|
||||
let fits = bytes <= byte_limit
|
||||
&& items <= item_limit
|
||||
&& queue_data_bytes.saturating_add(state.websocket_bytes)
|
||||
<= byte_limit.saturating_sub(bytes)
|
||||
&& queue_data_items <= item_limit.saturating_sub(items);
|
||||
(fits, byte_pressure)
|
||||
};
|
||||
if !fits {
|
||||
if websocket_byte_pressure {
|
||||
self.pressured.store(true, Ordering::Release);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
state.queue_bytes += bytes;
|
||||
state.queue_items += items;
|
||||
if control {
|
||||
state.queue_control_bytes += bytes;
|
||||
state.queue_control_items += items;
|
||||
}
|
||||
add_owner(&mut state.owner_bytes, owner, bytes);
|
||||
update_high_water(&mut state);
|
||||
true
|
||||
}
|
||||
|
||||
pub(super) fn release_queue(
|
||||
&self,
|
||||
owner: ProfileKey,
|
||||
bytes: usize,
|
||||
items: usize,
|
||||
control: bool,
|
||||
) {
|
||||
let mut state = self.state.lock();
|
||||
state.queue_bytes = state.queue_bytes.saturating_sub(bytes);
|
||||
state.queue_items = state.queue_items.saturating_sub(items);
|
||||
if control {
|
||||
state.queue_control_bytes = state.queue_control_bytes.saturating_sub(bytes);
|
||||
state.queue_control_items = state.queue_control_items.saturating_sub(items);
|
||||
}
|
||||
remove_owner(&mut state.owner_bytes, owner, bytes);
|
||||
drop(state);
|
||||
self.pressured.store(false, Ordering::Release);
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
|
||||
pub(super) fn try_reserve_websocket(
|
||||
self: &Arc<Self>,
|
||||
owner: ProfileKey,
|
||||
bytes: usize,
|
||||
class: WebSocketBudgetClass,
|
||||
) -> Option<WebSocketBudgetLease> {
|
||||
let mut state = self.state.lock();
|
||||
if state.closed || bytes == 0 {
|
||||
return None;
|
||||
}
|
||||
let websocket_limit = match class {
|
||||
WebSocketBudgetClass::Base => watermark(
|
||||
self.limits.websocket_bytes_global,
|
||||
self.limits.websocket_admission_watermark_pct,
|
||||
),
|
||||
WebSocketBudgetClass::Data => watermark(
|
||||
self.limits.websocket_bytes_global,
|
||||
self.limits.websocket_eviction_watermark_pct,
|
||||
),
|
||||
};
|
||||
let data_byte_limit = self
|
||||
.limits
|
||||
.pending_bytes_global
|
||||
.saturating_sub(self.limits.control_bytes_global);
|
||||
let queue_data_bytes = state.queue_bytes.saturating_sub(state.queue_control_bytes);
|
||||
if state.websocket_bytes > websocket_limit.saturating_sub(bytes)
|
||||
|| queue_data_bytes.saturating_add(state.websocket_bytes)
|
||||
> data_byte_limit.saturating_sub(bytes)
|
||||
{
|
||||
self.pressured.store(true, Ordering::Release);
|
||||
return None;
|
||||
}
|
||||
state.websocket_bytes += bytes;
|
||||
add_owner(&mut state.owner_bytes, owner, bytes);
|
||||
update_high_water(&mut state);
|
||||
Some(WebSocketBudgetLease {
|
||||
budget: Arc::clone(self),
|
||||
owner,
|
||||
bytes,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn notify(&self) -> Arc<Notify> {
|
||||
Arc::clone(&self.notify)
|
||||
}
|
||||
|
||||
pub(super) fn take_pressure(&self) -> bool {
|
||||
self.pressured.swap(false, Ordering::AcqRel)
|
||||
}
|
||||
|
||||
pub(super) fn owner_usage(&self, owner: ProfileKey) -> usize {
|
||||
self.state
|
||||
.lock()
|
||||
.owner_bytes
|
||||
.get(&owner)
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub(super) fn fair_share(&self, additional_owner: Option<ProfileKey>) -> usize {
|
||||
let state = self.state.lock();
|
||||
let mut owners = state.owner_bytes.len();
|
||||
if additional_owner.is_some_and(|owner| !state.owner_bytes.contains_key(&owner)) {
|
||||
owners += 1;
|
||||
}
|
||||
let admission = watermark(
|
||||
self.limits.websocket_bytes_global,
|
||||
self.limits.websocket_admission_watermark_pct,
|
||||
);
|
||||
admission / owners.max(1)
|
||||
}
|
||||
|
||||
pub(super) fn snapshot(&self) -> WebDataBudgetSnapshot {
|
||||
let state = self.state.lock();
|
||||
WebDataBudgetSnapshot {
|
||||
queue_bytes: state.queue_bytes,
|
||||
queue_items: state.queue_items,
|
||||
websocket_bytes: state.websocket_bytes,
|
||||
high_water_bytes: state.high_water_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn close(&self) {
|
||||
self.state.lock().closed = true;
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
|
||||
fn release_websocket(&self, owner: ProfileKey, bytes: usize) {
|
||||
let mut state = self.state.lock();
|
||||
state.websocket_bytes = state.websocket_bytes.saturating_sub(bytes);
|
||||
remove_owner(&mut state.owner_bytes, owner, bytes);
|
||||
drop(state);
|
||||
self.pressured.store(false, Ordering::Release);
|
||||
self.notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
impl WebSocketBudgetLease {
|
||||
/// Releases unused worst-case capacity after one message is assembled.
|
||||
pub(crate) fn shrink_to(&mut self, bytes: usize) {
|
||||
let bytes = bytes.min(self.bytes);
|
||||
let released = self.bytes - bytes;
|
||||
if released == 0 {
|
||||
return;
|
||||
}
|
||||
self.bytes = bytes;
|
||||
self.budget.release_websocket(self.owner, released);
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WebSocketBudgetLease {
|
||||
fn drop(&mut self) {
|
||||
self.budget.release_websocket(self.owner, self.bytes);
|
||||
}
|
||||
}
|
||||
|
||||
fn watermark(limit: usize, percentage: u8) -> usize {
|
||||
limit.saturating_mul(usize::from(percentage)) / 100
|
||||
}
|
||||
|
||||
pub(super) fn control_item_reserve(limits: &WebLimitsConfig) -> usize {
|
||||
limits
|
||||
.max_sessions_global
|
||||
.saturating_mul(16usize.saturating_add(limits.max_streams_per_session.saturating_mul(3)))
|
||||
}
|
||||
|
||||
fn add_owner(values: &mut HashMap<ProfileKey, usize>, owner: ProfileKey, bytes: usize) {
|
||||
*values.entry(owner).or_insert(0) += bytes;
|
||||
}
|
||||
|
||||
fn remove_owner(values: &mut HashMap<ProfileKey, usize>, owner: ProfileKey, bytes: usize) {
|
||||
let remove = if let Some(value) = values.get_mut(&owner) {
|
||||
*value = value.saturating_sub(bytes);
|
||||
*value == 0
|
||||
} else {
|
||||
false
|
||||
};
|
||||
if remove {
|
||||
values.remove(&owner);
|
||||
}
|
||||
}
|
||||
|
||||
fn update_high_water(state: &mut BudgetState) {
|
||||
state.high_water_bytes = state
|
||||
.high_water_bytes
|
||||
.max(state.queue_bytes.saturating_add(state.websocket_bytes));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn downlink_reservation_preserves_one_uplink_and_websocket_batch() {
|
||||
let limits = WebLimitsConfig::default();
|
||||
let uplink_bytes = limits
|
||||
.max_body_bytes
|
||||
.saturating_add(limits.max_frames_per_body.saturating_mul(QUEUE_ITEM_COST));
|
||||
let downlink_bytes = limits
|
||||
.pending_bytes_global
|
||||
.saturating_sub(limits.control_bytes_global)
|
||||
.saturating_sub(uplink_bytes)
|
||||
.saturating_sub(limits.carrier_batch_bytes);
|
||||
let budget = WebDataBudget::new(limits);
|
||||
|
||||
assert!(budget.try_reserve_queue([1; 32], downlink_bytes, 1, false, true));
|
||||
assert!(!budget.try_reserve_queue([1; 32], 1, 1, false, true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn item_limit_rejection_does_not_request_websocket_eviction() {
|
||||
let limits = WebLimitsConfig::default();
|
||||
let rejected_items = limits.pending_items_global.saturating_add(1);
|
||||
let budget = WebDataBudget::new(limits);
|
||||
let _websocket = budget
|
||||
.try_reserve_websocket([1; 32], 1, WebSocketBudgetClass::Data)
|
||||
.unwrap();
|
||||
|
||||
assert!(!budget.try_reserve_queue([2; 32], 1, rejected_items, false, false));
|
||||
assert!(!budget.take_pressure());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_byte_conflict_requests_pressure_eviction() {
|
||||
let limits = WebLimitsConfig::default();
|
||||
let data_bytes = limits
|
||||
.pending_bytes_global
|
||||
.saturating_sub(limits.control_bytes_global);
|
||||
let budget = WebDataBudget::new(limits);
|
||||
let _websocket = budget
|
||||
.try_reserve_websocket([1; 32], 1, WebSocketBudgetClass::Data)
|
||||
.unwrap();
|
||||
|
||||
assert!(!budget.try_reserve_queue([2; 32], data_bytes, 1, false, false));
|
||||
assert!(budget.take_pressure());
|
||||
}
|
||||
}
|
||||
@@ -69,6 +69,8 @@ impl WebProcessRuntime {
|
||||
/// Stops issuance, closes all sessions, and joins bounded child work.
|
||||
pub(crate) async fn shutdown(&self) {
|
||||
self.shutdown.cancel();
|
||||
self.close_websockets();
|
||||
self.data_budget.close();
|
||||
let sessions = {
|
||||
let mut state = self.state.lock();
|
||||
state.closed = true;
|
||||
@@ -94,15 +96,11 @@ impl WebProcessRuntime {
|
||||
let _ = tokio::time::timeout(Duration::from_secs(timeout_secs), waits).await;
|
||||
self.tasks.close();
|
||||
let _ = tokio::time::timeout(Duration::from_secs(timeout_secs), self.tasks.wait()).await;
|
||||
let (sessions_live, streams_live, pending_bytes, pending_items) = {
|
||||
let (sessions_live, streams_live) = {
|
||||
let state = self.state.lock();
|
||||
(
|
||||
state.sessions.len(),
|
||||
state.streams_live,
|
||||
state.pending_bytes,
|
||||
state.pending_items,
|
||||
)
|
||||
(state.sessions.len(), state.streams_live)
|
||||
};
|
||||
let budget = self.data_budget.snapshot();
|
||||
info!(
|
||||
target: "telemt::web",
|
||||
sessions_created = self.sessions_created.load(Ordering::Relaxed),
|
||||
@@ -111,8 +109,10 @@ impl WebProcessRuntime {
|
||||
streams_opened = self.streams_opened.load(Ordering::Relaxed),
|
||||
streams_rejected = self.streams_rejected.load(Ordering::Relaxed),
|
||||
streams_live,
|
||||
pending_bytes,
|
||||
pending_items,
|
||||
pending_bytes = budget.queue_bytes,
|
||||
pending_items = budget.queue_items,
|
||||
websocket_bytes = budget.websocket_bytes,
|
||||
data_high_water_bytes = budget.high_water_bytes,
|
||||
bytes_up = self.bytes_up.load(Ordering::Relaxed),
|
||||
bytes_down = self.bytes_down.load(Ordering::Relaxed),
|
||||
limit_hits = self.limit_hits.load(Ordering::Relaxed),
|
||||
@@ -122,6 +122,7 @@ impl WebProcessRuntime {
|
||||
|
||||
/// Expires credentials and closes idle sessions without holding locks across callbacks.
|
||||
pub(super) fn cleanup(&self) {
|
||||
self.cleanup_websockets();
|
||||
let now = Instant::now();
|
||||
let sessions = {
|
||||
let mut state = self.state.lock();
|
||||
|
||||
+15
-18
@@ -8,10 +8,12 @@ use sha2::{Digest, Sha256};
|
||||
use zeroize::Zeroizing;
|
||||
|
||||
use super::{ProfileKey, TOKEN_BYTES, TokenHash};
|
||||
use crate::config::{WebLimitsConfig, WebRuntimeConfig, WebRuntimeProfile};
|
||||
use crate::config::{WebRuntimeConfig, WebRuntimeProfile};
|
||||
use crate::maestro::generation::RuntimeGeneration;
|
||||
use crate::web::session::WebSession;
|
||||
|
||||
const WEB_PROFILE_OWNER_CONTEXT: &[u8] = b"telemt-web-profile-owner-v1\0";
|
||||
|
||||
/// One issued bootstrap and optional idempotent session-creation replay state.
|
||||
pub(super) struct Bootstrap {
|
||||
/// Credential and replay-state expiry deadline.
|
||||
@@ -74,14 +76,6 @@ pub(super) struct ManagerState {
|
||||
/// Process-wide live relay-task count.
|
||||
pub(super) streams_live: usize,
|
||||
stream_ports: HashMap<(IpAddr, SocketAddr), StreamPortState>,
|
||||
/// Total process-wide queued byte reservation.
|
||||
pub(super) pending_bytes: usize,
|
||||
/// Total process-wide queued item reservation.
|
||||
pub(super) pending_items: usize,
|
||||
/// Portion of queued bytes charged to the control reserve.
|
||||
pub(super) pending_control_bytes: usize,
|
||||
/// Portion of queued items charged to the control reserve.
|
||||
pub(super) pending_control_items: usize,
|
||||
/// Bootstrap issuance rate limiter.
|
||||
pub(super) bootstrap_rate: RateState,
|
||||
/// Session creation rate limiter.
|
||||
@@ -112,9 +106,19 @@ pub(super) fn new_unique_token(
|
||||
None
|
||||
}
|
||||
|
||||
/// Returns the precomputed capability as the stable process profile key.
|
||||
/// Derives a secret-independent quota owner stable across capability rotation.
|
||||
pub(super) fn profile_key(profile: &WebRuntimeProfile) -> ProfileKey {
|
||||
profile.capability
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(WEB_PROFILE_OWNER_CONTEXT);
|
||||
digest.update((profile.host.len() as u64).to_be_bytes());
|
||||
digest.update(profile.host.as_bytes());
|
||||
digest.update((profile.user.len() as u64).to_be_bytes());
|
||||
digest.update(profile.user.as_bytes());
|
||||
digest.update([match profile.secret_mode {
|
||||
crate::config::WebSecretMode::Plain => 0,
|
||||
crate::config::WebSecretMode::Dd => 1,
|
||||
}]);
|
||||
digest.finalize().into()
|
||||
}
|
||||
|
||||
/// Re-resolves an issued profile against the active generation without weakening identity.
|
||||
@@ -211,13 +215,6 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
/// Computes the process-wide item reserve required for session control progress.
|
||||
pub(super) fn control_item_reserve(limits: &WebLimitsConfig) -> usize {
|
||||
limits
|
||||
.max_sessions_global
|
||||
.saturating_mul(16usize.saturating_add(limits.max_streams_per_session.saturating_mul(3)))
|
||||
}
|
||||
|
||||
/// Allocates a non-zero source port unique among live streams for one KDF route.
|
||||
pub(super) fn allocate_stream_port(
|
||||
state: &mut ManagerState,
|
||||
|
||||
@@ -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;
|
||||
@@ -0,0 +1,40 @@
|
||||
use super::*;
|
||||
|
||||
fn entry(kind: WebSocketKind, opened: bool, peer_tick: u64) -> WebSocketEntry {
|
||||
WebSocketEntry {
|
||||
id: 1,
|
||||
owner: [0; 32],
|
||||
session_id: 1,
|
||||
client_ip: "192.0.2.10".parse().unwrap(),
|
||||
kind,
|
||||
liveness_interval_ms: 10,
|
||||
created_tick: 1,
|
||||
last_peer_tick: AtomicU64::new(peer_tick),
|
||||
last_progress_tick: AtomicU64::new(peer_tick),
|
||||
opened: AtomicBool::new(opened),
|
||||
cancel: CancellationToken::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preopen_and_dead_entries_precede_live_lane_and_multiplex_victims() {
|
||||
let preopen = entry(WebSocketKind::Multiplex, false, 90);
|
||||
let dead = entry(WebSocketKind::Multiplex, true, 1);
|
||||
let lane = entry(WebSocketKind::Lane(7), true, 90);
|
||||
let multiplex = entry(WebSocketKind::Multiplex, true, 90);
|
||||
|
||||
assert_eq!(entry_priority(&preopen, 100), 0);
|
||||
assert_eq!(entry_priority(&dead, 100), 0);
|
||||
assert_eq!(entry_priority(&lane, 100), 1);
|
||||
assert_eq!(entry_priority(&multiplex, 100), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dead_classification_keeps_each_connections_creation_time_interval() {
|
||||
let short_interval = entry(WebSocketKind::Multiplex, true, 80);
|
||||
let mut long_interval = entry(WebSocketKind::Multiplex, true, 80);
|
||||
long_interval.liveness_interval_ms = 100;
|
||||
|
||||
assert_eq!(entry_priority(&short_interval, 100), 0);
|
||||
assert_eq!(entry_priority(&long_interval, 100), 2);
|
||||
}
|
||||
Reference in New Issue
Block a user