Overhaul query cancellation

Multi-host support means we can't simply take the old approach - we need
to know which of the hosts we actually connected to. It's also nice to
move this from the connection to the client since that's what you'd
normally have access to.

Author: sfackler

tokio-postgres/src/config.rs

@@ -17,10 +17,10 @@ use tokio_io::{AsyncRead, AsyncWrite};
 
 #[cfg(feature = "runtime")]
 use crate::proto::ConnectFuture;
-use crate::proto::{CancelQueryRawFuture, HandshakeFuture};
-use crate::{CancelData, CancelQueryRaw, Error, Handshake, TlsMode};
+use crate::proto::HandshakeFuture;
 #[cfg(feature = "runtime")]
 use crate::{Connect, MakeTlsMode, Socket};
+use crate::{Error, Handshake, TlsMode};
 
 #[cfg(feature = "runtime")]
 #[derive(Debug, Copy, Clone, PartialEq)]
@@ -267,7 +267,7 @@ impl Config {
         S: AsyncRead + AsyncWrite,
         T: TlsMode<S>,
     {
-        Handshake(HandshakeFuture::new(stream, tls_mode, self.clone()))
+        Handshake(HandshakeFuture::new(stream, tls_mode, self.clone(), None))
     }
 
     #[cfg(feature = "runtime")]
@@ -277,19 +277,6 @@ impl Config {
     {
         Connect(ConnectFuture::new(make_tls_mode, Ok(self.clone())))
     }
-
-    pub fn cancel_query_raw<S, T>(
-        &self,
-        stream: S,
-        tls_mode: T,
-        cancel_data: CancelData,
-    ) -> CancelQueryRaw<S, T>
-    where
-        S: AsyncRead + AsyncWrite,
-        T: TlsMode<S>,
-    {
-        CancelQueryRaw(CancelQueryRawFuture::new(stream, tls_mode, cancel_data))
-    }
 }
 
 impl FromStr for Config {

tokio-postgres/src/lib.rs

@@ -206,6 +206,22 @@ impl Client {
         BatchExecute(self.0.batch_execute(query))
     }
 
+    #[cfg(feature = "runtime")]
+    pub fn cancel_query<T>(&mut self, make_tls_mode: T) -> CancelQuery<T>
+    where
+        T: MakeTlsMode<Socket>,
+    {
+        CancelQuery(self.0.cancel_query(make_tls_mode))
+    }
+
+    pub fn cancel_query_raw<S, T>(&mut self, stream: S, tls_mode: T) -> CancelQueryRaw<S, T>
+    where
+        S: AsyncRead + AsyncWrite,
+        T: TlsMode<S>,
+    {
+        CancelQueryRaw(self.0.cancel_query_raw(stream, tls_mode))
+    }
+
     pub fn is_closed(&self) -> bool {
         self.0.is_closed()
     }
@@ -222,10 +238,6 @@ impl<S> Connection<S>
 where
     S: AsyncRead + AsyncWrite,
 {
-    pub fn cancel_data(&self) -> CancelData {
-        self.0.cancel_data()
-    }
-
     pub fn parameter(&self, name: &str) -> Option<&str> {
         self.0.parameter(name)
     }
@@ -274,6 +286,25 @@ where
     }
 }
 
+#[cfg(feature = "runtime")]
+#[must_use = "futures do nothing unless polled"]
+pub struct CancelQuery<T>(proto::CancelQueryFuture<T>)
+where
+    T: MakeTlsMode<Socket>;
+
+#[cfg(feature = "runtime")]
+impl<T> Future for CancelQuery<T>
+where
+    T: MakeTlsMode<Socket>,
+{
+    type Item = ();
+    type Error = Error;
+
+    fn poll(&mut self) -> Poll<(), Error> {
+        self.0.poll()
+    }
+}
+
 #[must_use = "futures do nothing unless polled"]
 pub struct Handshake<S, T>(proto::HandshakeFuture<S, T>)
 where
@@ -478,15 +509,6 @@ impl Future for BatchExecute {
     }
 }
 
