PostgreSQL: cancelling Pool::begin_with can return an open transaction to the pool (0.9.0)

[slvnlrt]

Summary

With SQLx 0.9.0/PostgreSQL, cancelling Pool::begin_with after the server has executed BEGIN but before the driver consumes ReadyForQuery can return the same backend to the pool with an open transaction. A later plain pooled query runs in that transaction.

A deterministic reproduction was executed on PostgreSQL 17.11 and Rust 1.96.0. It observes the blocked backend through pg_stat_activity before cancellation, then reuses a single-connection pool. Actual output (backend PID is incidental):

blocked_pid=664027 reused_pid=664027 state=idle in transaction

Expected: rollback before reuse, or discard the uncertain connection. Actual: the next successful pooled SELECT leaves the backend idle in transaction.

Impact

This can cause a lock hang in otherwise autocommit application operations. More seriously, subsequent writes that callers believe are autocommitted can remain in the leaked transaction and later be rolled back on connection closure or idle_in_transaction_session_timeout. The reproduction below demonstrates the leaked transaction; the latter write-loss consequence follows from PostgreSQL transaction semantics rather than a separate write-loss experiment.

Reproduction

Create a small Rust project with:

[dependencies]
sqlx = { version = "=0.9.0", features = ["postgres", "runtime-tokio"] }
tokio = { version = "1", features = ["full"] }

Run against an isolated database with DATABASE_URL set. The observer holds an advisory lock; a multi-statement custom BEGIN waits at pg_advisory_xact_lock after BEGIN. This deliberately widens the same cancellation window before wait_until_ready() completes, without relying on a network proxy or guessing timing.

//! Run in an isolated PostgreSQL database: DATABASE_URL=... cargo run
//! Requires an isolated PostgreSQL database with pg_stat_activity visibility.

use std::time::{Duration, Instant};

use sqlx::{Connection, PgConnection, PgPool, postgres::PgPoolOptions};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let url = std::env::var("DATABASE_URL")?;
    let mut observer = PgConnection::connect(&url).await?;
    let pool: PgPool = PgPoolOptions::new().max_connections(1).connect(&url).await?;
    let observer_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()")
        .fetch_one(&mut observer)
        .await?;
    let key = i64::from(observer_pid);
    sqlx::query("SELECT pg_advisory_lock($1)")
        .bind(key)
        .execute(&mut observer)
        .await?;
    let marker = format!("sqlx_cancel_begin_{observer_pid}");
    let statement = format!("BEGIN; SELECT pg_advisory_xact_lock({key}) /* {marker} */");
    let pool_for_begin = pool.clone();
    let begin = tokio::spawn(async move {
        let _transaction = pool_for_begin
            .begin_with(sqlx::AssertSqlSafe(statement))
            .await?;
        Ok::<(), sqlx::Error>(())
    });

    let pattern = format!("%{marker}%");
    let deadline = Instant::now() + Duration::from_secs(5);
    let blocked_pid = loop {
        let pid: Option<i32> = sqlx::query_scalar(
            "SELECT pid FROM pg_stat_activity WHERE datname = current_database() \
             AND query LIKE $1 AND wait_event_type = 'Lock' AND wait_event = 'advisory'",
        )
        .bind(&pattern)
        .fetch_optional(&mut observer)
        .await?;
        if let Some(pid) = pid {
            break pid;
        }
        if Instant::now() >= deadline {
            return Err("BEGIN did not reach the blocked advisory lock".into());
        }
        tokio::time::sleep(Duration::from_millis(10)).await;
    };
    begin.abort();
    assert!(begin.await.unwrap_err().is_cancelled(), "begin task was not cancelled");
    let unlocked: bool = sqlx::query_scalar("SELECT pg_advisory_unlock($1)")
        .bind(key)
        .fetch_one(&mut observer)
        .await?;
    assert!(unlocked);

    let reused_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()")
        .fetch_one(&pool)
        .await?;
    let state: String = sqlx::query_scalar("SELECT state FROM pg_stat_activity WHERE pid = $1")
        .bind(reused_pid)
        .fetch_one(&mut observer)
        .await?;
    println!("blocked_pid={blocked_pid} reused_pid={reused_pid} state={state}");
    assert_eq!(blocked_pid, reused_pid, "the single-connection pool was not reused");
    assert_eq!(state, "idle in transaction", "this SQLx version did not reproduce the bug");
    pool.close().await;
    observer.close().await?;
    Ok(())
}

Source analysis

In sqlx-postgres 0.9.0/src/transaction.rs, PgTransactionManager::begin constructs a rollback guard, queues the BEGIN query, awaits wait_until_ready(), and only then increments transaction_depth. The guard's Drop calls start_rollback, which queues nothing while transaction_depth == 0. sqlx-core::Transaction::begin also constructs an open transaction guard before the await, but its drop reaches the same depth-zero no-op.

Pool return pings the connection, which can consume the outstanding ReadyForQuery and accept the connection as healthy without ending that server-side transaction. The application workaround we are using is cooperative cancellation outside owned database units, not resetting connections on checkout.

Could the driver account for a pending BEGIN before awaiting the response, so cancellation queues ROLLBACK (or marks the connection unsafe to reuse)?