mirror of
https://github.com/telemt/telemt.git
synced 2026-09-05 18:16:06 +03:00
Alexey
639 lines
21 KiB
Rust
639 lines
21 KiB
Rust
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
|
|
use bytes::Bytes;
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use hyper_util::rt::TokioIo;
|
|
use tokio_tungstenite::WebSocketStream;
|
|
use tokio_tungstenite::tungstenite::protocol::{Message, Role, WebSocketConfig};
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
use super::ConnectionIo;
|
|
use crate::web::manager::{
|
|
ManagerError, WebProcessRuntime, WebSocketBudgetLease, WebSocketConnection,
|
|
};
|
|
use crate::web::session::{WebSession, WebSocketLaneReservation};
|
|
use crate::web::trace::{TraceDirection, TraceWebSocketContext};
|
|
|
|
const READ_BUFFER_BYTES: usize = 64 * 1024;
|
|
const WRITE_BUFFER_BYTES: usize = 64 * 1024;
|
|
|
|
pub(super) async fn run_upgraded(
|
|
on_upgrade: hyper::upgrade::OnUpgrade,
|
|
runtime: Arc<WebProcessRuntime>,
|
|
session: Arc<WebSession>,
|
|
connection: WebSocketConnection,
|
|
mut lane_reservation: Option<WebSocketLaneReservation>,
|
|
trace: Option<TraceWebSocketContext>,
|
|
) {
|
|
let Ok(upgraded) = on_upgrade.await else {
|
|
return;
|
|
};
|
|
let Ok(parts) = upgraded.downcast::<TokioIo<ConnectionIo>>() else {
|
|
return;
|
|
};
|
|
let mut io = parts.io.into_inner();
|
|
io.enable_websocket(parts.read_buf);
|
|
let limits = runtime.active_generation().config().web.limits.clone();
|
|
let config = WebSocketConfig::default()
|
|
.read_buffer_size(READ_BUFFER_BYTES)
|
|
.write_buffer_size(WRITE_BUFFER_BYTES)
|
|
.max_write_buffer_size(
|
|
WRITE_BUFFER_BYTES
|
|
.saturating_add(limits.carrier_batch_bytes)
|
|
.saturating_add(1024),
|
|
)
|
|
.max_message_size(Some(limits.carrier_batch_bytes))
|
|
.max_frame_size(Some(limits.carrier_batch_bytes));
|
|
let mut socket = WebSocketStream::from_raw_socket(io, Role::Server, Some(config)).await;
|
|
connection.mark_opened();
|
|
let cancellation = connection.cancellation();
|
|
if let Some(reservation) = lane_reservation.as_mut() {
|
|
let _ = run_lane(
|
|
&mut socket,
|
|
&runtime,
|
|
&session,
|
|
&connection,
|
|
reservation,
|
|
cancellation.clone(),
|
|
trace.as_ref(),
|
|
)
|
|
.await;
|
|
} else {
|
|
let _ = run_multiplex(
|
|
&mut socket,
|
|
&runtime,
|
|
&session,
|
|
&connection,
|
|
cancellation.clone(),
|
|
trace.as_ref(),
|
|
)
|
|
.await;
|
|
}
|
|
let eviction = Duration::from_secs(
|
|
runtime
|
|
.active_generation()
|
|
.config()
|
|
.web
|
|
.timeouts
|
|
.websocket_eviction_secs,
|
|
);
|
|
let _ = tokio::time::timeout(eviction, socket.close(None)).await;
|
|
if let Some(reservation) = lane_reservation {
|
|
session.close_websocket_lane(reservation.lane_id());
|
|
drop(reservation);
|
|
} else {
|
|
session.close();
|
|
}
|
|
}
|
|
|
|
type CarrierSocket = WebSocketStream<ConnectionIo>;
|
|
|
|
async fn run_multiplex(
|
|
socket: &mut CarrierSocket,
|
|
runtime: &Arc<WebProcessRuntime>,
|
|
session: &Arc<WebSession>,
|
|
connection: &WebSocketConnection,
|
|
cancellation: CancellationToken,
|
|
trace: Option<&TraceWebSocketContext>,
|
|
) -> Result<(), ()> {
|
|
let mut sequence = 1u64;
|
|
let mut cursor = 0u64;
|
|
// The lease survives cancelled select branches and control frames interleaved
|
|
// inside one fragmented data message.
|
|
let mut read_budget = None;
|
|
let liveness_interval = connection.liveness_interval();
|
|
let mut next_ping = Instant::now() + liveness_interval;
|
|
loop {
|
|
let down = session.poll_down(cursor);
|
|
tokio::pin!(down);
|
|
let event = tokio::select! {
|
|
_ = cancellation.cancelled() => return Err(()),
|
|
_ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness,
|
|
incoming = read_message(
|
|
socket,
|
|
runtime,
|
|
session.profile_key(),
|
|
&cancellation,
|
|
&mut read_budget,
|
|
) => {
|
|
DriverEvent::Incoming(incoming?)
|
|
}
|
|
down = &mut down => DriverEvent::Down(down.map_err(|_| ())?),
|
|
};
|
|
match event {
|
|
DriverEvent::Incoming((message, _budget)) => match message {
|
|
Message::Binary(body) => {
|
|
let started = Instant::now();
|
|
let result =
|
|
process_multiplex(runtime, session, sequence, &body, &cancellation).await;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"binary",
|
|
&body,
|
|
started,
|
|
);
|
|
result?;
|
|
sequence = sequence.checked_add(1).ok_or(())?;
|
|
connection.mark_peer_activity();
|
|
next_ping = Instant::now() + liveness_interval;
|
|
}
|
|
Message::Pong(payload) => {
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"pong",
|
|
&payload,
|
|
Instant::now(),
|
|
);
|
|
connection.mark_peer_activity();
|
|
next_ping = Instant::now() + liveness_interval;
|
|
}
|
|
Message::Ping(payload) => {
|
|
let started = Instant::now();
|
|
flush(socket, runtime).await?;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"ping",
|
|
&payload,
|
|
started,
|
|
);
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Response,
|
|
"pong",
|
|
&payload,
|
|
started,
|
|
);
|
|
connection.mark_peer_activity();
|
|
next_ping = Instant::now() + liveness_interval;
|
|
}
|
|
Message::Close(_) => {
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"close",
|
|
&[],
|
|
Instant::now(),
|
|
);
|
|
return Ok(());
|
|
}
|
|
Message::Text(text) => {
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"text",
|
|
text.as_bytes(),
|
|
Instant::now(),
|
|
);
|
|
return Err(());
|
|
}
|
|
Message::Frame(_) => return Err(()),
|
|
},
|
|
DriverEvent::Down(result) => {
|
|
if result.body.is_empty() {
|
|
let started = Instant::now();
|
|
send(socket, runtime, Message::Ping(Bytes::new())).await?;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Response,
|
|
"ping",
|
|
&[],
|
|
started,
|
|
);
|
|
next_ping = Instant::now() + liveness_interval;
|
|
} else {
|
|
let _budget = reserve_data(
|
|
runtime,
|
|
session.profile_key(),
|
|
result.body.len(),
|
|
&cancellation,
|
|
)
|
|
.await?;
|
|
let body = result.body;
|
|
let started = Instant::now();
|
|
if trace.is_some() {
|
|
send(socket, runtime, Message::Binary(body.clone())).await?;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Response,
|
|
"binary",
|
|
&body,
|
|
started,
|
|
);
|
|
} else {
|
|
send(socket, runtime, Message::Binary(body)).await?;
|
|
}
|
|
connection.mark_progress();
|
|
}
|
|
cursor = result.next_cursor;
|
|
}
|
|
DriverEvent::Liveness => {
|
|
let started = Instant::now();
|
|
send(socket, runtime, Message::Ping(Bytes::new())).await?;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Response,
|
|
"ping",
|
|
&[],
|
|
started,
|
|
);
|
|
next_ping = Instant::now() + liveness_interval;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn run_lane(
|
|
socket: &mut CarrierSocket,
|
|
runtime: &Arc<WebProcessRuntime>,
|
|
session: &Arc<WebSession>,
|
|
connection: &WebSocketConnection,
|
|
reservation: &mut WebSocketLaneReservation,
|
|
cancellation: CancellationToken,
|
|
trace: Option<&TraceWebSocketContext>,
|
|
) -> Result<(), ()> {
|
|
let mut sequence = 1u64;
|
|
let mut cursor = 0u64;
|
|
// Lane reads use the same cancellation-safe fragmented-message ownership.
|
|
let mut read_budget = None;
|
|
let liveness_interval = connection.liveness_interval();
|
|
let mut next_ping = Instant::now() + liveness_interval;
|
|
loop {
|
|
let down = session.poll_down_lane(reservation.lane_id(), cursor);
|
|
tokio::pin!(down);
|
|
let event = tokio::select! {
|
|
_ = cancellation.cancelled() => return Err(()),
|
|
_ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness,
|
|
incoming = read_message(
|
|
socket,
|
|
runtime,
|
|
session.profile_key(),
|
|
&cancellation,
|
|
&mut read_budget,
|
|
) => {
|
|
DriverEvent::Incoming(incoming?)
|
|
}
|
|
down = &mut down => DriverEvent::Down(down.map_err(|_| ())?),
|
|
};
|
|
match event {
|
|
DriverEvent::Incoming((message, _budget)) => match message {
|
|
Message::Binary(body) => {
|
|
let started = Instant::now();
|
|
let result = process_lane(
|
|
runtime,
|
|
session,
|
|
reservation,
|
|
sequence,
|
|
&body,
|
|
&cancellation,
|
|
)
|
|
.await;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"binary",
|
|
&body,
|
|
started,
|
|
);
|
|
result?;
|
|
sequence = sequence.checked_add(1).ok_or(())?;
|
|
connection.mark_peer_activity();
|
|
next_ping = Instant::now() + liveness_interval;
|
|
}
|
|
Message::Pong(payload) => {
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"pong",
|
|
&payload,
|
|
Instant::now(),
|
|
);
|
|
connection.mark_peer_activity();
|
|
next_ping = Instant::now() + liveness_interval;
|
|
}
|
|
Message::Ping(payload) => {
|
|
let started = Instant::now();
|
|
flush(socket, runtime).await?;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"ping",
|
|
&payload,
|
|
started,
|
|
);
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Response,
|
|
"pong",
|
|
&payload,
|
|
started,
|
|
);
|
|
connection.mark_peer_activity();
|
|
next_ping = Instant::now() + liveness_interval;
|
|
}
|
|
Message::Close(_) => {
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"close",
|
|
&[],
|
|
Instant::now(),
|
|
);
|
|
return Ok(());
|
|
}
|
|
Message::Text(text) => {
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Request,
|
|
"text",
|
|
text.as_bytes(),
|
|
Instant::now(),
|
|
);
|
|
return Err(());
|
|
}
|
|
Message::Frame(_) => return Err(()),
|
|
},
|
|
DriverEvent::Down(result) => {
|
|
if result.lane_closed {
|
|
return Ok(());
|
|
}
|
|
if result.body.is_empty() {
|
|
let started = Instant::now();
|
|
send(socket, runtime, Message::Ping(Bytes::new())).await?;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Response,
|
|
"ping",
|
|
&[],
|
|
started,
|
|
);
|
|
next_ping = Instant::now() + liveness_interval;
|
|
} else {
|
|
let _budget = reserve_data(
|
|
runtime,
|
|
session.profile_key(),
|
|
result.body.len(),
|
|
&cancellation,
|
|
)
|
|
.await?;
|
|
let body = result.body;
|
|
let started = Instant::now();
|
|
if trace.is_some() {
|
|
send(socket, runtime, Message::Binary(body.clone())).await?;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Response,
|
|
"binary",
|
|
&body,
|
|
started,
|
|
);
|
|
} else {
|
|
send(socket, runtime, Message::Binary(body)).await?;
|
|
}
|
|
connection.mark_progress();
|
|
}
|
|
cursor = result.next_cursor;
|
|
}
|
|
DriverEvent::Liveness => {
|
|
let started = Instant::now();
|
|
send(socket, runtime, Message::Ping(Bytes::new())).await?;
|
|
record_message(
|
|
runtime,
|
|
trace,
|
|
TraceDirection::Response,
|
|
"ping",
|
|
&[],
|
|
started,
|
|
);
|
|
next_ping = Instant::now() + liveness_interval;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
enum DriverEvent {
|
|
Incoming((Message, Option<WebSocketBudgetLease>)),
|
|
Down(crate::web::session::PollResult),
|
|
Liveness,
|
|
}
|
|
|
|
async fn read_message(
|
|
socket: &mut CarrierSocket,
|
|
runtime: &Arc<WebProcessRuntime>,
|
|
owner: crate::web::manager::ProfileKey,
|
|
cancellation: &CancellationToken,
|
|
retained_budget: &mut Option<WebSocketBudgetLease>,
|
|
) -> Result<(Message, Option<WebSocketBudgetLease>), ()> {
|
|
tokio::select! {
|
|
_ = cancellation.cancelled() => return Err(()),
|
|
ready = socket.get_ref().readable() => ready.map_err(|_| ())?,
|
|
}
|
|
if retained_budget.is_none() {
|
|
let maximum = runtime
|
|
.active_generation()
|
|
.config()
|
|
.web
|
|
.limits
|
|
.carrier_batch_bytes;
|
|
*retained_budget = Some(reserve_data(runtime, owner, maximum, cancellation).await?);
|
|
}
|
|
let message = tokio::select! {
|
|
_ = cancellation.cancelled() => return Err(()),
|
|
message = socket.next() => message.ok_or(())?.map_err(|_| ())?,
|
|
};
|
|
if socket.get_ref().websocket_fragmented_message() {
|
|
return Ok((message, None));
|
|
}
|
|
let mut budget = retained_budget.take().ok_or(())?;
|
|
budget.shrink_to(message.len());
|
|
Ok((message, Some(budget)))
|
|
}
|
|
|
|
async fn reserve_data(
|
|
runtime: &Arc<WebProcessRuntime>,
|
|
owner: crate::web::manager::ProfileKey,
|
|
bytes: usize,
|
|
cancellation: &CancellationToken,
|
|
) -> Result<WebSocketBudgetLease, ()> {
|
|
let timeout = Duration::from_secs(
|
|
runtime
|
|
.active_generation()
|
|
.config()
|
|
.web
|
|
.timeouts
|
|
.websocket_backpressure_secs,
|
|
);
|
|
tokio::time::timeout(timeout, async {
|
|
loop {
|
|
let notify = runtime.budget_notify();
|
|
let notified = notify.notified();
|
|
if let Some(budget) = runtime.try_websocket_data_budget(owner, bytes.max(1)) {
|
|
return Ok(budget);
|
|
}
|
|
tokio::select! {
|
|
_ = cancellation.cancelled() => return Err(()),
|
|
_ = notified => {}
|
|
}
|
|
}
|
|
})
|
|
.await
|
|
.map_err(|_| ())?
|
|
}
|
|
|
|
async fn process_multiplex(
|
|
runtime: &Arc<WebProcessRuntime>,
|
|
session: &Arc<WebSession>,
|
|
sequence: u64,
|
|
body: &[u8],
|
|
cancellation: &CancellationToken,
|
|
) -> Result<(), ()> {
|
|
retry_backpressure(runtime, cancellation, || {
|
|
session.process_up(sequence, body).map(|_| ())
|
|
})
|
|
.await
|
|
}
|
|
|
|
async fn process_lane(
|
|
runtime: &Arc<WebProcessRuntime>,
|
|
session: &Arc<WebSession>,
|
|
reservation: &mut WebSocketLaneReservation,
|
|
sequence: u64,
|
|
body: &[u8],
|
|
cancellation: &CancellationToken,
|
|
) -> Result<(), ()> {
|
|
let timeout = Duration::from_secs(
|
|
runtime
|
|
.active_generation()
|
|
.config()
|
|
.web
|
|
.timeouts
|
|
.websocket_backpressure_secs,
|
|
);
|
|
tokio::time::timeout(timeout, async {
|
|
loop {
|
|
let notify = runtime.budget_notify();
|
|
let notified = notify.notified();
|
|
match session.process_websocket_lane(reservation, sequence, body) {
|
|
Ok(()) => return Ok(()),
|
|
Err(ManagerError::Backpressure) => {}
|
|
Err(_) => return Err(()),
|
|
}
|
|
tokio::select! {
|
|
_ = cancellation.cancelled() => return Err(()),
|
|
_ = notified => {}
|
|
}
|
|
}
|
|
})
|
|
.await
|
|
.map_err(|_| ())?
|
|
}
|
|
|
|
async fn retry_backpressure<F>(
|
|
runtime: &Arc<WebProcessRuntime>,
|
|
cancellation: &CancellationToken,
|
|
mut operation: F,
|
|
) -> Result<(), ()>
|
|
where
|
|
F: FnMut() -> Result<(), ManagerError>,
|
|
{
|
|
let timeout = Duration::from_secs(
|
|
runtime
|
|
.active_generation()
|
|
.config()
|
|
.web
|
|
.timeouts
|
|
.websocket_backpressure_secs,
|
|
);
|
|
tokio::time::timeout(timeout, async {
|
|
loop {
|
|
let notify = runtime.budget_notify();
|
|
let notified = notify.notified();
|
|
match operation() {
|
|
Ok(()) => return Ok(()),
|
|
Err(ManagerError::Backpressure) => {}
|
|
Err(_) => return Err(()),
|
|
}
|
|
tokio::select! {
|
|
_ = cancellation.cancelled() => return Err(()),
|
|
_ = notified => {}
|
|
}
|
|
}
|
|
})
|
|
.await
|
|
.map_err(|_| ())?
|
|
}
|
|
|
|
async fn send(
|
|
socket: &mut CarrierSocket,
|
|
runtime: &WebProcessRuntime,
|
|
message: Message,
|
|
) -> Result<(), ()> {
|
|
let timeout = Duration::from_secs(
|
|
runtime
|
|
.active_generation()
|
|
.config()
|
|
.web
|
|
.timeouts
|
|
.websocket_write_secs,
|
|
);
|
|
tokio::time::timeout(timeout, socket.send(message))
|
|
.await
|
|
.map_err(|_| ())?
|
|
.map_err(|_| ())
|
|
}
|
|
|
|
async fn flush(socket: &mut CarrierSocket, runtime: &WebProcessRuntime) -> Result<(), ()> {
|
|
let timeout = Duration::from_secs(
|
|
runtime
|
|
.active_generation()
|
|
.config()
|
|
.web
|
|
.timeouts
|
|
.websocket_write_secs,
|
|
);
|
|
tokio::time::timeout(timeout, socket.flush())
|
|
.await
|
|
.map_err(|_| ())?
|
|
.map_err(|_| ())
|
|
}
|
|
|
|
fn record_message(
|
|
runtime: &WebProcessRuntime,
|
|
trace: Option<&TraceWebSocketContext>,
|
|
direction: TraceDirection,
|
|
message_type: &'static str,
|
|
payload: &[u8],
|
|
started: Instant,
|
|
) {
|
|
let Some(trace) = trace else {
|
|
return;
|
|
};
|
|
runtime.trace().record_websocket_message(
|
|
trace,
|
|
direction,
|
|
message_type,
|
|
payload,
|
|
started.elapsed().as_micros().min(u128::from(u64::MAX)) as u64,
|
|
);
|
|
}
|