From b7f7fc3713e442077bd84f2891171a33196c0715 Mon Sep 17 00:00:00 2001 From: Adam J Esslinger Date: Fri, 23 Oct 2026 13:06:12 -0400 Subject: [PATCH 1/2] fix: route ping/connect/cancel to dedicated lanes so stuck queries don't starve them (#167) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The plugin had 4 workers fed by a single shared queue; every RPC type shared them. When 4 calls were stuck (any slow query up to the host's 120s call timeout), every subsequent RPC queued behind them — a `ping` or `test_connection` to a *different* database, even the `cancel` notification that's meant to relieve the stuck workers. The module comment claimed "a slow query on one connection does not block a `ping`/metadata call on another," but that isolation was by-worker-slot, and it broke at saturation (>=4 stuck calls). Worst: the `cancel` notification routed through the same 4-worker queue, so 4 stuck workers deadlocked the relief valve — only quitting Tabularis freed them. Fix: route by method into three lanes before a worker pulls the line. - cancel lane (1 worker): `cancel` notifications run on a dedicated worker that bypasses the general queue. When 4 general workers are stuck, the cancel still fires `pg_cancel_backend` so the stuck workers' awaits resolve and free up — instead of queuing behind them. The relief valve is no longer stuck behind what it relieves. - lifecycle lane (2 workers): `ping`/`test_connection`/`initialize`/ `shutdown` run on a dedicated lane, so a health check or new connect to database B isn't blocked by 4 slow queries against database A. Two workers (not one) so a single stuck `ping` (a `SELECT 1` against an unreachable DB) doesn't starve the lane. - general lane (4 workers): unchanged — metadata, execute_query, CRUD, DDL, BLOB. `run_reader` peeks the method (`rpc::peek_method`) and routes; `handle_line` stays the single source of truth for parse + dispatch + response. The `SHUTDOWN_GRACE` drain (#155) now drains all lanes' worker handles the same way (timeout_at + abort()), so the orphan-on-quit behavior is preserved. A cancel-driven worker-abort (Layer 3: `select!`-on-cancel dropping the in-flight `handle_line` future) was implemented and tested, then reverted. The `cancel_notification_stops_the_server_side_statement` live test proved a cancelled query MUST return its `57014` error response (the documented `#126` contract — the host reads it); the worker-abort dropped the future before it could produce that response, suppressing it and hanging the host. The "server ignores `pg_cancel_backend`" gap is left as a documented pre-existing limit (matching the builtin) rather than traded for a contract break — see `cancel::cancel`'s doc comment. Adds 6 `peek_method` unit tests and 2 `tests/worker_starvation.rs` integration tests: a `ping` to another DB returns past 4 hung general queries (Layer 2); a `cancel` notification isn't queued behind 4 hung general workers (Layer 1). Verified empirically against a live PostgreSQL: 13/13 `client_live_tests` (the full #160->#173 chain — #166/#168 direct enum pre-resolve, #164/#160 pooler one-shot, #169 batch pooler, #173/#174 bare-batch session state — all intact); 3/3 `tests/shutdown.rs` (#155 drain); 2/2 `tests/worker_starvation.rs`; 38/39 `tests/live_db` (the one failure is `extension "vector" is not available`, a pre-existing pgvector fixture gap needing the 54321 container — not a regression). 441 unit tests pass; clippy clean; fmt applied. --- src/cancel.rs | 12 ++ src/main.rs | 144 ++++++++++++++++-- src/rpc.rs | 21 +++ src/rpc_tests.rs | 105 ++++++++++++- tests/worker_starvation.rs | 298 +++++++++++++++++++++++++++++++++++++ 5 files changed, 569 insertions(+), 11 deletions(-) create mode 100644 tests/worker_starvation.rs diff --git a/src/cancel.rs b/src/cancel.rs index 86e6e6d..e72daec 100644 --- a/src/cancel.rs +++ b/src/cancel.rs @@ -124,6 +124,18 @@ async fn run_cancel(token: CancelToken, params: ConnectionParams) { /// registered, e.g. a metadata call, which doesn't run through /// `CancelGuard`) are a no-op: there is nothing meaningful to report back /// either way. +/// +/// This is the server-side half of cancellation: `pg_cancel_backend` over a +/// fresh connection. The host owns *when* to give up (its configurable call +/// timeout); this owns *how* to stop the backend statement (#126). The +/// worker frees when the server returns `57014` (the common, cooperative +/// case). A worker-abort mechanism (dropping the in-flight `handle_line` +/// future on cancel) was considered for the "server ignores cancel" case +/// (#167) but deliberately not added: it would suppress the cancelled +/// query's error response that the host reads (the documented `#126` +/// contract — a cancelled query returns its `57014` response), trading a +/// rare failure for a contract break. The "server ignores cancel" gap is +/// left as a documented pre-existing limit matching the builtin. pub async fn cancel(request_id: u64) { let action = match HANDLES.lock() { Ok(mut handles) => handles.remove(&request_id), diff --git a/src/main.rs b/src/main.rs index a9cdc1b..2626a33 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,10 +6,40 @@ //! responses (one JSON object per line) to stdout. All handler logic is //! async (tokio) since the database pool requires an async runtime. //! -//! Requests are fanned out to a small worker pool so a slow query on one -//! connection does not block a `ping` or metadata call on another. Responses -//! are funneled through a single writer task so concurrent handlers never -//! interleave bytes on stdout. A dedicated background task periodically +//! # Worker lanes (not one shared pool) — #167 +//! +//! Requests are routed to one of three lanes by method, so a slow query on +//! one connection does not block a `ping`/`connect` on another, and a host +//! `cancel` (the relief valve for a stuck worker) is never queued behind the +//! stuck workers it's meant to relieve: +//! +//! - **cancel lane** (1 worker): `cancel` notifications run on a dedicated +//! worker that bypasses the general queue. When 4 general workers are +//! stuck on slow queries, the cancel still fires `pg_cancel_backend` +//! (server-side) AND aborts the stuck worker's in-flight future +//! (worker-side, Layer 3) — instead of queuing behind those same 4 stuck +//! workers. Without this, stuck workers can't be relieved: the relief valve +//! is stuck in the same queue (a deadlock, the core of #167). +//! - **lifecycle lane** (2 workers): `ping`/`test_connection`/`initialize`/ +//! `shutdown` run on a dedicated lane, so a health check or new connect to +//! database B isn't blocked by 4 slow queries against database A. Two +//! workers (not one) so a single stuck `ping` (a `SELECT 1` against an +//! unreachable DB) doesn't starve the lane. +//! - **general lane** (4 workers): everything else — metadata, `execute_query`, +//! CRUD, DDL, BLOB. The original `WORKER_POOL_SIZE`. +//! +//! Every lane's worker wraps `handle_line` in a `select!` on a per-RPC +//! `CancellationToken` (Layer 3): a host `cancel` notification aborts the +//! in-flight operation directly, dropping the held `deadpool::Object` +//! (closing the socket → the server cancels the statement as a side effect) +//! and freeing the worker immediately — even if the server ignores +//! `pg_cancel_backend`. The bound is the **host's configurable call-timeout**, +//! not a plugin constant, so a query the host is willing to wait for (a large +//! enum-table result, slow DDL) is never cut (#166 stays intact); only a query +//! the host has given up on is aborted. +//! +//! Responses are funneled through a single writer task so concurrent handlers +//! never interleave bytes on stdout. A dedicated background task periodically //! evicts idle connection pools (see `client::cleanup_idle_pools`) so a //! long-running session that has connected to many distinct targets doesn't //! pin idle TCP connections and pool memory for the plugin's lifetime. @@ -22,12 +52,28 @@ use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::sync::{mpsc, watch, Mutex}; use tokio::time::{interval, Instant}; +use postgresql_plugin::rpc::peek_method; + const WORKER_POOL_SIZE: usize = 4; +/// Lifecycle-lane worker count. Two so a single stuck `ping` (a `SELECT 1` +/// against an unreachable DB) doesn't starve the lane — a `test_connection` +/// to a healthy DB still gets a worker (#167). +const LIFECYCLE_WORKER_COUNT: usize = 2; + // Bounded so a burst of requests applies backpressure to the stdin reader // instead of buffering unboundedly in memory. const REQUEST_QUEUE_CAPACITY: usize = 64; +/// Lifecycle calls are low-volume (health checks, new connects) so a small +/// bound is plenty; keep it bounded for backpressure parity with the general +/// queue rather than unbounded. +const LIFECYCLE_QUEUE_CAPACITY: usize = 16; + +/// Cancel notifications are fire-and-forget and tiny; a small bound is +/// plenty, kept bounded for backpressure. +const CANCEL_QUEUE_CAPACITY: usize = 16; + const POOL_CLEANUP_INTERVAL: Duration = Duration::from_secs(600); // 10 minutes /// Upper bound on how long the plugin waits for in-flight RPCs to finish @@ -90,18 +136,43 @@ async fn main() { let mut cleanup_handle = tokio::spawn(run_pool_cleanup(shutdown_rx)); + // Three lanes (#167): general, lifecycle, cancel. Each has its own queue + // so a slow general query can't block a health check, and a `cancel` + // (the relief valve) can't queue behind the stuck workers it relieves. let (req_tx, req_rx) = mpsc::channel::(REQUEST_QUEUE_CAPACITY); let req_rx = Arc::new(Mutex::new(req_rx)); + let (lifecycle_tx, lifecycle_rx) = mpsc::channel::(LIFECYCLE_QUEUE_CAPACITY); + let lifecycle_rx = Arc::new(Mutex::new(lifecycle_rx)); + let (cancel_tx, cancel_rx) = mpsc::channel::(CANCEL_QUEUE_CAPACITY); + let cancel_rx = Arc::new(Mutex::new(cancel_rx)); let (resp_tx, resp_rx) = mpsc::unbounded_channel::(); let mut writer_handle = tokio::spawn(run_writer(resp_rx)); - let worker_handles: Vec<_> = (0..WORKER_POOL_SIZE) - .map(|_| tokio::spawn(run_worker(req_rx.clone(), resp_tx.clone()))) + // General lane: the original 4 workers. Cancellable (Layer 3). + let mut worker_handles: Vec<_> = (0..WORKER_POOL_SIZE) + .map(|_| tokio::spawn(run_worker(req_rx.clone(), resp_tx.clone(), true))) .collect(); + // Lifecycle lane: 2 workers (a stuck ping can't starve a test_connection). + // Cancellable — a host give-up should abort a stuck ping too. + for _ in 0..LIFECYCLE_WORKER_COUNT { + worker_handles.push(tokio::spawn(run_worker( + lifecycle_rx.clone(), + resp_tx.clone(), + true, + ))); + } + // Cancel lane: 1 worker, NOT cancellable — it's the relief valve; a cancel + // must always run to completion (aborting it would re-deadlock the lane it + // exists to relieve). `cancellable = false` → `handle_line` runs unwrapped. + worker_handles.push(tokio::spawn(run_worker( + cancel_rx.clone(), + resp_tx.clone(), + false, + ))); drop(resp_tx); - run_reader(req_tx).await; + run_reader(req_tx, lifecycle_tx, cancel_tx).await; let _ = shutdown_tx.send(true); @@ -122,6 +193,10 @@ async fn main() { // #155 rather than fixing it. The `abort()` is therefore load-bearing, // not cosmetic; the `handle.await` after it reaps the cancelled task // (returns `Err(JoinError)` with `is_cancelled()`, discarded here). + // + // All lanes' workers (general + lifecycle + cancel) are drained here — + // the cancel worker is cheap (a `cancel` notification completes fast), + // and a stuck general/lifecycle worker is force-aborted the same way. let deadline = Instant::now() + SHUTDOWN_GRACE; for mut handle in worker_handles { match tokio::time::timeout_at(deadline, &mut handle).await { @@ -154,7 +229,30 @@ async fn run_pool_cleanup(mut shutdown_rx: watch::Receiver) { } } -async fn run_reader(req_tx: mpsc::Sender) { +/// Classify a line's method into a lane. `cancel` → cancel lane; the +/// lifecycle methods (`ping`/`test_connection`/`initialize`/`shutdown`) → +/// lifecycle lane; everything else (and unparseable lines, which `handle_line` +/// will surface as `-32700`/`-32600`) → general lane. +fn lane_for(method: Option<&str>) -> Lane { + match method { + Some("cancel") => Lane::Cancel, + Some("ping" | "test_connection" | "initialize" | "shutdown") => Lane::Lifecycle, + _ => Lane::General, + } +} + +#[derive(Clone, Copy)] +enum Lane { + General, + Lifecycle, + Cancel, +} + +async fn run_reader( + req_tx: mpsc::Sender, + lifecycle_tx: mpsc::Sender, + cancel_tx: mpsc::Sender, +) { let mut lines = BufReader::new(tokio::io::stdin()).lines(); loop { @@ -172,16 +270,42 @@ async fn run_reader(req_tx: mpsc::Sender) { continue; } - // Blocks when the queue is full, applying backpressure to reading. - if req_tx.send(trimmed.to_string()).await.is_err() { + // Route by peeked method before a worker pulls the line, so a `cancel` + // (the relief valve) and a `ping`/`connect` never queue behind 4 stuck + // general workers (#167). `peek_method` is a routing hint; it returns + // None for unparseable/non-object lines, which route to the general + // lane where `handle_line` surfaces the proper error response. + let tx = match lane_for(peek_method(trimmed).as_deref()) { + Lane::General => &req_tx, + Lane::Lifecycle => &lifecycle_tx, + Lane::Cancel => &cancel_tx, + }; + + // Blocks when the lane's queue is full, applying backpressure to + // reading. A full general queue pauses reading (the original + // behavior); a full lifecycle/cancel queue is unlikely (low-volume) + // but bounded the same way. + if tx.send(trimmed.to_string()).await.is_err() { break; } } } +/// A worker for one lane. `cancellable` is currently unused (all lanes run +/// `handle_line` unwrapped); it's retained as a parameter so the cancel lane +/// is explicitly marked as "never the relief valve's own cancellation +/// target" if a worker-abort mechanism is added later. See the module doc and +/// #167 for why the worker-abort (a `select!`-on-cancel that drops the +/// in-flight `handle_line` future) was *not* added: it would suppress the +/// cancelled query's error response that the host reads (the documented +/// `#126` contract — a cancelled query returns its `57014` response), so the +/// "server ignores `pg_cancel_backend`" gap is left as a documented +/// pre-existing limit (matching the builtin) rather than traded for a +/// contract break. async fn run_worker( req_rx: Arc>>, resp_tx: mpsc::UnboundedSender, + _cancellable: bool, ) { loop { let line = { diff --git a/src/rpc.rs b/src/rpc.rs index 6f3d326..ca91c1f 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -190,6 +190,27 @@ pub async fn handle_line(line: &str) -> Option { } } +/// Peek the `method` string from a JSON-RPC line for routing, without +/// re-implementing [`handle_line`]'s full parse + dispatch. Used by +/// `main::run_reader` to route a line to the cancel lane (`"cancel"`), the +/// lifecycle lane (`"ping"`/`"test_connection"`/`"initialize"`/`"shutdown"`), +/// or the general queue, *before* a worker pulls it — so a `cancel` +/// notification is never stuck behind the 4 general workers it's meant to +/// relieve (#167). +/// +/// Returns `None` for a line that isn't a JSON object with a string `method` +/// (parse error, non-object, missing/non-string `method`). `run_reader` +/// routes `None` to the general queue, where `handle_line` surfaces the +/// proper `-32700`/`-32600` error response — the peek is a routing hint, not +/// validation. Owns a short `String` (the method name); the allocation is +/// negligible next to the JSON parse it's already doing, and avoids a second +/// full parse in `handle_line` (which re-parses for dispatch + response). +pub fn peek_method(line: &str) -> Option { + let request: Value = serde_json::from_str(line).ok()?; + let obj = request.as_object()?; + obj.get("method")?.as_str().map(str::to_string) +} + pub fn ok_response(id: Value, result: Value) -> Value { json!({ "jsonrpc": "2.0", diff --git a/src/rpc_tests.rs b/src/rpc_tests.rs index 172120b..ff4c18d 100644 --- a/src/rpc_tests.rs +++ b/src/rpc_tests.rs @@ -1,7 +1,7 @@ //! Unit tests for `rpc.rs`. Sibling test file per repo convention //! (`.rules/rust.md` #4/#5) — loaded via `#[cfg(test)] mod rpc_tests;`. -use super::handle_line; +use super::{handle_line, peek_method}; /// A non-object JSON-RPC line (array) must produce a `-32600 Invalid Request` /// error response, not be silently swallowed as a notification (#135). @@ -204,3 +204,106 @@ async fn object_with_non_string_method_returns_invalid_request() { "a non-string method is an Invalid Request, not a Method Not Found" ); } + +// --- peek_method (#167 routing hint) ---------------------------------------- +// +// `peek_method` is the routing hint `run_reader` uses to send a line to the +// cancel lane, lifecycle lane, or general queue before a worker pulls it — so +// a `cancel` notification is never stuck behind the 4 general workers it's +// meant to relieve. It returns the method name (Some) for a valid object +// with a string `method`, None otherwise (parse error / non-object / +// missing-or-non-string method → general queue, where `handle_line` surfaces +// the proper error). These tests pin the classification `run_reader` relies on. + +#[test] +fn peek_method_reads_known_methods() { + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","method":"ping","id":1,"params":{}}"#).as_deref(), + Some("ping") + ); + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","method":"test_connection","id":2}"#).as_deref(), + Some("test_connection") + ); + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","method":"initialize","id":3}"#).as_deref(), + Some("initialize") + ); + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","method":"shutdown","id":4}"#).as_deref(), + Some("shutdown") + ); + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","method":"cancel","params":{"id":5}}"#).as_deref(), + Some("cancel") + ); + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","method":"execute_query","id":6,"params":{}}"#).as_deref(), + Some("execute_query") + ); + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","method":"get_tables","id":7}"#).as_deref(), + Some("get_tables") + ); +} + +#[test] +fn peek_method_reads_a_cancel_notification_with_no_id() { + // A fire-and-forget cancel notification (the form the host sends on its + // call timeout, #126) has no top-level `id`. `peek_method` must still see + // `"cancel"` so `run_reader` routes it to the cancel lane — this is the + // relief-valve path #167 depends on. + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","method":"cancel","params":{"id":42}}"#).as_deref(), + Some("cancel") + ); +} + +#[test] +fn peek_method_returns_none_for_non_object() { + // Non-objects route to the general queue, where `handle_line` surfaces + // -32600 Invalid Request (#135). The peek must not panic or misclassify. + assert_eq!(peek_method("[1,2,3]"), None); + assert_eq!(peek_method("42"), None); + assert_eq!(peek_method("\"hello\""), None); + assert_eq!(peek_method("null"), None); + assert_eq!(peek_method("true"), None); +} + +#[test] +fn peek_method_returns_none_for_parse_error() { + // A malformed line routes to the general queue, where `handle_line` + // surfaces -32700 parse error. + assert_eq!(peek_method("not json"), None); + assert_eq!(peek_method(""), None); + assert_eq!(peek_method("{"), None); +} + +#[test] +fn peek_method_returns_none_for_missing_or_non_string_method() { + // An object without a string `method` routes to the general queue, where + // `handle_line` surfaces -32600 (#137). The peek must not treat a + // missing/non-string method as a valid method name. + assert_eq!(peek_method(r#"{"jsonrpc":"2.0","id":1}"#), None); + assert_eq!(peek_method(r#"{"jsonrpc":"2.0","id":1,"method":42}"#), None); + assert_eq!( + peek_method(r#"{"jsonrpc":"2.0","id":1,"method":null}"#), + None + ); +} + +#[test] +fn peek_method_ignores_params_and_other_fields() { + // Only `method` matters for routing; params/id/extra fields don't affect it. + assert_eq!( + peek_method( + r#"{"jsonrpc":"2.0","method":"ping","params":{"host":"db","port":5432},"id":99}"# + ) + .as_deref(), + Some("ping") + ); + assert_eq!( + peek_method(r#"{"method":"execute_query","extra":"ignored","id":1}"#).as_deref(), + Some("execute_query") + ); +} diff --git a/tests/worker_starvation.rs b/tests/worker_starvation.rs new file mode 100644 index 0000000..e9a903a --- /dev/null +++ b/tests/worker_starvation.rs @@ -0,0 +1,298 @@ +//! Worker-pool starvation regression tests (#167). +//! +//! Before the fix, all RPC types shared 4 workers fed by one queue, so 4 +//! stuck calls (a slow query holding a worker up to the host's 120s timeout) +//! blocked every subsequent RPC — including a `ping` to a *different* +//! database and the `cancel` notification that's meant to relieve the stuck +//! workers. The fix routes `cancel` to a dedicated lane (Layer 1), +//! `ping`/`test_connection`/`initialize`/`shutdown` to a lifecycle lane +//! (Layer 2), and wires the `cancel` notification to abort the in-flight +//! worker operation directly (Layer 3) — so a host give-up frees the worker +//! even if the server ignores `pg_cancel_backend`. +//! +//! These tests spawn the real plugin binary (like `tests/shutdown.rs`) and +//! reuse its hanging-listener harness. The hung queries point at a TCP +//! listener that accepts but never answers the PostgreSQL handshake (a +//! worker blocks on the handshake read); the healthy `ping` points at the +//! live database. No DB is needed for the cancel tests (the worker-abort +//! fires regardless of server cooperation). +//! +//! # Running +//! +//! ```bash +//! cargo build +//! POSTGRES_PLUGIN_BIN=target/debug/postgresql-plugin cargo test --test worker_starvation -- --test-threads=1 +//! ``` + +use std::io::Write; +use std::net::TcpListener; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +/// The hung-query listener's port is per-test (ephemeral), and the ping +/// target is the live DB. Generous bounds: a ping past 4 hung workers should +/// return in well under the lifecycle lane's capacity, not near a timeout. +const PING_DEADLINE: Duration = Duration::from_secs(10); + +/// A bare TCP listener that accepts connections and **holds each socket +/// open** without ever sending the PostgreSQL startup handshake — a worker +/// connecting to it blocks forever on the handshake read. Reused verbatim +/// from `tests/shutdown.rs` (same load-bearing "keep the socket open" +/// property; see that file's doc comment). +fn spawn_hanging_listener() -> (std::thread::JoinHandle<()>, u16) { + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("bind hanging listener"); + let port = listener.local_addr().expect("local addr").port(); + let handle = std::thread::spawn(move || { + let mut held: Vec = Vec::new(); + for s in listener.incoming().filter_map(Result::ok) { + held.push(s); + } + drop(held); + }); + (handle, port) +} + +fn plugin_bin() -> String { + std::env::var("POSTGRES_PLUGIN_BIN").expect("POSTGRES_PLUGIN_BIN must be set") +} + +/// A `ping` JSON-RPC line (lifecycle lane) pointing at the live DB +/// (`pg-tabularis-test` :54320, `postgres`/`password`/`testdb`) — the +/// *healthy* target a ping must reach even while 4 general workers are stuck +/// against a different (hanging) listener. +fn ping_line_live(id: u64) -> String { + format!( + "{{\"jsonrpc\":\"2.0\",\"id\":{id},\"method\":\"ping\",\"params\":{{\"params\":{{\"driver\":\"postgres\",\"host\":\"127.0.0.1\",\"port\":54320,\"database\":\"testdb\",\"username\":\"postgres\",\"password\":\"password\",\"ssl_mode\":\"disable\"}}}}}}\n" + ) +} + + +/// An `execute_query` line (general lane) pointing at `port`. Against the +/// hanging listener this blocks on the connect handshake, holding a general +/// worker indefinitely. +fn execute_query_line(id: u64, port: u16) -> String { + format!( + "{{\"jsonrpc\":\"2.0\",\"id\":{id},\"method\":\"execute_query\",\"params\":{{\"params\":{{\"driver\":\"postgres\",\"host\":\"127.0.0.1\",\"port\":{port},\"database\":\"x\",\"username\":\"u\",\"password\":\"p\",\"ssl_mode\":\"disable\"}},\"query\":\"SELECT 1\"}}}}\n" + ) +} + +/// A fire-and-forget `cancel` notification (no top-level `id`) carrying the +/// request id to cancel. This is the form the host sends on its call timeout +/// (#126); it routes to the cancel lane (Layer 1). +fn cancel_notification(request_id: u64) -> String { + format!("{{\"jsonrpc\":\"2.0\",\"method\":\"cancel\",\"params\":{{\"id\":{request_id}}}}}\n") +} + +/// Layer 2: a `ping`/`test_connection` to a *different* database must NOT +/// queue behind 4 stuck general queries. Before the fix all RPCs shared 4 +/// workers, so 4 hung `execute_query` calls blocked a `ping` to another DB +/// until the host's 120s timeout. With the lifecycle lane, the `ping` runs +/// on its own worker and returns within `PING_DEADLINE`. +/// +/// The 4 hung queries point at a hanging listener (block on the handshake +/// read); the `ping` points at the live DB (`pg-tabularis-test` :54320), +/// which responds fast. Skipped if the live DB isn't reachable. +#[test] +fn ping_to_another_db_returns_past_four_hung_general_queries() { + if !live_db_up() { + eprintln!("skipped: live DB at 127.0.0.1:54320 not reachable"); + return; + } + let (_listener, hung_port) = spawn_hanging_listener(); + + let mut child = Command::new(plugin_bin()) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .expect("spawn plugin binary"); + let mut stdin = child.stdin.take().expect("plugin stdin"); + let stdout = child.stdout.take().expect("plugin stdout"); + + // Saturate the 4 general workers with hung execute_query calls (ids + // 1-4), each blocking on the hanging listener's handshake read. + for id in 1..=4 { + stdin + .write_all(execute_query_line(id, hung_port).as_bytes()) + .expect("write hung execute_query"); + } + stdin.flush().expect("flush"); + + // Give the workers a moment to pick up the hung calls and block. + std::thread::sleep(Duration::from_millis(300)); + + // A ping to the live DB (a different database) — lifecycle lane. Before + // the fix this queued behind the 4 stuck general workers and timed out; + // with the fix it runs on the lifecycle lane and returns fast. + let ping = ping_line_live(100); + let start = Instant::now(); + stdin.write_all(ping.as_bytes()).expect("write ping"); + stdin.flush().expect("flush"); + + let response = + read_stdout_until_marker(stdout, "\"id\":100", PING_DEADLINE).unwrap_or_else(|| { + // Cleanup: kill the stuck process so the test doesn't leave it hanging. + let _ = child.kill(); + panic!( + "ping to another DB did not return within {PING_DEADLINE:?} — the \ + lifecycle lane (#167 Layer 2) regressed: the ping queued behind \ + the 4 stuck general workers" + ) + }); + let elapsed = start.elapsed(); + + // The ping's response carries id 100 and a result (no error against the + // live DB). Assert it came back promptly, not near the deadline. + assert!( + response.contains("\"result\""), + "expected the ping's response (id 100, result), got: {response}" + ); + assert!( + elapsed < PING_DEADLINE, + "ping took {elapsed:?} — expected a prompt lifecycle-lane return, not \ + near the {PING_DEADLINE:?} deadline (the lane regressed)" + ); + + // Cleanup: close stdin so the process drains and exits. + drop(stdin); + let _ = wait_with_deadline(&mut child, Duration::from_secs(20)); +} + +/// Layer 1: a `cancel` notification must NOT queue behind the stuck general +/// workers it's meant to relieve. Before the fix the cancel shared the same +/// 4-worker queue, so 4 stuck queries deadlocked the relief valve. With the +/// cancel lane, the cancel runs on its own worker regardless. (Layer 3 — the +/// worker-abort — is asserted in the next test.) +#[test] +fn cancel_notification_is_not_queued_behind_hung_general_workers() { + let (_listener, hung_port) = spawn_hanging_listener(); + + let mut child = Command::new(plugin_bin()) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .expect("spawn plugin binary"); + let mut stdin = child.stdin.take().expect("plugin stdin"); + + // Saturate the 4 general workers with hung execute_query calls. + for id in 1..=4 { + stdin + .write_all(execute_query_line(id, hung_port).as_bytes()) + .expect("write hung execute_query"); + } + stdin.flush().expect("flush"); + std::thread::sleep(Duration::from_millis(300)); + + // A fire-and-forget cancel for one of the hung queries. Before the fix this + // queued behind the 4 stuck workers and never fired; with the cancel lane + // it runs immediately. The cancel is a notification (no id) → no stdout + // response, so we assert the cancel lane didn't wedge the process: a + // wedged cancel lane would backpressure the reader, so a second line + // couldn't be written / the process couldn't drain. (The cancel itself is a + // no-op for these hung queries — they're stuck on the connect handshake + // with no registered `CancelGuard`, so server-side `pg_cancel_backend` + // can't fire; only the stdin-EOF drain aborts them. That's the documented + // "server ignores cancel" limit, not a regression. The point here is the + // cancel *ran* on its own lane without queueing behind the stuck workers.) + let cancel = cancel_notification(1); + stdin.write_all(cancel.as_bytes()).expect("write cancel"); + stdin.flush().expect("flush"); + + // The cancel lane must have processed it promptly. Give it a moment, then + // confirm the process is still accepting input (not wedged). A wedged + // cancel lane would backpressure the reader; a second cancel for a + // different id should still be writable (the channel isn't full). + std::thread::sleep(Duration::from_millis(300)); + let cancel2 = cancel_notification(2); + stdin + .write_all(cancel2.as_bytes()) + .expect("write second cancel — process should still be reading"); + stdin.flush().expect("flush second cancel"); + + // The process must still be draining (the reader didn't wedge). A 5th + // execute_query (id 5) writes fine — it queues (the 4 general workers are + // still stuck on the handshake, so it won't run), but the write succeeding + // proves the reader and cancel lane are healthy. The stuck workers are + // cleaned up by the EOF drain below, not by the cancel. + stdin + .write_all(execute_query_line(5, hung_port).as_bytes()) + .expect("write 5th execute_query — reader must keep draining"); + stdin.flush().expect("flush"); + + // Cleanup: the process has stuck workers; close stdin to trigger the + // bounded drain (#155) and exit. + drop(stdin); + let status = wait_with_deadline(&mut child, Duration::from_secs(20)) + .expect("plugin did not exit after stdin EOF — the shutdown drain (#155) regressed"); + assert!( + status.code().is_some(), + "process killed by harness, not a clean drain exit; status: {status:?}" + ); +} + +fn live_db_up() -> bool { + std::net::TcpStream::connect_timeout( + &std::net::SocketAddr::from(([127, 0, 0, 1], 54320)), + Duration::from_millis(500), + ) + .is_ok() +} + +/// Read stdout line-by-line until a line containing `marker` arrives, or +/// `deadline` elapses — on a blocking thread so a regressed hang surfaces as +/// a test failure (timeout) rather than a hung test. Unlike `shutdown.rs`'s +/// read-to-EOF helper, this returns as soon as the target response line +/// arrives, so a process still running (held open by stuck workers) doesn't +/// block the assertion on EOF that never comes. +fn read_stdout_until_marker( + mut stdout: std::process::ChildStdout, + marker: &str, + deadline: Duration, +) -> Option { + let marker = marker.to_string(); + let handle = std::thread::spawn(move || { + let mut reader = std::io::BufReader::new(&mut stdout); + loop { + let mut line = String::new(); + match std::io::BufRead::read_line(&mut reader, &mut line) { + Ok(0) => return None, // EOF + Ok(_) => { + if line.contains(&marker) { + return Some(line); + } + } + Err(_) => return None, + } + } + }); + let start = Instant::now(); + loop { + if handle.is_finished() { + return handle.join().ok().flatten(); + } + if start.elapsed() > deadline { + return None; + } + std::thread::sleep(Duration::from_millis(20)); + } +} + +/// Wait for a child with a deadline, polling `try_wait`. Reused from +/// `tests/shutdown.rs`. +fn wait_with_deadline( + child: &mut std::process::Child, + deadline: Duration, +) -> Option { + let start = Instant::now(); + loop { + if let Ok(Some(status)) = child.try_wait() { + return Some(status); + } + if start.elapsed() > deadline { + let _ = child.kill(); + return None; + } + std::thread::sleep(Duration::from_millis(50)); + } +} From c292342784542fd115346a6c60bdb6efaef9b9cd Mon Sep 17 00:00:00 2001 From: Adam J Esslinger Date: Fri, 23 Oct 2026 13:45:09 -0400 Subject: [PATCH 2/2] fix: remove stale Layer 3 worker-abort docs, make cancel-lane test non-vacuous (#167) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A pre-handoff ensemble review (11 dimensions + /code-review) found the module doc and test-file doc still described the reverted Layer 3 worker-abort (select!/CancellationToken) as live and tested, when run_worker runs handle_line unwrapped with no such mechanism. Five findings fixed: - src/main.rs module doc + spawn-site comments: rewrite to describe Layers 1-2 only, with an honest 'worker-abort was considered but not added' paragraph pointing to cancel::cancel's doc for the rationale. Remove 'Cancellable (Layer 3)' inline comments. - tests/worker_starvation.rs module doc: rewrite to describe Layers 1-2 only; remove the dangling 'asserted in the next test' comment (there was no next test). - Cancel test rewritten as cancel_runs_on_its_own_lane_*: send cancel as a request (id 200, not a notification) and assert a stdout response within 10s — a misrouted cancel queues behind 4 stuck workers and never responds. Mutation-tested: misrouting cancel->General fails the test with the right message; the old test passed on that regression. - Add 5 direct lane_for unit tests to main.rs (the routing decision the whole fix depends on had no direct test). Add PartialEq/Eq/Debug to Lane for assert_eq!. - Remove dead _cancellable: bool param on run_worker and the unused cancel_notification test helper (YAGNI — their doc comments were the stale ones). - Document the cancel-lane single-worker tradeoff at the spawn site (no connect_timeout on cancel_query; a stuck cancel blocks the lane). - Reword peek_method doc: the double parse is an accepted cost of keeping handle_line's signature unchanged, not 'avoids a second parse'. Verified: 441 unit tests pass (incl. 5 new lane_for), 2/2 worker_starvation integration tests pass, clippy -D warnings clean, fmt clean. --- src/main.rs | 120 ++++++++++++++++++++++++------------ src/rpc.rs | 6 +- tests/worker_starvation.rs | 123 ++++++++++++++++++++++--------------- 3 files changed, 160 insertions(+), 89 deletions(-) diff --git a/src/main.rs b/src/main.rs index 2626a33..f9033e3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,10 +16,9 @@ //! - **cancel lane** (1 worker): `cancel` notifications run on a dedicated //! worker that bypasses the general queue. When 4 general workers are //! stuck on slow queries, the cancel still fires `pg_cancel_backend` -//! (server-side) AND aborts the stuck worker's in-flight future -//! (worker-side, Layer 3) — instead of queuing behind those same 4 stuck -//! workers. Without this, stuck workers can't be relieved: the relief valve -//! is stuck in the same queue (a deadlock, the core of #167). +//! (server-side) — instead of queuing behind those same 4 stuck workers. +//! Without this, stuck workers can't be relieved: the relief valve is +//! stuck in the same queue (a deadlock, the core of #167). //! - **lifecycle lane** (2 workers): `ping`/`test_connection`/`initialize`/ //! `shutdown` run on a dedicated lane, so a health check or new connect to //! database B isn't blocked by 4 slow queries against database A. Two @@ -28,15 +27,16 @@ //! - **general lane** (4 workers): everything else — metadata, `execute_query`, //! CRUD, DDL, BLOB. The original `WORKER_POOL_SIZE`. //! -//! Every lane's worker wraps `handle_line` in a `select!` on a per-RPC -//! `CancellationToken` (Layer 3): a host `cancel` notification aborts the -//! in-flight operation directly, dropping the held `deadpool::Object` -//! (closing the socket → the server cancels the statement as a side effect) -//! and freeing the worker immediately — even if the server ignores -//! `pg_cancel_backend`. The bound is the **host's configurable call-timeout**, -//! not a plugin constant, so a query the host is willing to wait for (a large -//! enum-table result, slow DDL) is never cut (#166 stays intact); only a query -//! the host has given up on is aborted. +//! A worker stuck on a slow query frees when the server returns `57014` +//! (the common, cooperative case after `pg_cancel_backend`) or when the +//! host's per-call timeout elapses. A worker-abort mechanism (dropping the +//! in-flight `handle_line` future on cancel, to free the worker even when +//! the server ignores `pg_cancel_backend`) was considered but deliberately +//! not added: it would suppress the cancelled query's error response that +//! the host reads (the `#126` contract — a cancelled query returns its +//! `57014` response). See `cancel::cancel`'s doc for the full rationale. The +//! "server ignores `pg_cancel_backend`" gap is a documented pre-existing +//! limit matching the builtin. //! //! Responses are funneled through a single writer task so concurrent handlers //! never interleave bytes on stdout. A dedicated background task periodically @@ -149,27 +149,24 @@ async fn main() { let (resp_tx, resp_rx) = mpsc::unbounded_channel::(); let mut writer_handle = tokio::spawn(run_writer(resp_rx)); - // General lane: the original 4 workers. Cancellable (Layer 3). + // General lane: the original 4 workers. let mut worker_handles: Vec<_> = (0..WORKER_POOL_SIZE) - .map(|_| tokio::spawn(run_worker(req_rx.clone(), resp_tx.clone(), true))) + .map(|_| tokio::spawn(run_worker(req_rx.clone(), resp_tx.clone()))) .collect(); // Lifecycle lane: 2 workers (a stuck ping can't starve a test_connection). - // Cancellable — a host give-up should abort a stuck ping too. for _ in 0..LIFECYCLE_WORKER_COUNT { worker_handles.push(tokio::spawn(run_worker( lifecycle_rx.clone(), resp_tx.clone(), - true, ))); } - // Cancel lane: 1 worker, NOT cancellable — it's the relief valve; a cancel - // must always run to completion (aborting it would re-deadlock the lane it - // exists to relieve). `cancellable = false` → `handle_line` runs unwrapped. - worker_handles.push(tokio::spawn(run_worker( - cancel_rx.clone(), - resp_tx.clone(), - false, - ))); + // Cancel lane: 1 worker — the relief valve. A cancel notification + // (fire-and-forget, fast) runs here so it never queues behind the 4 general + // workers it's meant to relieve (#167). One worker is a deliberate + // tradeoff: cancel is low-volume, but a `cancel_query` that hangs on a + // non-responding handshake blocks the whole lane (see `cancel::run_cancel` + // — there is no `connect_timeout` on the fresh cancel connection). + worker_handles.push(tokio::spawn(run_worker(cancel_rx.clone(), resp_tx.clone()))); drop(resp_tx); run_reader(req_tx, lifecycle_tx, cancel_tx).await; @@ -241,7 +238,7 @@ fn lane_for(method: Option<&str>) -> Lane { } } -#[derive(Clone, Copy)] +#[derive(Clone, Copy, PartialEq, Eq, Debug)] enum Lane { General, Lifecycle, @@ -291,21 +288,18 @@ async fn run_reader( } } -/// A worker for one lane. `cancellable` is currently unused (all lanes run -/// `handle_line` unwrapped); it's retained as a parameter so the cancel lane -/// is explicitly marked as "never the relief valve's own cancellation -/// target" if a worker-abort mechanism is added later. See the module doc and -/// #167 for why the worker-abort (a `select!`-on-cancel that drops the -/// in-flight `handle_line` future) was *not* added: it would suppress the -/// cancelled query's error response that the host reads (the documented -/// `#126` contract — a cancelled query returns its `57014` response), so the -/// "server ignores `pg_cancel_backend`" gap is left as a documented -/// pre-existing limit (matching the builtin) rather than traded for a -/// contract break. +/// A worker for one lane. All lanes run `handle_line` unwrapped — a +/// worker-abort mechanism (a `select!`-on-cancel that drops the in-flight +/// `handle_line` future to free the worker even when the server ignores +/// `pg_cancel_backend`) was considered for #167 but deliberately not added: +/// it would suppress the cancelled query's error response that the host +/// reads (the `#126` contract — a cancelled query returns its `57014` +/// response), trading a rare failure for a contract break. See the module +/// doc and `cancel::cancel`'s doc. The "server ignores `pg_cancel_backend`" +/// gap is left as a documented pre-existing limit (matching the builtin). async fn run_worker( req_rx: Arc>>, resp_tx: mpsc::UnboundedSender, - _cancellable: bool, ) { loop { let line = { @@ -344,3 +338,53 @@ async fn run_writer(mut resp_rx: mpsc::UnboundedReceiver) { let _ = stdout.flush().await; } } + +#[cfg(test)] +mod tests { + use super::{lane_for, Lane}; + + /// `lane_for` is the routing decision the whole #167 fix depends on, so + /// pin each classification directly — the integration tests in + /// `tests/worker_starvation.rs` exercise it end-to-end, but a direct unit + /// test catches a misrouted method arm that an end-to-end test (especially + /// the DB-gated ping test) might skip over in CI. + + #[test] + fn cancel_routes_to_cancel_lane() { + assert_eq!(lane_for(Some("cancel")), Lane::Cancel); + } + + #[test] + fn lifecycle_methods_route_to_lifecycle_lane() { + assert_eq!(lane_for(Some("ping")), Lane::Lifecycle); + assert_eq!(lane_for(Some("test_connection")), Lane::Lifecycle); + assert_eq!(lane_for(Some("initialize")), Lane::Lifecycle); + assert_eq!(lane_for(Some("shutdown")), Lane::Lifecycle); + } + + #[test] + fn general_methods_route_to_general_lane() { + assert_eq!(lane_for(Some("execute_query")), Lane::General); + assert_eq!(lane_for(Some("get_tables")), Lane::General); + assert_eq!(lane_for(Some("insert_row")), Lane::General); + assert_eq!(lane_for(Some("create_table")), Lane::General); + assert_eq!(lane_for(Some("get_blob")), Lane::General); + } + + #[test] + fn unparseable_lines_route_to_general_lane() { + // `peek_method` returns None for unparseable/non-object/missing-method + // lines, so `run_reader` passes `None` to `lane_for` — the `_` arm + // routes them to the general lane where `handle_line` surfaces the + // proper `-32700`/`-32600` error. + assert_eq!(lane_for(None), Lane::General); + } + + #[test] + fn unknown_method_routes_to_general_lane() { + // A method `handle_line` doesn't dispatch (e.g. a typo) routes to + // general, where `handle_line` returns `-32601 Method not found`. + assert_eq!(lane_for(Some("execute")), Lane::General); + assert_eq!(lane_for(Some("get_foo")), Lane::General); + } +} diff --git a/src/rpc.rs b/src/rpc.rs index ca91c1f..0970c03 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -203,8 +203,10 @@ pub async fn handle_line(line: &str) -> Option { /// routes `None` to the general queue, where `handle_line` surfaces the /// proper `-32700`/`-32600` error response — the peek is a routing hint, not /// validation. Owns a short `String` (the method name); the allocation is -/// negligible next to the JSON parse it's already doing, and avoids a second -/// full parse in `handle_line` (which re-parses for dispatch + response). +/// negligible next to the JSON parse it's already doing. `handle_line` then +/// re-parses the line for dispatch + response — the double parse is an +/// accepted cost of keeping `handle_line`'s signature unchanged (it stays +/// the single source of truth for parse + dispatch + response). pub fn peek_method(line: &str) -> Option { let request: Value = serde_json::from_str(line).ok()?; let obj = request.as_object()?; diff --git a/tests/worker_starvation.rs b/tests/worker_starvation.rs index e9a903a..cf04cd8 100644 --- a/tests/worker_starvation.rs +++ b/tests/worker_starvation.rs @@ -4,18 +4,24 @@ //! stuck calls (a slow query holding a worker up to the host's 120s timeout) //! blocked every subsequent RPC — including a `ping` to a *different* //! database and the `cancel` notification that's meant to relieve the stuck -//! workers. The fix routes `cancel` to a dedicated lane (Layer 1), +//! workers. The fix routes `cancel` to a dedicated lane (Layer 1) and //! `ping`/`test_connection`/`initialize`/`shutdown` to a lifecycle lane -//! (Layer 2), and wires the `cancel` notification to abort the in-flight -//! worker operation directly (Layer 3) — so a host give-up frees the worker -//! even if the server ignores `pg_cancel_backend`. +//! (Layer 2) — so neither queues behind 4 stuck general workers. +//! +//! A worker-abort (Layer 3 — dropping the in-flight `handle_line` future on +//! cancel, to free the worker even if the server ignores `pg_cancel_backend`) +//! was considered but NOT added: it suppresses the cancelled query's `57014` +//! error response that the host reads (the `#126` contract). The +//! "server ignores `pg_cancel_backend`" gap is a documented pre-existing +//! limit matching the builtin — see `cancel::cancel`'s doc. No test here +//! exercises a worker-abort; a stuck worker frees only when the server +//! returns `57014` or the host's call timeout elapses. //! //! These tests spawn the real plugin binary (like `tests/shutdown.rs`) and //! reuse its hanging-listener harness. The hung queries point at a TCP //! listener that accepts but never answers the PostgreSQL handshake (a //! worker blocks on the handshake read); the healthy `ping` points at the -//! live database. No DB is needed for the cancel tests (the worker-abort -//! fires regardless of server cooperation). +//! live database. //! //! # Running //! @@ -66,7 +72,6 @@ fn ping_line_live(id: u64) -> String { ) } - /// An `execute_query` line (general lane) pointing at `port`. Against the /// hanging listener this blocks on the connect handshake, holding a general /// worker indefinitely. @@ -76,11 +81,19 @@ fn execute_query_line(id: u64, port: u16) -> String { ) } -/// A fire-and-forget `cancel` notification (no top-level `id`) carrying the -/// request id to cancel. This is the form the host sends on its call timeout -/// (#126); it routes to the cancel lane (Layer 1). -fn cancel_notification(request_id: u64) -> String { - format!("{{\"jsonrpc\":\"2.0\",\"method\":\"cancel\",\"params\":{{\"id\":{request_id}}}}}\n") +/// A `cancel` sent as a regular JSON-RPC *request* (with a top-level `id`). +/// `handle_line` dispatches it to the same `cancel` handler the host's +/// fire-and-forget notification hits, which returns `ok_response(id, Null)` — +/// a stdout response the test can read. This is what makes the cancel-lane +/// test observable: a misrouted cancel (sent to the general lane) would queue +/// behind 4 stuck workers and never respond, where a cancel-lane cancel +/// responds immediately. No live DB is needed (the cancel is a no-op for +/// the hung queries — they have no registered `CancelGuard` — but the +/// handler still returns its response). +fn cancel_request(id: u64, request_id: u64) -> String { + format!( + "{{\"jsonrpc\":\"2.0\",\"id\":{id},\"method\":\"cancel\",\"params\":{{\"id\":{request_id}}}}}\n" + ) } /// Layer 2: a `ping`/`test_connection` to a *different* database must NOT @@ -158,13 +171,22 @@ fn ping_to_another_db_returns_past_four_hung_general_queries() { let _ = wait_with_deadline(&mut child, Duration::from_secs(20)); } -/// Layer 1: a `cancel` notification must NOT queue behind the stuck general -/// workers it's meant to relieve. Before the fix the cancel shared the same -/// 4-worker queue, so 4 stuck queries deadlocked the relief valve. With the -/// cancel lane, the cancel runs on its own worker regardless. (Layer 3 — the -/// worker-abort — is asserted in the next test.) +/// Layer 1: a `cancel` must NOT queue behind the stuck general workers it's +/// meant to relieve. Before the fix the cancel shared the same 4-worker queue, +/// so 4 stuck queries deadlocked the relief valve. With the cancel lane, the +/// cancel runs on its own worker regardless. +/// +/// To make the routing observable, the cancel is sent as a regular JSON-RPC +/// *request* (with a top-level `id`), not a fire-and-forget notification: the +/// `cancel` handler returns `ok_response(id, Null)`, a stdout line the test +/// reads. If the cancel were misrouted to the general lane (the #167 +/// regression), it would queue behind the 4 stuck workers and never respond — +/// the read would time out and the test would panic. A cancel-lane cancel +/// responds immediately. No live DB is needed: the cancel targets a hung +/// query with no registered `CancelGuard` (a no-op), but the handler still +/// returns its response on the cancel lane. #[test] -fn cancel_notification_is_not_queued_behind_hung_general_workers() { +fn cancel_runs_on_its_own_lane_not_behind_hung_general_workers() { let (_listener, hung_port) = spawn_hanging_listener(); let mut child = Command::new(plugin_bin()) @@ -174,6 +196,7 @@ fn cancel_notification_is_not_queued_behind_hung_general_workers() { .spawn() .expect("spawn plugin binary"); let mut stdin = child.stdin.take().expect("plugin stdin"); + let stdout = child.stdout.take().expect("plugin stdout"); // Saturate the 4 general workers with hung execute_query calls. for id in 1..=4 { @@ -184,41 +207,43 @@ fn cancel_notification_is_not_queued_behind_hung_general_workers() { stdin.flush().expect("flush"); std::thread::sleep(Duration::from_millis(300)); - // A fire-and-forget cancel for one of the hung queries. Before the fix this - // queued behind the 4 stuck workers and never fired; with the cancel lane - // it runs immediately. The cancel is a notification (no id) → no stdout - // response, so we assert the cancel lane didn't wedge the process: a - // wedged cancel lane would backpressure the reader, so a second line - // couldn't be written / the process couldn't drain. (The cancel itself is a - // no-op for these hung queries — they're stuck on the connect handshake - // with no registered `CancelGuard`, so server-side `pg_cancel_backend` - // can't fire; only the stdin-EOF drain aborts them. That's the documented - // "server ignores cancel" limit, not a regression. The point here is the - // cancel *ran* on its own lane without queueing behind the stuck workers.) - let cancel = cancel_notification(1); + // A cancel sent as a request (id 200) for one of the hung queries. Before + // the fix this queued behind the 4 stuck workers and never got a worker; + // with the cancel lane it runs immediately and returns its response. The + // cancel itself is a no-op for these hung queries (stuck on the connect + // handshake with no registered `CancelGuard`, so server-side + // `pg_cancel_backend` can't fire — the documented "server ignores cancel" + // limit, not a regression), but the handler still returns + // `ok_response(200, Null)` on the cancel lane, which is the observable + // signal the cancel ran without queueing behind the stuck workers. + let cancel = cancel_request(200, 1); + let start = Instant::now(); stdin.write_all(cancel.as_bytes()).expect("write cancel"); stdin.flush().expect("flush"); - // The cancel lane must have processed it promptly. Give it a moment, then - // confirm the process is still accepting input (not wedged). A wedged - // cancel lane would backpressure the reader; a second cancel for a - // different id should still be writable (the channel isn't full). - std::thread::sleep(Duration::from_millis(300)); - let cancel2 = cancel_notification(2); - stdin - .write_all(cancel2.as_bytes()) - .expect("write second cancel — process should still be reading"); - stdin.flush().expect("flush second cancel"); + let response = + read_stdout_until_marker(stdout, "\"id\":200", PING_DEADLINE).unwrap_or_else(|| { + let _ = child.kill(); + panic!( + "cancel (id 200) did not respond within {PING_DEADLINE:?} — the \ + cancel lane (#167 Layer 1) regressed: the cancel queued behind \ + the 4 stuck general workers instead of running on its own lane" + ) + }); + let elapsed = start.elapsed(); - // The process must still be draining (the reader didn't wedge). A 5th - // execute_query (id 5) writes fine — it queues (the 4 general workers are - // still stuck on the handshake, so it won't run), but the write succeeding - // proves the reader and cancel lane are healthy. The stuck workers are - // cleaned up by the EOF drain below, not by the cancel. - stdin - .write_all(execute_query_line(5, hung_port).as_bytes()) - .expect("write 5th execute_query — reader must keep draining"); - stdin.flush().expect("flush"); + // The cancel's response carries id 200 and a null result (no error — an + // unknown/unregistered cancel id is a no-op, never an error). Assert it + // came back promptly, not near the deadline. + assert!( + response.contains("\"result\""), + "expected the cancel's response (id 200, result), got: {response}" + ); + assert!( + elapsed < PING_DEADLINE, + "cancel took {elapsed:?} — expected a prompt cancel-lane return, not \ + near the {PING_DEADLINE:?} deadline (the lane regressed)" + ); // Cleanup: the process has stuck workers; close stdin to trigger the // bounded drain (#155) and exit.