diff --git a/src/test_handler.rs b/src/test_handler.rs index a69060b..d3dd86a 100644 --- a/src/test_handler.rs +++ b/src/test_handler.rs @@ -6,7 +6,9 @@ use http_body_util::{BodyExt, Empty, Full}; use hyper::{body::Bytes, Request}; use tokio::time::Duration; -use crate::commands::{MjCommand, PrintCommand, StaticCommand, ToSse}; +use crate::commands::{ + BusPubCommand, BusSubCommand, MjCommand, PrintCommand, StaticCommand, ToSse, +}; use crate::handler::{handle, AppConfig}; fn default_config() -> Arc { @@ -260,6 +262,197 @@ async fn test_sse_brotli_cancel_signals_error_to_client() { } } +/// `.bus sub` blocks in a signal-aware `recv_timeout` loop with nothing +/// published to it -- genuinely idle, the same shape as an idle `.cat +/// --follow`. Used instead of a timer-driven generator (`1.. | each { sleep +/// ...}`) because a generator that produces on its own eventually discovers +/// a dropped receiver via the existing `tx.send` failure regardless of this +/// fix, silently hiding the leak the same way live traffic does on the real +/// harness (see task doc's "method note that cost an hour"). +const IDLE_SSE_SCRIPT: &str = r#"{|req| + .bus sub "task4b" | each {|e| {data: $e.value} } | to sse +}"#; + +/// `test_engine` does not register `.bus sub`/`.bus pub` (no existing test +/// needs the bus), so build the engine directly rather than extending that +/// shared helper's command list for every other test. +fn idle_sse_engine() -> crate::Engine { + let mut engine = crate::Engine::new().unwrap(); + engine + .add_commands(vec![ + Box::new(ToSse {}), + Box::new(BusSubCommand::new(engine.bus.clone())), + Box::new(BusPubCommand::new(engine.bus.clone())), + ]) + .unwrap(); + engine + .set_http_nu_const(&crate::engine::HttpNuOptions::default()) + .unwrap(); + // Matches production (main.rs's setup_ctrlc_handler wires a real + // interrupt onto the engine): a non-empty Signals here means + // `.trigger()` actually does something, so a test that shares this + // process-wide flag across requests (the pre-fix bug) has a real signal + // to cross-trip, not a no-op `Signals::empty()`. + engine.set_signals(Arc::new(std::sync::atomic::AtomicBool::new(false))); + engine.parse_closure(IDLE_SSE_SCRIPT, None).unwrap(); + engine +} + +/// Publish until `body` yields a data frame, retrying the publish each time: +/// `.bus sub` subscribes to a broadcast channel with no backlog, on a thread +/// spun up inside the request's own eval thread, so a publish made before +/// that subscription exists is simply missed. +async fn publish_until_received( + bus: &crate::bus::Bus, + body: &mut http_body_util::combinators::BoxBody< + Bytes, + Box, + >, +) -> hyper::body::Frame { + loop { + bus.publish( + "task4b", + nu_protocol::Value::string("hi", nu_protocol::Span::unknown()), + ); + match tokio::time::timeout(Duration::from_millis(50), body.frame()).await { + Ok(Some(Ok(frame))) if frame.is_data() => return frame, + Ok(other) => panic!("expected SSE data frame, got: {other:?}"), + Err(_) => continue, // no subscriber registered yet; publish again + } + } +} + +/// A client that goes away (body dropped without a graceful close -- what a +/// killed curl looks like server-side) must free the request's job. Before +/// the per-connection Signals fix, nothing ever cancelled it: the job sat in +/// the table, its thread parked in a `.cat`/`.last`-equivalent blocking +/// read, for the life of the server. +#[tokio::test] +async fn test_client_disconnect_kills_job() { + let engine = idle_sse_engine(); + let jobs = engine.state.jobs.clone(); + let bus = engine.bus.clone(); + + let req = Request::builder() + .method("GET") + .uri("/sse") + .body(Empty::::new()) + .unwrap(); + + let resp = handle( + Arc::new(ArcSwap::from_pointee(engine)), + None, + default_config(), + req, + ) + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let mut body = resp.into_body(); + + // Confirm the job is live and streaming, not just registered. + publish_until_received(&bus, &mut body).await; + + assert_eq!( + jobs.lock().unwrap().iter().count(), + 1, + "job should be registered while the request streams" + ); + + // Simulate the client going away: drop the body without draining to + // completion or closing gracefully. Nothing more is published, so + // `.bus sub` is now genuinely idle -- parked in its own signal-aware + // wait, same as an idle `.cat --follow`. + drop(body); + + // The disconnect watcher notices `stream_tx.closed()` and kills the job + // asynchronously (a tokio task, not synchronous with the drop above); + // poll briefly rather than assume a fixed delay. + for _ in 0..50 { + if jobs.lock().unwrap().iter().count() == 0 { + return; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("job was not removed from the jobs table after the client disconnected"); +} + +/// The per-request Signals must not cross-trip: killing one request's job on +/// disconnect must not interrupt a different, still-connected request. Before +/// the fix, every request's `ThreadJob` shared `engine.state.signals()`, so +/// this would have stopped both. +#[tokio::test] +async fn test_client_disconnect_does_not_interrupt_other_requests() { + let engine = idle_sse_engine(); + let bus = engine.bus.clone(); + let engine = Arc::new(ArcSwap::from_pointee(engine)); + + let req_a = Request::builder() + .method("GET") + .uri("/sse") + .body(Empty::::new()) + .unwrap(); + let req_b = Request::builder() + .method("GET") + .uri("/sse") + .body(Empty::::new()) + .unwrap(); + + let resp_a = handle(engine.clone(), None, default_config(), req_a) + .await + .unwrap(); + let resp_b = handle(engine.clone(), None, default_config(), req_b) + .await + .unwrap(); + + let mut body_a = resp_a.into_body(); + let mut body_b = resp_b.into_body(); + + // Both requests subscribe to the same topic and both must receive the + // confirmation publish before either is considered "live". + let mut a_ready = false; + let mut b_ready = false; + while !(a_ready && b_ready) { + bus.publish( + "task4b", + nu_protocol::Value::string("hi", nu_protocol::Span::unknown()), + ); + if !a_ready { + if let Ok(Some(Ok(frame))) = + tokio::time::timeout(Duration::from_millis(50), body_a.frame()).await + { + assert!(frame.is_data()); + a_ready = true; + } + } + if !b_ready { + if let Ok(Some(Ok(frame))) = + tokio::time::timeout(Duration::from_millis(50), body_b.frame()).await + { + assert!(frame.is_data()); + b_ready = true; + } + } + } + + // A disconnects; B stays connected and idle (nothing published) in + // between, same as the leak scenario. + drop(body_a); + tokio::time::sleep(Duration::from_millis(200)).await; + + // B must still be alive: publish once more and confirm B receives it, + // not an error or a clean end from A's disconnect crossing over. + bus.publish( + "task4b", + nu_protocol::Value::string("still alive", nu_protocol::Span::unknown()), + ); + match tokio::time::timeout(Duration::from_secs(1), body_b.frame()).await { + Ok(Some(Ok(frame))) if frame.is_data() => {} + other => panic!("request B was interrupted by A's disconnect: {other:?}"), + } +} + fn assert_timing_sequence(timings: &[(String, Duration)]) { // Check values arrive in sequence for (i, (value, _)) in timings.iter().enumerate() { diff --git a/src/worker.rs b/src/worker.rs index 1e48d0b..10d3097 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -7,9 +7,10 @@ use crate::response::{ }; use nu_protocol::{ engine::{Job, StateWorkingSet, ThreadJob}, - format_cli_error, PipelineData, Value, + format_cli_error, PipelineData, Signals, Value, }; use std::io::Read; +use std::sync::atomic::AtomicBool; use std::sync::{mpsc, Arc}; use tokio::sync::{mpsc as tokio_mpsc, oneshot}; @@ -18,6 +19,29 @@ fn is_jsonl_record(value: &Value) -> bool { matches!(value, Value::Record { val, .. } if val.get("__html").is_none()) } +/// Kill this request's job once its response channel has no receiver left -- +/// the client disconnected, or the response finished and hyper dropped the +/// body. Either way there is nothing left to protect; `job.kill()` after a +/// normal finish is a harmless no-op (the job is already done, nothing +/// downstream inspects `signals.interrupted()` after the fact). +/// +/// `job.kill()` triggers this request's own `Signals` (see spawn_eval_thread: +/// each request gets its own, not a clone of the engine's process-wide one), +/// so a `.cat --follow`/`.last --follow` blocked in `Store::blocking_recv` +/// notices and returns instead of leaking its thread for the life of the +/// server. It also reaps any external-command PIDs this request's closure +/// spawned, via the same path Ctrl-C already used. +fn watch_disconnect( + runtime: &tokio::runtime::Handle, + tx: tokio_mpsc::Sender>, + job: ThreadJob, +) { + runtime.spawn(async move { + tx.closed().await; + let _ = job.kill(); + }); +} + type BoxError = Box; /// Result of pipeline evaluation containing content-type, HTTP response metadata, and body @@ -40,6 +64,8 @@ pub fn spawn_eval_thread( stream: nu_protocol::ByteStream, meta_tx: oneshot::Sender, body_tx: oneshot::Sender, + runtime: tokio::runtime::Handle, + job: ThreadJob, ) -> Result<(), BoxError> { RESPONSE_TX.with(|tx| { *tx.borrow_mut() = Some(meta_tx); @@ -121,6 +147,7 @@ pub fn spawn_eval_thread( PipelineData::ListStream(stream, meta) => { let http_meta = extract_http_response_meta(meta.as_ref()); let (stream_tx, stream_rx) = tokio_mpsc::channel(32); + watch_disconnect(&runtime, stream_tx.clone(), job.clone()); let mut iter = stream.into_inner(); // Peek first value to determine mode @@ -179,6 +206,7 @@ pub fn spawn_eval_thread( PipelineData::ByteStream(stream, meta) => { let http_meta = extract_http_response_meta(meta.as_ref()); let (stream_tx, stream_rx) = tokio_mpsc::channel(32); + watch_disconnect(&runtime, stream_tx.clone(), job.clone()); let content_type = meta .as_ref() .and_then(|m| m.content_type.clone()) @@ -220,10 +248,19 @@ pub fn spawn_eval_thread( } } + // Each request gets its own interrupt flag, not a clone of the engine's + // process-wide `engine.state.signals()`. With the process-wide one, + // killing any single request's job -- including the per-connection + // disconnect kill added below -- trips every other concurrent request's + // signals too. Ctrl-C still stops everything: it kills every job in the + // table (main.rs's setup_ctrlc_handler), one `ThreadJob::kill()` call + // per job, each tripping its own request's flag. + let interrupt = Arc::new(AtomicBool::new(false)); + let signals = Signals::new(interrupt); + // Create a thread job for this evaluation let (sender, _receiver) = mpsc::channel(); - let signals = engine.state.signals().clone(); - let job = ThreadJob::new(signals, Some("HTTP Request".to_string()), sender); + let job = ThreadJob::new(signals.clone(), Some("HTTP Request".to_string()), sender); // Add the job to the engine's job table let job_id = { @@ -231,6 +268,11 @@ pub fn spawn_eval_thread( jobs.add_job(Job::Thread(job.clone())) }; + // Captured here, on the async caller's thread, so it's usable from + // `inner`'s plain std::thread (which has no tokio context of its own) to + // spawn the disconnect watcher below. + let runtime = tokio::runtime::Handle::current(); + std::thread::spawn(move || -> Result<(), std::convert::Infallible> { let mut meta_tx_opt = Some(meta_tx); let mut body_tx_opt = Some(body_tx); @@ -239,7 +281,12 @@ pub fn spawn_eval_thread( // async runtime and we can still send a response back to the caller. let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { let mut local_engine = (*engine).clone(); - local_engine.state.current_job.background_thread_job = Some(job); + local_engine.state.current_job.background_thread_job = Some(job.clone()); + // Route this request's own Signals to the engine the closure + // actually runs on, not just the ThreadJob in the jobs table -- + // .cat/.last and any other signal-aware command check + // engine_state.signals(), not the job table. + local_engine.state.set_signals(signals.clone()); // Take the senders for the inner call. If the evaluation completes // successfully, these senders will have been consumed. Otherwise we @@ -250,6 +297,8 @@ pub fn spawn_eval_thread( stream, meta_tx_opt.take().unwrap(), body_tx_opt.take().unwrap(), + runtime, + job, ) }));