fix: integrate bounded Worker protocol transport with reviewed deadline fixes
This commit is contained in:
@@ -2537,16 +2537,24 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
method => {
|
method => {
|
||||||
let result = self.send_method(
|
// The connected protocol owns this bounded admission wait. The old
|
||||||
WorkerExecutionOperation::ProtocolMethod,
|
// generic adapter waiter could time out while leaving a detached send
|
||||||
execution.handle.clone(),
|
// task alive; it might deliver later after the connection had ended.
|
||||||
method,
|
// Join cancellation before releasing endpoint exclusion. A send which
|
||||||
);
|
// raced cancellation remains delivery-unknown, never automatically retried.
|
||||||
if result.is_accepted() {
|
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())
|
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());
|
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]
|
#[test]
|
||||||
fn absent_execution_has_typed_input_rejection_and_safe_stop() {
|
fn absent_execution_has_typed_input_rejection_and_safe_stop() {
|
||||||
let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap();
|
let backend = WorkerRuntimeExecutionBackend::new(FailingFactory).unwrap();
|
||||||
|
|||||||
@@ -2,9 +2,12 @@ use subjektiv::api::*;
|
|||||||
mod backend_job_profile;
|
mod backend_job_profile;
|
||||||
mod subjektiv_job;
|
mod subjektiv_job;
|
||||||
mod worker_operations;
|
mod worker_operations;
|
||||||
|
mod worker_protocol;
|
||||||
pub(crate) use worker_operations::{
|
pub(crate) use worker_operations::{
|
||||||
WorkerOperationContext, WorkspaceWorker, WorkspaceWorkerMethodSender,
|
WorkerOperationContext, WorkspaceWorker, WorkspaceWorkerMethodSender,
|
||||||
};
|
};
|
||||||
|
pub(crate) use worker_protocol::WorkerProtocolEnd;
|
||||||
|
use worker_protocol::*;
|
||||||
|
|
||||||
#[path = "server_drive.rs"]
|
#[path = "server_drive.rs"]
|
||||||
mod drive;
|
mod drive;
|
||||||
@@ -29083,51 +29086,19 @@ async fn worker_protocol_ws(
|
|||||||
match worker.connect_protocol(&context).await {
|
match worker.connect_protocol(&context).await {
|
||||||
Ok(connection) => workspace_worker_protocol_session(connection, socket, context).await,
|
Ok(connection) => workspace_worker_protocol_session(connection, socket, context).await,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
let _ = send_protocol_event(&mut socket, &protocol_error_event(error.to_string()))
|
let _ = tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, async {
|
||||||
.await;
|
let _ =
|
||||||
let _ = socket.send(WsMessage::Close(None)).await;
|
send_protocol_event(&mut socket, &protocol_error_event(error.to_string()))
|
||||||
|
.await;
|
||||||
|
let _ = socket.send(WsMessage::Close(None)).await;
|
||||||
|
})
|
||||||
|
.await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.into_response()
|
.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(
|
pub(crate) async fn connect_workspace_worker_protocol(
|
||||||
api: &WorkspaceApi,
|
api: &WorkspaceApi,
|
||||||
worker: &RuntimeWorkerRef,
|
worker: &RuntimeWorkerRef,
|
||||||
@@ -29199,39 +29170,21 @@ async fn connect_remote_worker_protocol(
|
|||||||
}
|
}
|
||||||
insert_authenticated_input_source_header(request.headers_mut(), input_source)?;
|
insert_authenticated_input_source_header(request.headers_mut(), input_source)?;
|
||||||
let (socket, _) =
|
let (socket, _) =
|
||||||
connect_async(request)
|
tokio::time::timeout(WORKER_PROTOCOL_TRANSPORT_TIMEOUT, connect_async(request))
|
||||||
.await
|
.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 {
|
.map_err(|error| Error::RuntimeOperationFailed {
|
||||||
runtime_id: config.worker.runtime_id.clone(),
|
runtime_id: config.worker.runtime_id.clone(),
|
||||||
code: "worker_protocol_connect_failed".to_string(),
|
code: "worker_protocol_connect_failed".into(),
|
||||||
message: error.to_string(),
|
message: error.to_string(),
|
||||||
})?;
|
})?;
|
||||||
let (mut sink, mut stream) = socket.split();
|
let (sink, stream) = socket.split();
|
||||||
let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256);
|
Ok(remote_worker_protocol_bridge(sink, stream))
|
||||||
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 })
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn connect_embedded_worker_protocol(
|
async fn connect_embedded_worker_protocol(
|
||||||
@@ -29249,47 +29202,48 @@ async fn connect_embedded_worker_protocol(
|
|||||||
code: "embedded_worker_protocol_connect_failed".to_string(),
|
code: "embedded_worker_protocol_connect_failed".to_string(),
|
||||||
message: error.to_string(),
|
message: error.to_string(),
|
||||||
})?;
|
})?;
|
||||||
let (methods, mut method_receiver) = tokio::sync::mpsc::channel(256);
|
Ok(spawn_worker_protocol_pump(
|
||||||
let (event_sender, events) = tokio::sync::mpsc::channel(512);
|
move |mut methods, events| async move {
|
||||||
tokio::spawn(async move {
|
if let Err(reason) = enqueue_worker_protocol_event(&events, transport.snapshot.clone())
|
||||||
if event_sender.send(transport.snapshot.clone()).await.is_err() {
|
{
|
||||||
return;
|
return reason;
|
||||||
}
|
}
|
||||||
loop {
|
loop {
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
_ = event_sender.closed() => break,
|
method = methods.recv() => {
|
||||||
method = method_receiver.recv() => {
|
let Some(method) = method else { return WorkerProtocolEnd::Disconnected; };
|
||||||
let Some(method) = method else { break };
|
let shutdown = matches!(&method, protocol::Method::Shutdown { .. });
|
||||||
let shutdown = matches!(&method, protocol::Method::Shutdown { .. });
|
match source.runtime.send_connected_protocol_method_scoped(
|
||||||
match source.runtime.send_connected_protocol_method_scoped(
|
&scope, &source.worker_ref, &transport, method,
|
||||||
&scope, &source.worker_ref, &transport, method,
|
) {
|
||||||
) {
|
Ok(direct_events) => {
|
||||||
Ok(direct_events) => {
|
for event in direct_events {
|
||||||
for event in direct_events {
|
let terminal = matches!(event, protocol::Event::Shutdown);
|
||||||
let terminal = matches!(event, protocol::Event::Shutdown);
|
if let Err(reason) = enqueue_worker_protocol_event(&events, event) { return reason; }
|
||||||
if event_sender.send(event).await.is_err() || terminal { return; }
|
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) => {
|
if shutdown { return WorkerProtocolEnd::Disconnected; }
|
||||||
let terminal = embedded_protocol_dispatch_is_terminal(&error, &transport);
|
}
|
||||||
if event_sender.send(protocol_error_event(error.to_string())).await.is_err() || terminal { return; }
|
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(
|
fn embedded_protocol_dispatch_is_terminal(
|
||||||
@@ -32502,6 +32456,7 @@ mod tests {
|
|||||||
mod drive_tests;
|
mod drive_tests;
|
||||||
mod drive_web_roundtrip_tests;
|
mod drive_web_roundtrip_tests;
|
||||||
mod external_checkout_tests;
|
mod external_checkout_tests;
|
||||||
|
mod protocol_backpressure_tests;
|
||||||
mod subject_spawn_tests;
|
mod subject_spawn_tests;
|
||||||
mod subjektiv_jobs_tests;
|
mod subjektiv_jobs_tests;
|
||||||
mod ticket_evidence_tests;
|
mod ticket_evidence_tests;
|
||||||
@@ -32511,7 +32466,9 @@ mod tests {
|
|||||||
mod workdir_delegation_tests;
|
mod workdir_delegation_tests;
|
||||||
mod worker_drive_tests;
|
mod worker_drive_tests;
|
||||||
mod worker_operations_tests;
|
mod worker_operations_tests;
|
||||||
|
mod worker_protocol_tests;
|
||||||
mod workspace_admission_tests;
|
mod workspace_admission_tests;
|
||||||
|
mod workspace_protocol_backpressure_tests;
|
||||||
include!("server_workspace_config_tests.rs");
|
include!("server_workspace_config_tests.rs");
|
||||||
use super::*;
|
use super::*;
|
||||||
use axum::body::{Body, to_bytes};
|
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 =
|
let connection =
|
||||||
connect_workspace_worker_protocol(&self.api, &self.identity, source).await?;
|
connect_workspace_worker_protocol(&self.api, &self.identity, source).await?;
|
||||||
Ok(WorkspaceWorkerConnection {
|
Ok(WorkspaceWorkerConnection {
|
||||||
sender: WorkspaceWorkerMethodSender {
|
sender: WorkspaceWorkerMethodSender::new(
|
||||||
worker: self.clone(),
|
self.clone(),
|
||||||
methods: connection.methods,
|
connection.methods,
|
||||||
},
|
connection.termination,
|
||||||
|
),
|
||||||
events: connection.events,
|
events: connection.events,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -485,12 +486,27 @@ impl WorkspaceWorker {
|
|||||||
pub(crate) struct WorkspaceWorkerMethodSender {
|
pub(crate) struct WorkspaceWorkerMethodSender {
|
||||||
worker: WorkspaceWorker,
|
worker: WorkspaceWorker,
|
||||||
methods: tokio::sync::mpsc::Sender<protocol::Method>,
|
methods: tokio::sync::mpsc::Sender<protocol::Method>,
|
||||||
|
pub(super) termination: WorkerProtocolTermination,
|
||||||
}
|
}
|
||||||
pub(crate) struct WorkspaceWorkerConnection {
|
pub(crate) struct WorkspaceWorkerConnection {
|
||||||
pub(crate) sender: WorkspaceWorkerMethodSender,
|
pub(crate) sender: WorkspaceWorkerMethodSender,
|
||||||
pub(crate) events: tokio::sync::mpsc::Receiver<protocol::Event>,
|
pub(crate) events: tokio::sync::mpsc::Receiver<protocol::Event>,
|
||||||
}
|
}
|
||||||
impl WorkspaceWorkerMethodSender {
|
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(
|
pub(crate) async fn send(
|
||||||
&self,
|
&self,
|
||||||
context: &WorkerOperationContext,
|
context: &WorkerOperationContext,
|
||||||
@@ -604,15 +620,25 @@ impl WorkspaceWorkerMethodSender {
|
|||||||
};
|
};
|
||||||
let method = authorize_browser_worker_method(method, source)
|
let method = authorize_browser_worker_method(method, source)
|
||||||
.map_err(|message| Error::InvalidInput(message.into()))?;
|
.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 {
|
Error::RuntimeOperationFailed {
|
||||||
runtime_id: self.worker.identity.runtime_id.clone(),
|
runtime_id: self.worker.identity.runtime_id.clone(),
|
||||||
code: "worker_protocol_closed".into(),
|
code: code.into(), message: message.into(),
|
||||||
message:
|
}.into()
|
||||||
"Worker execution connection closed; no implicit Restore or retry was performed"
|
|
||||||
.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 std::collections::HashMap;
|
||||||
|
|
||||||
use axum::extract::ws::{Message as WsMessage, WebSocket};
|
use axum::extract::ws::Message as WsMessage;
|
||||||
use futures::{SinkExt, StreamExt};
|
use futures::{SinkExt, StreamExt};
|
||||||
use protocol::subscription::{
|
use protocol::subscription::{
|
||||||
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
|
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
|
||||||
@@ -12,27 +12,43 @@ use worker_runtime::identity::RuntimeWorkerRef;
|
|||||||
|
|
||||||
use crate::runtime_subscription::RuntimeSubscriptionBroker;
|
use crate::runtime_subscription::RuntimeSubscriptionBroker;
|
||||||
use crate::server::{
|
use crate::server::{
|
||||||
WorkerOperationContext, WorkspaceApi, WorkspaceWorker, WorkspaceWorkerMethodSender,
|
WorkerOperationContext, WorkerProtocolEnd, WorkspaceApi, WorkspaceWorker,
|
||||||
|
WorkspaceWorkerMethodSender,
|
||||||
};
|
};
|
||||||
|
|
||||||
const OUTBOUND_CAPACITY: usize = 256;
|
const OUTBOUND_CAPACITY: usize = 256;
|
||||||
const WORKER_PROTOCOL_SNAPSHOT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
|
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 {
|
struct ActiveSubscription {
|
||||||
task: tokio::task::JoinHandle<()>,
|
task: SubscriptionTask,
|
||||||
methods: Option<WorkspaceWorkerMethodSender>,
|
methods: Option<WorkspaceWorkerMethodSender>,
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn serve_workspace_subscription(
|
pub(crate) async fn serve_workspace_subscription<S>(
|
||||||
api: WorkspaceApi,
|
api: WorkspaceApi,
|
||||||
socket: WebSocket,
|
socket: S,
|
||||||
context: WorkerOperationContext,
|
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 broker = api.runtime_subscription_broker().clone();
|
||||||
let (mut socket_sender, mut socket_receiver) = socket.split();
|
let (mut socket_sender, mut socket_receiver) = socket.split();
|
||||||
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
|
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
|
||||||
let (protocol_outbound, mut protocol_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 {
|
loop {
|
||||||
let message = tokio::select! {
|
let message = tokio::select! {
|
||||||
biased;
|
biased;
|
||||||
@@ -40,160 +56,179 @@ pub(crate) async fn serve_workspace_subscription(
|
|||||||
message = protocol_receiver.recv() => message,
|
message = protocol_receiver.recv() => message,
|
||||||
};
|
};
|
||||||
let Some(message) = message else { break };
|
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;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
});
|
let _ = writer_done.send(());
|
||||||
|
}));
|
||||||
let mut next_subscription_id = 1_u64;
|
let mut next_subscription_id = 1_u64;
|
||||||
let mut subscriptions = HashMap::<SubscriptionId, ActiveSubscription>::new();
|
let mut subscriptions = HashMap::<SubscriptionId, ActiveSubscription>::new();
|
||||||
|
|
||||||
while let Some(message) = socket_receiver.next().await {
|
let receiving = async {
|
||||||
let Ok(message) = message else { break };
|
while let Some(message) = socket_receiver.next().await {
|
||||||
match message {
|
let Ok(message) = message else { break };
|
||||||
WsMessage::Text(text) => {
|
match message {
|
||||||
let Ok(frame) = serde_json::from_str::<SubscriptionFrame>(text.as_str()) else {
|
WsMessage::Text(text) => {
|
||||||
break;
|
let Ok(frame) = serde_json::from_str::<SubscriptionFrame>(text.as_str()) else {
|
||||||
};
|
break;
|
||||||
if frame.validate().is_err() {
|
};
|
||||||
break;
|
if frame.validate().is_err() {
|
||||||
}
|
break;
|
||||||
subscriptions.retain(|_, subscription| !subscription.task.is_finished());
|
}
|
||||||
match frame.payload {
|
subscriptions.retain(|_, subscription| !subscription.task.0.is_finished());
|
||||||
SubscriptionFramePayload::Request(SubscriptionRequest::SubscribeEvents {
|
match frame.payload {
|
||||||
request_id,
|
SubscriptionFramePayload::Request(
|
||||||
selector,
|
SubscriptionRequest::SubscribeEvents {
|
||||||
}) => {
|
request_id,
|
||||||
let subscription_id = SubscriptionId::new(format!(
|
selector,
|
||||||
"workspace-subscription-{next_subscription_id}"
|
},
|
||||||
))
|
) => {
|
||||||
.expect("generated Workspace subscription id is valid");
|
let subscription_id = SubscriptionId::new(format!(
|
||||||
next_subscription_id = next_subscription_id.saturating_add(1);
|
"workspace-subscription-{next_subscription_id}"
|
||||||
match selector {
|
))
|
||||||
EventSubscriptionSelector::WorkspaceWorkers => {
|
.expect("generated Workspace subscription id is valid");
|
||||||
let task = tokio::spawn(run_workspace_workers(
|
next_subscription_id = next_subscription_id.saturating_add(1);
|
||||||
api.clone(),
|
match selector {
|
||||||
broker.clone(),
|
EventSubscriptionSelector::WorkspaceWorkers => {
|
||||||
request_id,
|
let task = tokio::spawn(run_workspace_workers(
|
||||||
subscription_id.clone(),
|
api.clone(),
|
||||||
control_outbound.clone(),
|
broker.clone(),
|
||||||
));
|
request_id,
|
||||||
subscriptions.insert(
|
subscription_id.clone(),
|
||||||
subscription_id,
|
control_outbound.clone(),
|
||||||
ActiveSubscription {
|
));
|
||||||
task,
|
subscriptions.insert(
|
||||||
methods: None,
|
subscription_id,
|
||||||
},
|
ActiveSubscription {
|
||||||
);
|
task: SubscriptionTask(task),
|
||||||
}
|
methods: None,
|
||||||
EventSubscriptionSelector::WorkerProtocol {
|
},
|
||||||
worker_id,
|
);
|
||||||
runtime_id: Some(runtime_id),
|
}
|
||||||
} => {
|
EventSubscriptionSelector::WorkerProtocol {
|
||||||
let connection = match WorkspaceWorker::resolve(
|
worker_id,
|
||||||
&api,
|
runtime_id: Some(runtime_id),
|
||||||
&runtime_id,
|
} => {
|
||||||
worker_id.as_str(),
|
let connection = match WorkspaceWorker::resolve(
|
||||||
) {
|
&api,
|
||||||
Ok(worker) => worker
|
&runtime_id,
|
||||||
.connect_protocol(&context)
|
worker_id.as_str(),
|
||||||
.await
|
) {
|
||||||
.map_err(|error| error.to_string()),
|
Ok(worker) => worker
|
||||||
Err(error) => Err(error.to_string()),
|
.connect_protocol(&context)
|
||||||
};
|
.await
|
||||||
match connection {
|
.map_err(|error| error.to_string()),
|
||||||
Ok(connection) => {
|
Err(error) => Err(error.to_string()),
|
||||||
let methods = connection.sender.clone();
|
};
|
||||||
let task = tokio::spawn(run_worker_protocol(
|
match connection {
|
||||||
request_id,
|
Ok(connection) => {
|
||||||
subscription_id.clone(),
|
let methods = connection.sender.clone();
|
||||||
runtime_id,
|
let task = tokio::spawn(run_worker_protocol(
|
||||||
worker_id,
|
request_id,
|
||||||
connection.events,
|
subscription_id.clone(),
|
||||||
control_outbound.clone(),
|
runtime_id,
|
||||||
protocol_outbound.clone(),
|
worker_id,
|
||||||
));
|
connection.events,
|
||||||
subscriptions.insert(
|
methods.end_status(),
|
||||||
subscription_id,
|
control_outbound.clone(),
|
||||||
ActiveSubscription {
|
protocol_outbound.clone(),
|
||||||
task,
|
));
|
||||||
methods: Some(methods),
|
subscriptions.insert(
|
||||||
},
|
subscription_id,
|
||||||
);
|
ActiveSubscription {
|
||||||
}
|
task: SubscriptionTask(task),
|
||||||
Err(error) => {
|
methods: Some(methods),
|
||||||
let _ = send_rejected(
|
},
|
||||||
&control_outbound,
|
);
|
||||||
request_id,
|
}
|
||||||
SubscriptionRejectionCode::ResourceNotFound,
|
Err(error) => {
|
||||||
error.to_string(),
|
let _ = send_rejected(
|
||||||
)
|
&control_outbound,
|
||||||
.await;
|
request_id,
|
||||||
|
SubscriptionRejectionCode::ResourceNotFound,
|
||||||
|
error.to_string(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
_ => {
|
||||||
_ => {
|
let _ = send_rejected(
|
||||||
let _ = send_rejected(
|
|
||||||
&control_outbound, request_id, SubscriptionRejectionCode::UnsupportedSelector,
|
&control_outbound, request_id, SubscriptionRejectionCode::UnsupportedSelector,
|
||||||
"Workspace clients may subscribe only to workspace_workers or a runtime-scoped worker_protocol selector".to_string(),
|
"Workspace clients may subscribe only to workspace_workers or a runtime-scoped worker_protocol selector".to_string(),
|
||||||
).await;
|
).await;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
SubscriptionFramePayload::Request(
|
||||||
SubscriptionFramePayload::Request(SubscriptionRequest::UnsubscribeEvents {
|
SubscriptionRequest::UnsubscribeEvents {
|
||||||
request_id,
|
request_id,
|
||||||
subscription_id,
|
subscription_id,
|
||||||
}) => {
|
},
|
||||||
if let Some(subscription) = subscriptions.remove(&subscription_id) {
|
) => {
|
||||||
subscription.task.abort();
|
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(
|
SubscriptionFramePayload::WorkerProtocol(message) => {
|
||||||
&control_outbound,
|
let Some(methods) = subscriptions
|
||||||
SubscriptionFrame::new(SubscriptionFramePayload::Response(
|
.get(&message.subscription_id)
|
||||||
SubscriptionResponse::Unsubscribed {
|
.and_then(|value| value.methods.clone())
|
||||||
request_id,
|
else {
|
||||||
subscription_id,
|
break;
|
||||||
},
|
};
|
||||||
)),
|
if methods.send(&context, message.method).await.is_err() {
|
||||||
)
|
break;
|
||||||
.await
|
}
|
||||||
.is_err()
|
}
|
||||||
{
|
SubscriptionFramePayload::Response(_)
|
||||||
|
| SubscriptionFramePayload::Event(_) => {
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
SubscriptionFramePayload::WorkerProtocol(message) => {
|
}
|
||||||
let Some(methods) = subscriptions
|
WsMessage::Ping(value) => {
|
||||||
.get(&message.subscription_id)
|
if control_outbound.send(WsMessage::Pong(value)).await.is_err() {
|
||||||
.and_then(|value| value.methods.clone())
|
|
||||||
else {
|
|
||||||
break;
|
|
||||||
};
|
|
||||||
if methods.send(&context, message.method).await.is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
SubscriptionFramePayload::Response(_) | SubscriptionFramePayload::Event(_) => {
|
|
||||||
break;
|
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;
|
||||||
}
|
}
|
||||||
|
// Never wait for a congested writer to drain after disconnect or admission failure.
|
||||||
for (_, subscription) in subscriptions {
|
writer.0.abort();
|
||||||
subscription.task.abort();
|
let _ = (&mut writer.0).await;
|
||||||
}
|
|
||||||
drop(control_outbound);
|
|
||||||
drop(protocol_outbound);
|
|
||||||
let _ = writer.await;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn send_rejected(
|
async fn send_rejected(
|
||||||
@@ -222,6 +257,7 @@ async fn run_worker_protocol(
|
|||||||
runtime_id: String,
|
runtime_id: String,
|
||||||
worker_id: protocol::subscription::SubscriptionWorkerId,
|
worker_id: protocol::subscription::SubscriptionWorkerId,
|
||||||
mut events: mpsc::Receiver<protocol::Event>,
|
mut events: mpsc::Receiver<protocol::Event>,
|
||||||
|
end_status: tokio::sync::watch::Receiver<Option<WorkerProtocolEnd>>,
|
||||||
control_outbound: mpsc::Sender<WsMessage>,
|
control_outbound: mpsc::Sender<WsMessage>,
|
||||||
protocol_outbound: mpsc::Sender<WsMessage>,
|
protocol_outbound: mpsc::Sender<WsMessage>,
|
||||||
) {
|
) {
|
||||||
@@ -313,8 +349,8 @@ async fn run_worker_protocol(
|
|||||||
SubscriptionFrame::new(SubscriptionFramePayload::Event(
|
SubscriptionFrame::new(SubscriptionFramePayload::Event(
|
||||||
SubscriptionEvent::SubscriptionClosed {
|
SubscriptionEvent::SubscriptionClosed {
|
||||||
subscription_id,
|
subscription_id,
|
||||||
code: SubscriptionTerminationCode::ResourceGone,
|
code: if matches!(*end_status.borrow(), Some(WorkerProtocolEnd::EventOverflow | WorkerProtocolEnd::MethodOverflow)) { SubscriptionTerminationCode::Lagged } else { SubscriptionTerminationCode::ResourceGone },
|
||||||
message: "Worker protocol stream closed".to_string(),
|
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(),
|
"runtime-1".to_string(),
|
||||||
worker_id(),
|
worker_id(),
|
||||||
events,
|
events,
|
||||||
|
tokio::sync::watch::channel(None).1,
|
||||||
control_outbound,
|
control_outbound,
|
||||||
protocol_outbound,
|
protocol_outbound,
|
||||||
)
|
)
|
||||||
@@ -660,6 +697,81 @@ mod tests {
|
|||||||
response
|
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]
|
#[tokio::test]
|
||||||
async fn worker_protocol_rejects_a_non_snapshot_initial_event() {
|
async fn worker_protocol_rejects_a_non_snapshot_initial_event() {
|
||||||
let (event_sender, events) = mpsc::channel(1);
|
let (event_sender, events) = mpsc::channel(1);
|
||||||
|
|||||||
@@ -0,0 +1,111 @@
|
|||||||
|
# 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