From 78e115bee5a573eb6e0abfbdccc4a0e677a80085 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 27 Jul 2026 12:32:44 +0200 Subject: [PATCH] fix: datatable full schema hangs behind a transaction-pooling postgres proxy (#10352) * fix: datatable full schema hangs behind a transaction-pooling postgres proxy * test: pause the clock in the pg connection shutdown test --- backend/windmill-api-settings/src/lib.rs | 7 +- .../windmill-api-workspaces/src/workspaces.rs | 17 ++--- .../src/workspaces_extra.rs | 4 +- backend/windmill-common/Cargo.toml | 4 ++ backend/windmill-common/src/lib.rs | 50 ++++++++++++-- backend/windmill-common/src/query_builders.rs | 67 +++++++++++++------ backend/windmill-common/src/workspaces.rs | 4 +- 7 files changed, 106 insertions(+), 47 deletions(-) diff --git a/backend/windmill-api-settings/src/lib.rs b/backend/windmill-api-settings/src/lib.rs index 8351ec0d2f..3d474ff4b1 100644 --- a/backend/windmill-api-settings/src/lib.rs +++ b/backend/windmill-api-settings/src/lib.rs @@ -1727,12 +1727,7 @@ async fn setup_custom_instance_pg_database_inner( logs.grant_permissions = "OK".to_string(); drop(client); // /!\ Drop before joining to avoid deadlock - join_handle - .await - .map_err(|e| error::Error::ExecutionErr(format!("join error: {}", e.to_string())))? - .map_err(|e| { - error::Error::ExecutionErr(format!("tokio_postgres error: {}", e.to_string())) - })?; + windmill_common::shutdown_pg_connection(join_handle).await?; Ok(()) } diff --git a/backend/windmill-api-workspaces/src/workspaces.rs b/backend/windmill-api-workspaces/src/workspaces.rs index 9089b6e8b6..c8f67c9d2e 100644 --- a/backend/windmill-api-workspaces/src/workspaces.rs +++ b/backend/windmill-api-workspaces/src/workspaces.rs @@ -2677,7 +2677,7 @@ async fn create_pg_database( if db_exists { drop(client); - let _ = join_handle.await; + let _ = windmill_common::shutdown_pg_connection(join_handle).await; return Err(Error::BadRequest(format!( "Database '{}' already exists on the resource server", req.target_dbname @@ -2695,10 +2695,7 @@ async fn create_pg_database( })?; drop(client); - join_handle - .await - .map_err(|e| Error::internal_err(format!("join error: {}", e)))? - .map_err(|e| Error::internal_err(format!("tokio_postgres error: {}", e)))?; + windmill_common::shutdown_pg_connection(join_handle).await?; } Ok(format!("Created database '{}'", req.target_dbname)) @@ -2801,10 +2798,7 @@ async fn get_datatable_full_schema( .map_err(Error::internal_err)?; drop(client); - join_handle - .await - .map_err(|e| Error::internal_err(format!("join error: {}", e)))? - .map_err(|e| Error::internal_err(format!("tokio_postgres error: {}", e)))?; + windmill_common::shutdown_pg_connection(join_handle).await?; Ok(Json(result)) } @@ -6389,10 +6383,7 @@ async fn snapshot_datatable_schema( .map_err(Error::internal_err)?; drop(client); - join_handle - .await - .map_err(|e| Error::internal_err(format!("join error: {}", e)))? - .map_err(|e| Error::internal_err(format!("tokio_postgres error: {}", e)))?; + windmill_common::shutdown_pg_connection(join_handle).await?; serde_json::to_value(schema) .map_err(|e| Error::internal_err(format!("Failed to serialize schema: {}", e))) diff --git a/backend/windmill-api-workspaces/src/workspaces_extra.rs b/backend/windmill-api-workspaces/src/workspaces_extra.rs index 19494c41fc..8755bc0a05 100644 --- a/backend/windmill-api-workspaces/src/workspaces_extra.rs +++ b/backend/windmill-api-workspaces/src/workspaces_extra.rs @@ -1352,7 +1352,7 @@ pub async fn drop_forked_datatable_databases( )); } drop(client); - let _ = join_handle.await; + let _ = windmill_common::shutdown_pg_connection(join_handle).await; } Err(e) => { errors.push(format!( @@ -1681,7 +1681,7 @@ async fn drop_fork_ducklake_metadata_schema( ) .await; drop(client); - let _ = join_handle.await; + let _ = windmill_common::shutdown_pg_connection(join_handle).await; res.map_err(|e| Error::internal_err(format!("{e:#}")))?; Ok(()) } diff --git a/backend/windmill-common/Cargo.toml b/backend/windmill-common/Cargo.toml index e45bba05c3..e2c3233183 100644 --- a/backend/windmill-common/Cargo.toml +++ b/backend/windmill-common/Cargo.toml @@ -130,6 +130,10 @@ opentelemetry-appender-tracing = { workspace = true, optional = true } tonic = { workspace = true, optional = true } equivalent = "1.0.2" +[dev-dependencies] +# `test-util` is not part of tokio's `full`; it is what lets tests pause the clock. +tokio = { workspace = true, features = ["test-util"] } + [target.'cfg(not(target_env = "msvc"))'.dependencies] tikv-jemalloc-ctl = { optional = true, workspace = true } diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 08a4ea93ed..1eda7503af 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -35,6 +35,7 @@ pub mod auth; pub mod bench; pub mod cache; pub mod client; +pub mod data_metrics; pub mod db; #[cfg(all(feature = "enterprise", feature = "private"))] mod db_entra_ee; @@ -62,7 +63,6 @@ pub mod instance_config; pub mod job_metrics; pub mod log_context; pub mod materialization; -pub mod data_metrics; pub mod min_version; pub mod notify_events; pub mod runtime_assets; @@ -1068,6 +1068,49 @@ impl PgDatabase { } } +/// How long a `tokio_postgres` connection task gets to wind down once its `Client` is dropped. +const PG_CONNECTION_SHUTDOWN_GRACE: std::time::Duration = std::time::Duration::from_secs(5); + +/// Wind down the task driving a `tokio_postgres` connection after its `Client` has been dropped, +/// surfacing whatever error the connection ended with. A teardown that has to be aborted is +/// reported as success — the work the client did is already done and complete. +/// +/// The task only finishes once the exchange the client left behind (its Terminate, and any +/// still-unanswered request) has been settled by the peer. A connection proxy that stops +/// replying leaves that pending forever, so waiting on the task without a deadline pins the +/// caller and the socket for the lifetime of the process. Aborting past the grace period drops +/// the stream, which is the only cleanup the task owes. +pub async fn shutdown_pg_connection( + join_handle: tokio::task::JoinHandle>, +) -> error::Result<()> { + let abort_handle = join_handle.abort_handle(); + match tokio::time::timeout(PG_CONNECTION_SHUTDOWN_GRACE, join_handle).await { + Ok(Ok(Ok(()))) => Ok(()), + Ok(Ok(Err(e))) => Err(error::Error::internal_err(format!( + "tokio_postgres error: {}", + e + ))), + Ok(Err(e)) => Err(error::Error::internal_err(format!("join error: {}", e))), + Err(_) => { + tracing::warn!( + "Postgres connection did not close within {}s of its client being dropped, aborting it", + PG_CONNECTION_SHUTDOWN_GRACE.as_secs() + ); + abort_handle.abort(); + Ok(()) + } + } +} + +#[cfg(test)] +mod pg_connection_shutdown_tests { + #[tokio::test(start_paused = true)] + async fn gives_up_on_a_connection_task_that_never_finishes() { + let never_finishes = tokio::spawn(std::future::pending()); + assert!(super::shutdown_pg_connection(never_finishes).await.is_ok()); + } +} + /// Validate a database name to prevent SQL injection. /// Must start with a letter, contain only alphanumeric characters, underscores, or hyphens, and be <= 63 chars. pub fn validate_dbname(dbname: &str) -> error::Result<()> { @@ -1219,10 +1262,7 @@ pub async fn create_custom_instance_database( } drop(client); - join_handle - .await - .map_err(|e| error::Error::internal_err(format!("join error: {}", e)))? - .map_err(|e| error::Error::internal_err(format!("tokio_postgres error: {}", e)))?; + shutdown_pg_connection(join_handle).await?; // Register in global_settings let status_json = serde_json::json!({ diff --git a/backend/windmill-common/src/query_builders.rs b/backend/windmill-common/src/query_builders.rs index 063d3958de..c0eedadbc4 100644 --- a/backend/windmill-common/src/query_builders.rs +++ b/backend/windmill-common/src/query_builders.rs @@ -4830,8 +4830,37 @@ fn pg_action_to_string(action: &str) -> String { } } +/// Rows of a simple-protocol result, dropping the framing messages. +fn simple_query_rows( + messages: Vec, +) -> Vec { + messages + .into_iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => Some(row), + _ => None, + }) + .collect() +} + +fn required_str<'a>( + row: &'a tokio_postgres::SimpleQueryRow, + column: &str, +) -> Result<&'a str, String> { + row.try_get(column) + .map_err(|e| format!("Failed to read column {}: {}", column, e))? + .ok_or_else(|| format!("Unexpected NULL in column {}", column)) +} + /// Introspect a PostgreSQL database and return the full schema. /// Takes a connected tokio_postgres Client. +/// +/// Both statements go through the simple query protocol. The extended protocol allocates a +/// named prepared statement per call and closes it when the statement handle drops; behind a +/// transaction-pooling proxy those names are shared with, and outlive, other sessions on the +/// same backend, and the exchange then stalls with no reply — the connection never becomes +/// idle again and the request hangs. Neither statement takes parameters, so nothing here +/// needs the extended protocol. pub async fn pg_get_full_schema( client: &tokio_postgres::Client, ) -> Result { @@ -4840,7 +4869,7 @@ pub async fn pg_get_full_schema( // per-column correlated subqueries — on large catalogs those subqueries run // once per column and make the introspection time out. let column_rows = client - .query( + .simple_query( "SELECT ns.nspname AS schema_name, c.relname AS table_name, @@ -4862,13 +4891,13 @@ pub async fn pg_get_full_schema( AND NOT a.attisdropped AND ns.nspname NOT IN ('pg_catalog', 'information_schema') ORDER BY ns.nspname, c.relname, a.attnum", - &[], ) .await + .map(simple_query_rows) .map_err(|e| format!("Failed to query columns: {}", e))?; let fk_rows = client - .query( + .simple_query( "SELECT ns.nspname AS schema_name, c.relname AS table_name, @@ -4890,21 +4919,21 @@ pub async fn pg_get_full_schema( WHERE con.contype = 'f' AND ns.nspname NOT IN ('pg_catalog', 'information_schema') ORDER BY ns.nspname, c.relname, con.conname, u.ord", - &[], ) .await + .map(simple_query_rows) .map_err(|e| format!("Failed to query foreign keys: {}", e))?; let mut result: FullDatabaseSchema = std::collections::HashMap::new(); for row in &column_rows { - let schema_name: &str = row.get("schema_name"); - let table_name: &str = row.get("table_name"); - let column_name: &str = row.get("column_name"); - let datatype: &str = row.get("datatype"); + let schema_name = required_str(row, "schema_name")?; + let table_name = required_str(row, "table_name")?; + let column_name = required_str(row, "column_name")?; + let datatype = required_str(row, "datatype")?; let default_value: Option<&str> = row.get("default_value"); - let nullable: bool = row.get("nullable"); - let is_primary_key: bool = row.get("is_primary_key"); + let nullable = required_str(row, "nullable")? == "t"; + let is_primary_key = required_str(row, "is_primary_key")? == "t"; let pk_constraint_name: Option<&str> = row.get("pk_constraint_name"); let schema_tables = result.entry(schema_name.to_string()).or_default(); @@ -4937,15 +4966,15 @@ pub async fn pg_get_full_schema( > = std::collections::HashMap::new(); for row in &fk_rows { - let schema_name: &str = row.get("schema_name"); - let table_name: &str = row.get("table_name"); - let fk_name: &str = row.get("fk_constraint_name"); - let source_column: &str = row.get("source_column"); - let ref_schema: &str = row.get("ref_schema"); - let ref_table: &str = row.get("ref_table"); - let ref_column: &str = row.get("ref_column"); - let on_delete: &str = row.get("on_delete"); - let on_update: &str = row.get("on_update"); + let schema_name = required_str(row, "schema_name")?; + let table_name = required_str(row, "table_name")?; + let fk_name = required_str(row, "fk_constraint_name")?; + let source_column = required_str(row, "source_column")?; + let ref_schema = required_str(row, "ref_schema")?; + let ref_table = required_str(row, "ref_table")?; + let ref_column = required_str(row, "ref_column")?; + let on_delete = required_str(row, "on_delete")?; + let on_update = required_str(row, "on_update")?; let target_table = if ref_schema == schema_name { ref_table.to_string() diff --git a/backend/windmill-common/src/workspaces.rs b/backend/windmill-common/src/workspaces.rs index 135b8b9558..fb73ba7e2c 100644 --- a/backend/windmill-common/src/workspaces.rs +++ b/backend/windmill-common/src/workspaces.rs @@ -1866,7 +1866,7 @@ async fn inspect_fork_catalog( ); let table_rows = client.query(&qt, &[]).await; drop(client); - let _ = join_handle.await; + let _ = crate::shutdown_pg_connection(join_handle).await; existing_schemas.extend( same_catalog_res @@ -1920,7 +1920,7 @@ async fn inspect_fork_catalog( let join_handle = tokio::spawn(async move { connection.await }); let res = query_schemas(&client, schemas).await; drop(client); - let _ = join_handle.await; + let _ = crate::shutdown_pg_connection(join_handle).await; res.map_err(|e| Error::internal_err(format!("{e}"))) } .await;