Add unwind safety

src/lib.rs

@@ -50,6 +50,7 @@ use std::error;
 use std::fmt;
 use std::mem;
 use std::ops::{Deref, DerefMut};
+use std::panic::{RefUnwindSafe, UnwindSafe};
 use std::sync::{Arc, Weak};
 use std::time::{Duration, Instant};
 
@@ -175,6 +176,13 @@ where
     cond: Condvar,
 }
 
+// Condvar isn't UnwindSafe, but our pool is. The only thing you could do
+// with it which causes problems in the face of unwinding is check out a
+// connection. Since we do not check a connection back into the pool if it
+// is dropped during a panic, we can safely implement UnwindSafe.
+impl<M: ManageConnection + UnwindSafe> UnwindSafe for SharedPool<M> {}
+impl<M: ManageConnection + RefUnwindSafe> RefUnwindSafe for SharedPool<M> {}
+
 fn drop_conns<M>(
     shared: &Arc<SharedPool<M>>,
     mut internals: MutexGuard<PoolInternals<M::Connection>>,
@@ -466,6 +474,11 @@ where
         }
     }
 
+    fn drop_conn(&self, conn: Conn<M::Connection>) {
+        let internals = self.0.internals.lock();
+        drop_conns(&self.0, internals, vec![conn.conn]);
+    }
+
     /// Returns information about the current state of the pool.
     pub fn state(&self) -> State {
         let internals = self.0.internals.lock();
@@ -569,7 +582,13 @@ where
     M: ManageConnection,
 {
     fn drop(&mut self) {
-        self.pool.put_back(self.conn.take().unwrap());
+        use std::thread::panicking;
+
+        if panicking() {
+            self.pool.drop_conn(self.conn.take().unwrap());
+        } else {
+            self.pool.put_back(self.conn.take().unwrap());
+        }
     }
 }
 

src/test.rs

@@ -1,4 +1,5 @@
 use antidote::Mutex;
+use std::panic::catch_unwind;
 use std::sync::atomic::{
     AtomicBool, AtomicIsize, AtomicUsize, Ordering, ATOMIC_BOOL_INIT, ATOMIC_USIZE_INIT,
 };
@@ -545,3 +546,51 @@ fn conns_drop_on_pool_drop() {
     }
     panic!("timed out waiting for connections to drop");
 }
+
+#[test]
+fn connections_are_not_returned_to_pool_on_panic() {
+    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;
+
+    struct Connection;
+
+    impl Drop for Connection {
+        fn drop(&mut self) {
+            DROPPED.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    struct Handler;
+
+    impl ManageConnection for Handler {
+        type Connection = Connection;
+        type Error = Error;
+
+        fn connect(&self) -> Result<Connection, Error> {
+            Ok(Connection)
+        }
+
+        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
+            Ok(())
+        }
+
+        fn has_broken(&self, _: &mut Connection) -> bool {
+            false
+        }
+    }
+
+    let pool = Pool::builder()
+        .max_lifetime(Some(Duration::from_secs(1)))
+        .max_size(10)
+        .min_idle(Some(0))
+        .build(Handler)
+        .unwrap();
+
+    let _ = catch_unwind(|| {
+        let _conn = pool.get().unwrap();
+        panic!();
+    });
+
+    assert_eq!(DROPPED.load(Ordering::SeqCst), 1);
+    assert_eq!(pool.state().connections, 0);
+    assert_eq!(pool.state().idle_connections, 0);
+}