diff --git a/sqlx-postgres/src/transaction.rs b/sqlx-postgres/src/transaction.rs index 3f4122ea82..6ae0da53fc 100644 --- a/sqlx-postgres/src/transaction.rs +++ b/sqlx-postgres/src/transaction.rs @@ -27,11 +27,18 @@ impl TransactionManager for PgTransactionManager { let rollback = Rollback::new(conn); rollback.conn.queue_simple_query(statement.as_str())?; + // Claim the depth before the round trip, not after. `start_rollback` -- which both + // this guard and `Transaction`'s own drop guard call -- is a no-op while the depth is + // zero, so a future cancelled during the await below would otherwise leave the server + // in a transaction with nothing queued to end it. `Pool` then hands that connection + // to the next borrower, whose statements silently run inside it. Unwound just below + // if the BEGIN did not take. + rollback.conn.inner.transaction_depth += 1; rollback.conn.wait_until_ready().await?; if !rollback.conn.in_transaction() { + rollback.conn.inner.transaction_depth -= 1; return Err(Error::BeginFailed); } - rollback.conn.inner.transaction_depth += 1; rollback.defuse(); Ok(()) diff --git a/tests/postgres/postgres.rs b/tests/postgres/postgres.rs index 26f827b837..0e5f2829c9 100644 --- a/tests/postgres/postgres.rs +++ b/tests/postgres/postgres.rs @@ -2244,3 +2244,54 @@ async fn it_can_recover_from_copy_in_invalid_params() -> anyhow::Result<()> { ) .await } + +// Regression: a future cancelled while `BEGIN`'s round trip is in flight used to leave the +// session inside a transaction. `start_rollback` is a no-op while `transaction_depth` is +// zero, and the depth was raised only after the await, so neither drop guard queued a +// `ROLLBACK` -- and `return_to_pool` validates with a bare `wait_until_ready` that never +// looks at the `ReadyForQuery` transaction-status byte, so the connection was handed to the +// next borrower with the transaction still open. +#[sqlx_macros::test] +async fn it_rolls_back_a_transaction_cancelled_during_begin() -> anyhow::Result<()> { + let pool = PgPoolOptions::new() + .max_connections(1) + .min_connections(0) + .connect(&dotenvy::var("DATABASE_URL")?) + .await?; + + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&pool) + .await?; + + // A plain `BEGIN` answers too quickly to cancel reliably; the sleep widens the same + // round trip so the cancellation lands inside it. + let cancelled = sqlx_core::rt::timeout( + Duration::from_millis(300), + pool.begin_with(AssertSqlSafe("BEGIN; SELECT pg_sleep(2);".to_string())), + ) + .await; + assert!(cancelled.is_err(), "the begin should not have completed"); + + // Outlast the sleep: the queued `ROLLBACK` is only flushed once the abandoned statement + // has answered and the connection is on its way back to the pool. + sqlx_core::rt::sleep(Duration::from_millis(3500)).await; + + let mut conn = new::().await?; + let state: Option = + sqlx::query_scalar("SELECT state FROM pg_stat_activity WHERE pid = $1") + .bind(pid) + .fetch_optional(&mut conn) + .await?; + + assert_eq!( + state.as_deref(), + Some("idle"), + "connection was returned to the pool still inside a transaction" + ); + + // and the pooled connection is still usable + let one: i32 = sqlx::query_scalar("SELECT 1").fetch_one(&pool).await?; + assert_eq!(one, 1); + + Ok(()) +}