fix: cancel an abandoned pg query and close failed connections cleanly (#11511)

* fix: cancel the query of a timed out or cancelled postgresql job Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019cpWeGy8jRMtWqvuRBpgu4 * test: fail the cancelled-query test unless the cancel happened Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019cpWeGy8jRMtWqvuRBpgu4 * fix: leave a failed pg job's connection clean for a pooler before closing it Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019cpWeGy8jRMtWqvuRBpgu4 * fix: cancel a failed pg job's unread results and roll back only inside a transaction Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019cpWeGy8jRMtWqvuRBpgu4 * fix: redo a pg connection's closing reset when a late cancel lands on it Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019cpWeGy8jRMtWqvuRBpgu4 * fix: let a late cancel land on the settle query without skipping the reset Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019cpWeGy8jRMtWqvuRBpgu4 --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-05 08:02:28 +00:00 · 2026-10-04 14:04:24 +02:00
parent 33f6de0cd2
commit 26856b8061
3 changed files with 258 additions and 21 deletions
@@ -1875,6 +1875,87 @@ async fn test_postgresql_open_transaction_not_reused(db: Pool<Postgres>) -> 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<Postgres>) -> 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)]
@@ -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(
@@ -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<tokio::task::JoinHandle<()>>,
    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.