Don't hold strong references in statements

There's no need for the connection to stay open until statements drop -
they'll be cleaned up anyway once the connection dies.

Author: sfackler

tokio-postgres/src/proto/client.rs

@@ -4,7 +4,7 @@ use postgres_protocol;
 use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 use disconnected;
 use error::{self, Error};
@@ -15,67 +15,120 @@ use proto::query::QueryStream;
 use proto::statement::Statement;
 use types::{IsNull, Oid, ToSql, Type};
 
-pub struct PendingRequest {
-    sender: mpsc::UnboundedSender<Request>,
-    messages: Result<Vec<u8>, Error>,
-}
+pub struct PendingRequest(Result<Vec<u8>, Error>);
 
-impl PendingRequest {
-    pub fn send(self) -> Result<mpsc::Receiver<Message>, Error> {
-        let messages = self.messages?;
-        let (sender, receiver) = mpsc::channel(0);
-        self.sender
-            .unbounded_send(Request { messages, sender })
-            .map(|_| receiver)
-            .map_err(|_| disconnected())
+pub struct WeakClient(Weak<Inner>);
+
+impl WeakClient {
+    pub fn upgrade(&self) -> Option<Client> {
+        self.0.upgrade().map(Client)
     }
 }
 
-pub struct State {
-    pub types: HashMap<Oid, Type>,
-    pub typeinfo_query: Option<Statement>,
-    pub typeinfo_enum_query: Option<Statement>,
-    pub typeinfo_composite_query: Option<Statement>,
+struct State {
+    types: HashMap<Oid, Type>,
+    typeinfo_query: Option<Statement>,
+    typeinfo_enum_query: Option<Statement>,
+    typeinfo_composite_query: Option<Statement>,
 }
 
-#[derive(Clone)]
-pub struct Client {
-    pub state: Arc<Mutex<State>>,
+struct Inner {
+    state: Mutex<State>,
     sender: mpsc::UnboundedSender<Request>,
 }
 
+#[derive(Clone)]
+pub struct Client(Arc<Inner>);
+
 impl Client {
     pub fn new(sender: mpsc::UnboundedSender<Request>) -> Client {
-        Client {
-            state: Arc::new(Mutex::new(State {
+        Client(Arc::new(Inner {
+            state: Mutex::new(State {
                 types: HashMap::new(),
                 typeinfo_query: None,
                 typeinfo_enum_query: None,
                 typeinfo_composite_query: None,
-            })),
+            }),
             sender,
-        }
+        }))
+    }
+
+    pub fn downgrade(&self) -> WeakClient {
+        WeakClient(Arc::downgrade(&self.0))
+    }
+
+    pub fn cached_type(&self, oid: Oid) -> Option<Type> {
+        self.0.state.lock().types.get(&oid).cloned()
+    }
+
+    pub fn cache_type(&self, ty: &Type) {
+        self.0.state.lock().types.insert(ty.oid(), ty.clone());
+    }
+
+    pub fn typeinfo_query(&self) -> Option<Statement> {
+        self.0.state.lock().typeinfo_query.clone()
+    }
+
+    pub fn set_typeinfo_query(&self, statement: &Statement) {
+        self.0.state.lock().typeinfo_query = Some(statement.clone());
     }
 
-    pub fn prepare(&mut self, name: String, query: &str, param_types: &[Type]) -> PrepareFuture {
+    pub fn typeinfo_enum_query(&self) -> Option<Statement> {
+        self.0.state.lock().typeinfo_enum_query.clone()
+    }
+
+    pub fn set_typeinfo_enum_query(&self, statement: &Statement) {
+        self.0.state.lock().typeinfo_enum_query = Some(statement.clone());
+    }
+
+    pub fn typeinfo_composite_query(&self) -> Option<Statement> {
+        self.0.state.lock().typeinfo_composite_query.clone()
+    }
+
+    pub fn set_typeinfo_composite_query(&self, statement: &Statement) {
+        self.0.state.lock().typeinfo_composite_query = Some(statement.clone());
+    }
+
+    pub fn send(&self, request: PendingRequest) -> Result<mpsc::Receiver<Message>, Error> {
+        let messages = request.0?;
+        let (sender, receiver) = mpsc::channel(0);
+        self.0
+            .sender
+            .unbounded_send(Request { messages, sender })
+            .map(|_| receiver)
+            .map_err(|_| disconnected())
+    }
+
+    pub fn prepare(&self, name: String, query: &str, param_types: &[Type]) -> PrepareFuture {
         let pending = self.pending(|buf| {
             frontend::parse(&name, query, param_types.iter().map(|t| t.oid()), buf)?;
             frontend::describe(b'S', &name, buf)?;
             frontend::sync(buf);
             Ok(())
         });
 
-        PrepareFuture::new(pending, self.sender.clone(), name, self.clone())
+        PrepareFuture::new(self.clone(), pending, name)
     }
 
-    pub fn execute(&mut self, statement: &Statement, params: &[&ToSql]) -> ExecuteFuture {
+    pub fn execute(&self, statement: &Statement, params: &[&ToSql]) -> ExecuteFuture {
         let pending = self.pending_execute(statement, params);
-        ExecuteFuture::new(pending, statement.clone())
+        ExecuteFuture::new(self.clone(), pending, statement.clone())
     }
 
-    pub fn query(&mut self, statement: &Statement, params: &[&ToSql]) -> QueryStream {
+    pub fn query(&self, statement: &Statement, params: &[&ToSql]) -> QueryStream {
         let pending = self.pending_execute(statement, params);
-        QueryStream::new(pending, statement.clone())
+        QueryStream::new(self.clone(), pending, statement.clone())
+    }
+
+    pub fn close_statement(&self, name: &str) {
+        let mut buf = vec![];
+        frontend::close(b'S', name, &mut buf).expect("statement name not valid");
+        frontend::sync(&mut buf);
+        let (sender, _) = mpsc::channel(0);
+        let _ = self.0.sender.unbounded_send(Request {
+            messages: buf,
+            sender,
+        });
     }
 
     fn pending_execute(&self, statement: &Statement, params: &[&ToSql]) -> PendingRequest {
@@ -109,9 +162,6 @@ impl Client {
         F: FnOnce(&mut Vec<u8>) -> Result<(), Error>,
     {
         let mut buf = vec![];
-        PendingRequest {
-            sender: self.sender.clone(),
-            messages: messages(&mut buf).map(|()| buf),
-        }
+        PendingRequest(messages(&mut buf).map(|()| buf))
     }
 }

tokio-postgres/src/proto/execute.rs

@@ -4,14 +4,15 @@ use postgres_protocol::message::backend::Message;
 use state_machine_future::RentToOwn;
 
 use error::{self, Error};
-use proto::client::PendingRequest;
+use proto::client::{Client, PendingRequest};
 use proto::statement::Statement;
 use {bad_response, disconnected};
 
 #[derive(StateMachineFuture)]
 pub enum Execute {
     #[state_machine_future(start, transitions(ReadResponse))]
     Start {
+        client: Client,
         request: PendingRequest,
         statement: Statement,
     },
@@ -31,7 +32,7 @@ pub enum Execute {
 impl PollExecute for Execute {
     fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
         let state = state.take();
-        let receiver = state.request.send()?;
+        let receiver = state.client.send(state.request)?;
 
         // the statement can drop after this point, since its close will queue up after the execution
         transition!(ReadResponse { receiver })
@@ -82,7 +83,7 @@ impl PollExecute for Execute {
 }
 
 impl ExecuteFuture {
-    pub fn new(request: PendingRequest, statement: Statement) -> ExecuteFuture {
-        Execute::start(request, statement)
+    pub fn new(client: Client, request: PendingRequest, statement: Statement) -> ExecuteFuture {
+        Execute::start(client, request, statement)
     }
 }

tokio-postgres/src/proto/prepare.rs

@@ -8,7 +8,6 @@ use std::vec;
 
 use error::{self, Error};
 use proto::client::{Client, PendingRequest};
-use proto::connection::Request;
 use proto::statement::Statement;
 use proto::typeinfo::TypeinfoFuture;
 use types::{Oid, Type};
@@ -19,47 +18,41 @@ use {bad_response, disconnected};
 pub enum Prepare {
     #[state_machine_future(start, transitions(ReadParseComplete))]
     Start {
+        client: Client,
         request: PendingRequest,
-        sender: mpsc::UnboundedSender<Request>,
         name: String,
-        client: Client,
     },
     #[state_machine_future(transitions(ReadParameterDescription))]
     ReadParseComplete {
-        sender: mpsc::UnboundedSender<Request>,
+        client: Client,
         receiver: mpsc::Receiver<Message>,
         name: String,
-        client: Client,
     },
     #[state_machine_future(transitions(ReadRowDescription))]
     ReadParameterDescription {
-        sender: mpsc::UnboundedSender<Request>,
+        client: Client,
         receiver: mpsc::Receiver<Message>,
         name: String,
-        client: Client,
     },
     #[state_machine_future(transitions(ReadReadyForQuery))]
     ReadRowDescription {
-        sender: mpsc::UnboundedSender<Request>,
+        client: Client,
         receiver: mpsc::Receiver<Message>,
         name: String,
         parameters: Vec<Oid>,
-        client: Client,
     },
     #[state_machine_future(transitions(GetParameterTypes, GetColumnTypes, Finished))]
     ReadReadyForQuery {
-        sender: mpsc::UnboundedSender<Request>,
+        client: Client,
         receiver: mpsc::Receiver<Message>,
         name: String,
         parameters: Vec<Oid>,
         columns: Vec<(String, Oid)>,
-        client: Client,
     },
     #[state_machine_future(transitions(GetColumnTypes, Finished))]
     GetParameterTypes {
         future: TypeinfoFuture,
         remaining_parameters: vec::IntoIter<Oid>,
-        sender: mpsc::UnboundedSender<Request>,
         name: String,
         parameters: Vec<Type>,
         columns: Vec<(String, Oid)>,
@@ -69,7 +62,6 @@ pub enum Prepare {
         future: TypeinfoFuture,
         cur_column_name: String,
         remaining_columns: vec::IntoIter<(String, Oid)>,
-        sender: mpsc::UnboundedSender<Request>,
         name: String,
         parameters: Vec<Type>,
         columns: Vec<Column>,
@@ -83,10 +75,9 @@ pub enum Prepare {
 impl PollPrepare for Prepare {
     fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
         let state = state.take();
-        let receiver = state.request.send()?;
+        let receiver = state.client.send(state.request)?;
 
         transition!(ReadParseComplete {
-            sender: state.sender,
             receiver,
             name: state.name,
             client: state.client,
@@ -101,7 +92,6 @@ impl PollPrepare for Prepare {
 
         match message {
             Some(Message::ParseComplete) => transition!(ReadParameterDescription {
-                sender: state.sender,
                 receiver: state.receiver,
                 name: state.name,
                 client: state.client,
@@ -120,7 +110,6 @@ impl PollPrepare for Prepare {
 
         match message {
             Some(Message::ParameterDescription(body)) => transition!(ReadRowDescription {
-                sender: state.sender,
                 receiver: state.receiver,
                 name: state.name,
                 parameters: body.parameters().collect()?,
@@ -148,7 +137,6 @@ impl PollPrepare for Prepare {
         };
 
         transition!(ReadReadyForQuery {
-            sender: state.sender,
             receiver: state.receiver,
             name: state.name,
             parameters: state.parameters,
@@ -174,7 +162,6 @@ impl PollPrepare for Prepare {
             transition!(GetParameterTypes {
                 future: TypeinfoFuture::new(oid, state.client),
                 remaining_parameters: parameters,
-                sender: state.sender,
                 name: state.name,
                 parameters: vec![],
                 columns: state.columns,
@@ -187,15 +174,14 @@ impl PollPrepare for Prepare {
                 future: TypeinfoFuture::new(oid, state.client),
                 cur_column_name: name,
                 remaining_columns: columns,
-                sender: state.sender,
                 name: state.name,
                 parameters: vec![],
                 columns: vec![],
             });
         }
 
         transition!(Finished(Statement::new(
-            state.sender,
+            state.client.downgrade(),
             state.name,
             vec![],
             vec![]
@@ -222,15 +208,14 @@ impl PollPrepare for Prepare {
                 future: TypeinfoFuture::new(oid, client),
                 cur_column_name: name,
                 remaining_columns: columns,
-                sender: state.sender,
                 name: state.name,
                 parameters: state.parameters,
                 columns: vec![],
             })
         }
 
         transition!(Finished(Statement::new(
-            state.sender,
+            client.downgrade(),
             state.name,
             state.parameters,
             vec![],
@@ -240,7 +225,7 @@ impl PollPrepare for Prepare {
     fn poll_get_column_types<'a>(
         state: &'a mut RentToOwn<'a, GetColumnTypes>,
     ) -> Poll<AfterGetColumnTypes, Error> {
-        loop {
+        let client = loop {
             let (ty, client) = try_ready!(state.future.poll());
             let name = mem::replace(&mut state.cur_column_name, String::new());
             state.columns.push(Column::new(name, ty));
@@ -250,13 +235,13 @@ impl PollPrepare for Prepare {
                     state.cur_column_name = name;
                     state.future = TypeinfoFuture::new(oid, client);
                 }
-                None => break,
+                None => break client,
             }
-        }
+        };
         let state = state.take();
 
         transition!(Finished(Statement::new(
-            state.sender,
+            client.downgrade(),
             state.name,
             state.parameters,
             state.columns,
@@ -265,12 +250,7 @@ impl PollPrepare for Prepare {
 }
 
 impl PrepareFuture {
-    pub fn new(
-        request: PendingRequest,
-        sender: mpsc::UnboundedSender<Request>,
-        name: String,
-        client: Client,
-    ) -> PrepareFuture {
-        Prepare::start(request, sender, name, client)
+    pub fn new(client: Client, request: PendingRequest, name: String) -> PrepareFuture {
+        Prepare::start(client, request, name)
     }
 }