Fix transaction not being rolled back on TransactionBuilder::start() Future dropped before completion

tokio-postgres/src/client.rs

@@ -1,4 +1,4 @@
-use crate::codec::{BackendMessages, FrontendMessage};
+use crate::codec::BackendMessages;
 use crate::config::SslMode;
 use crate::connection::{Request, RequestMessages};
 use crate::copy_out::CopyOutStream;
@@ -21,7 +21,7 @@ use fallible_iterator::FallibleIterator;
 use futures_channel::mpsc;
 use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
-use postgres_protocol::message::{backend::Message, frontend};
+use postgres_protocol::message::backend::Message;
 use postgres_types::BorrowToSql;
 use std::collections::HashMap;
 use std::fmt;
@@ -469,43 +469,7 @@ impl Client {
     ///
     /// The transaction will roll back by default - use the `commit` method to commit it.
     pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
-        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))
+        Transaction::begin(self).await
     }
 
     /// Returns a builder for a transaction with custom settings.

tokio-postgres/src/transaction.rs

@@ -9,7 +9,7 @@ use crate::types::{BorrowToSql, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
 use crate::{
-    bind, query, slice_iter, CancelToken, Client, CopyInSink, Error, Portal, Row,
+    bind, query, slice_iter, CancelToken, Client, CopyInSink, Error, IsolationLevel, Portal, Row,
     SimpleQueryMessage, Statement, ToStatement,
 };
 use bytes::Buf;
@@ -56,14 +56,78 @@ impl<'a> Drop for Transaction<'a> {
 }
 
 impl<'a> Transaction<'a> {
-    pub(crate) fn new(client: &'a mut Client) -> Transaction<'a> {
+    fn new(client: &'a mut Client) -> Transaction<'a> {
         Transaction {
             client,
             savepoint: None,
             done: false,
         }
     }
 
+    pub(crate) async fn begin(client: &'a mut Client) -> Result<Transaction<'a>, Error> {
+        let transaction = Transaction::new(client);
+
+        transaction.client.batch_execute("BEGIN").await?;
+
+        Ok(transaction)
+    }
+
+    pub(crate) async fn start(
+        client: &'a mut Client,
+        isolation_level: Option<IsolationLevel>,
+        read_only: Option<bool>,
+        deferrable: Option<bool>,
+    ) -> Result<Transaction<'a>, Error> {
+        let transaction = Transaction::new(client);
+
+        let mut query = "START TRANSACTION".to_string();
+        let mut first = true;
+
+        if let Some(level) = isolation_level {
+            first = false;
+
+            query.push_str(" ISOLATION LEVEL ");
+            let level = match level {
+                IsolationLevel::ReadUncommitted => "READ UNCOMMITTED",
+                IsolationLevel::ReadCommitted => "READ COMMITTED",
+                IsolationLevel::RepeatableRead => "REPEATABLE READ",
+                IsolationLevel::Serializable => "SERIALIZABLE",
+            };
+            query.push_str(level);
+        }
+
+        if let Some(read_only) = read_only {
+            if !first {
+                query.push(',');
+            }
+            first = false;
+
+            let s = if read_only {
+                " READ ONLY"
+            } else {
+                " READ WRITE"
+            };
+            query.push_str(s);
+        }
+
+        if let Some(deferrable) = deferrable {
+            if !first {
+                query.push(',');
+            }
+
+            let s = if deferrable {
+                " DEFERRABLE"
+            } else {
+                " NOT DEFERRABLE"
+            };
+            query.push_str(s);
+        }
+
+        transaction.client.batch_execute(&query).await?;
+
+        Ok(transaction)
+    }
+
     /// Consumes the transaction, committing all changes made within it.
     pub async fn commit(mut self) -> Result<(), Error> {
         self.done = true;

tokio-postgres/src/transaction_builder.rs

@@ -63,51 +63,12 @@ impl<'a> TransactionBuilder<'a> {
     ///
     /// The transaction will roll back by default - use the `commit` method to commit it.
     pub async fn start(self) -> Result<Transaction<'a>, Error> {
-        let mut query = "START TRANSACTION".to_string();
-        let mut first = true;
-
-        if let Some(level) = self.isolation_level {
-            first = false;
-
-            query.push_str(" ISOLATION LEVEL ");
-            let level = match level {
-                IsolationLevel::ReadUncommitted => "READ UNCOMMITTED",
-                IsolationLevel::ReadCommitted => "READ COMMITTED",
-                IsolationLevel::RepeatableRead => "REPEATABLE READ",
-                IsolationLevel::Serializable => "SERIALIZABLE",
-            };
-            query.push_str(level);
-        }
-
-        if let Some(read_only) = self.read_only {
-            if !first {
-                query.push(',');
-            }
-            first = false;
-
-            let s = if read_only {
-                " READ ONLY"
-            } else {
-                " READ WRITE"
-            };
-            query.push_str(s);
-        }
-
-        if let Some(deferrable) = self.deferrable {
-            if !first {
-                query.push(',');
-            }
-
-            let s = if deferrable {
-                " DEFERRABLE"
-            } else {
-                " NOT DEFERRABLE"
-            };
-            query.push_str(s);
-        }
-
-        self.client.batch_execute(&query).await?;
-
-        Ok(Transaction::new(self.client))
+        Transaction::start(
+            self.client,
+            self.isolation_level,
+            self.read_only,
+            self.deferrable,
+        )
+        .await
     }
 }

tokio-postgres/tests/test/main.rs

@@ -448,6 +448,30 @@ async fn transaction_future_cancellation() {
     }
 }
 
+#[tokio::test]
+async fn start_transaction_future_cancellation() {
+    let mut client = connect("user=postgres").await;
+
+    for i in 0.. {
+        let done = {
+            let txn = client.build_transaction().start();
+            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;