WEB: websocket + websocket-lanes as Carrier

This commit is contained in:
Alexey
2026-08-26 09:16:06 +03:00
parent 8e577ec5ca
commit d2edd90479
50 changed files with 3642 additions and 350 deletions
+16 -6
View File
@@ -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;
}
}
+366
View File
@@ -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());
}
}
+10 -9
View File
@@ -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
View File
@@ -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,
+313
View File
@@ -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;
+40
View File
@@ -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);
}