diff --git a/crates/base/src/server.rs b/crates/base/src/server.rs index 0b33426d..ba0eb85e 100644 --- a/crates/base/src/server.rs +++ b/crates/base/src/server.rs @@ -235,9 +235,10 @@ impl Service> for WorkerService { } }; - // If the token has already been canceled, return 503 instead of - // dropping the socket connection without a response. - if cancel.is_cancelled() { + // User-worker teardown also cancels this shared token. The main worker + // can recover from that failure and return a valid response, so only use + // the cancellation fallback when it could not produce a response. + if res.is_err() && cancel.is_cancelled() { error!("connection aborted (uri: {:?})", req_uri.to_string()); return Ok( Response::builder() @@ -1005,3 +1006,91 @@ fn accept_stream( } }); } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn cancelled_request_preserves_main_worker_response() { + let (tx, mut rx) = mpsc::unbounded_channel::(); + let (mut service, connection_token) = + WorkerService::new(SharedMetricSource::default(), tx); + + // Force teardown to win the race, without relying on scheduling or CPU + // timing. The main worker can still recover and provide a complete response. + let main_worker = tokio::spawn(async move { + let msg = rx.recv().await.unwrap(); + msg.conn_token.unwrap().cancel(); + let (mut body_tx, body) = Body::channel(); + msg + .res_tx + .send(Ok( + Response::builder() + .status(http_v02::StatusCode::INTERNAL_SERVER_ERROR) + .header("content-type", "application/json") + .header("x-main-worker", "recovered") + .body(body) + .unwrap(), + )) + .unwrap(); + body_tx.send_data("{\"msg\":".into()).await.unwrap(); + body_tx + .send_data("\"worker failed\"}".into()) + .await + .unwrap(); + }); + + let response = service.call(Request::new(Body::empty())).await.unwrap(); + assert_eq!( + response.status(), + http_v02::StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!(response.headers()["x-main-worker"], "recovered"); + assert_eq!(response.headers()["content-type"], "application/json"); + assert!(!response.headers().contains_key("x-served-by")); + let body = hyper_v014::body::to_bytes(response.into_body()) + .await + .unwrap(); + assert_eq!(body.as_ref(), br#"{"msg":"worker failed"}"#); + main_worker.await.unwrap(); + assert!(!connection_token.is_cancelled()); + } + + #[tokio::test] + async fn response_drop_cancels_request_but_not_keep_alive_connection() { + let (tx, mut rx) = mpsc::unbounded_channel::(); + let (mut service, connection_token) = + WorkerService::new(SharedMetricSource::default(), tx); + let main_worker = tokio::spawn(async move { + let msg = rx.recv().await.unwrap(); + let token = msg.conn_token.unwrap(); + msg + .res_tx + .send(Ok(Response::new(Body::from("ok")))) + .unwrap(); + token + }); + let response = service.call(Request::new(Body::empty())).await.unwrap(); + let request_token = main_worker.await.unwrap(); + assert!(!request_token.is_cancelled()); + drop(response); + assert!(request_token.is_cancelled()); + assert!(!connection_token.is_cancelled()); + } + + #[tokio::test] + async fn connection_close_still_cancels_pending_request() { + let (tx, mut rx) = mpsc::unbounded_channel::(); + let (mut service, connection_token) = + WorkerService::new(SharedMetricSource::default(), tx); + let response = tokio::spawn(service.call(Request::new(Body::empty()))); + let msg = rx.recv().await.unwrap(); + let request_token = msg.conn_token.unwrap(); + assert!(!request_token.is_cancelled()); + connection_token.cancel(); + assert!(request_token.is_cancelled()); + drop(msg.res_tx); + response.await.unwrap().unwrap(); + } +} diff --git a/crates/base/test_cases/main_worker_recovery/index.ts b/crates/base/test_cases/main_worker_recovery/index.ts new file mode 100644 index 00000000..a92aa9b1 --- /dev/null +++ b/crates/base/test_cases/main_worker_recovery/index.ts @@ -0,0 +1,30 @@ +// Keep a failed handle to exercise dispatch after teardown as well as a fresh +// worker. HTTP keep-alive is controlled independently by the integration test. +let failedWorker; + +Deno.serve(async (req: Request) => { + const scenario = new URL(req.url).pathname.slice(1); + try { + const worker = scenario === 'retired' ? failedWorker : await EdgeRuntime.userWorkers.create({ + servicePath: scenario === 'cpu' + ? '../../examples/cpu' + : scenario === 'rejection' + ? './test_cases/unhandled-rejection' + : './test_cases/module-scope-throw', + memoryLimitMb: 150, + workerTimeoutMs: 10_000, + cpuTimeSoftLimitMs: 100, + cpuTimeHardLimitMs: 200, + forceCreate: true, + noModuleCache: false, + envVars: [], + }); + failedWorker = worker; + return await worker.fetch(req); + } catch (e) { + return Response.json({ msg: e.toString() }, { + status: 500, + headers: { 'x-main-worker': 'recovered' }, + }); + } +}); diff --git a/crates/base/test_cases/module-scope-throw/index.ts b/crates/base/test_cases/module-scope-throw/index.ts new file mode 100644 index 00000000..d40f3b13 --- /dev/null +++ b/crates/base/test_cases/module-scope-throw/index.ts @@ -0,0 +1 @@ +throw new Error('module-scope regression fixture'); diff --git a/crates/base/test_cases/unhandled-rejection/index.ts b/crates/base/test_cases/unhandled-rejection/index.ts new file mode 100644 index 00000000..8022e8dc --- /dev/null +++ b/crates/base/test_cases/unhandled-rejection/index.ts @@ -0,0 +1,4 @@ +Deno.serve(() => { + Promise.reject(new Error('unhandled rejection regression fixture')); + return new Promise(() => {}); +}); diff --git a/crates/base/tests/integration_tests.rs b/crates/base/tests/integration_tests.rs index b54d36f1..a5fbb5be 100644 --- a/crates/base/tests/integration_tests.rs +++ b/crates/base/tests/integration_tests.rs @@ -1263,6 +1263,75 @@ async fn req_failure_case_intentional_peer_reset_secure() { req_failure_case_intentional_peer_reset(new_localhost_tls(true)).await; } +#[tokio::test] +#[serial] +async fn req_failure_case_preserves_main_worker_response() { + integration_test!( + "./test_cases/main_worker_recovery", + NON_SECURE_PORT, + "cpu", + None, + None, + None, + ( + |(port, _, _, _, ..)| async move { + for keep_alive in [false, true] { + let client = Client::builder() + .pool_max_idle_per_host(if keep_alive { 1 } else { 0 }) + .timeout(Duration::from_secs(10)) + .build() + .unwrap(); + for scenario in + ["cpu", "retired", "throw", "retired", "rejection", "retired"] + { + let res = client + .post(format!("http://localhost:{port}/{scenario}")) + .header("content-type", "application/json") + .body(r#"{"base":100}"#) + .send() + .await + .unwrap(); + assert_eq!(res.status().as_u16(), 500, "{scenario}"); + assert_eq!(res.headers()["x-main-worker"], "recovered"); + assert!(!res.headers().contains_key("x-served-by")); + let payload = res.json::().await.unwrap(); + let expected = match scenario { + "cpu" => payload.msg.starts_with("WorkerRequestCancelled:"), + // The pool may have removed the worker, marked it retired, or + // still retain its original failure. Preserve the main response + // in each case, independently of pool cleanup timing. + "retired" => { + payload.msg.starts_with("WorkerAlreadyRetired:") + || payload.msg.starts_with("InvalidWorkerResponse:") + } + _ => payload.msg.starts_with("InvalidWorkerResponse:"), + }; + assert!(expected, "{scenario}: {}", payload.msg); + } + } + Some( + Client::new() + .get(format!("http://localhost:{port}/throw")) + .send() + .await, + ) + }, + |resp| async { + let res = resp.unwrap(); + assert_eq!(res.status().as_u16(), 500); + assert_eq!(res.headers()["x-main-worker"], "recovered"); + assert!(res + .json::() + .await + .unwrap() + .msg + .starts_with("InvalidWorkerResponse:")); + } + ), + TerminationToken::new() + ); +} + #[tokio::test] #[serial] async fn req_failure_case_op_cancel_from_server_due_to_cpu_resource_limit() { @@ -1291,58 +1360,21 @@ async fn req_failure_case_op_cancel_from_server_due_to_cpu_resource_limit_2() { .await; } -/// When the supervisor tears down a user worker that exceeded its CPU limit, -/// two paths race each other and both outcomes are correct: -/// -/// 1. The connection token is canceled before the main worker's response is -/// handed back to the server, so the server serves 503 on its own. -/// 2. The main worker's `WorkerRequestCancelled` handler wins the race, and its -/// own 500 payload is relayed to the client untouched. -/// -/// Which one wins depends on scheduling alone, so accept either, but keep -/// asserting the shape of the response so that unrelated failures (most -/// notably the request body being detached from its receiver) are still caught. +/// Worker teardown can cancel the shared token before the main worker replies. +/// A complete main-worker response must still reach the client in that order. async fn assert_op_cancel_from_server_response( resp: Result, ) { let res = resp.unwrap(); - let status = res.status().as_u16(); - let served_by = res - .headers() - .get("x-served-by") - .map(|v| v.to_str().unwrap().to_owned()); - - match status { - 503 => { - assert_eq!( - served_by.as_deref(), - Some(concat!(env!("CARGO_PKG_NAME"), "/server")) - ); - } - - 500 => { - assert_eq!(served_by, None); - - let payload = res.json::().await; + assert_eq!(res.status().as_u16(), 500); + assert!(!res.headers().contains_key("x-served-by")); - assert!(payload.is_ok()); - - let msg = payload.unwrap().msg; - - assert!( - !msg.starts_with("TypeError: request body receiver not connected"), - "unexpected error message: {msg}" - ); - assert!( - msg - == "WorkerRequestCancelled: request has been cancelled by supervisor" - || msg == "broken pipe", - "unexpected error message: {msg}" - ); - } - - _ => panic!("unexpected status code: {status}"), - } + let msg = res.json::().await.unwrap().msg; + assert!( + msg == "WorkerRequestCancelled: request has been cancelled by supervisor" + || msg == "broken pipe", + "unexpected error message: {msg}" + ); } async fn test_oak_file_upload(