fix connection leak

suharev7 · Oct 6, 2019 · 6168956 · 6168956
1 parent b8cd03b
commit 6168956
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 2 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -363,6 +363,8 @@ impl ClientHandle {
     {
         let context = self.context.clone();
         let pool = self.pool.clone();
+        let release_pool = self.pool.clone();
+
         let query = Query::from(sql);
         self.wrap_future(|mut c| {
             info!("[execute]    {}", query.get_sql());
@@ -387,6 +389,10 @@ impl ClientHandle {
                     }
                     _ => future::err::<_, Error>(Error::Driver(DriverError::UnexpectedPacket)),
                 })
+                .map_err(move |err| {
+                    release_pool.release_conn();
+                    err.into()
+                })
                 .map(Option::unwrap)
         })
     }
@@ -409,6 +415,7 @@ impl ClientHandle {
 
         let context = self.context.clone();
         let pool = self.pool.clone();
+        let release_pool = self.pool.clone();
 
         self.wrap_future(|mut c| {
             info!("[insert]     {}", query.get_sql());
@@ -439,6 +446,10 @@ impl ClientHandle {
                             .map(|(c, _)| c),
                     )
                 })
+                .map_err(move |err| {
+                    release_pool.release_conn();
+                    err.into()
+                })
         })
     }
 

diff --git a/src/pool/mod.rs b/src/pool/mod.rs
@@ -300,10 +300,9 @@ mod test {
 
     use tokio::prelude::*;
 
-    use crate::{errors::Error, io::BoxFuture, test_misc::DATABASE_URL, types::Options};
+    use crate::{ClientHandle, errors::Error, io::BoxFuture, test_misc::DATABASE_URL, types::{Block, Options}};
 
     use super::Pool;
-    use crate::ClientHandle;
 
     /// Same as `tokio::run`, but will panic if future panics and will return the result
     /// of future execution.
@@ -480,5 +479,44 @@ mod test {
 
         let info = pool.info();
         assert_eq!(info.ongoing, 0);
+        assert_eq!(info.tasks_len, 0);
+        assert_eq!(info.idle_len, 0);
+    }
+
+    #[test]
+    fn test_wrong_insert() {
+        let pool = Pool::new(DATABASE_URL.as_str());
+
+        let done = pool
+            .get_handle()
+            .and_then(|c| {
+                let block = Block::new();
+                c.insert("unexisting", block)
+            });
+
+        run(done).unwrap_err();
+
+        let info = pool.info();
+        assert_eq!(info.ongoing, 0);
+        assert_eq!(info.tasks_len, 0);
+        assert_eq!(info.idle_len, 0);
+    }
+
+    #[test]
+    fn test_wrong_execute() {
+        let pool = Pool::new(DATABASE_URL.as_str());
+
+        let done = pool
+            .get_handle()
+            .and_then(|c| {
+                c.execute("DROP TABLE unexisting")
+            });
+
+        run(done).unwrap_err();
+
+        let info = pool.info();
+        assert_eq!(info.ongoing, 0);
+        assert_eq!(info.tasks_len, 0);
+        assert_eq!(info.idle_len, 0);
     }
 }