-/// Contains information necessary to cancel queries for a session.
-#[derive(Copy, Clone, Debug)]
-pub struct CancelData {
-    /// The process ID of the session.
-    pub process_id: i32,
-    /// The secret key for the session.
-    pub secret_key: i32,
-}
-
 /// An asynchronous notification.
 #[derive(Clone, Debug)]
 pub struct Notification {

tokio-postgres/src/proto/cancel_query.rs

@@ -0,0 +1,105 @@
+use futures::{try_ready, Future, Poll};
+use state_machine_future::{transition, RentToOwn, StateMachineFuture};
+use std::io;
+
+use crate::proto::{CancelQueryRawFuture, ConnectSocketFuture};
+use crate::{Config, Error, Host, MakeTlsMode, Socket};
+
+#[derive(StateMachineFuture)]
+pub enum CancelQuery<T>
+where
+    T: MakeTlsMode<Socket>,
+{
+    #[state_machine_future(start, transitions(ConnectingSocket))]
+    Start {
+        make_tls_mode: T,
+        idx: Option<usize>,
+        config: Config,
+        process_id: i32,
+        secret_key: i32,
+    },
+    #[state_machine_future(transitions(Canceling))]
+    ConnectingSocket {
+        future: ConnectSocketFuture,
+        tls_mode: T::TlsMode,
+        process_id: i32,
+        secret_key: i32,
+    },
+    #[state_machine_future(transitions(Finished))]
+    Canceling {
+        future: CancelQueryRawFuture<Socket, T::TlsMode>,
+    },
+    #[state_machine_future(ready)]
+    Finished(()),
+    #[state_machine_future(error)]
+    Failed(Error),
+}
+
+impl<T> PollCancelQuery<T> for CancelQuery<T>
+where
+    T: MakeTlsMode<Socket>,
+{
+    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<T>>) -> Poll<AfterStart<T>, Error> {
+        let mut state = state.take();
+
+        let idx = state.idx.ok_or_else(|| {
+            Error::connect(io::Error::new(io::ErrorKind::InvalidInput, "unknown host"))
+        })?;
+
+        let hostname = match &state.config.0.host[idx] {
+            Host::Tcp(host) => &**host,
+            // postgres doesn't support TLS over unix sockets, so the choice here doesn't matter
+            #[cfg(unix)]
+            Host::Unix(_) => "",
+        };
+        let tls_mode = state
+            .make_tls_mode
+            .make_tls_mode(hostname)
+            .map_err(|e| Error::tls(e.into()))?;
+
+        transition!(ConnectingSocket {
+            future: ConnectSocketFuture::new(state.config, idx),
+            tls_mode,
+            process_id: state.process_id,
+            secret_key: state.secret_key,
+        })
+    }
+
+    fn poll_connecting_socket<'a>(
+        state: &'a mut RentToOwn<'a, ConnectingSocket<T>>,
+    ) -> Poll<AfterConnectingSocket<T>, Error> {
+        let socket = try_ready!(state.future.poll());
+        let state = state.take();
+
+        transition!(Canceling {
+            future: CancelQueryRawFuture::new(
+                socket,
+                state.tls_mode,
+                state.process_id,
+                state.secret_key
+            ),
+        })
+    }
+
+    fn poll_canceling<'a>(
+        state: &'a mut RentToOwn<'a, Canceling<T>>,
+    ) -> Poll<AfterCanceling, Error> {
+        try_ready!(state.future.poll());
+        transition!(Finished(()))
+    }
+}
+
+impl<T> CancelQueryFuture<T>
+where
+    T: MakeTlsMode<Socket>,
+{
+    pub fn new(
+        make_tls_mode: T,
+        idx: Option<usize>,
+        config: Config,
+        process_id: i32,
+        secret_key: i32,
+    ) -> CancelQueryFuture<T> {
+        CancelQuery::start(make_tls_mode, idx, config, process_id, secret_key)
+    }
+}

tokio-postgres/src/proto/cancel_query_raw.rs

@@ -6,7 +6,7 @@ use tokio_io::{AsyncRead, AsyncWrite};
 
 use crate::error::Error;
 use crate::proto::TlsFuture;
-use crate::{CancelData, TlsMode};
+use crate::TlsMode;
 
 #[derive(StateMachineFuture)]
 pub enum CancelQueryRaw<S, T>
