mirror of
https://github.com/telemt/telemt.git
synced 2026-09-14 14:34:11 +03:00
WEB: websocket + websocket-lanes as Carrier
This commit is contained in:
Alexey
@@ -0,0 +1,638 @@
|
||||
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,
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user