Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: bound Worker protocol backpressure and connection ownership

Browse Source
This commit is contained in:
Hare
2026-10-12 02:56:42 +09:00
parent 98023b9872
commit 21df210a7f
9 changed files with 2036 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/server.rs
+59 -102
View File
@@ -2,9 +2,12 @@ use subjektiv::api::*;
mod backend_job_profile;
mod subjektiv_job;
mod worker_operations;
mod worker_protocol;
pub(crate) use worker_operations::{
WorkerOperationContext, WorkspaceWorker, WorkspaceWorkerMethodSender,
};
pub(crate) use worker_protocol::WorkerProtocolEnd;
use worker_protocol::*;
#[path = "server_drive.rs"]
mod drive;
@@ -29083,51 +29086,19 @@ async fn worker_protocol_ws(
match worker.connect_protocol(&context).await {
Ok(connection) => workspace_worker_protocol_session(connection, socket, context).await,
Err(error) => {
let _ = send_protocol_event(&mut socket, &protocol_error_event(error.to_string()))
.await;
let _ = socket.send(WsMessage::Close(None)).await;
let _ = tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, async {
let _ =
send_protocol_event(&mut socket, &protocol_error_event(error.to_string()))
.await;
let _ = socket.send(WsMessage::Close(None)).await;
})
.await;
}
}
})
.into_response()
}
async fn workspace_worker_protocol_session(
mut connection: worker_operations::WorkspaceWorkerConnection,
mut socket: WebSocket,
context: WorkerOperationContext,
) {
loop {
tokio::select! {
inbound = socket.next() => match inbound {
Some(Ok(WsMessage::Text(text))) => {
let result = match decode_method(&text) {
Ok(method) => connection.sender.send(&context, method).await.map_err(|error| error.to_string()),
Err(error) => Err(error.to_string()),
};
if let Err(error) = result {
if !send_protocol_event(&mut socket, &protocol_error_event(error)).await { break; }
}
}
Some(Ok(WsMessage::Ping(value))) => { if socket.send(WsMessage::Pong(value)).await.is_err() { break; } }
Some(Ok(WsMessage::Pong(_))) => {}
_ => break,
},
event = connection.events.recv() => {
let Some(event) = event else { break; };
let shutdown = matches!(event, protocol::Event::Shutdown);
if !send_protocol_event(&mut socket, &event).await || shutdown { break; }
}
}
}
let _ = socket.send(WsMessage::Close(None)).await;
}
pub(crate) struct WorkspaceWorkerProtocolConnection {
pub(crate) methods: tokio::sync::mpsc::Sender<protocol::Method>,
pub(crate) events: tokio::sync::mpsc::Receiver<protocol::Event>,
}
pub(crate) async fn connect_workspace_worker_protocol(
api: &WorkspaceApi,
worker: &RuntimeWorkerRef,
@@ -29199,39 +29170,21 @@ async fn connect_remote_worker_protocol(
}
insert_authenticated_input_source_header(request.headers_mut(), input_source)?;
let (socket, _) =
connect_async(request)
tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, connect_async(request))
.await
.map_err(|_| Error::RuntimeOperationFailed {
runtime_id: config.worker.runtime_id.clone(),
code: "worker_protocol_connect_timeout".into(),
message: "Worker protocol connection handshake timed out; no method was admitted"
.into(),
})?
.map_err(|error| Error::RuntimeOperationFailed {
runtime_id: config.worker.runtime_id.clone(),
code: "worker_protocol_connect_failed".to_string(),
code: "worker_protocol_connect_failed".into(),
message: error.to_string(),
})?;
let (mut sink, mut stream) = socket.split();
let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256);
let (event_sender, events) = tokio::sync::mpsc::channel(512);
tokio::spawn(async move {
loop {
tokio::select! {
method = method_receiver.recv() => {
let Some(method) = method else { break };
let Ok(text) = protocol::stream::encode_method(&method) else { break };
if sink.send(TungsteniteMessage::Text(text.into())).await.is_err() { break; }
}
message = stream.next() => match message {
Some(Ok(TungsteniteMessage::Text(text))) => {
let Ok(event) = protocol::stream::decode_event(text.as_ref()) else { break; };
if event_sender.send(event).await.is_err() { break; }
}
Some(Ok(TungsteniteMessage::Ping(value))) => {
if sink.send(TungsteniteMessage::Pong(value)).await.is_err() { break; }
}
Some(Ok(TungsteniteMessage::Pong(_))) => {}
_ => break,
}
}
}
});
Ok(WorkspaceWorkerProtocolConnection { methods, events })
let (sink, stream) = socket.split();
Ok(remote_worker_protocol_bridge(sink, stream))
}
async fn connect_embedded_worker_protocol(
@@ -29249,47 +29202,48 @@ async fn connect_embedded_worker_protocol(
code: "embedded_worker_protocol_connect_failed".to_string(),
message: error.to_string(),
})?;
let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256);
let (event_sender, events) = tokio::sync::mpsc::channel(512);
tokio::spawn(async move {
if event_sender.send(transport.snapshot.clone()).await.is_err() {
return;
}
loop {
tokio::select! {
_ = event_sender.closed() => break,
method = method_receiver.recv() => {
let Some(method) = method else { break };
let shutdown = matches!(&method, protocol::Method::Shutdown { .. });
match source.runtime.send_connected_protocol_method_scoped(
&scope, &source.worker_ref, &transport, method,
) {
Ok(direct_events) => {
for event in direct_events {
let terminal = matches!(event, protocol::Event::Shutdown);
if event_sender.send(event).await.is_err() || terminal { return; }
Ok(spawn_worker_protocol_pump(
move |mut methods, events| async move {
if let Err(reason) = enqueue_worker_protocol_event(&events, transport.snapshot.clone())
{
return reason;
}
loop {
tokio::select! {
method = methods.recv() => {
let Some(method) = method else { return WorkerProtocolEnd::Disconnected; };
let shutdown = matches!(&method, protocol::Method::Shutdown { .. });
match source.runtime.send_connected_protocol_method_scoped(
&scope, &source.worker_ref, &transport, method,
) {
Ok(direct_events) => {
for event in direct_events {
let terminal = matches!(event, protocol::Event::Shutdown);
if let Err(reason) = enqueue_worker_protocol_event(&events, event) { return reason; }
if terminal { return WorkerProtocolEnd::Disconnected; }
}
}
Err(error) => {
let terminal = embedded_protocol_dispatch_is_terminal(&error, &transport);
if let Err(reason) = enqueue_worker_protocol_event(&events, protocol_error_event(error.to_string())) { return reason; }
if terminal { return WorkerProtocolEnd::Disconnected; }
}
}
Err(error) => {
let terminal = embedded_protocol_dispatch_is_terminal(&error, &transport);
if event_sender.send(protocol_error_event(error.to_string())).await.is_err() || terminal { return; }
if shutdown { return WorkerProtocolEnd::Disconnected; }
}
event = transport.events.recv() => match event {
Some(event) => {
let terminal = matches!(event, protocol::Event::Shutdown);
if let Err(reason) = enqueue_worker_protocol_event(&events, event) { return reason; }
if terminal { return WorkerProtocolEnd::Disconnected; }
}
// Stop, disconnect and upstream lag never cause an implicit reattach.
None => return WorkerProtocolEnd::Disconnected,
}
if shutdown { break; }
}
event = transport.events.recv() => match event {
Some(event) => {
let terminal = matches!(event, protocol::Event::Shutdown);
if event_sender.send(event).await.is_err() || terminal { break; }
}
// The execution relay closes on Stop, disconnect or lag.
// Uncertain streams are terminal, never grounds to reattach.
None => break,
}
}
}
});
Ok(WorkspaceWorkerProtocolConnection { methods, events })
},
))
}
fn embedded_protocol_dispatch_is_terminal(
@@ -32502,6 +32456,7 @@ mod tests {
mod drive_tests;
mod drive_web_roundtrip_tests;
mod external_checkout_tests;
mod protocol_backpressure_tests;
mod subject_spawn_tests;
mod subjektiv_jobs_tests;
mod ticket_evidence_tests;
@@ -32511,7 +32466,9 @@ mod tests {
mod workdir_delegation_tests;
mod worker_drive_tests;
mod worker_operations_tests;
mod worker_protocol_tests;
mod workspace_admission_tests;
mod workspace_protocol_backpressure_tests;
include!("server_workspace_config_tests.rs");
use super::*;
use axum::body::{Body, to_bytes};
crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs
+920
View File
@@ -0,0 +1,920 @@
use super::*;
const PROTOCOL_TEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(2);
// Completion replies are produced synchronously by the actual attached execution.
// Unlike flooding the fixture's 32-entry broadcast bus, this reaches the Server's
// event admission boundary without first causing unrelated Runtime relay lag.
#[tokio::test]
async fn embedded_event_overflow_terminates_without_losing_the_admitted_prefix() {
let (runtime, worker_ref, execution) = runtime_with_worker_backend();
let mut connection = connect_embedded_worker_protocol(
embedded_protocol_test_source(&runtime, &worker_ref),
embedded_protocol_test_scope(),
)
.await
.unwrap();
// Leave the snapshot queued. Exactly capacity - 1 completion replies fit;
// the next reply must end this connection rather than block its bridge.
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, async {
for index in 0..WORKER_PROTOCOL_EVENT_CAPACITY {
connection
.methods
.send(protocol_test_completion(&format!("overflow-{index}")))
.await
.unwrap();
}
assert_eq!(
connection.termination.ended().await,
WorkerProtocolEnd::EventOverflow
);
connection.methods.closed().await;
})
.await
.expect("full embedded event queue did not terminate the connection");
assert_embedded_protocol_snapshot(
next_embedded_protocol_event(&mut connection).await,
&worker_ref,
);
for index in 0..WORKER_PROTOCOL_EVENT_CAPACITY - 1 {
assert!(matches!(
next_embedded_protocol_event(&mut connection).await,
protocol::Event::Completions { request_id: Some(id), .. }
if id == format!("overflow-{index}")
));
}
assert!(
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, connection.events.recv())
.await
.expect("overflowed event receiver did not close")
.is_none()
);
assert!(
connection
.methods
.try_send(protocol_test_completion("after-overflow"))
.is_err()
);
assert_eq!(
execution.protocol_methods().len(),
WORKER_PROTOCOL_EVENT_CAPACITY
);
}
#[tokio::test]
async fn embedded_protocol_delivers_snapshot_then_ordered_replies_and_live_events() {
let (runtime, worker_ref, execution) = runtime_with_worker_backend();
let mut connection = connect_embedded_worker_protocol(
embedded_protocol_test_source(&runtime, &worker_ref),
embedded_protocol_test_scope(),
)
.await
.unwrap();
for id in ["first", "second", "third"] {
connection
.methods
.try_send(protocol_test_completion(id))
.unwrap();
}
assert_embedded_protocol_snapshot(
next_embedded_protocol_event(&mut connection).await,
&worker_ref,
);
for expected in ["first", "second", "third"] {
assert!(matches!(
next_embedded_protocol_event(&mut connection).await,
protocol::Event::Completions { request_id: Some(id), .. } if id == expected
));
}
execution.emit_worker_event(
&worker_ref,
protocol::Event::TextDelta {
text: "live-after-replies".into(),
},
);
assert!(matches!(
next_embedded_protocol_event(&mut connection).await,
protocol::Event::TextDelta { text } if text == "live-after-replies"
));
assert_eq!(connection.termination.reason(), None);
}
#[tokio::test]
async fn embedded_overflow_requires_explicit_reconnect_and_stop_ends_the_new_connection() {
let (runtime, worker_ref, execution) = runtime_with_worker_backend();
let source = embedded_protocol_test_source(&runtime, &worker_ref);
let mut old = connect_embedded_worker_protocol(source.clone(), embedded_protocol_test_scope())
.await
.unwrap();
assert_embedded_protocol_snapshot(next_embedded_protocol_event(&mut old).await, &worker_ref);
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, async {
for index in 0..=WORKER_PROTOCOL_EVENT_CAPACITY {
old.methods
.send(protocol_test_completion(&format!("old-{index}")))
.await
.unwrap();
}
assert_eq!(
old.termination.ended().await,
WorkerProtocolEnd::EventOverflow
);
old.methods.closed().await;
})
.await
.expect("overflowed connection remained live");
// Overflow ends only the attached connection, not its Worker. Reconnection
// is explicit and starts from a fresh execution snapshot, not old backlog.
let mut fresh = connect_embedded_worker_protocol(source, embedded_protocol_test_scope())
.await
.unwrap();
assert_embedded_protocol_snapshot(next_embedded_protocol_event(&mut fresh).await, &worker_ref);
fresh
.methods
.try_send(protocol_test_completion("explicit-reconnect"))
.unwrap();
assert!(matches!(
next_embedded_protocol_event(&mut fresh).await,
protocol::Event::Completions { request_id: Some(id), .. } if id == "explicit-reconnect"
));
assert!(
old.methods
.try_send(protocol_test_completion("must-not-reattach"))
.is_err()
);
assert_eq!(
old.termination.reason(),
Some(WorkerProtocolEnd::EventOverflow)
);
assert_eq!(
execution.protocol_methods().len(),
WORKER_PROTOCOL_EVENT_CAPACITY + 2
);
runtime.stop_worker(&worker_ref, None).unwrap();
assert_eq!(
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, fresh.termination.ended())
.await
.expect("Stop did not end the attached protocol connection"),
WorkerProtocolEnd::Disconnected
);
assert_embedded_protocol_closed(&mut fresh).await;
}
async fn backpressure_browser_context(api: &WorkspaceApi) -> (WorkerOperationContext, String) {
seed_test_api_token(api.store.as_ref(), "backpressure-browser");
let token = "backpressure-browser-session";
let token_hash = crate::auth::token_hash(token);
api.store
.create_browser_session(&BrowserSessionRecord {
token_hash: token_hash.clone(),
session_id: "backpressure-session".into(),
user_id: "user-backpressure-browser".into(),
created_at: "2026-01-01T00:00:00Z".into(),
expires_at: "2099-01-01T00:00:00Z".into(),
revoked_at: None,
})
.unwrap();
let AuthConfig::Passkey {
cookie_name,
origin,
..
} = &api.config.auth;
let mut headers = HeaderMap::new();
headers.insert(ORIGIN, origin.parse().unwrap());
headers.insert(
axum::http::header::COOKIE,
format!("{cookie_name}={token}").parse().unwrap(),
);
let actor = resolve_actor(&ServerAuthApi::from(api), &headers)
.await
.unwrap()
.expect("fixture browser session must authenticate");
(
WorkerOperationContext::Browser {
headers,
source: authenticated_browser_input_source(&actor),
},
token_hash,
)
}
async fn undrained_protocol_connection() -> (
WorkspaceWorkerProtocolConnection,
tokio::sync::mpsc::Receiver<protocol::Method>,
) {
let (receiver_sender, receiver) = tokio::sync::oneshot::channel();
let connection = spawn_worker_protocol_pump(move |methods, events| async move {
receiver_sender.send(methods).unwrap();
// Keep the event endpoint live until common termination cancels this
// fixture pump. Method consumption is controlled solely by the test.
let _events = events;
std::future::pending::<WorkerProtocolEnd>().await
});
let receiver = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, receiver)
.await
.expect("fixture pump did not expose its undrained method endpoint")
.unwrap();
(connection, receiver)
}
async fn full_protocol_method_queue() -> (
WorkspaceWorkerProtocolConnection,
tokio::sync::mpsc::Receiver<protocol::Method>,
) {
let (connection, receiver) = undrained_protocol_connection().await;
for index in 0..WORKER_PROTOCOL_METHOD_CAPACITY {
connection
.methods
.try_send(protocol_test_completion(&format!("queued-{index}")))
.unwrap();
}
(connection, receiver)
}
#[tokio::test]
async fn browser_method_queue_overflow_fails_fast_and_does_not_pin_session_cleanup() {
let workspace = tempfile::tempdir().unwrap();
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
let identity = spawn_ticket_check_source(&api, "method-overflow");
let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
let (context, _) = backpressure_browser_context(&api).await;
let (connection, mut receiver) = full_protocol_method_queue().await;
let termination = connection.termination.clone();
let sender =
WorkspaceWorkerMethodSender::new(worker.clone(), connection.methods, termination.clone());
let _events = connection.events;
let result = tokio::time::timeout(
PROTOCOL_TEST_TIMEOUT,
sender.send(
&context,
protocol::Method::Cancel {
command: protocol::WorkerCommandEnvelope::new(801),
},
),
)
.await
.expect("method admission blocked while holding the session guard");
assert!(result.is_err(), "full method queue must reject Cancel too");
assert_eq!(
termination.reason(),
Some(WorkerProtocolEnd::MethodOverflow)
);
assert_eq!(receiver.len(), WORKER_PROTOCOL_METHOD_CAPACITY);
assert!(matches!(
receiver.try_recv().unwrap(),
protocol::Method::ListCompletions { request_id: Some(id), .. } if id == "queued-0"
));
assert!(
sender
.send(
&context,
protocol_test_completion("must-not-retry-after-drain")
)
.await
.is_err()
);
assert_eq!(receiver.len(), WORKER_PROTOCOL_METHOD_CAPACITY - 1);
let stopped = tokio::time::timeout(
PROTOCOL_TEST_TIMEOUT,
worker.stop(
&WorkerOperationContext::Backend,
WorkerLifecycleRequest {
reason: Some("cleanup-after-method-overflow".into()),
ticket_assignment: None,
},
),
)
.await
.expect("full method queue pinned Stop's session cleanup")
.unwrap();
assert_eq!(stopped.state, InternalWorkerOperationState::Accepted);
assert!(execution.protocol_methods().is_empty());
for index in 1..WORKER_PROTOCOL_METHOD_CAPACITY {
assert!(matches!(
receiver.try_recv().unwrap(),
protocol::Method::ListCompletions { request_id: Some(id), .. }
if id == format!("queued-{index}")
));
}
assert!(receiver.try_recv().is_err());
}
#[tokio::test]
async fn revoked_browser_credentials_reject_before_method_overflow_admission() {
let workspace = tempfile::tempdir().unwrap();
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
let identity = spawn_ticket_check_source(&api, "revoked-full-queue");
let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
let (context, token_hash) = backpressure_browser_context(&api).await;
let (connection, receiver) = full_protocol_method_queue().await;
let termination = connection.termination.clone();
let sender = WorkspaceWorkerMethodSender::new(worker, connection.methods, termination.clone());
let _events = connection.events;
assert!(
api.store
.revoke_browser_session(&token_hash, &now_registry_timestamp())
.unwrap()
);
let result = tokio::time::timeout(
PROTOCOL_TEST_TIMEOUT,
sender.send(&context, protocol_test_completion("unauthorized")),
)
.await
.expect("revoked credentials waited for method capacity");
assert!(matches!(result, Err(Error::WorkspacePermissionDenied(_))));
assert_eq!(termination.reason(), None);
assert_eq!(receiver.len(), WORKER_PROTOCOL_METHOD_CAPACITY);
assert!(execution.protocol_methods().is_empty());
}
#[tokio::test]
async fn cancelling_pending_browser_admission_does_not_enqueue_or_hold_the_session_guard() {
let workspace = tempfile::tempdir().unwrap();
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
let identity = spawn_ticket_check_source(&api, "cancel-pending-admission");
let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
let (context, _) = backpressure_browser_context(&api).await;
let (connection, mut receiver) = undrained_protocol_connection().await;
let termination = connection.termination.clone();
let sender = WorkspaceWorkerMethodSender::new(worker, connection.methods, termination.clone());
let _events = connection.events;
let session_lock = current_worker_session_lock(&api, &identity);
let guard = session_lock.lock().await;
{
let admission = sender.send(&context, protocol_test_completion("cancelled-admission"));
tokio::pin!(admission);
assert!(futures::poll!(admission.as_mut()).is_pending());
// Dropping a blocked admission future must not schedule a detached retry.
}
drop(guard);
assert!(receiver.try_recv().is_err());
let guard = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session_lock.lock())
.await
.expect("cancelled admission retained the session guard");
close_current_worker_session_locked(&api, &identity)
.await
.unwrap();
drop(guard);
let command = protocol::WorkerCommandEnvelope::new(802);
sender
.send(&context, protocol::Method::Cancel { command })
.await
.unwrap();
assert!(matches!(
receiver.try_recv().unwrap(),
protocol::Method::Cancel { command: admitted } if admitted == command
));
assert!(receiver.try_recv().is_err());
assert_eq!(termination.reason(), None);
assert!(execution.protocol_methods().is_empty());
}
// In-process browser transport: one real JSON method frame, then an idle
// connection. Its writes are immediately ready, isolating authorization waits
// from transport congestion and without crossing a complete product flow.
struct AdmissionCancellationSocket {
inbound: Option<WsMessage>,
written: Vec<WsMessage>,
}
impl futures::Stream for AdmissionCancellationSocket {
type Item = std::result::Result<WsMessage, axum::Error>;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
match self.inbound.take() {
Some(message) => std::task::Poll::Ready(Some(Ok(message))),
None => std::task::Poll::Pending,
}
}
}
impl futures::Sink<WsMessage> for AdmissionCancellationSocket {
type Error = std::convert::Infallible;
fn poll_ready(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::result::Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
fn start_send(
mut self: std::pin::Pin<&mut Self>,
message: WsMessage,
) -> std::result::Result<(), Self::Error> {
self.written.push(message);
Ok(())
}
fn poll_flush(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::result::Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
fn poll_close(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::result::Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn browser_protocol_termination_cancels_admission_waiting_for_session_authority() {
let workspace = tempfile::tempdir().unwrap();
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
let identity = spawn_ticket_check_source(&api, "browser-authorization-cancellation");
let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
let (context, _) = backpressure_browser_context(&api).await;
let (connection, mut receiver) = undrained_protocol_connection().await;
let termination = connection.termination.clone();
let sender = WorkspaceWorkerMethodSender::new(worker, connection.methods, termination.clone());
let connection = worker_operations::WorkspaceWorkerConnection {
sender,
events: connection.events,
};
let session_lock = current_worker_session_lock(&api, &identity);
let guard = session_lock.lock().await;
let method = protocol::Method::Cancel {
command: protocol::WorkerCommandEnvelope::new(803),
};
let mut socket = AdmissionCancellationSocket {
inbound: Some(WsMessage::Text(
protocol::stream::encode_method(&method).unwrap().into(),
)),
written: Vec::new(),
};
{
let session = workspace_worker_protocol_session(connection, &mut socket, context);
tokio::pin!(session);
// SQLite browser authorization is synchronous; this poll consumes the
// frame and reaches the deliberately held session lock.
assert!(futures::poll!(session.as_mut()).is_pending());
termination.terminate(WorkerProtocolEnd::Disconnected);
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session)
.await
.expect("terminal connection did not cancel pending browser admission");
}
assert!(matches!(socket.written.last(), Some(WsMessage::Close(_))));
assert!(receiver.try_recv().is_err());
drop(guard);
let _guard = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session_lock.lock())
.await
.expect("cancelled browser authorization pinned session cleanup");
assert!(execution.protocol_methods().is_empty());
}
struct BackpressureTestServer(tokio::task::JoinHandle<()>);
impl Drop for BackpressureTestServer {
fn drop(&mut self) {
self.0.abort();
}
}
#[tokio::test]
async fn browser_protocol_session_ends_on_full_method_queue_without_blocking_cleanup() {
let workspace = tempfile::tempdir().unwrap();
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
let identity = spawn_ticket_check_source(&api, "browser-session-overflow");
let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
let (context, _) = backpressure_browser_context(&api).await;
let (connection, receiver) = full_protocol_method_queue().await;
let termination = connection.termination.clone();
let sender = WorkspaceWorkerMethodSender::new(worker, connection.methods, termination.clone());
let events = connection.events;
let (session_done, finished) = tokio::sync::oneshot::channel();
let pending = Arc::new(std::sync::Mutex::new(Some((
worker_operations::WorkspaceWorkerConnection { sender, events },
session_done,
))));
let app = Router::new().route(
"/protocol/ws",
get(move |ws: WebSocketUpgrade| {
let pending = pending.clone();
let context = context.clone();
async move {
ws.on_upgrade(move |socket| async move {
let (connection, session_done) = pending.lock().unwrap().take().unwrap();
workspace_worker_protocol_session(connection, socket, context).await;
let _ = session_done.send(());
})
}
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let _server = BackpressureTestServer(tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
}));
let (mut socket, _) = connect_async(format!("ws://{address}/protocol/ws"))
.await
.unwrap();
tokio::time::timeout(
PROTOCOL_TEST_TIMEOUT,
send_protocol_test_ws_method(&mut socket, &protocol_test_completion("browser-overflow")),
)
.await
.expect("browser method frame did not reach the loopback transport");
assert_eq!(
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, termination.ended())
.await
.expect("browser session did not report method overflow"),
WorkerProtocolEnd::MethodOverflow
);
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, finished)
.await
.expect("browser protocol session did not finish after overflow")
.expect("browser protocol session dropped its completion signal");
let session_lock = current_worker_session_lock(&api, &identity);
let _guard = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session_lock.lock())
.await
.expect("browser method overflow pinned session cleanup");
close_current_worker_session_locked(&api, &identity)
.await
.unwrap();
assert_eq!(receiver.len(), WORKER_PROTOCOL_METHOD_CAPACITY);
assert!(execution.protocol_methods().is_empty());
}
// Accepting an output frame is not proof of delivery: this browser transport
// never finishes flushing, including the final best-effort Close frame.
struct StalledBrowserOutput {
entered: Option<tokio::sync::oneshot::Sender<()>>,
dropped: Option<tokio::sync::oneshot::Sender<()>>,
written: Arc<std::sync::Mutex<Vec<WsMessage>>>,
}
impl futures::Stream for StalledBrowserOutput {
type Item = std::result::Result<WsMessage, axum::Error>;
fn poll_next(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
std::task::Poll::Pending
}
}
impl futures::Sink<WsMessage> for StalledBrowserOutput {
type Error = std::convert::Infallible;
fn poll_ready(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::result::Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
fn start_send(
mut self: std::pin::Pin<&mut Self>,
message: WsMessage,
) -> std::result::Result<(), Self::Error> {
self.written.lock().unwrap().push(message);
if let Some(entered) = self.entered.take() {
let _ = entered.send(());
}
Ok(())
}
fn poll_flush(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::result::Result<(), Self::Error>> {
std::task::Poll::Pending
}
fn poll_close(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::result::Result<(), Self::Error>> {
std::task::Poll::Pending
}
}
impl Drop for StalledBrowserOutput {
fn drop(&mut self) {
if let Some(dropped) = self.dropped.take() {
let _ = dropped.send(());
}
}
}
fn assert_bounded_unknown_delivery_close(messages: &[WsMessage]) {
let Some(WsMessage::Close(Some(frame))) = messages.last() else {
panic!("terminal browser session must attempt an explicit Close frame: {messages:?}");
};
assert_eq!(frame.code, 1011);
assert!(frame.reason.contains("delivery unknown"), "{frame:?}");
assert!(frame.reason.contains("no retry"), "{frame:?}");
assert!(
frame.reason.len() < 123,
"Close reason exceeds its wire bound: {frame:?}"
);
}
async fn assert_same_worker_reconnect_and_stop(
worker: &WorkspaceWorker,
context: &WorkerOperationContext,
) {
let mut fresh = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, worker.connect_protocol(context))
.await
.expect("terminated transport retained same Worker's connection authority")
.unwrap();
assert!(matches!(
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, fresh.events.recv())
.await
.expect("same Worker reconnect did not deliver a fresh snapshot"),
Some(protocol::Event::Snapshot { .. })
));
let stopped = tokio::time::timeout(
PROTOCOL_TEST_TIMEOUT,
worker.stop(
&WorkerOperationContext::Backend,
WorkerLifecycleRequest {
reason: Some("cleanup-after-stalled-protocol".into()),
ticket_assignment: None,
},
),
)
.await
.expect("terminated transport pinned same Worker's Stop/session cleanup")
.unwrap();
assert_eq!(stopped.state, InternalWorkerOperationState::Accepted);
assert_eq!(
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, fresh.sender.termination.ended())
.await
.expect("Stop did not terminate the fresh same Worker connection"),
WorkerProtocolEnd::Disconnected
);
}
async fn assert_stalled_browser_output_reclaimed(explicit_termination: bool) {
let workspace = tempfile::tempdir().unwrap();
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
let identity = spawn_ticket_check_source(&api, "stalled-browser-output");
let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
let (context, _) = backpressure_browser_context(&api).await;
// Use the real attached embedded execution snapshot as browser output.
let raw = connect_workspace_worker_protocol(&api, &identity, None)
.await
.unwrap();
let methods = raw.methods.clone();
let termination = raw.termination;
let connection = worker_operations::WorkspaceWorkerConnection {
sender: WorkspaceWorkerMethodSender::new(worker.clone(), raw.methods, termination.clone()),
events: raw.events,
};
let (entered_sender, entered) = tokio::sync::oneshot::channel();
let (dropped_sender, dropped) = tokio::sync::oneshot::channel();
let written = Arc::new(std::sync::Mutex::new(Vec::new()));
let socket = StalledBrowserOutput {
entered: Some(entered_sender),
dropped: Some(dropped_sender),
written: written.clone(),
};
let started = tokio::time::Instant::now();
{
let session = workspace_worker_protocol_session(connection, socket, context.clone());
tokio::pin!(session);
assert!(futures::poll!(session.as_mut()).is_pending());
// Poll the real session alongside its output-acceptance oneshot rather
// than assuming how many scheduler turns the attached pump needs.
let output_accepted = async {
tokio::select! {
result = entered => result.expect("browser sink dropped before receiving output"),
_ = session.as_mut() => panic!("browser session ended before accepting output"),
}
};
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, output_accepted)
.await
.expect("attached snapshot never reached the stalled browser output");
{
let messages = written.lock().unwrap();
assert_eq!(messages.len(), 1);
let WsMessage::Text(text) = &messages[0] else {
panic!("first browser output must be the attached snapshot");
};
assert!(matches!(
protocol::stream::decode_event(text.as_str()).unwrap(),
protocol::Event::Snapshot { .. }
));
}
if explicit_termination {
termination.terminate(WorkerProtocolEnd::Disconnected);
} else {
tokio::time::advance(std::time::Duration::from_secs(9)).await;
assert!(futures::poll!(session.as_mut()).is_pending());
assert_eq!(termination.reason(), None);
tokio::time::advance(std::time::Duration::from_secs(1)).await;
}
assert!(futures::poll!(session.as_mut()).is_pending());
assert_eq!(
termination.reason(),
Some(if explicit_termination {
WorkerProtocolEnd::Disconnected
} else {
WorkerProtocolEnd::TransportTimeout
}),
"browser end must publish the concrete common termination cause"
);
{
let messages = written.lock().unwrap();
assert_eq!(messages.len(), 2, "must not retry uncertain browser output");
assert_bounded_unknown_delivery_close(&messages);
}
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, methods.closed())
.await
.expect("browser termination did not cancel its actual attached pump");
let session_lock = current_worker_session_lock(&api, &identity);
let guard = tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session_lock.lock())
.await
.expect("blocked browser Close pinned session ownership");
close_current_worker_session_locked(&api, &identity)
.await
.unwrap();
drop(guard);
// The terminal Close itself is stalled. It must get a separate bounded
// attempt, not become a second unbounded output wait.
tokio::time::advance(std::time::Duration::from_secs(9)).await;
assert!(futures::poll!(session.as_mut()).is_pending());
tokio::time::advance(std::time::Duration::from_secs(1)).await;
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, session.as_mut())
.await
.expect("stalled terminal Close did not reach its deadline");
}
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, dropped)
.await
.expect("completed browser session retained its stalled sink")
.expect("browser sink dropped without its ownership signal");
assert_eq!(written.lock().unwrap().len(), 2);
assert_eq!(
tokio::time::Instant::now() - started,
std::time::Duration::from_secs(if explicit_termination { 10 } else { 20 })
);
assert!(execution.protocol_methods().is_empty());
assert_same_worker_reconnect_and_stop(&worker, &context).await;
}
#[tokio::test(start_paused = true)]
async fn browser_blocked_output_deadline_ends_pump_and_bounds_the_terminal_close() {
assert_stalled_browser_output_reclaimed(false).await;
}
#[tokio::test(start_paused = true)]
async fn browser_common_termination_preempts_blocked_output_and_bounds_the_terminal_close() {
assert_stalled_browser_output_reclaimed(true).await;
}
#[tokio::test]
async fn embedded_disconnect_and_last_owner_drop_release_the_execution_event_subscription() {
for explicit_termination in [false, true] {
let (runtime, worker_ref, execution) = runtime_with_worker_backend();
let endpoint = execution
.contexts
.lock()
.unwrap()
.get(&worker_ref)
.cloned()
.unwrap();
let mut connection = connect_embedded_worker_protocol(
embedded_protocol_test_source(&runtime, &worker_ref),
embedded_protocol_test_scope(),
)
.await
.unwrap();
assert_embedded_protocol_snapshot(
next_embedded_protocol_event(&mut connection).await,
&worker_ref,
);
assert_eq!(endpoint.events.receiver_count(), 1);
let WorkspaceWorkerProtocolConnection {
methods,
mut events,
termination,
} = connection;
let last_owner = termination.clone();
drop(termination);
assert!(!methods.is_closed());
if explicit_termination {
last_owner.terminate(WorkerProtocolEnd::Disconnected);
// Keep the owner alive until watch cancellation is observed so
// abort-on-drop cannot hide a broken explicit-disconnect path.
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, methods.closed())
.await
.expect("embedded explicit disconnect did not reclaim its pump");
}
// Both channel endpoints remain held. Without explicit termination,
// only dropping the last client owner can cancel this attached pump.
drop(last_owner);
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, methods.closed())
.await
.expect("embedded last-owner drop did not reclaim its method pump");
assert!(
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, events.recv())
.await
.expect("reclaimed embedded event endpoint stayed open")
.is_none()
);
// The fixture's execution relay notices its dropped receiver on its
// next event. Wake that relay once, then observe subscription closure.
let _ = endpoint.events.send(protocol::Event::TextDelta {
text: "execution-relay-close-probe".into(),
});
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, endpoint.events.closed())
.await
.expect("dropped embedded transport retained its execution event subscription");
assert_eq!(endpoint.events.receiver_count(), 0);
assert!(execution.protocol_methods().is_empty());
runtime.stop_worker(&worker_ref, None).unwrap();
}
}
#[tokio::test(start_paused = true)]
async fn stalled_remote_bridge_browser_cancel_overflow_reclaims_wire_and_allows_same_worker_progress()
{
let workspace = tempfile::tempdir().unwrap();
let (api, execution) = test_api_with_recording_backend(workspace.path()).await;
let identity = spawn_ticket_check_source(&api, "remote-browser-overflow");
let worker = WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
let (context, _) = backpressure_browser_context(&api).await;
let (raw, entered, dropped, sent) = super::worker_protocol_tests::stalled_remote_bridge();
raw.methods
.try_send(protocol_test_completion("wire-delivery-unknown"))
.unwrap();
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, entered)
.await
.expect("actual remote pump did not enter its stalled wire send")
.unwrap();
for index in 0..WORKER_PROTOCOL_METHOD_CAPACITY {
raw.methods
.try_send(protocol_test_completion(&format!("remote-queued-{index}")))
.unwrap();
}
let methods = raw.methods.clone();
let termination = raw.termination;
let connection = worker_operations::WorkspaceWorkerConnection {
sender: WorkspaceWorkerMethodSender::new(worker.clone(), raw.methods, termination.clone()),
events: raw.events,
};
let cancel = protocol::Method::Cancel {
command: protocol::WorkerCommandEnvelope::new(804),
};
let mut socket = AdmissionCancellationSocket {
inbound: Some(WsMessage::Text(
protocol::stream::encode_method(&cancel).unwrap().into(),
)),
written: Vec::new(),
};
// This real browser session decodes Cancel and passes the authenticated
// context through the actual method sender into the saturated remote pump.
tokio::time::timeout(
PROTOCOL_TEST_TIMEOUT,
workspace_worker_protocol_session(connection, &mut socket, context.clone()),
)
.await
.expect("browser Cancel waited for the stalled remote wire instead of failing admission");
assert_eq!(
termination.reason(),
Some(WorkerProtocolEnd::MethodOverflow)
);
assert!(socket.inbound.is_none());
assert_bounded_unknown_delivery_close(&socket.written);
tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, dropped)
.await
.expect("method overflow did not drop the stalled remote sink")
.unwrap();
assert!(methods.is_closed());
{
let messages = sent.lock().unwrap();
assert_eq!(
messages.len(),
1,
"must not retry or drain unknown remote delivery"
);
let TungsteniteMessage::Text(text) = &messages[0] else {
panic!("remote pump must send an actual encoded method frame");
};
assert!(matches!(
decode_method(text.as_str()).unwrap(),
protocol::Method::ListCompletions { request_id: Some(id), .. }
if id == "wire-delivery-unknown"
));
}
assert!(execution.protocol_methods().is_empty());
assert_same_worker_reconnect_and_stop(&worker, &context).await;
}
crates/workspace-server/src/server/tests/worker_protocol_tests.rs
+210
View File
@@ -0,0 +1,210 @@
use super::*;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::sync::oneshot;
// The real bridge owns this sink exactly as it owns a TCP WebSocket sink. It
// accepts bytes but never finishes flushing: delivery is deliberately unknown.
struct StalledWire {
entered: Option<oneshot::Sender<()>>,
dropped: Option<oneshot::Sender<()>>,
sent: std::sync::Arc<std::sync::Mutex<Vec<TungsteniteMessage>>>,
}
impl futures::Sink<TungsteniteMessage> for StalledWire {
type Error = tokio_tungstenite::tungstenite::Error;
fn poll_ready(
self: Pin<&mut Self>,
_: &mut Context<'_>,
) -> Poll<std::result::Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(
mut self: Pin<&mut Self>,
message: TungsteniteMessage,
) -> std::result::Result<(), Self::Error> {
self.sent.lock().unwrap().push(message);
if let Some(entered) = self.entered.take() {
let _ = entered.send(());
}
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_: &mut Context<'_>,
) -> Poll<std::result::Result<(), Self::Error>> {
Poll::Pending
}
fn poll_close(
self: Pin<&mut Self>,
_: &mut Context<'_>,
) -> Poll<std::result::Result<(), Self::Error>> {
Poll::Pending
}
}
impl Drop for StalledWire {
fn drop(&mut self) {
if let Some(dropped) = self.dropped.take() {
let _ = dropped.send(());
}
}
}
pub(super) fn stalled_remote_bridge() -> (
WorkspaceWorkerProtocolConnection,
oneshot::Receiver<()>,
oneshot::Receiver<()>,
std::sync::Arc<std::sync::Mutex<Vec<TungsteniteMessage>>>,
) {
let (entered_tx, entered) = oneshot::channel();
let (dropped_tx, dropped) = oneshot::channel();
let sent = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let sink = StalledWire {
entered: Some(entered_tx),
dropped: Some(dropped_tx),
sent: sent.clone(),
};
let connection = remote_worker_protocol_bridge(sink, futures::stream::pending());
(connection, entered, dropped, sent)
}
#[tokio::test(start_paused = true)]
async fn remote_pending_method_timeout_closes_both_queues_and_never_retries_unknown_delivery() {
let (mut connection, entered, dropped, sent) = stalled_remote_bridge();
connection
.methods
.try_send(protocol_test_completion("delivery-unknown"))
.unwrap();
entered.await.unwrap();
for _ in 0..WORKER_PROTOCOL_METHOD_CAPACITY {
connection
.methods
.try_send(protocol::Method::ListPendingSubmissions)
.unwrap();
}
// Queues are saturated while the actual bridge is stalled in a wire send.
assert!(
connection
.methods
.try_send(protocol::Method::ListRewindTargets)
.is_err()
);
tokio::time::advance(std::time::Duration::from_secs(10)).await;
assert_eq!(
connection.termination.ended().await,
WorkerProtocolEnd::TransportTimeout
);
dropped.await.unwrap();
assert!(connection.methods.is_closed());
assert!(connection.events.recv().await.is_none());
assert_eq!(
sent.lock().unwrap().len(),
1,
"must not retry or drain uncertain pending methods"
);
}
#[tokio::test]
async fn remote_disconnect_and_last_capability_drop_reclaim_a_pending_wire_send() {
for explicit in [false, true] {
let (connection, entered, dropped, sent) = stalled_remote_bridge();
connection
.methods
.try_send(protocol_test_completion("cancelled-wire"))
.unwrap();
entered.await.unwrap();
if explicit {
connection
.termination
.terminate(WorkerProtocolEnd::Disconnected);
}
drop(connection);
tokio::time::timeout(std::time::Duration::from_secs(2), dropped)
.await
.unwrap()
.unwrap();
assert_eq!(sent.lock().unwrap().len(), 1);
}
}
#[tokio::test]
async fn remote_actual_connector_overflow_ends_connection_and_preserves_delivered_event_order() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (release_tx, release) = oneshot::channel();
let (wire_closed_tx, wire_closed) = oneshot::channel();
let peer = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut socket = tokio_tungstenite::accept_hdr_async(
stream,
|request: &tokio_tungstenite::tungstenite::handshake::server::Request, response| {
assert_eq!(
request.headers()["authorization"],
"Bearer test-worker-protocol"
);
assert_eq!(
request.headers()[protocol::AUTHENTICATED_ACCOUNT_ID_HEADER],
"test-account"
);
Ok(response)
},
)
.await
.unwrap();
release.await.unwrap();
for i in 0..=WORKER_PROTOCOL_EVENT_CAPACITY {
let event = protocol::Event::TextDelta {
text: i.to_string(),
};
let text = protocol::stream::encode_event(&event).unwrap();
// TCP may observe closure before all writes flush; either way the
// consumer must reach the same explicit bounded-overflow end.
if socket
.send(TungsteniteMessage::Text(text.into()))
.await
.is_err()
{
break;
}
}
while let Some(Ok(message)) = socket.next().await {
if matches!(message, TungsteniteMessage::Close(_)) {
break;
}
}
let _ = wire_closed_tx.send(());
});
let config = RuntimeObservationSourceConfig {
worker: RuntimeWorkerRef::new("test-runtime", "test-worker"),
endpoint: format!("ws://{addr}/protocol/ws"),
bearer_token: Some("test-worker-protocol".into()),
};
let mut connection = connect_remote_worker_protocol(
config,
Some(&protocol::AuthenticatedInputSource::Account {
account_id: "test-account".into(),
}),
)
.await
.unwrap();
release_tx.send(()).unwrap();
assert_eq!(
tokio::time::timeout(
std::time::Duration::from_secs(2),
connection.termination.ended()
)
.await
.unwrap(),
WorkerProtocolEnd::EventOverflow
);
for i in 0..WORKER_PROTOCOL_EVENT_CAPACITY {
assert!(
matches!(connection.events.recv().await, Some(protocol::Event::TextDelta { text }) if text == i.to_string())
);
}
assert!(connection.events.recv().await.is_none());
assert!(connection.methods.is_closed());
tokio::time::timeout(std::time::Duration::from_secs(2), wire_closed)
.await
.unwrap()
.unwrap();
peer.await.unwrap();
}
crates/workspace-server/src/server/tests/workspace_protocol_backpressure_tests.rs
+140
View File
@@ -0,0 +1,140 @@
use super::*;
use protocol::subscription::{
EventSubscriptionSelector, SubscriptionFrame, SubscriptionFramePayload, SubscriptionRequest,
SubscriptionRequestId,
};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::sync::{mpsc, oneshot};
// A controlled browser wire. Receiving remains independent of a blocked writer,
// reproducing the real multiplexer split without filling an OS socket by timing.
struct CongestedBrowser {
inbound: mpsc::Receiver<WsMessage>,
entered: Option<oneshot::Sender<()>>,
dropped: Option<oneshot::Sender<()>>,
}
impl futures::Stream for CongestedBrowser {
type Item = std::result::Result<WsMessage, axum::Error>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.inbound.poll_recv(cx).map(|message| message.map(Ok))
}
}
impl futures::Sink<WsMessage> for CongestedBrowser {
type Error = std::convert::Infallible;
fn poll_ready(
self: Pin<&mut Self>,
_: &mut Context<'_>,
) -> Poll<std::result::Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(mut self: Pin<&mut Self>, _: WsMessage) -> std::result::Result<(), Self::Error> {
if let Some(entered) = self.entered.take() {
let _ = entered.send(());
}
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_: &mut Context<'_>,
) -> Poll<std::result::Result<(), Self::Error>> {
Poll::Pending
}
fn poll_close(
self: Pin<&mut Self>,
_: &mut Context<'_>,
) -> Poll<std::result::Result<(), Self::Error>> {
Poll::Pending
}
}
impl Drop for CongestedBrowser {
fn drop(&mut self) {
if let Some(dropped) = self.dropped.take() {
let _ = dropped.send(());
}
}
}
#[tokio::test(start_paused = true)]
async fn actual_workspace_multiplexer_stalled_writer_timeout_or_disconnect_reclaims_connection_tasks()
{
for exit in ["timeout", "disconnect", "caller-cancel"] {
let workspace = tempfile::tempdir().unwrap();
let (api, _) = test_api_with_recording_backend(workspace.path()).await;
let identity = spawn_ticket_check_source(&api, "multiplexer-stalled-writer");
let (inbound, receiver) = mpsc::channel(300);
let (entered_tx, entered) = oneshot::channel();
let (dropped_tx, dropped) = oneshot::channel();
let socket = CongestedBrowser {
inbound: receiver,
entered: Some(entered_tx),
dropped: Some(dropped_tx),
};
let frame = SubscriptionFrame::new(SubscriptionFramePayload::Request(
SubscriptionRequest::SubscribeEvents {
request_id: SubscriptionRequestId::new("saturated-browser").unwrap(),
selector: EventSubscriptionSelector::WorkspaceWorkers,
},
));
inbound
.try_send(WsMessage::Text(
serde_json::to_string(&frame).unwrap().into(),
))
.unwrap();
let task = tokio::spawn(crate::workspace_subscription::serve_workspace_subscription(
api.clone(),
socket,
WorkerOperationContext::Backend,
));
entered.await.unwrap(); // Actual subscription snapshot reached the blocked wire.
match exit {
"timeout" => {
// Fill the control queue as well; the receiving side may now be
// waiting for Pong capacity. Writer failure must preempt that wait.
for _ in 0..258 {
inbound
.try_send(WsMessage::Ping(Vec::new().into()))
.unwrap();
}
tokio::time::advance(std::time::Duration::from_secs(10)).await;
task.await.unwrap();
}
"disconnect" => {
drop(inbound);
task.await.unwrap();
}
"caller-cancel" => {
task.abort();
assert!(task.await.unwrap_err().is_cancelled());
}
_ => unreachable!(),
}
dropped.await.unwrap();
let worker =
WorkspaceWorker::resolve(&api, &identity.runtime_id, &identity.worker_id).unwrap();
let mut fresh = worker
.connect_protocol(&WorkerOperationContext::Backend)
.await
.unwrap();
assert!(matches!(
fresh.events.recv().await,
Some(protocol::Event::Snapshot { .. })
));
let stopped = worker
.stop(
&WorkerOperationContext::Backend,
WorkerLifecycleRequest {
reason: Some("cleanup-after-multiplexer-exit".into()),
ticket_assignment: None,
},
)
.await
.unwrap();
assert_eq!(stopped.state, InternalWorkerOperationState::Accepted);
assert!(
current_worker_session_lock(&api, &identity)
.try_lock()
.is_ok()
);
}
}
crates/workspace-server/src/server/worker_operations.rs
+37 -11
View File
@@ -470,10 +470,11 @@ impl WorkspaceWorker {
let connection =
connect_workspace_worker_protocol(&self.api, &self.identity, source).await?;
Ok(WorkspaceWorkerConnection {
sender: WorkspaceWorkerMethodSender {
worker: self.clone(),
methods: connection.methods,
},
sender: WorkspaceWorkerMethodSender::new(
self.clone(),
connection.methods,
connection.termination,
),
events: connection.events,
})
}
@@ -485,12 +486,27 @@ impl WorkspaceWorker {
pub(crate) struct WorkspaceWorkerMethodSender {
worker: WorkspaceWorker,
methods: tokio::sync::mpsc::Sender<protocol::Method>,
pub(super) termination: WorkerProtocolTermination,
}
pub(crate) struct WorkspaceWorkerConnection {
pub(crate) sender: WorkspaceWorkerMethodSender,
pub(crate) events: tokio::sync::mpsc::Receiver<protocol::Event>,
}
impl WorkspaceWorkerMethodSender {
pub(crate) fn end_status(&self) -> tokio::sync::watch::Receiver<Option<WorkerProtocolEnd>> {
self.termination.subscribe()
}
pub(super) fn new(
worker: WorkspaceWorker,
methods: tokio::sync::mpsc::Sender<protocol::Method>,
termination: WorkerProtocolTermination,
) -> Self {
Self {
worker,
methods,
termination,
}
}
pub(crate) async fn send(
&self,
context: &WorkerOperationContext,
@@ -604,15 +620,25 @@ impl WorkspaceWorkerMethodSender {
};
let method = authorize_browser_worker_method(method, source)
.map_err(|message| Error::InvalidInput(message.into()))?;
self.methods.send(method).await.map_err(|_| {
// Session exclusion fences attach/stop only through admission, never delivery.
let result = if self.termination.reason().is_some() {
Err(tokio::sync::mpsc::error::TrySendError::Closed(method))
} else {
self.methods.try_send(method)
};
result.map_err(|error| {
let (code, message) = match error {
tokio::sync::mpsc::error::TrySendError::Full(_) => {
self.termination.terminate(WorkerProtocolEnd::MethodOverflow);
("worker_protocol_overflow", "Worker method queue overflow; this method was not admitted. Earlier pending methods may have unknown delivery; reconnect with a fresh snapshot, no automatic retry")
}
tokio::sync::mpsc::error::TrySendError::Closed(_) =>
("worker_protocol_closed", "Worker connection closed; pending delivery may be unknown; no implicit Restore or retry was performed"),
};
Error::RuntimeOperationFailed {
runtime_id: self.worker.identity.runtime_id.clone(),
code: "worker_protocol_closed".into(),
message:
"Worker execution connection closed; no implicit Restore or retry was performed"
.into(),
}
.into()
code: code.into(), message: message.into(),
}.into()
})
}
}
crates/workspace-server/src/server/worker_protocol.rs
+247
View File
@@ -0,0 +1,247 @@
//! Bounded operational protocol connections. Queue admission never waits for delivery.
//! Overflow ends the connection: reconnect/resubscribe with a fresh snapshot, never replay
//! a method whose transport delivery or execution outcome may be unknown.
use super::*;
use futures::{Sink, Stream};
use std::future::Future;
use tokio::sync::{mpsc, watch};
pub(super) const WORKER_PROTOCOL_METHOD_CAPACITY: usize = 256;
pub(super) const WORKER_PROTOCOL_EVENT_CAPACITY: usize = 512;
pub(super) const WORKER_PROTOCOL_TRANSPORT_TIMEOUT: std::time::Duration =
std::time::Duration::from_secs(10);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum WorkerProtocolEnd {
MethodOverflow,
EventOverflow,
TransportTimeout,
Disconnected,
}
impl WorkerProtocolEnd {
fn message(self) -> &'static str {
match self {
Self::MethodOverflow => "Worker protocol method queue overflow",
Self::EventOverflow => "Worker protocol event queue overflow/lag",
Self::TransportTimeout => "Worker protocol transport timed out",
Self::Disconnected => "Worker protocol disconnected",
}
}
}
struct WorkerProtocolPump(tokio::task::AbortHandle);
impl Drop for WorkerProtocolPump {
fn drop(&mut self) {
self.0.abort();
}
}
/// The last client capability owns the pump. The pump holds only the watch sender,
/// not this owner, so dropping a cancelled connection cannot leave an orphan task.
#[derive(Clone)]
pub(crate) struct WorkerProtocolTermination {
state: watch::Sender<Option<WorkerProtocolEnd>>,
_pump: std::sync::Arc<WorkerProtocolPump>,
}
impl WorkerProtocolTermination {
pub(crate) fn subscribe(&self) -> watch::Receiver<Option<WorkerProtocolEnd>> {
self.state.subscribe()
}
pub(crate) fn reason(&self) -> Option<WorkerProtocolEnd> {
*self.state.borrow()
}
pub(crate) fn terminate(&self, reason: WorkerProtocolEnd) {
finish(&self.state, reason);
}
pub(crate) async fn ended(&self) -> WorkerProtocolEnd {
ended(&mut self.state.subscribe()).await
}
}
fn finish(state: &watch::Sender<Option<WorkerProtocolEnd>>, reason: WorkerProtocolEnd) {
state.send_if_modified(|current| {
if current.is_some() {
return false;
}
*current = Some(reason);
true
});
}
async fn ended(state: &mut watch::Receiver<Option<WorkerProtocolEnd>>) -> WorkerProtocolEnd {
loop {
if let Some(reason) = *state.borrow_and_update() {
return reason;
}
if state.changed().await.is_err() {
return WorkerProtocolEnd::Disconnected;
}
}
}
pub(crate) struct WorkspaceWorkerProtocolConnection {
pub(crate) methods: mpsc::Sender<protocol::Method>,
pub(crate) events: mpsc::Receiver<protocol::Event>,
pub(crate) termination: WorkerProtocolTermination,
}
pub(super) fn spawn_worker_protocol_pump<F, Fut>(pump: F) -> WorkspaceWorkerProtocolConnection
where
F: FnOnce(mpsc::Receiver<protocol::Method>, mpsc::Sender<protocol::Event>) -> Fut,
Fut: Future<Output = WorkerProtocolEnd> + Send + 'static,
{
let (methods, receiver) = mpsc::channel(WORKER_PROTOCOL_METHOD_CAPACITY);
let (sender, events) = mpsc::channel(WORKER_PROTOCOL_EVENT_CAPACITY);
let (state, mut shutdown) = watch::channel(None);
let task_state = state.clone();
let closed = sender.clone();
let future = pump(receiver, sender);
let task = tokio::spawn(async move {
let reason = tokio::select! {
biased;
reason = ended(&mut shutdown) => reason,
_ = closed.closed() => WorkerProtocolEnd::Disconnected,
reason = future => reason,
};
finish(&task_state, reason);
// Dropping the future releases both queues and the captured execution/WS endpoint.
});
WorkspaceWorkerProtocolConnection {
methods,
events,
termination: WorkerProtocolTermination {
state,
_pump: std::sync::Arc::new(WorkerProtocolPump(task.abort_handle())),
},
}
}
pub(super) fn enqueue_worker_protocol_event(
sender: &mpsc::Sender<protocol::Event>,
event: protocol::Event,
) -> std::result::Result<(), WorkerProtocolEnd> {
sender.try_send(event).map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => WorkerProtocolEnd::EventOverflow,
mpsc::error::TrySendError::Closed(_) => WorkerProtocolEnd::Disconnected,
})
}
/// Actual remote bridge pump, also usable with a bounded in-process wire fixture.
/// A transport send may have partially succeeded when interrupted: no retry or success
/// acknowledgement is synthesized. All transport waits have a deadline; cancellation
/// of the common owner drops the pending send and both halves of this one connection.
pub(super) fn remote_worker_protocol_bridge<S, R>(
mut sink: S,
mut stream: R,
) -> WorkspaceWorkerProtocolConnection
where
S: Sink<TungsteniteMessage> + Unpin + Send + 'static,
R: Stream<
Item = std::result::Result<TungsteniteMessage, tokio_tungstenite::tungstenite::Error>,
> + Unpin
+ Send
+ 'static,
S::Error: Send,
{
spawn_worker_protocol_pump(move |mut methods, events| async move {
loop {
tokio::select! {
method = methods.recv() => {
let Some(method) = method else { return WorkerProtocolEnd::Disconnected; };
let Ok(text) = protocol::stream::encode_method(&method) else { return WorkerProtocolEnd::Disconnected; };
match tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, sink.send(TungsteniteMessage::Text(text.into()))).await {
Ok(Ok(())) => {},
Ok(Err(_)) => return WorkerProtocolEnd::Disconnected,
Err(_) => return WorkerProtocolEnd::TransportTimeout,
}
},
message = stream.next() => match message {
Some(Ok(TungsteniteMessage::Text(text))) => {
let Ok(event) = protocol::stream::decode_event(text.as_ref()) else { return WorkerProtocolEnd::Disconnected; };
if let Err(reason) = enqueue_worker_protocol_event(&events, event) { return reason; }
}
Some(Ok(TungsteniteMessage::Ping(value))) => {
match tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, sink.send(TungsteniteMessage::Pong(value))).await {
Ok(Ok(())) => {},
Ok(Err(_)) => return WorkerProtocolEnd::Disconnected,
Err(_) => return WorkerProtocolEnd::TransportTimeout,
}
}
Some(Ok(TungsteniteMessage::Pong(_))) => {},
_ => return WorkerProtocolEnd::Disconnected,
}
}
}
})
}
/// Browser admission can await authorization/session ownership, but never method queue
/// capacity. Common termination preempts that wait. Browser writes are bounded too.
pub(super) async fn workspace_worker_protocol_session<S>(
mut connection: worker_operations::WorkspaceWorkerConnection,
mut socket: S,
context: WorkerOperationContext,
) where
S: Stream<Item = std::result::Result<WsMessage, axum::Error>> + Sink<WsMessage> + Unpin,
{
let termination = connection.sender.termination.clone();
let session = async {
loop {
tokio::select! {
inbound = socket.next() => match inbound {
Some(Ok(WsMessage::Text(text))) => {
let result = match decode_method(&text) {
Ok(method) => connection.sender.send(&context, method).await.map_err(|error| error.to_string()),
Err(error) => Err(error.to_string()),
};
if let Err(error) = result {
let event = protocol_error_event(error);
let Ok(text) = protocol::stream::encode_event(&event) else { return WorkerProtocolEnd::Disconnected; };
if let Err(reason) = browser_send(&mut socket, WsMessage::Text(text.into())).await { return reason; }
}
}
Some(Ok(WsMessage::Ping(value))) => {
if let Err(reason) = browser_send(&mut socket, WsMessage::Pong(value)).await { return reason; }
}
Some(Ok(WsMessage::Pong(_))) => {},
_ => return WorkerProtocolEnd::Disconnected,
},
event = connection.events.recv() => {
let Some(event) = event else { return WorkerProtocolEnd::Disconnected; };
let shutdown = matches!(event, protocol::Event::Shutdown);
let Ok(text) = protocol::stream::encode_event(&event) else { return WorkerProtocolEnd::Disconnected; };
if let Err(reason) = browser_send(&mut socket, WsMessage::Text(text.into())).await { return reason; }
if shutdown { return WorkerProtocolEnd::Disconnected; }
}
}
}
};
let reason = tokio::select! {
biased;
reason = termination.ended() => reason,
reason = session => reason,
};
termination.terminate(reason);
// Best-effort explicit close even if the downstream transport is also blocked.
// The close reason deliberately makes no claim that admitted methods did not run.
let _ = tokio::time::timeout(
WORKER_PROTOCOL_TRANSPORT_TIMEOUT,
socket.send(WsMessage::Close(Some(axum::extract::ws::CloseFrame {
code: 1011,
reason: format!(
"{}; pending delivery unknown; resync snapshot; no retry",
reason.message()
)
.into(),
}))),
)
.await;
}
async fn browser_send<S: Sink<WsMessage> + Unpin>(
socket: &mut S,
message: WsMessage,
) -> std::result::Result<(), WorkerProtocolEnd> {
match tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, socket.send(message)).await {
Ok(Ok(())) => Ok(()),
Ok(Err(_)) => Err(WorkerProtocolEnd::Disconnected),
Err(_) => Err(WorkerProtocolEnd::TransportTimeout),
}
}
crates/workspace-server/src/workspace_subscription.rs
+252 -140
View File
@@ -1,6 +1,6 @@
use std::collections::HashMap;
use axum::extract::ws::{Message as WsMessage, WebSocket};
use axum::extract::ws::Message as WsMessage;
use futures::{SinkExt, StreamExt};
use protocol::subscription::{
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
@@ -12,27 +12,43 @@ use worker_runtime::identity::RuntimeWorkerRef;
use crate::runtime_subscription::RuntimeSubscriptionBroker;
use crate::server::{
WorkerOperationContext, WorkspaceApi, WorkspaceWorker, WorkspaceWorkerMethodSender,
WorkerOperationContext, WorkerProtocolEnd, WorkspaceApi, WorkspaceWorker,
WorkspaceWorkerMethodSender,
};
const OUTBOUND_CAPACITY: usize = 256;
const WORKER_PROTOCOL_SNAPSHOT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
// Every spawned writer/subscription is connection-owned, including caller cancellation.
struct SubscriptionTask(tokio::task::JoinHandle<()>);
impl Drop for SubscriptionTask {
fn drop(&mut self) {
self.0.abort();
}
}
struct ActiveSubscription {
task: tokio::task::JoinHandle<()>,
task: SubscriptionTask,
methods: Option<WorkspaceWorkerMethodSender>,
}
pub(crate) async fn serve_workspace_subscription(
pub(crate) async fn serve_workspace_subscription<S>(
api: WorkspaceApi,
socket: WebSocket,
socket: S,
context: WorkerOperationContext,
) {
) where
S: futures::Stream<Item = Result<WsMessage, axum::Error>>
+ futures::Sink<WsMessage>
+ Unpin
+ Send
+ 'static,
S::Error: Send,
{
let broker = api.runtime_subscription_broker().clone();
let (mut socket_sender, mut socket_receiver) = socket.split();
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
let (protocol_outbound, mut protocol_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
let writer = tokio::spawn(async move {
let (writer_done, finished) = tokio::sync::oneshot::channel();
let mut writer = SubscriptionTask(tokio::spawn(async move {
loop {
let message = tokio::select! {
biased;
@@ -40,160 +56,179 @@ pub(crate) async fn serve_workspace_subscription(
message = protocol_receiver.recv() => message,
};
let Some(message) = message else { break };
if socket_sender.send(message).await.is_err() {
if !matches!(
tokio::time::timeout(
WORKER_PROTOCOL_SNAPSHOT_TIMEOUT,
socket_sender.send(message)
)
.await,
Ok(Ok(()))
) {
break;
}
}
});
let _ = writer_done.send(());
}));
let mut next_subscription_id = 1_u64;
let mut subscriptions = HashMap::<SubscriptionId, ActiveSubscription>::new();
while let Some(message) = socket_receiver.next().await {
let Ok(message) = message else { break };
match message {
WsMessage::Text(text) => {
let Ok(frame) = serde_json::from_str::<SubscriptionFrame>(text.as_str()) else {
break;
};
if frame.validate().is_err() {
break;
}
subscriptions.retain(|_, subscription| !subscription.task.is_finished());
match frame.payload {
SubscriptionFramePayload::Request(SubscriptionRequest::SubscribeEvents {
request_id,
selector,
}) => {
let subscription_id = SubscriptionId::new(format!(
"workspace-subscription-{next_subscription_id}"
))
.expect("generated Workspace subscription id is valid");
next_subscription_id = next_subscription_id.saturating_add(1);
match selector {
EventSubscriptionSelector::WorkspaceWorkers => {
let task = tokio::spawn(run_workspace_workers(
api.clone(),
broker.clone(),
request_id,
subscription_id.clone(),
control_outbound.clone(),
));
subscriptions.insert(
subscription_id,
ActiveSubscription {
task,
methods: None,
},
);
}
EventSubscriptionSelector::WorkerProtocol {
worker_id,
runtime_id: Some(runtime_id),
} => {
let connection = match WorkspaceWorker::resolve(
&api,
&runtime_id,
worker_id.as_str(),
) {
Ok(worker) => worker
.connect_protocol(&context)
.await
.map_err(|error| error.to_string()),
Err(error) => Err(error.to_string()),
};
match connection {
Ok(connection) => {
let methods = connection.sender.clone();
let task = tokio::spawn(run_worker_protocol(
request_id,
subscription_id.clone(),
runtime_id,
worker_id,
connection.events,
control_outbound.clone(),
protocol_outbound.clone(),
));
subscriptions.insert(
subscription_id,
ActiveSubscription {
task,
methods: Some(methods),
},
);
}
Err(error) => {
let _ = send_rejected(
&control_outbound,
request_id,
SubscriptionRejectionCode::ResourceNotFound,
error.to_string(),
)
.await;
let receiving = async {
while let Some(message) = socket_receiver.next().await {
let Ok(message) = message else { break };
match message {
WsMessage::Text(text) => {
let Ok(frame) = serde_json::from_str::<SubscriptionFrame>(text.as_str()) else {
break;
};
if frame.validate().is_err() {
break;
}
subscriptions.retain(|_, subscription| !subscription.task.0.is_finished());
match frame.payload {
SubscriptionFramePayload::Request(
SubscriptionRequest::SubscribeEvents {
request_id,
selector,
},
) => {
let subscription_id = SubscriptionId::new(format!(
"workspace-subscription-{next_subscription_id}"
))
.expect("generated Workspace subscription id is valid");
next_subscription_id = next_subscription_id.saturating_add(1);
match selector {
EventSubscriptionSelector::WorkspaceWorkers => {
let task = tokio::spawn(run_workspace_workers(
api.clone(),
broker.clone(),
request_id,
subscription_id.clone(),
control_outbound.clone(),
));
subscriptions.insert(
subscription_id,
ActiveSubscription {
task: SubscriptionTask(task),
methods: None,
},
);
}
EventSubscriptionSelector::WorkerProtocol {
worker_id,
runtime_id: Some(runtime_id),
} => {
let connection = match WorkspaceWorker::resolve(
&api,
&runtime_id,
worker_id.as_str(),
) {
Ok(worker) => worker
.connect_protocol(&context)
.await
.map_err(|error| error.to_string()),
Err(error) => Err(error.to_string()),
};
match connection {
Ok(connection) => {
let methods = connection.sender.clone();
let task = tokio::spawn(run_worker_protocol(
request_id,
subscription_id.clone(),
runtime_id,
worker_id,
connection.events,
methods.end_status(),
control_outbound.clone(),
protocol_outbound.clone(),
));
subscriptions.insert(
subscription_id,
ActiveSubscription {
task: SubscriptionTask(task),
methods: Some(methods),
},
);
}
Err(error) => {
let _ = send_rejected(
&control_outbound,
request_id,
SubscriptionRejectionCode::ResourceNotFound,
error.to_string(),
)
.await;
}
}
}
}
_ => {
let _ = send_rejected(
_ => {
let _ = send_rejected(
&control_outbound, request_id, SubscriptionRejectionCode::UnsupportedSelector,
"Workspace clients may subscribe only to workspace_workers or a runtime-scoped worker_protocol selector".to_string(),
).await;
}
}
}
}
SubscriptionFramePayload::Request(SubscriptionRequest::UnsubscribeEvents {
request_id,
subscription_id,
}) => {
if let Some(subscription) = subscriptions.remove(&subscription_id) {
subscription.task.abort();
SubscriptionFramePayload::Request(
SubscriptionRequest::UnsubscribeEvents {
request_id,
subscription_id,
},
) => {
if let Some(mut subscription) = subscriptions.remove(&subscription_id) {
subscription.task.0.abort();
let _ = (&mut subscription.task.0).await;
}
if send_frame(
&control_outbound,
SubscriptionFrame::new(SubscriptionFramePayload::Response(
SubscriptionResponse::Unsubscribed {
request_id,
subscription_id,
},
)),
)
.await
.is_err()
{
break;
}
}
if send_frame(
&control_outbound,
SubscriptionFrame::new(SubscriptionFramePayload::Response(
SubscriptionResponse::Unsubscribed {
request_id,
subscription_id,
},
)),
)
.await
.is_err()
{
SubscriptionFramePayload::WorkerProtocol(message) => {
let Some(methods) = subscriptions
.get(&message.subscription_id)
.and_then(|value| value.methods.clone())
else {
break;
};
if methods.send(&context, message.method).await.is_err() {
break;
}
}
SubscriptionFramePayload::Response(_)
| SubscriptionFramePayload::Event(_) => {
break;
}
}
SubscriptionFramePayload::WorkerProtocol(message) => {
let Some(methods) = subscriptions
.get(&message.subscription_id)
.and_then(|value| value.methods.clone())
else {
break;
};
if methods.send(&context, message.method).await.is_err() {
break;
}
}
SubscriptionFramePayload::Response(_) | SubscriptionFramePayload::Event(_) => {
}
WsMessage::Ping(value) => {
if control_outbound.send(WsMessage::Pong(value)).await.is_err() {
break;
}
}
WsMessage::Pong(_) => {}
WsMessage::Close(_) | WsMessage::Binary(_) => break,
}
WsMessage::Ping(value) => {
if control_outbound.send(WsMessage::Pong(value)).await.is_err() {
break;
}
}
WsMessage::Pong(_) => {}
WsMessage::Close(_) | WsMessage::Binary(_) => break,
}
};
// Writer termination also preempts method authorization/control queue waits.
tokio::select! { _ = receiving => {}, _ = finished => {} }
for (_, mut subscription) in subscriptions {
subscription.task.0.abort();
let _ = (&mut subscription.task.0).await;
}
for (_, subscription) in subscriptions {
subscription.task.abort();
}
drop(control_outbound);
drop(protocol_outbound);
let _ = writer.await;
// Never wait for a congested writer to drain after disconnect or admission failure.
writer.0.abort();
let _ = (&mut writer.0).await;
}
async fn send_rejected(
@@ -222,6 +257,7 @@ async fn run_worker_protocol(
runtime_id: String,
worker_id: protocol::subscription::SubscriptionWorkerId,
mut events: mpsc::Receiver<protocol::Event>,
end_status: tokio::sync::watch::Receiver<Option<WorkerProtocolEnd>>,
control_outbound: mpsc::Sender<WsMessage>,
protocol_outbound: mpsc::Sender<WsMessage>,
) {
@@ -313,8 +349,8 @@ async fn run_worker_protocol(
SubscriptionFrame::new(SubscriptionFramePayload::Event(
SubscriptionEvent::SubscriptionClosed {
subscription_id,
code: SubscriptionTerminationCode::ResourceGone,
message: "Worker protocol stream closed".to_string(),
code: if matches!(*end_status.borrow(), Some(WorkerProtocolEnd::EventOverflow | WorkerProtocolEnd::MethodOverflow)) { SubscriptionTerminationCode::Lagged } else { SubscriptionTerminationCode::ResourceGone },
message: "Worker protocol stream ended; pending method delivery may be unknown; resubscribe with a fresh snapshot, no automatic retry".to_string(),
},
)),
)
@@ -646,6 +682,7 @@ mod tests {
"runtime-1".to_string(),
worker_id(),
events,
tokio::sync::watch::channel(None).1,
control_outbound,
protocol_outbound,
)
@@ -660,6 +697,81 @@ mod tests {
response
}
#[tokio::test]
async fn worker_bridge_overflow_is_serialized_as_lagged_not_resource_disappearance() {
for (reason, expected) in [
(
WorkerProtocolEnd::EventOverflow,
SubscriptionTerminationCode::Lagged,
),
(
WorkerProtocolEnd::MethodOverflow,
SubscriptionTerminationCode::Lagged,
),
(
WorkerProtocolEnd::Disconnected,
SubscriptionTerminationCode::ResourceGone,
),
] {
let (event_sender, events) = mpsc::channel(1);
event_sender
.try_send(protocol::Event::Snapshot {
session: protocol::SessionSnapshot {
entries: vec![],
pending_submissions: Default::default(),
},
greeting: protocol::Greeting {
worker_name: "worker-1".into(),
cwd: String::new(),
provider: "fixture".into(),
model: "fixture".into(),
reasoning: None,
scope_summary: String::new(),
tools: vec![],
context_window: 0,
context_tokens: 0,
context_usage: None,
},
state: protocol::WorkerStateSnapshot::initial(),
in_flight: Default::default(),
internal_workers: vec![],
})
.unwrap();
drop(event_sender);
let (control, mut receiver) = mpsc::channel(2);
let (protocol, _events) = mpsc::channel(1);
run_worker_protocol(
request_id(),
subscription_id(),
"runtime-1".into(),
worker_id(),
events,
tokio::sync::watch::channel(Some(reason)).1,
control,
protocol,
)
.await;
let WsMessage::Text(snapshot) = receiver.recv().await.unwrap() else {
panic!("expected snapshot response");
};
let snapshot: SubscriptionFrame = serde_json::from_str(snapshot.as_str()).unwrap();
assert!(matches!(
snapshot.payload,
SubscriptionFramePayload::Response(SubscriptionResponse::Subscribed { .. })
));
let WsMessage::Text(closed) = receiver.recv().await.unwrap() else {
panic!("expected terminal control frame");
};
let frame: SubscriptionFrame = serde_json::from_str(closed.as_str()).unwrap();
frame.validate().unwrap();
assert!(
matches!(frame.payload, SubscriptionFramePayload::Event(SubscriptionEvent::SubscriptionClosed { code, message, .. })
if code == expected && message.contains("unknown") && message.contains("fresh snapshot")),
"{closed}"
);
}
}
#[tokio::test]
async fn worker_protocol_rejects_a_non_snapshot_initial_event() {
let (event_sender, events) = mpsc::channel(1);