Session Residence + Ownership

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-08-26 18:44:17 +03:00
parent 016ad247a4
commit 8b2b88f30c
10 changed files with 436 additions and 71 deletions
+2
View File
@@ -284,11 +284,13 @@ pub(super) async fn handle(
.admit_websocket(
session.profile_key(),
session.trace_session_id(),
session.token_hash(),
effective_ip,
kind,
BASE_BUDGET_BYTES,
Duration::from_secs(timeouts.long_poll_secs),
Duration::from_secs(timeouts.websocket_eviction_secs),
session.carrier_cancellation(),
)
.await
{
+72 -15
View File
@@ -30,7 +30,16 @@ pub(super) async fn run_upgraded(
trace: Option<TraceWebSocketContext>,
acknowledge_commit: bool,
) {
let Ok(upgraded) = on_upgrade.await else {
let cancellation = connection.cancellation();
let timeouts = runtime.active_generation().config().web.timeouts.clone();
let upgraded = tokio::select! {
_ = cancellation.cancelled() => return,
result = tokio::time::timeout(
Duration::from_secs(timeouts.websocket_upgrade_secs),
on_upgrade,
) => result,
};
let Ok(Ok(upgraded)) = upgraded else {
return;
};
let Ok(parts) = upgraded.downcast::<TokioIo<ConnectionIo>>() else {
@@ -51,7 +60,6 @@ pub(super) async fn run_upgraded(
.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,
@@ -84,7 +92,9 @@ pub(super) async fn run_upgraded(
.timeouts
.websocket_eviction_secs,
);
let _ = tokio::time::timeout(eviction, socket.close(None)).await;
if !cancellation.is_cancelled() {
let _ = tokio::time::timeout(eviction, socket.close(None)).await;
}
if let Some(reservation) = lane_reservation {
session.close_websocket_lane(reservation.lane_id());
drop(reservation);
@@ -111,11 +121,22 @@ 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 mut active = false;
loop {
let down = session.poll_down(cursor);
tokio::pin!(down);
let event = tokio::select! {
_ = cancellation.cancelled() => return Err(()),
_ = tokio::time::sleep_until(open_deadline.into()), if !active => return Err(()),
_ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness,
incoming = read_message(
socket,
@@ -145,7 +166,13 @@ async fn run_multiplex(
result?;
if acknowledge_commit && sequence == 1 && session.is_carrier_committed() {
let started = Instant::now();
send(socket, runtime, Message::Binary(Bytes::new())).await?;
send(
socket,
runtime,
Message::Binary(Bytes::new()),
&cancellation,
)
.await?;
record_message(
runtime,
trace,
@@ -155,6 +182,10 @@ async fn run_multiplex(
started,
);
}
if !active {
connection.mark_active();
active = true;
}
sequence = sequence.checked_add(1).ok_or(())?;
connection.mark_peer_activity();
next_ping = Instant::now() + liveness_interval;
@@ -173,7 +204,7 @@ async fn run_multiplex(
}
Message::Ping(payload) => {
let started = Instant::now();
flush(socket, runtime).await?;
flush(socket, runtime, &cancellation).await?;
record_message(
runtime,
trace,
@@ -220,7 +251,7 @@ async fn run_multiplex(
DriverEvent::Down(result) => {
if result.body.is_empty() {
let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new())).await?;
send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
record_message(
runtime,
trace,
@@ -241,7 +272,7 @@ async fn run_multiplex(
let body = result.body;
let started = Instant::now();
if trace.is_some() {
send(socket, runtime, Message::Binary(body.clone())).await?;
send(socket, runtime, Message::Binary(body.clone()), &cancellation).await?;
record_message(
runtime,
trace,
@@ -251,7 +282,7 @@ async fn run_multiplex(
started,
);
} else {
send(socket, runtime, Message::Binary(body)).await?;
send(socket, runtime, Message::Binary(body), &cancellation).await?;
}
connection.mark_progress();
}
@@ -259,7 +290,7 @@ async fn run_multiplex(
}
DriverEvent::Liveness => {
let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new())).await?;
send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
record_message(
runtime,
trace,
@@ -291,11 +322,22 @@ 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 mut active = false;
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(open_deadline.into()), if !active => return Err(()),
_ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness,
incoming = read_message(
socket,
@@ -332,7 +374,18 @@ async fn run_lane(
result?;
if acknowledge_commit && sequence == 1 && session.is_carrier_committed() {
let started = Instant::now();
send(socket, runtime, Message::Binary(Bytes::new())).await?;
if send(
socket,
runtime,
Message::Binary(Bytes::new()),
&cancellation,
)
.await
.is_err()
{
session.close();
return Err(());
}
record_message(
runtime,
trace,
@@ -342,6 +395,10 @@ async fn run_lane(
started,
);
}
if !active {
connection.mark_active();
active = true;
}
sequence = sequence.checked_add(1).ok_or(())?;
connection.mark_peer_activity();
next_ping = Instant::now() + liveness_interval;
@@ -360,7 +417,7 @@ async fn run_lane(
}
Message::Ping(payload) => {
let started = Instant::now();
flush(socket, runtime).await?;
flush(socket, runtime, &cancellation).await?;
record_message(
runtime,
trace,
@@ -410,7 +467,7 @@ async fn run_lane(
}
if result.body.is_empty() {
let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new())).await?;
send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
record_message(
runtime,
trace,
@@ -431,7 +488,7 @@ async fn run_lane(
let body = result.body;
let started = Instant::now();
if trace.is_some() {
send(socket, runtime, Message::Binary(body.clone())).await?;
send(socket, runtime, Message::Binary(body.clone()), &cancellation).await?;
record_message(
runtime,
trace,
@@ -441,7 +498,7 @@ async fn run_lane(
started,
);
} else {
send(socket, runtime, Message::Binary(body)).await?;
send(socket, runtime, Message::Binary(body), &cancellation).await?;
}
connection.mark_progress();
}
@@ -449,7 +506,7 @@ async fn run_lane(
}
DriverEvent::Liveness => {
let started = Instant::now();
send(socket, runtime, Message::Ping(Bytes::new())).await?;
send(socket, runtime, Message::Ping(Bytes::new()), &cancellation).await?;
record_message(
runtime,
trace,
+14 -8
View File
@@ -160,6 +160,7 @@ pub(super) async fn send(
socket: &mut CarrierSocket,
runtime: &WebProcessRuntime,
message: Message,
cancellation: &CancellationToken,
) -> Result<(), ()> {
let timeout = Duration::from_secs(
runtime
@@ -169,15 +170,18 @@ pub(super) async fn send(
.timeouts
.websocket_write_secs,
);
tokio::time::timeout(timeout, socket.send(message))
.await
.map_err(|_| ())?
.map_err(|_| ())
tokio::select! {
_ = cancellation.cancelled() => Err(()),
result = tokio::time::timeout(timeout, socket.send(message)) => {
result.map_err(|_| ())?.map_err(|_| ())
}
}
}
pub(super) async fn flush(
socket: &mut CarrierSocket,
runtime: &WebProcessRuntime,
cancellation: &CancellationToken,
) -> Result<(), ()> {
let timeout = Duration::from_secs(
runtime
@@ -187,10 +191,12 @@ pub(super) async fn flush(
.timeouts
.websocket_write_secs,
);
tokio::time::timeout(timeout, socket.flush())
.await
.map_err(|_| ())?
.map_err(|_| ())
tokio::select! {
_ = cancellation.cancelled() => Err(()),
result = tokio::time::timeout(timeout, socket.flush()) => {
result.map_err(|_| ())?.map_err(|_| ())
}
}
}
pub(super) fn record_message(