fix: route ping/connect/cancel to dedicated lanes so stuck queries don't starve them (#167)

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),

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) — 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`.
+//!
+//! 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
 //! 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,40 @@ 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::<String>(REQUEST_QUEUE_CAPACITY);
     let req_rx = Arc::new(Mutex::new(req_rx));
+    let (lifecycle_tx, lifecycle_rx) = mpsc::channel::<String>(LIFECYCLE_QUEUE_CAPACITY);
+    let lifecycle_rx = Arc::new(Mutex::new(lifecycle_rx));
+    let (cancel_tx, cancel_rx) = mpsc::channel::<String>(CANCEL_QUEUE_CAPACITY);
+    let cancel_rx = Arc::new(Mutex::new(cancel_rx));
 
     let (resp_tx, resp_rx) = mpsc::unbounded_channel::<String>();
     let mut writer_handle = tokio::spawn(run_writer(resp_rx));
 
-    let worker_handles: Vec<_> = (0..WORKER_POOL_SIZE)
+    // 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())))
         .collect();
+    // Lifecycle lane: 2 workers (a stuck ping can't starve a test_connection).
+    for _ in 0..LIFECYCLE_WORKER_COUNT {
+        worker_handles.push(tokio::spawn(run_worker(
+            lifecycle_rx.clone(),
+            resp_tx.clone(),
+        )));
+    }
+    // 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).await;
+    run_reader(req_tx, lifecycle_tx, cancel_tx).await;
 
     let _ = shutdown_tx.send(true);
 
@@ -122,6 +190,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 +226,30 @@ async fn run_pool_cleanup(mut shutdown_rx: watch::Receiver<bool>) {
     }
 }
 
-async fn run_reader(req_tx: mpsc::Sender<String>) {
+/// 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, PartialEq, Eq, Debug)]
+enum Lane {
+    General,
+    Lifecycle,
+    Cancel,
+}
+
+async fn run_reader(
+    req_tx: mpsc::Sender<String>,
+    lifecycle_tx: mpsc::Sender<String>,
+    cancel_tx: mpsc::Sender<String>,
+) {
     let mut lines = BufReader::new(tokio::io::stdin()).lines();
 
     loop {
@@ -172,13 +267,36 @@ async fn run_reader(req_tx: mpsc::Sender<String>) {
             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. 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<Mutex<mpsc::Receiver<String>>>,
     resp_tx: mpsc::UnboundedSender<String>,
@@ -220,3 +338,53 @@ async fn run_writer(mut resp_rx: mpsc::UnboundedReceiver<String>) {
         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);
+    }
+}

src/rpc.rs

@@ -190,6 +190,29 @@ pub async fn handle_line(line: &str) -> Option<Value> {
     }
 }
 
+/// 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. `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<String> {
+    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",