Repository: diesel-rs/diesel (language: rust)

Issue:
Performance regression when getting multiple connections with `r2d2::Pool`
## Setup

### Versions

- **Rust:** rustc 1.62.0-nightly (e85edd9a8 2022-04-28)
- **Diesel:** v2.0.0-rc0
- **Database:** postgresql
- **Operating System** Linux 5.17.3-arch1-1

### Feature Flags

- **diesel:** ["postgres", "r2d2"]

## Problem Description

As is reported in #3144, performance regression is observed when upgrading `diesel` from `v1.4.8` to `v2.0.0-rc0`. Aggregation queries (e.g. `count(*)` are particularly affected, while the performance hit on normal queries is minimal.

### What are you trying to accomplish?

### What is the expected output?

The performance of `v2.0.0-rc0` should be on par with `v1.4.8`.

### What is the actual output?

`cargo bench` output of `table.filter(id.eq(1)).count().get_result()`:

| hash | setup | perf |
|---|---|---|
| [d0bbda8](https://github.com/heyrict/diesel-aggr-sample/tree/d0bbda8a7066cf9b310e3beed2eece6418faeec9) |diesel 1.4 + r2d2| 185,936 ns/iter (+/- 15,282) |
| [599c492](https://github.com/heyrict/diesel-aggr-sample/tree/599c4923f1832fbafaaac33e6df75f6db8815e84) |diesel 2.0 + r2d2| 1,388,410 ns/iter (+/- 158,035) |
| [d3c23ba](https://github.com/heyrict/diesel-aggr-sample/tree/d3c23ba13a3202e59e28e0525a5cc7f0f40f461c) |diesel 1.4 - r2d2| 103,562 ns/iter (+/- 2,997) |
| [58d1210](https://github.com/heyrict/diesel-aggr-sample/tree/58d1210e65ac3ba2dc7e52b5b77853d323c23f9f) |diesel 2.0 - r2d2| 109,238 ns/iter (+/- 5,887) |

### Are you seeing any additional errors?


### Steps to reproduce

The minimal reproduction sample is [here](https://github.com/heyrict/diesel-aggr-sample).

<!--
Please include as much of your codebase as needed to reproduce the error.  If the relevant files are large, please consider linking to a public repository or a [Gist](https://gist.github.com/). This includes normally the following parts:

* The exact code where your hit the problem
* Relevant parts your schema, so any `table!` macro calls required for the
* Any other type definitions involved in the code, which produces your problem
-->

## Checklist

- [x] I have already looked over the [issue tracker](https://github.com/diesel-rs/diesel/issues)  and the [disussion forum](https://github.com/diesel-rs/diesel/discussions)  for similar possible closed issues.
<!--
If you are unsure if your issue is a duplicate of an existing issue please link the issue in question here
--> 
- [x] This issue can be reproduced on Rust's stable channel. (Your issue will be
  closed if this is not the case)
- [x] This issue can be reproduced without requiring a third party crate

<!--
Thank you for your submission!  You're helping make Diesel more robust 🎉

We'll try to respond as quickly as possible.
-->



Discussion / hints:

Fix patch:
diff --git a/diesel/src/mysql/connection/mod.rs b/diesel/src/mysql/connection/mod.rs
index 7e5bc4746690..7cc7a7140b90 100644
--- a/diesel/src/mysql/connection/mod.rs
+++ b/diesel/src/mysql/connection/mod.rs
@@ -144,11 +144,15 @@ impl crate::r2d2::R2D2Connection for MysqlConnection {
     }
 
     fn is_broken(&mut self) -> bool {
-        self.transaction_state
-            .status
-            .transaction_depth()
-            .map(|d| d.is_none())
-            .unwrap_or(true)
+        match self.transaction_state.status.transaction_depth() {
+            // all transactions are closed
+            // so we don't consider this connection broken
+            Ok(None) => false,
+            // The transaction manager is in an error state
+            // or contains an open transaction
+            // Therefore we consider this connection broken
+            Err(_) | Ok(Some(_)) => true,
+        }
     }
 }
 
diff --git a/diesel/src/pg/connection/mod.rs b/diesel/src/pg/connection/mod.rs
index c1a6cf3c8b44..fcd608b4ab6b 100644
--- a/diesel/src/pg/connection/mod.rs
+++ b/diesel/src/pg/connection/mod.rs
@@ -177,11 +177,15 @@ impl crate::r2d2::R2D2Connection for PgConnection {
     }
 
     fn is_broken(&mut self) -> bool {
-        self.transaction_state
-            .status
-            .transaction_depth()
-            .map(|d| d.is_none())
-            .unwrap_or(true)
+        match self.transaction_state.status.transaction_depth() {
+            // all transactions are closed
+            // so we don't consider this connection broken
+            Ok(None) => false,
+            // The transaction manager is in an error state
+            // or contains an open transaction
+            // Therefore we consider this connection broken
+            Err(_) | Ok(Some(_)) => true,
+        }
     }
 }
 
diff --git a/diesel/src/r2d2.rs b/diesel/src/r2d2.rs
index ddac7c96e818..62f5efb9eb10 100644
--- a/diesel/src/r2d2.rs
+++ b/diesel/src/r2d2.rs
@@ -327,6 +327,28 @@ where
     }
 }
 
+#[derive(QueryId)]
+pub(crate) struct CheckConnectionQuery;
+
+impl<DB> QueryFragment<DB> for CheckConnectionQuery
+where
+    DB: Backend,
+{
+    fn walk_ast<'b>(
+        &'b self,
+        mut pass: crate::query_builder::AstPass<'_, 'b, DB>,
+    ) -> QueryResult<()> {
+        pass.push_sql("SELECT 1");
+        Ok(())
+    }
+}
+
+impl Query for CheckConnectionQuery {
+    type SqlType = crate::sql_types::Integer;
+}
+
+impl<C> RunQueryDsl<C> for CheckConnectionQuery {}
+
 #[cfg(test)]
 mod tests {
     use std::sync::mpsc;
@@ -394,26 +416,110 @@ mod tests {
         let query = select("foo".into_sql::<Text>());
         assert_eq!("foo", query.get_result::<String>(&mut conn).unwrap());
     }
-}
 
-#[derive(QueryId)]
-pub(crate) struct CheckConnectionQuery;
+    #[test]
+    fn check_pool_does_actually_hold_connections() {
+        use std::sync::atomic::{AtomicU32, Ordering};
+
+        #[derive(Debug)]
+        struct TestEventHandler {
+            acquire_count: Arc<AtomicU32>,
+            release_count: Arc<AtomicU32>,
+            checkin_count: Arc<AtomicU32>,
+            checkout_count: Arc<AtomicU32>,
+        }
 
-impl<DB> QueryFragment<DB> for CheckConnectionQuery
-where
-    DB: Backend,
-{
-    fn walk_ast<'b>(
-        &'b self,
-        mut pass: crate::query_builder::AstPass<'_, 'b, DB>,
-    ) -> QueryResult<()> {
-        pass.push_sql("SELECT 1");
-        Ok(())
-    }
-}
+        impl r2d2::HandleEvent for TestEventHandler {
+            fn handle_acquire(&self, _event: r2d2::event::AcquireEvent) {
+                self.acquire_count.fetch_add(1, Ordering::Relaxed);
+            }
+            fn handle_release(&self, _event: r2d2::event::ReleaseEvent) {
+                self.release_count.fetch_add(1, Ordering::Relaxed);
+            }
+            fn handle_checkout(&self, _event: r2d2::event::CheckoutEvent) {
+                self.checkout_count.fetch_add(1, Ordering::Relaxed);
+            }
+            fn handle_checkin(&self, _event: r2d2::event::CheckinEvent) {
+                self.checkin_count.fetch_add(1, Ordering::Relaxed);
+            }
+        }
 
-impl Query for CheckConnectionQuery {
-    type SqlType = crate::sql_types::Integer;
-}
+        let acquire_count = Arc::new(AtomicU32::new(0));
+        let release_count = Arc::new(AtomicU32::new(0));
+        let checkin_count = Arc::new(AtomicU32::new(0));
+        let checkout_count = Arc::new(AtomicU32::new(0));
 
-impl<C> RunQueryDsl<C> for CheckConnectionQuery {}
+        let handler = Box::new(TestEventHandler {
+            acquire_count: acquire_count.clone(),
+            release_count: release_count.clone(),
+            checkin_count: checkin_count.clone(),
+            checkout_count: checkout_count.clone(),
+        });
+
+        let manager = ConnectionManager::<TestConnection>::new(database_url());
+        let pool = Pool::builder()
+            .max_size(1)
+            .test_on_check_out(true)
+            .event_handler(handler)
+            .build(manager)
+            .unwrap();
+
+        assert_eq!(acquire_count.load(Ordering::Relaxed), 1);
+        assert_eq!(release_count.load(Ordering::Relaxed), 0);
+        assert_eq!(checkin_count.load(Ordering::Relaxed), 0);
+        assert_eq!(checkout_count.load(Ordering::Relaxed), 0);
+
+        // check that we reuse connections with the pool
+        {
+            let conn = pool.get().unwrap();
+
+            assert_eq!(acquire_count.load(Ordering::Relaxed), 1);
+            assert_eq!(release_count.load(Ordering::Relaxed), 0);
+            assert_eq!(checkin_count.load(Ordering::Relaxed), 0);
+            assert_eq!(checkout_count.load(Ordering::Relaxed), 1);
+            std::mem::drop(conn);
+        }
+
+        assert_eq!(acquire_count.load(Ordering::Relaxed), 1);
+        assert_eq!(release_count.load(Ordering::Relaxed), 0);
+        assert_eq!(checkin_count.load(Ordering::Relaxed), 1);
+        assert_eq!(checkout_count.load(Ordering::Relaxed), 1);
+
+        // check that we remove a connection with open transactions from the pool
+        {
+            let mut conn = pool.get().unwrap();
+
+            assert_eq!(acquire_count.load(Ordering::Relaxed), 1);
+            assert_eq!(release_count.load(Ordering::Relaxed), 0);
+            assert_eq!(checkin_count.load(Ordering::Relaxed), 1);
+            assert_eq!(checkout_count.load(Ordering::Relaxed), 2);
+
+            <TestConnection as Connection>::TransactionManager::begin_transaction(&mut *conn)
+                .unwrap();
+        }
+
+        assert_eq!(acquire_count.load(Ordering::Relaxed), 1);
+        assert_eq!(release_count.load(Ordering::Relaxed), 1);
+        assert_eq!(checkin_count.load(Ordering::Relaxed), 2);
+        assert_eq!(checkout_count.load(Ordering::Relaxed), 2);
+
+        // check that we remove a connection from the pool that was
+        // open during panicing
+        #[allow(unreachable_code, unused_variables)]
+        std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+            let conn = pool.get();
+            assert_eq!(acquire_count.load(Ordering::Relaxed), 2);
+            assert_eq!(release_count.load(Ordering::Relaxed), 1);
+            assert_eq!(checkin_count.load(Ordering::Relaxed), 2);
+            assert_eq!(checkout_count.load(Ordering::Relaxed), 3);
+            panic!();
+            std::mem::drop(conn);
+        }))
+        .unwrap_err();
+
+        assert_eq!(acquire_count.load(Ordering::Relaxed), 2);
+        assert_eq!(release_count.load(Ordering::Relaxed), 2);
+        assert_eq!(checkin_count.load(Ordering::Relaxed), 3);
+        assert_eq!(checkout_count.load(Ordering::Relaxed), 3);
+    }
+}
diff --git a/diesel/src/sqlite/connection/mod.rs b/diesel/src/sqlite/connection/mod.rs
index 3d85958519ba..b6dffcb87f2b 100644
--- a/diesel/src/sqlite/connection/mod.rs
+++ b/diesel/src/sqlite/connection/mod.rs
@@ -141,11 +141,15 @@ impl crate::r2d2::R2D2Connection for crate::sqlite::SqliteConnection {
     }
 
     fn is_broken(&mut self) -> bool {
-        self.transaction_state
-            .status
-            .transaction_depth()
-            .map(|d| d.is_none())
-            .unwrap_or(true)
+        match self.transaction_state.status.transaction_depth() {
+            // all transactions are closed
+            // so we don't consider this connection broken
+            Ok(None) => false,
+            // The transaction manager is in an error state
+            // or contains an open transaction
+            // Therefore we consider this connection broken
+            Err(_) | Ok(Some(_)) => true,
+        }
     }
 }
 


Test patch: