From fbc37777698fd5d0dae50cc546352af0ae0d1372 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Fri, 4 Feb 2022 16:01:35 -0800 Subject: [PATCH] "query cancellation" --- src/client.rs | 36 +++++++++++++++++++++++++++--------- src/main.rs | 17 ++++++++++++++--- src/pool.rs | 13 ++++++++++++- src/server.rs | 34 ++++++++++++++++++++++++++++++++++ 4 files changed, 87 insertions(+), 13 deletions(-) diff --git a/src/client.rs b/src/client.rs index 5f0a2dc..90b7b86 100644 --- a/src/client.rs +++ b/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 { + pub async fn startup( + mut stream: TcpStream, + client_server_map: ClientServerMap, + ) -> Result { 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) -> 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)); + } } diff --git a/src/main.rs b/src/main.rs index 58aa2b2..b53dee4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,7 +16,7 @@ mod messages; mod pool; mod server; -type ClientServerMap = Arc>>; +type ClientServerMap = Arc>>; #[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(); } } } diff --git a/src/pool.rs b/src/pool.rs index 2683e17..65b37c6 100644 --- a/src/pool.rs +++ b/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?) } diff --git a/src/server.rs b/src/server.rs index 2c11efc..2c384ad 100644 --- a/src/server.rs +++ b/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 { 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`.