WEB Carriers Safe-matrix Refactored

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-08-26 20:33:33 +03:00
parent 8b2b88f30c
commit 34eeb2d856
49 changed files with 3406 additions and 936 deletions
+136 -54
View File
@@ -11,7 +11,9 @@ use super::ConnectionIo;
use crate::web::manager::{
WebProcessRuntime, WebSocketBudgetLease, WebSocketConnection,
};
use crate::web::session::{WebSession, WebSocketLaneReservation};
use crate::web::session::{
WebSession, WebSocketLaneReservation, WebSocketProbeReservation,
};
use crate::web::trace::{TraceDirection, TraceWebSocketContext};
const READ_BUFFER_BYTES: usize = 64 * 1024;
@@ -27,11 +29,12 @@ pub(super) async fn run_upgraded(
session: Arc<WebSession>,
connection: WebSocketConnection,
mut lane_reservation: Option<WebSocketLaneReservation>,
_probe_reservation: Option<WebSocketProbeReservation>,
trace: Option<TraceWebSocketContext>,
acknowledge_commit: bool,
) {
let cancellation = connection.cancellation();
let timeouts = runtime.active_generation().config().web.timeouts.clone();
let timeouts = session.timeouts().clone();
let upgraded = tokio::select! {
_ = cancellation.cancelled() => return,
result = tokio::time::timeout(
@@ -47,7 +50,7 @@ pub(super) async fn run_upgraded(
};
let mut io = parts.io.into_inner();
io.enable_websocket(parts.read_buf);
let limits = runtime.active_generation().config().web.limits.clone();
let limits = session.limits().clone();
let config = WebSocketConfig::default()
.read_buffer_size(READ_BUFFER_BYTES)
.write_buffer_size(WRITE_BUFFER_BYTES)
@@ -59,7 +62,9 @@ pub(super) async fn run_upgraded(
.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();
if !connection.mark_opened() {
return;
}
if let Some(reservation) = lane_reservation.as_mut() {
let _ = run_lane(
&mut socket,
@@ -84,16 +89,11 @@ pub(super) async fn run_upgraded(
)
.await;
}
let eviction = Duration::from_secs(
runtime
.active_generation()
.config()
.web
.timeouts
.websocket_eviction_secs,
);
if !cancellation.is_cancelled() {
let _ = tokio::time::timeout(eviction, socket.close(None)).await;
let eviction = Duration::from_secs(timeouts.websocket_eviction_secs);
tokio::select! {
biased;
_ = cancellation.cancelled() => {}
_ = tokio::time::timeout(eviction, socket.close(None)) => {}
}
if let Some(reservation) = lane_reservation {
session.close_websocket_lane(reservation.lane_id());
@@ -121,15 +121,12 @@ async fn run_multiplex(
let mut read_budget = None;
let liveness_interval = connection.liveness_interval();
let mut next_ping = Instant::now() + liveness_interval;
let open_deadline = Instant::now()
+ Duration::from_secs(
runtime
.active_generation()
.config()
.web
.timeouts
.websocket_open_secs,
);
let open_deadline =
Instant::now() + Duration::from_secs(session.timeouts().websocket_open_secs);
let backpressure_timeout =
Duration::from_secs(session.timeouts().websocket_backpressure_secs);
let write_timeout = Duration::from_secs(session.timeouts().websocket_write_secs);
let maximum_message = session.limits().carrier_batch_bytes;
let mut active = false;
loop {
let down = session.poll_down(cursor);
@@ -144,6 +141,8 @@ async fn run_multiplex(
session.profile_key(),
&cancellation,
&mut read_budget,
maximum_message,
backpressure_timeout,
) => {
DriverEvent::Incoming(incoming?)
}
@@ -153,8 +152,15 @@ async fn run_multiplex(
DriverEvent::Incoming((message, _budget)) => match message {
Message::Binary(body) => {
let started = Instant::now();
let result =
process_multiplex(runtime, session, sequence, &body, &cancellation).await;
let result = process_multiplex(
runtime,
session,
sequence,
&body,
&cancellation,
backpressure_timeout,
)
.await;
record_message(
runtime,
trace,
@@ -163,14 +169,17 @@ async fn run_multiplex(
&body,
started,
);
result?;
if acknowledge_commit && sequence == 1 && session.is_carrier_committed() {
let progressed = result?;
if acknowledge_commit && sequence == 1 {
if !session.needs_websocket_commit_ack(connection.id()) {
return Err(());
}
let started = Instant::now();
send(
socket,
runtime,
Message::Binary(Bytes::new()),
&cancellation,
write_timeout,
)
.await?;
record_message(
@@ -181,9 +190,19 @@ async fn run_multiplex(
&[],
started,
);
if !session.websocket_commit_ack_written(connection.id()) {
session.close();
return Err(());
}
} else if acknowledge_commit && sequence > 1 && progressed {
if !session.websocket_peer_after_commit_ack(connection.id()) {
return Err(());
}
}
if !active {
connection.mark_active();
if !active && progressed {
if !connection.mark_active() {
return Err(());
}
active = true;
}
sequence = sequence.checked_add(1).ok_or(())?;
@@ -204,7 +223,7 @@ async fn run_multiplex(
}
Message::Ping(payload) => {
let started = Instant::now();
flush(socket, runtime, &cancellation).await?;
flush(socket, &cancellation, write_timeout).await?;
record_message(
runtime,
trace,
@@ -251,7 +270,13 @@ async fn run_multiplex(
DriverEvent::Down(result) => {
if result.body.is_empty() {
let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
send(
socket,
Message::Ping(Bytes::new()),
&cancellation,
write_timeout,
)
.await?;
record_message(
runtime,
trace,
@@ -267,12 +292,19 @@ async fn run_multiplex(
session.profile_key(),
result.body.len(),
&cancellation,
backpressure_timeout,
)
.await?;
let body = result.body;
let started = Instant::now();
if trace.is_some() {
send(socket, runtime, Message::Binary(body.clone()), &cancellation).await?;
send(
socket,
Message::Binary(body.clone()),
&cancellation,
write_timeout,
)
.await?;
record_message(
runtime,
trace,
@@ -282,7 +314,13 @@ async fn run_multiplex(
started,
);
} else {
send(socket, runtime, Message::Binary(body), &cancellation).await?;
send(
socket,
Message::Binary(body),
&cancellation,
write_timeout,
)
.await?;
}
connection.mark_progress();
}
@@ -290,7 +328,13 @@ async fn run_multiplex(
}
DriverEvent::Liveness => {
let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
send(
socket,
Message::Ping(Bytes::new()),
&cancellation,
write_timeout,
)
.await?;
record_message(
runtime,
trace,
@@ -322,15 +366,12 @@ async fn run_lane(
let mut read_budget = None;
let liveness_interval = connection.liveness_interval();
let mut next_ping = Instant::now() + liveness_interval;
let open_deadline = Instant::now()
+ Duration::from_secs(
runtime
.active_generation()
.config()
.web
.timeouts
.websocket_open_secs,
);
let open_deadline =
Instant::now() + Duration::from_secs(session.timeouts().websocket_open_secs);
let backpressure_timeout =
Duration::from_secs(session.timeouts().websocket_backpressure_secs);
let write_timeout = Duration::from_secs(session.timeouts().websocket_write_secs);
let maximum_message = session.limits().carrier_batch_bytes;
let mut active = false;
loop {
let down = session.poll_down_lane(reservation.lane_id(), cursor);
@@ -345,6 +386,8 @@ async fn run_lane(
session.profile_key(),
&cancellation,
&mut read_budget,
maximum_message,
backpressure_timeout,
) => {
DriverEvent::Incoming(incoming?)
}
@@ -361,6 +404,7 @@ async fn run_lane(
sequence,
&body,
&cancellation,
backpressure_timeout,
)
.await;
record_message(
@@ -371,14 +415,17 @@ async fn run_lane(
&body,
started,
);
result?;
if acknowledge_commit && sequence == 1 && session.is_carrier_committed() {
let progressed = result?;
if acknowledge_commit && sequence == 1 {
if !session.needs_websocket_commit_ack(connection.id()) {
return Err(());
}
let started = Instant::now();
if send(
socket,
runtime,
Message::Binary(Bytes::new()),
&cancellation,
write_timeout,
)
.await
.is_err()
@@ -394,9 +441,19 @@ async fn run_lane(
&[],
started,
);
if !session.websocket_commit_ack_written(connection.id()) {
session.close();
return Err(());
}
} else if acknowledge_commit && sequence > 1 && progressed {
if !session.websocket_peer_after_commit_ack(connection.id()) {
return Err(());
}
}
if !active {
connection.mark_active();
if !active && progressed {
if !connection.mark_active() {
return Err(());
}
active = true;
}
sequence = sequence.checked_add(1).ok_or(())?;
@@ -417,7 +474,7 @@ async fn run_lane(
}
Message::Ping(payload) => {
let started = Instant::now();
flush(socket, runtime, &cancellation).await?;
flush(socket, &cancellation, write_timeout).await?;
record_message(
runtime,
trace,
@@ -467,7 +524,13 @@ async fn run_lane(
}
if result.body.is_empty() {
let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
send(
socket,
Message::Ping(Bytes::new()),
&cancellation,
write_timeout,
)
.await?;
record_message(
runtime,
trace,
@@ -483,12 +546,19 @@ async fn run_lane(
session.profile_key(),
result.body.len(),
&cancellation,
backpressure_timeout,
)
.await?;
let body = result.body;
let started = Instant::now();
if trace.is_some() {
send(socket, runtime, Message::Binary(body.clone()), &cancellation).await?;
send(
socket,
Message::Binary(body.clone()),
&cancellation,
write_timeout,
)
.await?;
record_message(
runtime,
trace,
@@ -498,7 +568,13 @@ async fn run_lane(
started,
);
} else {
send(socket, runtime, Message::Binary(body), &cancellation).await?;
send(
socket,
Message::Binary(body),
&cancellation,
write_timeout,
)
.await?;
}
connection.mark_progress();
}
@@ -506,7 +582,13 @@ async fn run_lane(
}
DriverEvent::Liveness => {
let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
send(
socket,
Message::Ping(Bytes::new()),
&cancellation,
write_timeout,
)
.await?;
record_message(
runtime,
trace,
+36 -58
View File
@@ -16,19 +16,24 @@ pub(super) async fn read_message(
owner: crate::web::manager::ProfileKey,
cancellation: &CancellationToken,
retained_budget: &mut Option<WebSocketBudgetLease>,
maximum: usize,
backpressure_timeout: Duration,
) -> 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?);
*retained_budget = Some(
reserve_data(
runtime,
owner,
maximum,
cancellation,
backpressure_timeout,
)
.await?,
);
}
let message = tokio::select! {
_ = cancellation.cancelled() => return Err(()),
@@ -47,17 +52,13 @@ pub(super) async fn reserve_data(
owner: crate::web::manager::ProfileKey,
bytes: usize,
cancellation: &CancellationToken,
timeout: Duration,
) -> Result<WebSocketBudgetLease, ()> {
let timeout = Duration::from_secs(
runtime
.active_generation()
.config()
.web
.timeouts
.websocket_backpressure_secs,
);
tokio::time::timeout(timeout, async {
loop {
if cancellation.is_cancelled() {
return Err(());
}
let notify = runtime.budget_notify();
let notified = notify.notified();
if let Some(budget) = runtime.try_websocket_data_budget(owner, bytes.max(1)) {
@@ -79,9 +80,10 @@ pub(super) async fn process_multiplex(
sequence: u64,
body: &[u8],
cancellation: &CancellationToken,
) -> Result<(), ()> {
retry_backpressure(runtime, cancellation, || {
session.process_up(sequence, body).map(|_| ())
timeout: Duration,
) -> Result<bool, ()> {
retry_backpressure(runtime, cancellation, timeout, || {
session.process_websocket_multiplex(sequence, body)
})
.await
}
@@ -93,21 +95,17 @@ pub(super) async fn process_lane(
sequence: u64,
body: &[u8],
cancellation: &CancellationToken,
) -> Result<(), ()> {
let timeout = Duration::from_secs(
runtime
.active_generation()
.config()
.web
.timeouts
.websocket_backpressure_secs,
);
timeout: Duration,
) -> Result<bool, ()> {
tokio::time::timeout(timeout, async {
loop {
if cancellation.is_cancelled() {
return Err(());
}
let notify = runtime.budget_notify();
let notified = notify.notified();
match session.process_websocket_lane(reservation, sequence, body) {
Ok(()) => return Ok(()),
Ok(progressed) => return Ok(progressed),
Err(ManagerError::Backpressure) => {}
Err(_) => return Err(()),
}
@@ -121,28 +119,24 @@ pub(super) async fn process_lane(
.map_err(|_| ())?
}
async fn retry_backpressure<F>(
async fn retry_backpressure<F, T>(
runtime: &Arc<WebProcessRuntime>,
cancellation: &CancellationToken,
timeout: Duration,
mut operation: F,
) -> Result<(), ()>
) -> Result<T, ()>
where
F: FnMut() -> Result<(), ManagerError>,
F: FnMut() -> Result<T, ManagerError>,
{
let timeout = Duration::from_secs(
runtime
.active_generation()
.config()
.web
.timeouts
.websocket_backpressure_secs,
);
tokio::time::timeout(timeout, async {
loop {
if cancellation.is_cancelled() {
return Err(());
}
let notify = runtime.budget_notify();
let notified = notify.notified();
match operation() {
Ok(()) => return Ok(()),
Ok(value) => return Ok(value),
Err(ManagerError::Backpressure) => {}
Err(_) => return Err(()),
}
@@ -158,18 +152,10 @@ where
pub(super) async fn send(
socket: &mut CarrierSocket,
runtime: &WebProcessRuntime,
message: Message,
cancellation: &CancellationToken,
timeout: Duration,
) -> Result<(), ()> {
let timeout = Duration::from_secs(
runtime
.active_generation()
.config()
.web
.timeouts
.websocket_write_secs,
);
tokio::select! {
_ = cancellation.cancelled() => Err(()),
result = tokio::time::timeout(timeout, socket.send(message)) => {
@@ -180,17 +166,9 @@ pub(super) async fn send(
pub(super) async fn flush(
socket: &mut CarrierSocket,
runtime: &WebProcessRuntime,
cancellation: &CancellationToken,
timeout: Duration,
) -> Result<(), ()> {
let timeout = Duration::from_secs(
runtime
.active_generation()
.config()
.web
.timeouts
.websocket_write_secs,
);
tokio::select! {
_ = cancellation.cancelled() => Err(()),
result = tokio::time::timeout(timeout, socket.flush()) => {
+32 -1
View File
@@ -160,6 +160,7 @@ fn create_automatic_session(
None,
[9; 32],
),
false,
)
.unwrap()
.token;
@@ -194,6 +195,7 @@ fn create_session(runtime: &Arc<WebProcessRuntime>) -> (String, TokenHash) {
client_ip,
&hello,
CarrierRequest::legacy([0; 32]),
false,
)
.unwrap()
.token;
@@ -382,7 +384,7 @@ async fn malformed_websocket_lane_closes_only_that_lane() {
}
#[tokio::test]
async fn automatic_websocket_carriers_ack_the_first_committing_message() {
async fn automatic_websocket_carriers_commit_after_acknowledged_peer_progress() {
for carrier in [WebCarrier::Websocket, WebCarrier::WebsocketLanes] {
let live = live_negotiation_runtime(carrier, Arc::from([carrier]));
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -409,6 +411,34 @@ async fn automatic_websocket_carriers_ack_the_first_committing_message() {
.unwrap()
.is_carrier_committed()
);
socket
.send(Message::Binary(frame::encode(
FrameType::Window,
7,
&frame::window_payload(1),
)))
.await
.unwrap();
socket
.send(Message::Ping(Bytes::from_static(b"commit")))
.await
.unwrap();
loop {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.unwrap()
.unwrap()
.unwrap();
if message == Message::Pong(Bytes::from_static(b"commit")) {
break;
}
}
assert!(
live.runtime
.get_session(session_hash, "proxy.example.com")
.unwrap()
.is_carrier_committed()
);
let _ = socket.close(None).await;
live.shutdown().await;
@@ -448,6 +478,7 @@ async fn failed_automatic_multiplex_socket_remains_supersedable() {
Some(CarrierFailure::Upgrade),
[9; 32],
),
false,
)
.unwrap();
assert_eq!(replacement.carrier, WebCarrier::Https);