feat(control): force-close idle pgwire connections via listener watchdog

An idle-in-transaction pgwire connection could hold its Data-Plane
staging overlay indefinitely, since process_socket owns the connection
loop and never times out on its own. The listener now races each
connection against a bounded re-check tick and drops the socket future
once it exceeds the configured idle or absolute session timeout, then
runs the existing connection-end teardown to reclaim the overlay.

An in-flight guard tracks per-connection statement execution so a
legitimately long-running query is never idle-killed, and the idle
window only starts once a statement has completed.

Author: farhan-syah

nodedb-test-support/src/pgwire_harness/start.rs

@@ -35,6 +35,14 @@ pub(super) struct StartConfig {
     /// Raft snapshot builder requires a routing table to resolve a group's
     /// vShards, so round-trip tests inject one here.
     pub routing: Option<nodedb_cluster::RoutingTable>,
+    /// Idle session timeout in seconds applied to the node's `SharedState`
+    /// before it is shared. `0` (the default) leaves the idle watchdog
+    /// disabled; a small value (e.g. `1`) lets tests exercise the pgwire
+    /// listener's idle-timeout force-close and overlay reclamation.
+    pub idle_timeout_secs: u64,
+    /// Absolute session lifetime in seconds applied to `SharedState` before it
+    /// is shared. `0` (the default) disables the absolute cap.
+    pub session_absolute_timeout_secs: u64,
 }
 
 impl Default for StartConfig {
@@ -44,6 +52,8 @@ impl Default for StartConfig {
             lockout: None,
             columnar_flush_threshold: None,
             routing: None,
+            idle_timeout_secs: 0,
+            session_absolute_timeout_secs: 0,
         }
     }
 }
@@ -84,6 +94,20 @@ impl TestServer {
         .await
     }
 
+    /// Spawn a single-core NodeDB server with a short idle session timeout so
+    /// tests can exercise the pgwire listener's idle-timeout watchdog: an
+    /// idle-in-transaction connection is force-closed after `idle_secs`,
+    /// triggering `on_connection_end` overlay reclamation. All other settings
+    /// stay at their defaults (trust-mode auth, lockout disabled, no absolute
+    /// cap).
+    pub async fn start_with_idle_timeout(idle_secs: u64) -> Self {
+        Self::start_with_config(StartConfig {
+            idle_timeout_secs: idle_secs,
+            ..Default::default()
+        })
+        .await
+    }
+
     /// Spawn a single-core NodeDB server with a cluster routing table
     /// installed on `SharedState::cluster_routing`.
     ///
@@ -146,6 +170,10 @@ impl TestServer {
             if let Some(routing) = cfg.routing {
                 s.cluster_routing = Some(std::sync::Arc::new(std::sync::RwLock::new(routing)));
             }
+            s.set_session_timeouts_for_test(
+                cfg.idle_timeout_secs,
+                cfg.session_absolute_timeout_secs,
+            );
         }
         let shared = shared;
 

nodedb/src/control/server/pgwire/factory.rs

@@ -232,6 +232,19 @@ impl NodeDbPgHandlerFactory {
         self.handler.reclaim_open_txn(addr).await;
         self.handler.sessions.remove(addr);
     }
