From f6189a95f24af45068ecfd6e3c4e6b71ac8e43fe Mon Sep 17 00:00:00 2001 From: ilslv Date: Thu, 28 Oct 2021 12:10:30 +0300 Subject: [PATCH] Fix transaction not being rolled back on `Client::transaction()` `Future` dropped before completion --- tokio-postgres/src/client.rs | 41 +++++++++- tokio-postgres/tests/test/main.rs | 122 +++++++++++++++++++++++++++++- 2 files changed, 159 insertions(+), 4 deletions(-) diff --git a/tokio-postgres/src/client.rs b/tokio-postgres/src/client.rs index 4a099d94..dea77da9 100644 --- a/tokio-postgres/src/client.rs +++ b/tokio-postgres/src/client.rs @@ -1,4 +1,4 @@ -use crate::codec::BackendMessages; +use crate::codec::{BackendMessages, FrontendMessage}; use crate::config::{Host, SslMode}; use crate::connection::{Request, RequestMessages}; use crate::copy_out::CopyOutStream; @@ -19,7 +19,7 @@ use fallible_iterator::FallibleIterator; use futures::channel::mpsc; use futures::{future, pin_mut, ready, StreamExt, TryStreamExt}; use parking_lot::Mutex; -use postgres_protocol::message::backend::Message; +use postgres_protocol::message::{backend::Message, frontend}; use postgres_types::BorrowToSql; use std::collections::HashMap; use std::fmt; @@ -488,7 +488,42 @@ impl Client { /// /// The transaction will roll back by default - use the `commit` method to commit it. pub async fn transaction(&mut self) -> Result, Error> { - self.batch_execute("BEGIN").await?; + struct RollbackIfNotDone<'me> { + client: &'me Client, + done: bool, + } + + impl<'a> Drop for RollbackIfNotDone<'a> { + fn drop(&mut self) { + if self.done { + return; + } + + let buf = self.client.inner().with_buf(|buf| { + frontend::query("ROLLBACK", buf).unwrap(); + buf.split().freeze() + }); + let _ = self + .client + .inner() + .send(RequestMessages::Single(FrontendMessage::Raw(buf))); + } + } + + // This is done, as `Future` created by this method can be dropped after + // `RequestMessages` is synchronously send to the `Connection` by + // `batch_execute()`, but before `Responses` is asynchronously polled to + // completion. In that case `Transaction` won't be created and thus + // won't be rolled back. + { + let mut cleaner = RollbackIfNotDone { + client: self, + done: false, + }; + self.batch_execute("BEGIN").await?; + cleaner.done = true; + } + Ok(Transaction::new(self)) } diff --git a/tokio-postgres/tests/test/main.rs b/tokio-postgres/tests/test/main.rs index 31d7fa29..dcfbc530 100644 --- a/tokio-postgres/tests/test/main.rs +++ b/tokio-postgres/tests/test/main.rs @@ -3,9 +3,12 @@ use bytes::{Bytes, BytesMut}; use futures::channel::mpsc; use futures::{ - future, join, pin_mut, stream, try_join, FutureExt, SinkExt, StreamExt, TryStreamExt, + future, join, pin_mut, stream, try_join, Future, FutureExt, SinkExt, StreamExt, TryStreamExt, }; +use pin_project_lite::pin_project; use std::fmt::Write; +use std::pin::Pin; +use std::task::{Context, Poll}; use std::time::Duration; use tokio::net::TcpStream; use tokio::time; @@ -22,6 +25,35 @@ mod parse; mod runtime; mod types; +pin_project! { + /// Polls `F` at most `polls_left` times returning `Some(F::Output)` if + /// [`Future`] returned [`Poll::Ready`] or [`None`] otherwise. + struct Cancellable { + #[pin] + fut: F, + polls_left: usize, + } +} + +impl Future for Cancellable { + type Output = Option; + + fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { + let this = self.project(); + match this.fut.poll(ctx) { + Poll::Ready(r) => Poll::Ready(Some(r)), + Poll::Pending => { + *this.polls_left = this.polls_left.saturating_sub(1); + if *this.polls_left == 0 { + Poll::Ready(None) + } else { + Poll::Pending + } + } + } + } +} + async fn connect_raw(s: &str) -> Result<(Client, Connection), Error> { let socket = TcpStream::connect("127.0.0.1:5433").await.unwrap(); let config = s.parse::().unwrap(); @@ -35,6 +67,20 @@ async fn connect(s: &str) -> Client { client } +async fn current_transaction_id(client: &Client) -> i64 { + client + .query("SELECT txid_current()", &[]) + .await + .unwrap() + .pop() + .unwrap() + .get::<_, i64>("txid_current") +} + +async fn in_transaction(client: &Client) -> bool { + current_transaction_id(client).await == current_transaction_id(client).await +} + #[tokio::test] async fn plain_password_missing() { connect_raw("user=pass_user dbname=postgres") @@ -377,6 +423,80 @@ async fn transaction_rollback() { assert_eq!(rows.len(), 0); } +#[tokio::test] +async fn transaction_future_cancellation() { + let mut client = connect("user=postgres").await; + + for i in 0.. { + let done = { + let txn = client.transaction(); + let fut = Cancellable { + fut: txn, + polls_left: i, + }; + fut.await + .map(|res| res.expect("transaction failed")) + .is_some() + }; + + assert!(!in_transaction(&client).await); + + if done { + break; + } + } +} + +#[tokio::test] +async fn transaction_commit_future_cancellation() { + let mut client = connect("user=postgres").await; + + for i in 0.. { + let done = { + let txn = client.transaction().await.unwrap(); + let commit = txn.commit(); + let fut = Cancellable { + fut: commit, + polls_left: i, + }; + fut.await + .map(|res| res.expect("transaction failed")) + .is_some() + }; + + assert!(!in_transaction(&client).await); + + if done { + break; + } + } +} + +#[tokio::test] +async fn transaction_rollback_future_cancellation() { + let mut client = connect("user=postgres").await; + + for i in 0.. { + let done = { + let txn = client.transaction().await.unwrap(); + let rollback = txn.rollback(); + let fut = Cancellable { + fut: rollback, + polls_left: i, + }; + fut.await + .map(|res| res.expect("transaction failed")) + .is_some() + }; + + assert!(!in_transaction(&client).await); + + if done { + break; + } + } +} + #[tokio::test] async fn transaction_rollback_drop() { let mut client = connect("user=postgres").await;