@@ -17,7 +17,8 @@ where
     #[state_machine_future(start, transitions(SendingCancel))]
     Start {
         future: TlsFuture<S, T>,
-        cancel_data: CancelData,
+        process_id: i32,
+        secret_key: i32,
     },
     #[state_machine_future(transitions(FlushingCancel))]
     SendingCancel {
@@ -40,11 +41,7 @@ where
         let (stream, _) = try_ready!(state.future.poll());
 
         let mut buf = vec![];
-        frontend::cancel_request(
-            state.cancel_data.process_id,
-            state.cancel_data.secret_key,
-            &mut buf,
-        );
+        frontend::cancel_request(state.process_id, state.secret_key, &mut buf);
 
         transition!(SendingCancel {
             future: io::write_all(stream, buf),
@@ -74,7 +71,12 @@ where
     S: AsyncRead + AsyncWrite,
     T: TlsMode<S>,
 {
-    pub fn new(stream: S, tls_mode: T, cancel_data: CancelData) -> CancelQueryRawFuture<S, T> {
-        CancelQueryRaw::start(TlsFuture::new(stream, tls_mode), cancel_data)
+    pub fn new(
+        stream: S,
+        tls_mode: T,
+        process_id: i32,
+        secret_key: i32,
+    ) -> CancelQueryRawFuture<S, T> {
+        CancelQueryRaw::start(TlsFuture::new(stream, tls_mode), process_id, secret_key)
     }
 }

tokio-postgres/src/proto/client.rs

@@ -8,6 +8,7 @@ use postgres_protocol::message::frontend;
 use std::collections::HashMap;
 use std::error::Error as StdError;
 use std::sync::{Arc, Weak};
+use tokio_io::{AsyncRead, AsyncWrite};
 
 use crate::proto::bind::BindFuture;
 use crate::proto::connection::{Request, RequestMessages};
@@ -20,8 +21,13 @@ use crate::proto::prepare::PrepareFuture;
 use crate::proto::query::QueryStream;
 use crate::proto::simple_query::SimpleQueryStream;
 use crate::proto::statement::Statement;
+#[cfg(feature = "runtime")]
+use crate::proto::CancelQueryFuture;
+use crate::proto::CancelQueryRawFuture;
 use crate::types::{IsNull, Oid, ToSql, Type};
-use crate::Error;
+use crate::{Config, Error, TlsMode};
+#[cfg(feature = "runtime")]
+use crate::{MakeTlsMode, Socket};
 
 pub struct PendingRequest(Result<(RequestMessages, IdleGuard), Error>);
 
@@ -44,13 +50,25 @@ struct Inner {
     state: Mutex<State>,
     idle: IdleState,
     sender: mpsc::UnboundedSender<Request>,
+    process_id: i32,
+    secret_key: i32,
+    #[cfg_attr(not(feature = "runtime"), allow(dead_code))]
+    config: Config,
+    #[cfg_attr(not(feature = "runtime"), allow(dead_code))]
+    idx: Option<usize>,
 }
 
 #[derive(Clone)]
 pub struct Client(Arc<Inner>);
 
 impl Client {
-    pub fn new(sender: mpsc::UnboundedSender<Request>) -> Client {
+    pub fn new(
+        sender: mpsc::UnboundedSender<Request>,
+        process_id: i32,
+        secret_key: i32,
+        config: Config,
+        idx: Option<usize>,
+    ) -> Client {
         Client(Arc::new(Inner {
             state: Mutex::new(State {
                 types: HashMap::new(),
@@ -60,6 +78,10 @@ impl Client {
             }),
             idle: IdleState::new(),
             sender,
+            process_id,
+            secret_key,
+            config,
+            idx,
         }))
     }
 
@@ -222,6 +244,28 @@ impl Client {
         self.close(b'P', name)
     }
 
+    #[cfg(feature = "runtime")]
+    pub fn cancel_query<T>(&self, make_tls_mode: T) -> CancelQueryFuture<T>
+    where
+        T: MakeTlsMode<Socket>,
+    {
+        CancelQueryFuture::new(
+            make_tls_mode,
+            self.0.idx,
+            self.0.config.clone(),
+            self.0.process_id,
+            self.0.secret_key,
+        )
+    }
+
+    pub fn cancel_query_raw<S, T>(&self, stream: S, tls_mode: T) -> CancelQueryRawFuture<S, T>
+    where
+        S: AsyncRead + AsyncWrite,
+        T: TlsMode<S>,
+    {
+        CancelQueryRawFuture::new(stream, tls_mode, self.0.process_id, self.0.secret_key)
+    }
+
     fn close(&self, ty: u8, name: &str) {
         let mut buf = vec![];
         frontend::close(ty, name, &mut buf).expect("statement name not valid");

tokio-postgres/src/proto/connect_once.rs

@@ -65,7 +65,7 @@ where
 
         transition!(Handshaking {
             target_session_attrs: state.config.0.target_session_attrs,
-            future: HandshakeFuture::new(socket, state.tls_mode, state.config),
+            future: HandshakeFuture::new(socket, state.tls_mode, state.config, Some(state.idx)),
         })
     }