Restructure conn pool to RAII style

Author: sfackler

src/lib.rs

@@ -158,7 +158,7 @@ pub struct PostgresConnection {
 }
 
 impl Drop for PostgresConnection {
-    fn drop(&self) {
+    fn drop(&mut self) {
         do io_error::cond.trap(|_| {}).inside {
             self.write_messages([&Terminate]);
         }
@@ -461,7 +461,7 @@ pub struct PostgresTransaction<'self> {
 
 #[unsafe_destructor]
 impl<'self> Drop for PostgresTransaction<'self> {
-    fn drop(&self) {
+    fn drop(&mut self) {
         do io_error::cond.trap(|_| {}).inside {
             if task::failing() || !self.commit.take() {
                 if self.nested {
@@ -568,7 +568,7 @@ impl ResultDescription {
 
 #[unsafe_destructor]
 impl<'self> Drop for NormalPostgresStatement<'self> {
-    fn drop(&self) {
+    fn drop(&mut self) {
         do io_error::cond.trap(|_| {}).inside {
             self.conn.write_messages([
                 &Close {
@@ -792,7 +792,7 @@ pub struct PostgresResult<'self> {
 
 #[unsafe_destructor]
 impl<'self> Drop for PostgresResult<'self> {
-    fn drop(&self) {
+    fn drop(&mut self) {
         do io_error::cond.trap(|_| {}).inside {
             self.stmt.conn.write_messages([
                 &Close {

src/pool/mod.rs

@@ -1,9 +1,13 @@
 extern mod extra;
 
 use extra::arc::MutexArc;
-use std::cell::Cell;
 
-use super::{PostgresConnection, PostgresConnectError};
+use super::{PostgresConnection,
+            NormalPostgresStatement,
+            PostgresDbError,
+            PostgresConnectError,
+            PostgresTransaction};
+use super::types::ToSql;
 
 pub struct PostgresConnectionPoolConfig {
     initial_size: uint,
@@ -75,37 +79,69 @@ impl PostgresConnectionPool {
         })
     }
 
-    pub fn try_with_connection<T>(&mut self,
-                                  blk: &fn(&PostgresConnection) -> T)
-            -> Result<T, PostgresConnectError> {
-        // Can't use access_cond here since PostgresConnection uses Cell
+    pub fn try_get_connection(&self) -> Result<PooledPostgresConnection,
+                                               PostgresConnectError> {
         let conn = unsafe {
-            do self.pool.unsafe_access_cond |pool, cond| {
+            do self.pool.unsafe_access_cond |pool, cvar| {
                 while pool.pool.is_empty() {
-                    cond.wait();
+                    cvar.wait();
                 }
 
                 pool.pool.pop()
             }
         };
 
-        let ret = blk(&conn);
+        Ok(PooledPostgresConnection {
+            pool: self.clone(),
+            conn: Some(conn)
+        })
+    }
 
-        let conn = Cell::new(conn);
+    pub fn get_connection(&self) -> PooledPostgresConnection {
+        match self.try_get_connection() {
+            Ok(conn) => conn,
+            Err(err) => fail!("Unable to get connection: %s", err.to_str())
+        }
+    }
+}
+
+// Should be a newtype
+pub struct PooledPostgresConnection {
+    priv pool: PostgresConnectionPool,
+    // Todo remove the Option wrapper when drop takes self by value
+    priv conn: Option<PostgresConnection>
+}
+
+impl Drop for PooledPostgresConnection {
+    fn drop(&mut self) {
         unsafe {
-            do self.pool.unsafe_access |pool| {
-                pool.pool.push(conn.take());
+            do self.pool.pool.unsafe_access |pool| {
+                pool.pool.push(self.conn.take_unwrap());
             }
         }
+    }
+}
 
-        Ok(ret)
+impl PooledPostgresConnection {
+    pub fn try_prepare<'a>(&'a self, query: &str)
+            -> Result<NormalPostgresStatement<'a>, PostgresDbError> {
+        self.conn.get_ref().try_prepare(query)
     }
 
-    pub fn with_connection<T>(&mut self, blk: &fn(&PostgresConnection) -> T)
-            -> T {
-        match self.try_with_connection(blk) {
-            Ok(ret) => ret,
-            Err(err) => fail!("Error getting connection: %s", err.to_str())
-        }
+    pub fn prepare<'a>(&'a self, query: &str) -> NormalPostgresStatement<'a> {
+        self.conn.get_ref().prepare(query)
+    }
+
+    pub fn try_update(&self, query: &str, params: &[&ToSql])
+            -> Result<uint, PostgresDbError> {
+        self.conn.get_ref().try_update(query, params)
+    }
+
+    pub fn update(&self, query: &str, params: &[&ToSql]) -> uint {
+        self.conn.get_ref().update(query, params)
+    }
+
+    pub fn in_transaction<T>(&self, blk: &fn(&PostgresTransaction) -> T) -> T {
+        self.conn.get_ref().in_transaction(blk)
     }
 }

src/pool/test.rs

@@ -15,33 +15,31 @@ fn test_pool() {
         min_size: 2,
         max_size: 2
     };
-    let mut pool = PostgresConnectionPool::new("postgres://postgres@localhost",
-                                               config).unwrap();
+    let pool = PostgresConnectionPool::new("postgres://postgres@localhost",
+                                           config).unwrap();
 
     let (stream1, stream2) = DuplexStream::<(), ()>();
 
     let pool1 = Cell::new(pool.clone());
     let mut fut1 = do future::spawn {
-        let mut pool = pool1.take();
+        let pool = pool1.take();
 
-        do pool.with_connection |_conn| {
-            stream1.send(());
-            stream1.recv();
-        }
+        let _conn = pool.get_connection();
+        stream1.send(());
+        stream1.recv();
     };
 
     let pool2 = Cell::new(pool.clone());
     let mut fut2 = do future::spawn {
-        let mut pool = pool2.take();
+        let pool = pool2.take();
 
-        do pool.with_connection |_conn| {
-            stream2.send(());
-            stream2.recv();
-        }
+        let _conn = pool.get_connection();
+        stream2.send(());
+        stream2.recv();
     };
 
     fut1.get();
     fut2.get();
 
-    do pool.with_connection |_conn| { }
+    pool.get_connection();
 }