+
+    /// Whether the connection at `addr` is eligible for idle timeout right now:
+    /// its session has zero statements in flight and has been silent for at
+    /// least `idle_ms`. Used by the pgwire listener watchdog, which owns the
+    /// per-connection task but cannot see inside pgwire's `process_socket`
+    /// loop. Returns `false` when the session is missing (nothing to time out).
+    pub fn session_idle_eligible(&self, addr: &std::net::SocketAddr, idle_ms: u64) -> bool {
+        self.handler.sessions.idle_eligible(
+            addr,
+            idle_ms,
+            crate::control::server::shared::session::now_unix_ms(),
+        )
+    }
 }
 
 impl PgWireServerHandlers for NodeDbPgHandlerFactory {

nodedb/src/control/server/pgwire/handler/core.rs

@@ -33,6 +33,7 @@ use crate::control::state::SharedState;
 use crate::types::RequestId;
 
 use super::super::types::notice_warning;
+use super::in_flight::InFlightGuard;
 use super::plan::extract_collection;
 use super::prepared::{NodeDbQueryParser, ParsedStatement};
 use crate::control::server::shared::session::{SessionStore, TransactionState};
@@ -314,6 +315,12 @@ impl SimpleQueryHandler for NodeDbPgHandler {
         let addr = client.socket_addr();
         self.sessions.ensure_session(addr);
 
+        // Mark this statement in flight for its whole duration so the idle
+        // watchdog never closes a connection mid-statement; the guard also
+        // stamps last-activity on drop, restarting the idle window at
+        // statement completion.
+        let _in_flight = InFlightGuard::new(&self.sessions, addr);
+
         let identity = self.resolve_identity(client, &addr)?;
         self.enforce_database_access(&identity, &addr)?;
 
@@ -443,10 +450,15 @@ impl ExtendedQueryHandler for NodeDbPgHandler {
         C::Error: Debug,
         PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
     {
+        let addr = client.socket_addr();
+        // Mark this statement in flight for the duration of execution so a
+        // long-running prepared statement is never idle-killed; the guard
+        // stamps last-activity on drop.
+        let _in_flight = InFlightGuard::new(&self.sessions, addr);
+
         let result = self.execute_prepared(client, portal, max_rows).await;
         // Mirror the simple-query path: surface any queued NOTICE messages
         // (e.g. `truncated_before_horizon`) before returning.
-        let addr = client.socket_addr();
         for message in self.sessions.drain_notices(&addr) {
             let notice = notice_warning(&message);
             let _ = client

nodedb/src/control/server/pgwire/handler/in_flight.rs

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: BUSL-1.1
+
+//! RAII in-flight guard for pgwire statement execution.
+//!
+//! Constructed at the top of a query handler and dropped when the statement
+//! finishes (including early return, error, or panic-unwind). On construction
+//! it bumps the connection's in-flight counter; on drop it decrements it and
+//! stamps last-activity to "now". The listener watchdog reads that state to
+//! decide idle eligibility, so this guard is what guarantees a long-running
+//! statement is never idle-killed and that the idle window only starts once a
+//! statement has actually completed.
+
+use std::net::SocketAddr;
+
+use crate::control::server::shared::session::SessionStore;
+
+/// Marks a statement as in flight for the lifetime of the guard.
+pub(crate) struct InFlightGuard<'a> {
+    sessions: &'a SessionStore,
+    addr: SocketAddr,
+}
+
+impl<'a> InFlightGuard<'a> {
+    /// Begin a request: increment the connection's in-flight counter.
+    pub(crate) fn new(sessions: &'a SessionStore, addr: SocketAddr) -> Self {
+        sessions.begin_request(&addr);
+        Self { sessions, addr }
+    }
+}
+
+impl Drop for InFlightGuard<'_> {
+    fn drop(&mut self) {
+        // Decrement in-flight and stamp last-activity on every exit path.
+        self.sessions.end_request(&self.addr);
+    }
+}

nodedb/src/control/server/pgwire/handler/mod.rs

@@ -7,6 +7,7 @@ mod cursor_cmds;
 mod cursor_query;
 mod dispatch;
 mod facet;
+mod in_flight;
 pub mod listen_notify;
 mod listen_notify_exec;
 mod live_select;

nodedb/src/control/server/pgwire/listener.rs

@@ -2,7 +2,7 @@
 
 use std::net::SocketAddr;
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use tokio::net::TcpListener;
 use tokio::sync::Semaphore;
@@ -60,6 +60,13 @@ impl PgListener {
         bus: crate::control::shutdown::ShutdownBus,
     ) -> crate::Result<()> {
         let conn_state = Arc::clone(&state);
+        // Session-timeout policy, read once at listener startup (config is fixed
+        // for the process lifetime). `idle` closes a connection that has been
+        // silent between statements for this many seconds — including an
+        // idle-in-transaction connection holding a staged-write overlay.
+        // `absolute` caps total connection lifetime. `0` disables either.
+        let idle_timeout_secs = conn_state.idle_timeout_secs();
+        let absolute_timeout_secs = conn_state.session_absolute_timeout_secs();
         let factory = Arc::new(NodeDbPgHandlerFactory::new(state, auth_mode));
 
         // Register with the shutdown bus so the sequencer waits for us to drain
@@ -122,15 +129,27 @@ impl PgListener {
                             info!(%peer_addr, "new pgwire connection");
                             let factory = Arc::clone(&factory);
                             let tls = tls_acceptor.clone();
+                            let idle = idle_timeout_secs;
+                            let absolute = absolute_timeout_secs;
                             connections.spawn(async move {
-                                if let Err(e) =
-                                    process_socket(stream, tls, Arc::clone(&factory)).await
-                                {
-                                    warn!(%peer_addr, error = %e, "pgwire session error");
+                                if idle == 0 && absolute == 0 {
+                                    // No session timeouts configured: run the
+                                    // plain socket loop with no watchdog wakeups.
+                                    if let Err(e) =
+                                        process_socket(stream, tls, Arc::clone(&factory)).await
+                                    {
+                                        warn!(%peer_addr, error = %e, "pgwire session error");
+                                    }
+                                } else {
+                                    run_with_watchdog(
+                                        stream, tls, &factory, peer_addr, idle, absolute,
+                                    )
+                                    .await;
                                 }
                                 // Reclaim any abandoned-transaction Data-Plane
                                 // overlays and drop the shared session entry now
-                                // that the connection has ended.
+                                // that the connection has ended (whether it ended
+                                // on its own or was force-closed by the watchdog).
                                 factory.on_connection_end(&peer_addr).await;
                                 drop(permit);
                                 peer_addr
@@ -191,3 +210,57 @@ impl PgListener {
         Ok(())
     }
 }
+
+/// Run `process_socket` under an idle + absolute session-timeout watchdog.
+///
+/// pgwire's `process_socket` owns the connection loop and is hard-typed to a
+/// `TcpStream`, so timeouts cannot be enforced inside it. Instead we race the
+/// socket future against a bounded re-check tick: on each wake, close the
+/// connection — by dropping the future, which drops the socket — if the
+/// absolute lifetime is exceeded or the connection is idle-eligible (zero
+/// in-flight statements AND silent past the idle window). A statement in
+/// flight keeps the connection non-idle, so a legitimately long-running query
+/// is never idle-killed. The caller runs `on_connection_end` afterwards to
+/// reclaim any staged overlay and drop the session entry.
+///
+/// Only invoked when at least one of `idle`/`absolute` is non-zero; the
+/// all-disabled path skips the watchdog entirely (no wakeups).
+async fn run_with_watchdog(
+    stream: tokio::net::TcpStream,
+    tls: Option<pgwire::tokio::TlsAcceptor>,
+    factory: &Arc<NodeDbPgHandlerFactory>,
+    peer_addr: SocketAddr,
+    idle: u64,
+    absolute: u64,
+) {
+    let started = Instant::now();
+    let mut fut = std::pin::pin!(process_socket(stream, tls, Arc::clone(factory)));
+    // Bounded re-check tick — never a busy loop. Activity may advance the idle
+    // deadline between wakes, which is recomputed via `session_idle_eligible`.
+    let tick = Duration::from_secs(1);
+    loop {
+        tokio::select! {
+            r = &mut fut => {
+                if let Err(e) = r {
+                    warn!(%peer_addr, error = %e, "pgwire session error");
+                }
+                break;
+            }
+            _ = tokio::time::sleep(tick) => {
+                if absolute > 0 && started.elapsed().as_secs() >= absolute {
+                    info!(%peer_addr, absolute, "pgwire absolute session timeout, closing");
+                    break;
+                }
+                if idle > 0
+                    && factory.session_idle_eligible(&peer_addr, idle.saturating_mul(1000))
+                {
+                    info!(%peer_addr, idle, "pgwire idle session timeout, closing");
+                    break;
+                }
+            }
+        }
+    }
+    // On a timeout break the loop exits with `fut` still pending; it is dropped
+    // here at scope end, which closes the socket. The caller then runs the
+    // connection-end teardown hook.
+}

nodedb/src/control/server/shared/session/mod.rs

@@ -40,6 +40,7 @@ pub use self::staging_gate::{
     DetachedTxnScope, DmlTxnCtx, InTxnRoute, StagedTagKind, StagedWriteOutcome, StagingGateError,
     route_in_tx_write,
 };
+pub(crate) use self::state::now_unix_ms;
 pub use self::state::{ConnSession, CursorState, TransactionState};
 pub use self::store::SessionStore;
 pub use self::temp_tables::TempTableEntry;

nodedb/src/control/server/shared/session/state.rs

@@ -3,6 +3,17 @@
 //! Per-connection session state types.
 
 use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::sync::atomic::{AtomicU32, AtomicU64};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+/// Milliseconds since the Unix epoch. Falls back to `0` instead of panicking
+/// if the system clock is set before the epoch (`duration_since` errors).
+pub(crate) fn now_unix_ms() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_millis() as u64)
+        .unwrap_or(0)
+}
 
 use crate::types::{DatabaseId, Lsn, TenantId, TxnId, VShardId};
 use nodedb_physical::physical_task::PhysicalTask;
@@ -150,6 +161,15 @@ pub struct ConnSession {
     /// GAP_FREE sequence reservations pending commit/rollback.
     /// On COMMIT: each reservation is finalized. On ROLLBACK: counter decremented.
     pub pending_sequence_reservations: Vec<crate::control::sequence::gap_free::ReservationHandle>,
+    /// Millis-since-epoch of the last statement COMPLETION on this connection
+    /// (also set to "now" at connection start). Read by the pgwire listener
+    /// watchdog to decide idle eligibility: a connection is idle only when it
+    /// has been silent (no statement completing) for the idle window.
+    pub last_activity_ms: AtomicU64,
+    /// Count of currently-executing statements on this connection. A connection
+    /// is idle-eligible only when this is zero — a legitimately long-running
+    /// statement (in flight) must never be idle-killed.
+    pub in_flight: AtomicU32,
 }
 
 impl ConnSession {
@@ -199,6 +219,8 @@ impl ConnSession {
             temp_tables: super::temp_tables::TempTableRegistry::new(),
             plan_cache: crate::control::server::shared::session::plan_cache::PlanCache::new(128),
             pending_sequence_reservations: Vec::new(),
+            last_activity_ms: AtomicU64::new(now_unix_ms()),
+            in_flight: AtomicU32::new(0),
         }
     }
 }

nodedb/src/control/server/shared/session/store.rs

@@ -5,12 +5,13 @@
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::RwLock;
+use std::sync::atomic::Ordering::Relaxed;
 
 use nodedb_types::DatabaseId;
 
 use crate::types::TenantId;
 
-use super::state::{ConnSession, TransactionState};
+use super::state::{ConnSession, TransactionState, now_unix_ms};
 
 /// Concurrent session store — keyed by socket address.
 pub struct SessionStore {
@@ -180,6 +181,51 @@ impl SessionStore {
             .identity = Some(identity);
     }
 
+    /// Record the start of a statement executing on this connection. Bumps the
+    /// in-flight counter so the idle watchdog never closes a connection with a
+    /// statement in progress. Creates the session entry if absent (mirrors
+    /// `set_identity`) so the extended-query path — which can begin execution
+    /// before `ensure_session` runs — still has its in-flight state tracked.
+    pub fn begin_request(&self, addr: &SocketAddr) {
+        let mut sessions = self.sessions.write().unwrap_or_else(|p| p.into_inner());
+        sessions
+            .entry(*addr)
+            .or_insert_with(ConnSession::new)
+            .in_flight
+            .fetch_add(1, Relaxed);
+    }
+
+    /// Record the completion of a statement on this connection: decrement the
+    /// in-flight counter (saturating — never underflows if the session was
+    /// already removed by a concurrent teardown) and stamp last-activity to
+    /// "now" so the idle window restarts from statement completion.
+    pub fn end_request(&self, addr: &SocketAddr) {
+        let mut sessions = self.sessions.write().unwrap_or_else(|p| p.into_inner());
+        if let Some(session) = sessions.get_mut(addr) {
+            // Exclusive write lock held: no concurrent mutator, so a
+            // load-check-store is a safe saturating decrement.
+            if session.in_flight.load(Relaxed) > 0 {
+                session.in_flight.fetch_sub(1, Relaxed);
+            }
+            session.last_activity_ms.store(now_unix_ms(), Relaxed);
+        }
+    }
+
+    /// Whether the connection at `addr` is eligible for idle timeout: the
+    /// session exists, has zero statements in flight, and its last activity is
+    /// at least `idle_ms` in the past relative to `now_ms`. Returns `false`
+    /// when the session is missing — nothing to time out.
+    pub fn idle_eligible(&self, addr: &SocketAddr, idle_ms: u64, now_ms: u64) -> bool {
+        let sessions = self.sessions.read().unwrap_or_else(|p| p.into_inner());
+        match sessions.get(addr) {
+            Some(session) => {
+                session.in_flight.load(Relaxed) == 0
+                    && now_ms.saturating_sub(session.last_activity_ms.load(Relaxed)) >= idle_ms
+            }
+            None => false,
+        }
+    }
+
     /// Reset per-session state for a `USE DATABASE` switch:
     ///   1. Aborts any open transaction (discards tx_buffer, resets state to Idle).
     ///   2. Clears all SQL-level prepared statements.