Add idle connection timeout to pgwire proxy

Author: jumpbetweenrows

.claude/CLAUDE.md

@@ -9,6 +9,7 @@ QueryProxy is a Rust PostgreSQL wire protocol proxy. See `README.md` for full de
 Key versions: pgwire 0.38, DataFusion 51, axum 0.8, SeaORM 1, tokio-postgres 0.7, argon2 0.5, aes-gcm 0.10, jsonwebtoken 9. Admin UI: React 19, Vite 6, Tailwind 4, TanStack Query 5, react-router-dom 7.
 
 ## Key Files
+- `proxy/src/server.rs` — `process_socket_with_idle_timeout()` (replaces pgwire's `process_socket`; adds idle + startup timeouts)
 - `proxy/src/admin/mod.rs` — `AdminState`, `ApiErr`, `admin_router()`
 - `proxy/src/admin/jwt.rs` — `AdminClaims` / `AuthClaims` extractors
 - `proxy/src/admin/datasource_types.rs` — `split_config`, `merge_config`, `get_type_defs`
@@ -51,6 +52,7 @@ Always get Arrow types from the library's `get_schema()` during discovery — th
 - Cache invalidation: `engine_cache.invalidate(name)` after catalog operations (keeps shared pool). `engine_cache.invalidate_all(name)` after datasource edit/delete (removes pool too). Never swap these — see README § Performance.
 - Discovery jobs: one active job per datasource enforced by `JobStore.active_by_ds`; cancellation via `CancellationToken` passed through all `DiscoveryProvider` methods
 - Catalog UUID v5 key format: `"{parent_uuid}:{child_name}"` — same natural key → same ID → re-discovery is a safe upsert
+- Idle timeout: `process_socket_with_idle_timeout` in `server.rs` replaces `pgwire::tokio::process_socket`. Env var `BR_IDLE_TIMEOUT_SECS` (default 900). Tests use `tokio::time::pause()` + `advance()` — do not add real `sleep()` calls in server tests.
 
 ## Pre-commit Hook
 A `.githooks/pre-commit` script runs `cargo fmt --check` and `cargo clippy -p proxy -- -D warnings`. Enable it once per clone:

Cargo.lock

@@ -155,6 +155,7 @@ betweenrows/
 │       └── types/                    TypeScript interfaces
 └── proxy/src/
     ├── main.rs                       entry point: CLI, DB init, EngineCache, servers
+    ├── server.rs                     process_socket_with_idle_timeout (idle + startup timeouts)
     ├── handler.rs                    pgwire StartupHandler + query handlers
     ├── auth.rs                       Argon2 auth, user creation
     ├── crypto.rs                     AES-256-GCM encrypt/decrypt
@@ -210,6 +211,7 @@ cargo run -p proxy -- user create --username alice --password secret --tenant ac
 | `BR_ADMIN_USER` | `admin` | Auto-seed username |
 | `BR_ADMIN_PASSWORD` | *(required on first boot)* | Auto-seed password. **Must be set** when no users exist in DB. |
 | `BR_ADMIN_TENANT` | `default` | Auto-seed tenant |
+| `BR_IDLE_TIMEOUT_SECS` | `900` (15 min) | Close pgwire connections that receive no messages for this many seconds. Enables Fly.io `auto_stop_machines` to work correctly with GUI clients (e.g. TablePlus) that hold idle connections indefinitely. Set to `0` to disable (not recommended on Fly.io). |
 | `BR_CORS_ALLOWED_ORIGINS` | *(empty, same-origin only)* | Comma-separated list of allowed CORS origins for the Admin API |
 | `RUST_LOG` | `info` | Log filter (standard Rust/tracing convention) |
 
@@ -225,6 +227,19 @@ Data sources are configured via the Admin UI (`/datasources`) or REST API — no
 
 ## Performance
 
+### Idle Connection Timeout
+
+pgwire 0.38 has no built-in idle timeout — `socket.next().await` blocks indefinitely after authentication. This prevents Fly.io `auto_stop_machines` from ever triggering when a GUI client like TablePlus is open, because the VM only stops when it has zero connections.
+
+`proxy/src/server.rs` replaces pgwire's `process_socket` with a custom message loop (`process_socket_with_idle_timeout`) that adds a `tokio::select!` branch racing each `socket.next()` against a `sleep(idle_timeout)`. The timer resets after every received message — a running query does not count as idle.
+
+Default timeout is 15 minutes (`BR_IDLE_TIMEOUT_SECS=900`). TCP keepalive (60 s time, 10 s interval) is also set on each accepted socket to detect dead connections from crashed clients or network failures.
+
+When idle timeout fires, a log line is emitted at `INFO` level:
+```
+Idle connection timed out after 900s
+```
+
 ### Arrow Type Alignment (query time)
 
 During catalog discovery, column types are captured using `datafusion-table-providers`' own `get_schema()` function rather than a manual PG-to-Arrow mapping. This guarantees that the stored schema matches exactly what the library produces at query time.
@@ -406,6 +421,6 @@ Catalog entity IDs (schemas, tables, columns) are deterministic UUID v5 fingerpr
 
 ```bash
 cargo build -p proxy          # compile
-cargo test -p proxy           # run tests (84 unit tests)
+cargo test -p proxy           # run tests (101 unit tests)
 cd admin-ui && npm run build  # production UI bundle
 ```

README.md

@@ -41,6 +41,9 @@ rand_core = { version = "0.6", features = ["std"] }
 aes-gcm = "0.10"
 base64 = "0.22"
 
+# TCP socket options (keepalive)
+socket2 = { version = "0.5", features = ["all"] }
+
 # Logging
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

proxy/Cargo.toml

@@ -11,4 +11,5 @@ pub mod engine;
 pub mod entity;
 pub mod handler;
 pub mod hooks;
+pub mod server;
 pub mod sql_rewrite;

proxy/src/lib.rs

@@ -1,13 +1,15 @@
 use clap::{Parser, Subcommand};
 use migration::{Migrator, MigratorTrait};
-use pgwire::tokio::process_socket;
 use proxy::admin::{AdminState, admin_router};
 use proxy::auth::Auth;
 use proxy::engine::EngineCache;
 use proxy::handler::ProxyHandler;
+use proxy::server::process_socket_with_idle_timeout;
 use rand_core::RngCore;
 use sea_orm::Database;
+use socket2::{SockRef, TcpKeepalive};
 use std::sync::Arc;
+use std::time::Duration;
 use tokio::net::TcpListener;
 
 #[derive(Parser)]
@@ -229,15 +231,36 @@ async fn serve(
 
     let bind_addr =
         std::env::var("BR_PROXY_BIND_ADDR").unwrap_or_else(|_| "127.0.0.1:5434".to_string());
+
+    let idle_timeout_secs: u64 = std::env::var("BR_IDLE_TIMEOUT_SECS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(900);
+    let idle_timeout = Duration::from_secs(idle_timeout_secs);
+    tracing::info!(
+        secs = idle_timeout_secs,
+        "Idle connection timeout configured"
+    );
+
+    let keepalive = TcpKeepalive::new()
+        .with_time(Duration::from_secs(60))
+        .with_interval(Duration::from_secs(10));
+
     let listener = TcpListener::bind(&bind_addr).await?;
     tracing::info!(addr = %bind_addr, "Proxy online");
 
     loop {
         let (incoming_socket, _) = listener.accept().await?;
-        let handler_clone = handler.clone();
 
+        if let Err(e) = SockRef::from(&incoming_socket).set_tcp_keepalive(&keepalive) {
+            tracing::warn!(error = %e, "Failed to set TCP keepalive");
+        }
+
+        let handler_clone = handler.clone();
         tokio::spawn(async move {
-            if let Err(e) = process_socket(incoming_socket, None, handler_clone).await {
+            if let Err(e) =
+                process_socket_with_idle_timeout(incoming_socket, handler_clone, idle_timeout).await
+            {
                 tracing::error!(error = %e, "Connection error");
             }
         });