diff --git a/backend/tests/worker.rs b/backend/tests/worker.rs index 6966b40e1f..f420dd2a12 100644 --- a/backend/tests/worker.rs +++ b/backend/tests/worker.rs @@ -1875,6 +1875,87 @@ async fn test_postgresql_open_transaction_not_reused(db: Pool) -> anyh Ok(()) } +/// Closing the connection of a cancelled job does not stop its query: the server +/// only notices once it next writes to the client. +#[sqlx::test(fixtures("base"))] +#[serial(pg_cache)] +async fn test_postgresql_cancelled_job_stops_its_query(db: Pool) -> anyhow::Result<()> { + use windmill_worker::pg_executor::clear_pg_cache; + + initialize_tracing().await; + clear_pg_cache().await; + + let server = ApiServer::start(db.clone()).await?; + let port = server.addr.port(); + let dbname: String = sqlx::query_scalar("SELECT current_database()") + .fetch_one(&db) + .await?; + let still_running = "SELECT count(*) FROM pg_stat_activity + WHERE datname = current_database() AND state = 'active' + AND pid <> pg_backend_pid() AND query LIKE '%pg_sleep(30)%'"; + + let cancelled = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let job = RunJob::from(JobPayload::Code(RawCode { + hash: None, + content: "SELECT pg_sleep(30);".to_string(), + path: None, + lock: None, + language: ScriptLang::Postgresql, + cache_ttl: None, + cache_ignore_s3_path: None, + dedicated_worker: None, + concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default() + .into(), + debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(), + modules: None, + tag: None, + })) + .arg( + "database", + json!({"host": "localhost", "port": 5432, "dbname": dbname, "user": "postgres", "password": "changeme"}), + ) + .run_until_complete_with(&db, false, port, |id| { + let (db, cancelled) = (db.clone(), cancelled.clone()); + async move { + tokio::spawn(async move { + while sqlx::query_scalar::<_, i64>(still_running) + .fetch_one(&db) + .await + .unwrap() + == 0 + { + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + sqlx::query("UPDATE v2_job_queue SET canceled_by = 'test-user' WHERE id = $1") + .bind(id) + .execute(&db) + .await + .unwrap(); + cancelled.store(true, std::sync::atomic::Ordering::Relaxed); + }); + } + }) + .await; + assert!(!job.success); + // Otherwise the job failed before its query ran, and there was nothing to stop. + assert!(cancelled.load(std::sync::atomic::Ordering::Relaxed)); + + let mut running = -1; + for _ in 0..50 { + running = sqlx::query_scalar::<_, i64>(still_running) + .fetch_one(&db) + .await?; + if running == 0 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + assert_eq!(running, 0); + + clear_pg_cache().await; + Ok(()) +} + /// A worker alternating between two databases keeps a connection for each. #[sqlx::test(fixtures("base"))] #[serial(pg_cache)] diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index d0e7c6d652..5754b8c8c1 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -1245,18 +1245,51 @@ impl PgDatabase { } } + fn sslmode_requires_tls(&self) -> bool { + matches!( + self.sslmode.as_deref(), + Some("require") | Some("verify-ca") | Some("verify-full") + ) + } + + /// Asks the server to cancel what the connection behind `token` is running. Dropping a + /// connection does not: the server works on until it next writes to the client. + /// + /// The request goes out the way the connection was made: `token_auth` is set for one + /// authenticated with an access token, which is over TLS whatever the sslmode. + pub async fn cancel_query( + &self, + token: tokio_postgres::CancelToken, + token_auth: bool, + ) -> Result<(), error::Error> { + if token_auth || self.sslmode_requires_tls() { + let mut connector = native_tls::TlsConnector::builder(); + Self::configure_pg_tls_verification( + &mut connector, + self.sslmode.as_deref(), + self.root_certificate_pem.as_deref(), + self.accept_invalid_certs, + )?; + let connector = + postgres_native_tls::MakeTlsConnector::new(connector.build().map_err(to_anyhow)?); + token.cancel_query(connector).await.map_err(to_anyhow)?; + } else { + token + .cancel_query(tokio_postgres::tls::NoTls) + .await + .map_err(to_anyhow)?; + } + Ok(()) + } + async fn connect_inner( &self, ) -> Result<(tokio_postgres::Client, TokioPgConnection), error::Error> { use native_tls::TlsConnector; use postgres_native_tls::MakeTlsConnector; use tokio_postgres::tls::NoTls; - let ssl_mode_is_require = matches!( - self.sslmode.as_deref(), - Some("require") | Some("verify-ca") | Some("verify-full") - ); - if ssl_mode_is_require { + if self.sslmode_requires_tls() { tracing::info!("Creating new connection"); let mut connector = TlsConnector::builder(); Self::configure_pg_tls_verification( diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index ac6a2ae79b..aac7d549db 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -79,6 +79,19 @@ const PG_OWN_CONNECTION_STALL: Duration = Duration::from_secs(5); /// statements. A server that vanished without closing the socket never answers /// it, and waiting for TCP to notice takes the keepalive budget (~2 min). const PG_RESET_PROBE_TIMEOUT: Duration = Duration::from_secs(5); +/// How long a closed connection gets to wind down. An idle one only has its Terminate to +/// send; one closed with results still arriving is cut off here. +const PG_CLOSE_GRACE: Duration = Duration::from_secs(2); +/// How long the connection of a failed job gets to end its query and reset its session. +/// Past this it is closed regardless. +const PG_CLOSE_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5); +/// How long a failed job's last query gets to be answered before it is taken to be still +/// sending results, and cancelled. +const PG_UNREAD_RESULT_GRACE: Duration = Duration::from_millis(250); +/// Clears the session state a job can leave behind. Each statement is explained where a +/// cached connection is probed with it. +const PG_SESSION_RESET: &str = "RESET ALL; RESET SESSION AUTHORIZATION; UNLISTEN *; CLOSE ALL; \ + SELECT pg_advisory_unlock_all();"; struct PgConnectionCache { /// Least recently used first. @@ -145,13 +158,32 @@ struct PgConnection { key: String, scope: String, client: Client, - connection_task: tokio::task::JoinHandle<()>, + connection_task: Option>, last_used: Instant, } impl Drop for PgConnection { + /// The connection task sends the protocol Terminate and ends by itself once the client is + /// gone, right after this returns. Aborting it straight away closes the socket without + /// one, and a transaction pooler may then hand the server connection to its next client + /// as the job left it, inside a failed transaction. fn drop(&mut self) { - self.connection_task.abort(); + let Some(mut task) = self.connection_task.take() else { + return; + }; + match tokio::runtime::Handle::try_current() { + Ok(runtime) => { + runtime.spawn(async move { + if tokio::time::timeout(PG_CLOSE_GRACE, &mut task) + .await + .is_err() + { + task.abort(); + } + }); + } + Err(_) => task.abort(), + } } } @@ -248,6 +280,83 @@ impl Drop for PgConnectionLease { } } +/// Closes the connection of a job that failed or left a transaction open, in the background +/// so the job ends without waiting. A pooler in transaction mode hands the server connection +/// to its next client the way this one leaves it, and one closed inside a transaction it may +/// not get back at all, so the transaction is ended and the session reset first. +/// +/// `abandoned` is set for a query the timeout or a cancellation cut short. Closing does not +/// stop it, the server works on until it next writes to the client, so it is cancelled. So +/// is one still sending results the job stopped reading, which would otherwise be drained. +fn close_unclean_connection( + lease: PgConnectionLease, + database: PgDatabase, + token_auth: bool, + abandoned: bool, +) { + if lease.conn.is_none() { + return; + } + tokio::spawn(async move { + let client = lease.client(); + let cleanup = async { + let cancel = || database.cancel_query(client.cancel_token(), token_auth); + if abandoned { + cancel().await?; + } + // Answered once the server is done with the job's last query. A cancel request + // is only sent by then, not acted on, and a pooler drops one whose client left. + let settled = client.batch_execute(""); + tokio::pin!(settled); + let answer = if abandoned { + settled.await + } else { + match tokio::time::timeout(PG_UNREAD_RESULT_GRACE, &mut settled).await { + Ok(answer) => answer, + Err(_) => { + cancel().await?; + settled.await + } + } + }; + // A cancel request that found the query already over lands on what comes next. + let cancelled = |e: &tokio_postgres::Error| { + e.code() == Some(&tokio_postgres::error::SqlState::QUERY_CANCELED) + }; + answer + .or_else(|e| if cancelled(&e) { Ok(()) } else { Err(e) }) + .map_err(to_anyhow)?; + let reset = || async { + let in_transaction = + client.transaction_status() != tokio_postgres::TransactionStatus::Idle; + let rollback = if in_transaction { "ROLLBACK; " } else { "" }; + client + .batch_execute(&format!("{rollback}{PG_SESSION_RESET}")) + .await + }; + let mut done = reset().await; + if done.as_ref().is_err_and(cancelled) { + done = reset().await; + } + done.map_err(to_anyhow)?; + Ok::<_, Error>(()) + }; + match tokio::time::timeout(PG_CLOSE_CLEANUP_TIMEOUT, cleanup).await { + Ok(Ok(())) => {} + // A connection that failed the job has nothing left to clean up. + Ok(Err(e)) if abandoned => { + tracing::warn!("could not cancel an abandoned postgres query: {e}") + } + Ok(Err(_)) => {} + Err(_) => tracing::warn!( + "a postgres connection did not settle within {}s of its job ending, closing it", + PG_CLOSE_CLEANUP_TIMEOUT.as_secs() + ), + } + drop(lease); + }); +} + async fn sweep_idle_pg_connections() { loop { tokio::time::sleep(Duration::from_secs(1)).await; @@ -1248,13 +1357,7 @@ pub async fn do_postgresql( // // Doubles as a liveness probe — if the connection is broken any // statement in the chain fails, or it never answers, and we replace it. - let reset = lease.client().batch_execute( - "RESET ALL; \ - RESET SESSION AUTHORIZATION; \ - UNLISTEN *; \ - CLOSE ALL; \ - SELECT pg_advisory_unlock_all();", - ); + let reset = lease.client().batch_execute(PG_SESSION_RESET); if matches!( tokio::time::timeout(PG_RESET_PROBE_TIMEOUT, reset).await, Ok(Ok(())) @@ -1287,7 +1390,7 @@ pub async fn do_postgresql( key: database_string.clone(), scope: lease.scope.clone(), client, - connection_task, + connection_task: Some(connection_task), last_used: Instant::now(), }); } @@ -1451,6 +1554,14 @@ pub async fn do_postgresql( result_f.await } }; + // Still unset once the job is over: the query was abandoned mid-flight, by the timeout or + // a cancellation, rather than ended by an error of its own. + let query_ended = std::sync::atomic::AtomicBool::new(false); + let result_f = async { + let result = result_f.await; + query_ended.store(true, Ordering::Relaxed); + result + }; let result = if run_inline { result_f.await @@ -1468,15 +1579,26 @@ pub async fn do_postgresql( Box::pin(futures::stream::once(async { 0 })), ) .await - } - .map_err(|e| map_s3object_jsonb_overflow(e, had_s3object_input))?; + }; + let result = match result { + Ok(result) => result, + Err(e) => { + close_unclean_connection( + lease, + database, + auth_mode != PgAuthMode::Password, + !query_ended.load(Ordering::Relaxed), + ); + return Err(map_s3object_jsonb_overflow(e, had_s3object_input)); + } + }; *mem_peak = size.load(Ordering::Relaxed) as i32; // A transaction the script left open would otherwise carry over into the - // next job that reuses the connection. Closing the connection rolls it - // back, as it always has without the cache. The status comes with the last - // reply, so a script that ends cleanly pays nothing for this check. + // next job that reuses the connection, so it is rolled back and the connection + // closed. The status comes with the last reply, so a script that ends cleanly + // pays nothing for this check. if lease.client().transaction_status() != tokio_postgres::TransactionStatus::Idle { if !run_inline { windmill_queue::append_logs( @@ -1488,14 +1610,15 @@ pub async fn do_postgresql( ) .await; } + close_unclean_connection(lease, database, auth_mode != PgAuthMode::Password, false); } else { // Keep the connection only after a query that read its results to the end. // An error may have stopped reading early (the result cap, a decode failure), // as a timeout, a dropped future or a first-row collection do: the server // keeps sending, and the next job's reset probe would wait on it. lease.keep = !*CLOUD_HOSTED && !collection_strategy.collect_first_row_only(); + drop(lease); } - drop(lease); *mem_peak = (result.get().len() / 1000) as i32; // And then check that we got back the same string we sent over.