fix: bound Worker protocol backpressure and connection ownership
This commit is contained in:
@@ -2537,16 +2537,24 @@ where
|
||||
}
|
||||
}
|
||||
method => {
|
||||
let result = self.send_method(
|
||||
WorkerExecutionOperation::ProtocolMethod,
|
||||
execution.handle.clone(),
|
||||
method,
|
||||
);
|
||||
if result.is_accepted() {
|
||||
// The connected protocol owns this bounded admission wait. The old
|
||||
// generic adapter waiter could time out while leaving a detached send
|
||||
// task alive; it might deliver later after the connection had ended.
|
||||
// Join cancellation before releasing endpoint exclusion. A send which
|
||||
// raced cancellation remains delivery-unknown, never automatically retried.
|
||||
let handle = execution.handle.clone();
|
||||
self.run_cancellable_on_adapter_runtime(USER_INPUT_TASK_TIMEOUT, async move {
|
||||
handle.send(method).await.map_err(|error| {
|
||||
format!("connected Worker method delivery failed or is unknown: {error}")
|
||||
})?;
|
||||
Ok(Vec::new())
|
||||
} else {
|
||||
Err(result)
|
||||
}
|
||||
})
|
||||
.map_err(|message| {
|
||||
WorkerExecutionResult::errored(
|
||||
WorkerExecutionOperation::ProtocolMethod,
|
||||
message,
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -7898,6 +7906,48 @@ mod tests {
|
||||
assert!(backend.worker_snapshot(&worker_ref).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connected_protocol_admission_timeout_joins_pending_send_before_capacity_reopens() {
|
||||
// This is the adapter's cancellation/join policy boundary used by the
|
||||
// captured protocol endpoint, not an execution-side-effect rollback.
|
||||
let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap();
|
||||
let (sender, mut receiver) = tokio::sync::mpsc::channel(1);
|
||||
sender.try_send(Method::ListPendingSubmissions).unwrap();
|
||||
let dropped = Arc::new(AtomicBool::new(false));
|
||||
struct PendingSendDrop(Arc<AtomicBool>);
|
||||
impl Drop for PendingSendDrop {
|
||||
fn drop(&mut self) {
|
||||
self.0.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
let owned = PendingSendDrop(dropped.clone());
|
||||
let error = backend
|
||||
.run_cancellable_on_adapter_runtime(Duration::from_millis(5), async move {
|
||||
let _owned = owned;
|
||||
sender
|
||||
.send(Method::ListRewindTargets)
|
||||
.await
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.unwrap_err();
|
||||
assert!(error.contains("was cancelled"), "{error}");
|
||||
assert!(
|
||||
dropped.load(Ordering::SeqCst),
|
||||
"must join task cancellation before returning"
|
||||
);
|
||||
assert!(matches!(
|
||||
receiver.try_recv(),
|
||||
Ok(Method::ListPendingSubmissions)
|
||||
));
|
||||
assert!(
|
||||
matches!(
|
||||
receiver.try_recv(),
|
||||
Err(tokio::sync::mpsc::error::TryRecvError::Disconnected)
|
||||
),
|
||||
"timeout must not leave an admission task capable of a late send"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_execution_has_typed_input_rejection_and_safe_stop() {
|
||||
let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap();
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
# Worker operational protocol backpressure
|
||||
|
||||
The Workspace Worker operational connection is separate from Workspace request
|
||||
admission (T-732). Both the dedicated Worker WebSocket and the Workspace
|
||||
multiplexer use the identity-bound `WorkspaceWorkerMethodSender` and the same
|
||||
remote/embedded bridge. This does not change Ticket assignment, general input
|
||||
routing, or database execution placement.
|
||||
|
||||
## Ownership and admission
|
||||
|
||||
- Each bridge owns one captured remote WebSocket or embedded execution endpoint.
|
||||
It never follows a restored execution under the same Worker identity.
|
||||
- Methods have 256 queue slots; events have 512. Browser authorization, input
|
||||
source attribution, interactive-Worker checks, and session exclusion remain at
|
||||
method admission. Queue admission is `try_send`, not a transport-capacity wait
|
||||
under the session guard. Success here means **queued**, not delivered/executed.
|
||||
- Event enqueue never waits for its consumer. Snapshot is first; direct replies
|
||||
and live events retain the existing bridge ordering. No additional event pump
|
||||
can reorder them. Overflow ends the connection rather than skipping an event
|
||||
and continuing an uncertain stream.
|
||||
- The first terminal cause wins: method overflow, event overflow/lag, transport
|
||||
timeout, or disconnect. Terminal state cannot reopen when capacity becomes
|
||||
available. It preempts pending browser admission and remote transport waits.
|
||||
- The last client capability owns an abort-on-drop pump handle. The pump does not
|
||||
retain that capability, avoiding an ownership cycle. Event-consumer closure
|
||||
also cancels the pump even when method capabilities remain. Cancellation drops
|
||||
the bounded channels and captured transport; it does not stop the Worker.
|
||||
|
||||
A rejected full-queue method is known **not admitted to that queue**. Previously
|
||||
queued methods, interrupted wire sends, and execution calls which have begun may
|
||||
have unknown delivery or side effects. They are never converted to success or a
|
||||
claim of non-execution, and are never automatically retried. Cancel is subject to
|
||||
this same rule: it either enters the bounded queue or the connection explicitly
|
||||
ends; there is no silent special-case drop or unbounded control lane.
|
||||
|
||||
## Transport deadlines and termination
|
||||
|
||||
Remote handshake and sends, browser writes (including Pong/error), and Workspace
|
||||
multiplexer writes have 10-second deadlines. The dedicated browser session makes
|
||||
one separately bounded Close attempt (at most another 10 seconds), with code 1011
|
||||
and a short reason requiring snapshot resynchronization and warning of unknown
|
||||
pending delivery. A completely blocked wire cannot guarantee delivery of the
|
||||
Close/error frame; releasing the socket is the terminal signal in that case.
|
||||
|
||||
The multiplexer retains separate bounded control and high-volume protocol lanes
|
||||
and control priority. Writer termination cancels the receive/admission loop, even
|
||||
if that loop is awaiting control capacity or session ownership. Receive exit
|
||||
aborts/joins the writer and subscription tasks instead of waiting indefinitely
|
||||
for a slow writer to drain. Caller cancellation has abort-on-drop ownership.
|
||||
Bridge overflow is exposed as `SubscriptionClosed { code: Lagged }`, not normal
|
||||
resource disappearance, when the terminal control frame can be delivered.
|
||||
|
||||
Embedded dispatch is an existing synchronous Runtime API, not a new database or
|
||||
executor redesign. Its connected-method adapter wait uses the existing
|
||||
10-second cancellation-and-join waiter rather than leaving a detached send after
|
||||
timeout. Existing completion and lifecycle ownership/fencing are retained. A
|
||||
synchronous embedded dispatch already begun is not immediately preemptible by
|
||||
browser disconnect; its bounded adapter wait settles before endpoint exclusion
|
||||
is released. Cancellation/join is resource reclamation, **not rollback of an
|
||||
already delivered method**.
|
||||
|
||||
## Recovery
|
||||
|
||||
Reconnect explicitly, resubscribe where using the multiplexer, and replace local
|
||||
state with the fresh snapshot before live delivery (T-415). Do not replay old
|
||||
method queues or transient protocol deltas. An ended connection never implicitly
|
||||
Restores, reattaches, or retargets the Worker. Slow connections do not hold Server
|
||||
session exclusion during transport waits, so other connections and Stop/session
|
||||
cleanup can proceed.
|
||||
|
||||
## Evidence and limits
|
||||
|
||||
The initial static audit used `develop@a9a67fd6`; the same wait edges were present
|
||||
at `develop@98023b98`. This is not evidence of a production incident. Controlled
|
||||
fixtures exercise the actual remote connector, embedded bridge, browser method
|
||||
admission/session, and Workspace multiplexer. Reinstating only the old bounded
|
||||
queue waits makes the remote/embedded overflow and browser session-exclusion
|
||||
fixtures time out; the corrected paths terminate without guessed sleeps.
|
||||
|
||||
Fixtures use channel/barrier ownership, controllable sinks, loopback WebSockets,
|
||||
and paused time for transport deadlines. They cover ordered admitted prefixes,
|
||||
auth rejection before admission, Cancel overflow, pending admission cancellation,
|
||||
unknown wire delivery without retry, same-Worker fresh connection/Stop, and
|
||||
transport/channel reclamation. They are not full browser/Runtime process E2E or a
|
||||
production load benchmark. Existing attach/stop/restore and serialization tests
|
||||
remain required alongside these regressions.
|
||||
|
||||
### Implementation validation (T-747)
|
||||
|
||||
- `cargo check --offline` at Workspace root: passed (normal TUI/Server closure).
|
||||
- `cargo test -p yoi-workspace-server --offline --lib protocol_`: 29 passed,
|
||||
including the real bridges, browser admission/ownership, and stalled multiplexer.
|
||||
- `cargo test -p yoi-workspace-server --offline`: 867 library tests passed, one
|
||||
existing ignored test, 11 binary tests passed, including the lag-code wire
|
||||
classification matrix (`--lib workspace_subscription::tests`: 4 passed).
|
||||
- `cargo test -p worker-runtime --offline`: 416 library and 19 tests in each of
|
||||
the two binaries passed, including connected admission cancellation/join and
|
||||
existing endpoint fencing/lifecycle regressions.
|
||||
- `cargo test -p protocol --offline`: 80 passed (wire contracts unchanged).
|
||||
- `cargo fmt --all -- --check`, `git diff --check HEAD`, and
|
||||
`env -u LD_LIBRARY_PATH deno run --allow-read --allow-run=git tools/check-state-contract-names.ts`:
|
||||
passed; naming checker reported 0 unapproved / 0 stale exceptions.
|
||||
|
||||
Additional Runtime compile probes with `--no-default-features`, with only
|
||||
`http-server,fs-store`, and with only `ws-server` failed in pre-existing feature
|
||||
boundaries (respectively missing gated request dependencies, missing futures
|
||||
imports/dependency, and missing filesystem-backed retention methods). Replacing
|
||||
only the changed Runtime file with its exact `98023b98` content produced the same
|
||||
error groups in all three probes. These configurations are not claimed validated;
|
||||
this patch does not repair their unrelated feature wiring. The normal
|
||||
`ws-server,fs-store` configuration is covered by root check and whole-crate tests.
|
||||
|
||||
Reference in New Issue
Block a user