"query cancellation"

src/client.rs

@@ -11,6 +11,8 @@ use tokio::net::TcpStream;
 use crate::errors::Error;
 use crate::messages::*;
 use crate::pool::ServerPool;
+use crate::server::Server;
+use crate::ClientServerMap;
 
 /// The client state.
 pub struct Client {
@@ -21,13 +23,17 @@ pub struct Client {
     cancel_mode: bool,
     process_id: i32,
     secret_key: i32,
+    client_server_map: ClientServerMap,
 }
 
 impl Client {
     /// Given a TCP socket, trick the client into thinking we are
     /// the Postgres server. Perform the authentication and place
     /// the client in query-ready mode.
-    pub async fn startup(mut stream: TcpStream) -> Result<Client, Error> {
+    pub async fn startup(
+        mut stream: TcpStream,
+        client_server_map: ClientServerMap,
+    ) -> Result<Client, Error> {
         loop {
             // Could be StartupMessage or SSLRequest
             // which makes this variable length.
@@ -86,6 +92,7 @@ impl Client {
                         cancel_mode: false,
                         process_id: process_id,
                         secret_key: secret_key,
+                        client_server_map: client_server_map,
                     });
                 }
 
@@ -104,6 +111,7 @@ impl Client {
                         cancel_mode: true,
                         process_id: process_id,
                         secret_key: secret_key,
+                        client_server_map: client_server_map,
                     });
                 }
 
@@ -118,12 +126,16 @@ impl Client {
     pub async fn handle(&mut self, pool: Pool<ServerPool>) -> Result<(), Error> {
         // Special: cancelling existing running query
         if self.cancel_mode {
-            // TODO: Implement this
-            println!(
-                ">> Query cancellation requested: {}, {}",
-                self.process_id, self.secret_key
-            );
-            return Ok(());
+            let (process_id, secret_key) = {
+                let guard = self.client_server_map.lock().unwrap();
+                match guard.get(&(self.process_id, self.secret_key)) {
+                    Some((process_id, secret_key)) => (process_id.clone(), secret_key.clone()),
+                    None => return Ok(()),
+                }
+            };
+
+            // TODO: pass actual server host and port somewhere.
+            return Ok(Server::cancel("127.0.0.1", "5432", process_id, secret_key).await?);
         }
 
         loop {
@@ -138,6 +150,7 @@ impl Client {
             let server = &mut *proxy;
 
             server.set_name(&self.name).await?;
+            server.claim(self.process_id, self.secret_key);
 
             loop {
                 let mut message = match read_message(&mut self.read).await {
@@ -180,7 +193,6 @@ impl Client {
 
                         // Release server
                         if !server.in_transaction() {
-                            drop(server);
                             break;
                         }
                     }
@@ -234,7 +246,6 @@ impl Client {
 
                         // Release server
                         if !server.in_transaction() {
-                            drop(server);
                             break;
                         }
                     }
@@ -264,6 +275,13 @@ impl Client {
                     }
                 }
             }
+
+            self.release();
         }
     }
+
+    pub fn release(&mut self) {
+        let mut guard = self.client_server_map.lock().unwrap();
+        guard.remove(&(self.process_id, self.secret_key));
+    }
 }

src/main.rs

@@ -16,7 +16,7 @@ mod messages;
 mod pool;
 mod server;
 
-type ClientServerMap = Arc<Mutex<HashMap<i32, i32>>>;
+type ClientServerMap = Arc<Mutex<HashMap<(i32, i32), (i32, i32)>>>;
 
 #[tokio::main]
 async fn main() {
@@ -30,11 +30,20 @@ async fn main() {
         }
     };
 
-    let manager = pool::ServerPool::new("127.0.0.1", "5432", "lev", "lev", "lev");
+    let client_server_map: ClientServerMap = Arc::new(Mutex::new(HashMap::new()));
+    let manager = pool::ServerPool::new(
+        "127.0.0.1",
+        "5432",
+        "lev",
+        "lev",
+        "lev",
+        client_server_map.clone(),
+    );
     let pool = Pool::builder().max_size(15).build(manager).await.unwrap();
 
     loop {
         let pool = pool.clone();
+        let client_server_map = client_server_map.clone();
 
         let (socket, addr) = match listener.accept().await {
             Ok((socket, addr)) => (socket, addr),
@@ -49,8 +58,9 @@ async fn main() {
             println!(">> Client {:?} connected.", addr);
 
             let pool = pool.clone();
+            let client_server_map = client_server_map.clone();
 
-            match client::Client::startup(socket).await {
+            match client::Client::startup(socket, client_server_map).await {
                 Ok(mut client) => {
                     println!(">> Client {:?} authenticated successfully!", addr);
 
@@ -61,6 +71,7 @@ async fn main() {
 
                         Err(err) => {
                             println!(">> Client disconnected with error: {:?}", err);
+                            client.release();
                         }
                     }
                 }

src/pool.rs

@@ -3,6 +3,7 @@ use bb8::{ManageConnection, PooledConnection};
 
 use crate::errors::Error;
 use crate::server::Server;
+use crate::ClientServerMap;
 
 pub struct ServerPool {
     host: String,
@@ -10,16 +11,25 @@ pub struct ServerPool {
     user: String,
     password: String,
     database: String,
+    client_server_map: ClientServerMap,
 }
 
 impl ServerPool {
-    pub fn new(host: &str, port: &str, user: &str, password: &str, database: &str) -> ServerPool {
+    pub fn new(
+        host: &str,
+        port: &str,
+        user: &str,
+        password: &str,
+        database: &str,
+        client_server_map: ClientServerMap,
+    ) -> ServerPool {
         ServerPool {
             host: host.to_string(),
             port: port.to_string(),
             user: user.to_string(),
             password: password.to_string(),
             database: database.to_string(),
+            client_server_map: client_server_map,
         }
     }
 }
@@ -42,6 +52,7 @@ impl ManageConnection for ServerPool {
             &self.user,
             &self.password,
             &self.database,
+            self.client_server_map.clone(),
         )
         .await?)
     }

src/server.rs

@@ -10,6 +10,7 @@ use tokio::net::TcpStream;
 
 use crate::errors::Error;
 use crate::messages::*;
+use crate::ClientServerMap;
 
 /// Server state.
 pub struct Server {
@@ -22,6 +23,7 @@ pub struct Server {
     in_transaction: bool,
     data_available: bool,
     bad: bool,
+    client_server_map: ClientServerMap,
 }
 
 impl Server {
@@ -33,6 +35,7 @@ impl Server {
         user: &str,
         password: &str,
         database: &str,
+        client_server_map: ClientServerMap,
     ) -> Result<Server, Error> {
         let mut stream = match TcpStream::connect(&format!("{}:{}", host, port)).await {
             Ok(stream) => stream,
@@ -144,6 +147,7 @@ impl Server {
                         in_transaction: false,
                         data_available: false,
                         bad: false,
+                        client_server_map: client_server_map,
                     });
                 }
 
@@ -155,6 +159,31 @@ impl Server {
         }
     }
 
+    /// Issue a cancellation request to the server.
+    /// Uses a separate connection that's not part of the connection pool.
+    pub async fn cancel(
+        host: &str,
+        port: &str,
+        process_id: i32,
+        secret_key: i32,
+    ) -> Result<(), Error> {
+        let mut stream = match TcpStream::connect(&format!("{}:{}", host, port)).await {
+            Ok(stream) => stream,
+            Err(err) => {
+                println!(">> Could not connect to server: {}", err);
+                return Err(Error::SocketError);
+            }
+        };
+
+        let mut bytes = BytesMut::with_capacity(16);
+        bytes.put_i32(16);
+        bytes.put_i32(80877102);
+        bytes.put_i32(process_id);
+        bytes.put_i32(secret_key);
+
+        Ok(write_all(&mut stream, bytes).await?)
+    }
+
     /// Send data to the server from the client.
     pub async fn send(&mut self, messages: BytesMut) -> Result<(), Error> {
         match write_all_half(&mut self.write, messages).await {
@@ -276,6 +305,11 @@ impl Server {
         self.bad = true;
     }
 
+    pub fn claim(&mut self, process_id: i32, secret_key: i32) {
+        let mut guard = self.client_server_map.lock().unwrap();
+        guard.insert((process_id, secret_key), (self.backend_id, self.secret_key));
+    }
+
     /// Execute an arbitrary query against the server.
     /// It will use the Simple query protocol.
     /// Result will not be returned, so this is useful for things like `SET` or `ROLLBACK`.