fix(operator): drop wall-clock query timeout, keep TCP keepalives (#56)

Author: passcod

AGENTS.md

@@ -22,3 +22,10 @@
 - Always work from a branch. If you're on `main`, create a new branch.
 - If you can at all do something from the operator, instead of in a script that runs in a job, do so. For example, if you need to query something in a database, you should be able to do so from the operator, instead of writing SQL code in bash logic.
 <!-- end rules -->
+
+## Lessons learned
+
+- Detect dead Postgres connections with **TCP keepalives** (`socket2::TcpKeepalive` on the raw `TcpStream` before handing it to `tokio_postgres::Config::connect_raw`), not with a wall-clock `tokio::time::timeout` around the query. Keepalives fire on the kernel's own timer regardless of in-flight queries, so a legitimately long-running statement (e.g. `DROP SCHEMA … CASCADE` on a populated dbt schema, which can run for tens of minutes) keeps working, while a NAT-evicted / silently-dead socket still errors within ~90s and the reconcile retries.
+- Wall-clock timeouts on a long-tail-distributed operation make a retry loop **never converge**: every attempt hits the cap, error policy retries, every retry hits the cap again. Don't add a timeout unless you have a real p99 number for the operation it wraps.
+- The operator's own libpq sockets are separate from the migration Job's libpq URI. Keepalive fixes in the Job (`keepalives=1&keepalives_idle=...`) do **not** cover the operator — set both.
+- The kube port-forward path is not a raw TCP socket; socket-level keepalives don't apply. That path keeps its own liveness via the API server WebSocket — only `connect_via_tcp` needs the socket2 setup.

src/controllers/postgres.rs

@@ -7,24 +7,20 @@ use kube::{
 };
 use socket2::{SockRef, TcpKeepalive};
 use tokio::net::TcpStream;
-use tokio_postgres::{NoTls, Row, ToStatement, types::ToSql};
+use tokio_postgres::NoTls;
 use tracing::{debug, info, warn};
 
 use crate::error::{Error, Result};
 
 pub const DEFAULT_PG_VERSION: i32 = 18;
 
-/// Per-query timeout. Backstop for hung connections that survive TCP
-/// keepalives (e.g. a server that ACKs keepalive probes but the query
-/// itself wedges). Generous because `DROP SCHEMA ... CASCADE` on a
-/// populated dbt schema can take minutes; intent is runaway detection,
-/// not a query SLA.
-const QUERY_TIMEOUT: Duration = Duration::from_secs(5 * 60);
-
 /// TCP keepalive parameters for operator-side libpq sockets. Matches the
 /// settings the migration Job uses in its libpq URI (see [[fix(migration)
 /// TCP keepalives]]). Detects NAT-evicted / silently-dead sockets within
 /// ~90s so reconciles can fail and retry instead of hanging forever.
+/// Keepalives fire on the kernel's own timer regardless of in-flight
+/// queries, so a legitimately long-running `DROP SCHEMA ... CASCADE` on
+/// a populated dbt schema is fine — only dead sockets are killed.
 const KEEPALIVE_IDLE: Duration = Duration::from_secs(60);
 const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(10);
 const KEEPALIVE_RETRIES: u32 = 3;
@@ -40,65 +36,6 @@ fn set_tcp_keepalive(stream: &TcpStream) -> Result<()> {
 	Ok(())
 }
 
-/// Run a parameterised query with the global [`QUERY_TIMEOUT`]. Returns a
-/// timeout-shaped error if the query doesn't complete in time, so the
-/// reconcile fails fast instead of hanging on a wedged connection.
-pub async fn query_timeout<T>(
-	pg: &tokio_postgres::Client,
-	stmt: &T,
-	params: &[&(dyn ToSql + Sync)],
-) -> Result<Vec<Row>>
-where
-	T: ?Sized + ToStatement,
-{
-	tokio::time::timeout(QUERY_TIMEOUT, pg.query(stmt, params))
-		.await
-		.map_err(|_| Error::MissingField(format!("query timed out after {QUERY_TIMEOUT:?}")))?
-		.map_err(Into::into)
-}
-
-pub async fn query_one_timeout<T>(
-	pg: &tokio_postgres::Client,
-	stmt: &T,
-	params: &[&(dyn ToSql + Sync)],
-) -> Result<Row>
-where
-	T: ?Sized + ToStatement,
-{
-	tokio::time::timeout(QUERY_TIMEOUT, pg.query_one(stmt, params))
-		.await
-		.map_err(|_| Error::MissingField(format!("query timed out after {QUERY_TIMEOUT:?}")))?
-		.map_err(Into::into)
-}
-
-pub async fn query_opt_timeout<T>(
-	pg: &tokio_postgres::Client,
-	stmt: &T,
-	params: &[&(dyn ToSql + Sync)],
-) -> Result<Option<Row>>
-where
-	T: ?Sized + ToStatement,
-{
-	tokio::time::timeout(QUERY_TIMEOUT, pg.query_opt(stmt, params))
-		.await
-		.map_err(|_| Error::MissingField(format!("query timed out after {QUERY_TIMEOUT:?}")))?
-		.map_err(Into::into)
-}
-
-pub async fn execute_timeout<T>(
-	pg: &tokio_postgres::Client,
-	stmt: &T,
-	params: &[&(dyn ToSql + Sync)],
-) -> Result<u64>
-where
-	T: ?Sized + ToStatement,
-{
-	tokio::time::timeout(QUERY_TIMEOUT, pg.execute(stmt, params))
-		.await
-		.map_err(|_| Error::MissingField(format!("execute timed out after {QUERY_TIMEOUT:?}")))?
-		.map_err(Into::into)
-}
-
 /// Holds a Postgres client and any resources that must stay alive for the
 /// duration of the connection (e.g. a port-forwarder).
 pub struct PgConnection {
@@ -282,15 +219,15 @@ pub async fn discover_restore_database(
 	.await?;
 	let pg = &conn.client;
 
-	let row = query_opt_timeout(
-		pg,
-		"SELECT datname FROM pg_database \
-		 WHERE datname NOT IN ('postgres', 'template0', 'template1') \
-		 ORDER BY pg_database_size(datname) DESC \
-		 LIMIT 1",
-		&[],
-	)
-	.await?;
+	let row = pg
+		.query_opt(
+			"SELECT datname FROM pg_database \
+			 WHERE datname NOT IN ('postgres', 'template0', 'template1') \
+			 ORDER BY pg_database_size(datname) DESC \
+			 LIMIT 1",
+			&[],
+		)
+		.await?;
 
 	match row {
 		Some(r) => {
@@ -314,7 +251,9 @@ pub async fn discover_restore_database(
 /// you're going to make multiple queries on the same database so the
 /// connection can be reused.
 pub async fn database_size_on(pg: &tokio_postgres::Client) -> Result<u64> {
-	let row = query_one_timeout(pg, "SELECT pg_database_size(current_database())", &[]).await?;
+	let row = pg
+		.query_one("SELECT pg_database_size(current_database())", &[])
+		.await?;
 	let size: i64 = row.get(0);
 	Ok(size as u64)
 }
@@ -356,12 +295,12 @@ pub async fn existing_schemas_on(
 	if candidates.is_empty() {
 		return Ok(Vec::new());
 	}
-	let rows = query_timeout(
-		pg,
-		"SELECT nspname FROM pg_catalog.pg_namespace WHERE nspname = ANY($1)",
-		&[&candidates],
-	)
-	.await?;
+	let rows = pg
+		.query(
+			"SELECT nspname FROM pg_catalog.pg_namespace WHERE nspname = ANY($1)",
+			&[&candidates],
+		)
+		.await?;
 	let found: HashSet<String> = rows.iter().map(|r| r.get::<_, String>(0)).collect();
 	Ok(candidates
 		.iter()
@@ -376,7 +315,7 @@ pub async fn drop_schemas_on(pg: &tokio_postgres::Client, schemas: &[String]) ->
 	for schema in schemas {
 		let stmt = format!("DROP SCHEMA IF EXISTS {} CASCADE", quote_ident(schema));
 		debug!(schema = schema, "dropping schema");
-		execute_timeout(pg, stmt.as_str(), &[]).await?;
+		pg.execute(stmt.as_str(), &[]).await?;
 	}
 	Ok(())
 }