name servers & dont leave open servers with bad state

src/client.rs

@@ -9,13 +9,15 @@ use bytes::{Buf, BufMut, BytesMut};
 use crate::errors::Error;
 use crate::messages::*;
 
-use bb8::Pool;
 use crate::pool::ServerPool;
+use bb8::Pool;
+use rand::{distributions::Alphanumeric, Rng};
 
 pub struct Client {
     read: BufReader<OwnedReadHalf>,
     write: OwnedWriteHalf,
     buffer: BytesMut,
+    name: String,
 }
 
 impl Client {
@@ -57,10 +59,17 @@ impl Client {
 
                     let (read, write) = stream.into_split();
 
+                    let name: String = rand::thread_rng()
+                        .sample_iter(&Alphanumeric)
+                        .take(7)
+                        .map(char::from)
+                        .collect();
+
                     return Ok(Client {
                         read: BufReader::new(read),
                         write: write,
                         buffer: BytesMut::with_capacity(8196),
+                        name: name,
                     });
                 }
 
@@ -72,7 +81,6 @@ impl Client {
     }
 
     pub async fn handle(&mut self, pool: Pool<ServerPool>) -> Result<(), Error> {
-        
         loop {
             // Only grab a connection once we have some traffic on the socket
             // TODO: this is not the most optimal way to share servers.
@@ -84,8 +92,20 @@ impl Client {
             let mut proxy = pool.get().await.unwrap();
             let server = &mut *proxy;
 
+            server.set_name(&self.name).await?;
+
             loop {
-                let mut message = read_message(&mut self.read).await?;
+                let mut message = match read_message(&mut self.read).await {
+                    Ok(message) => message,
+                    Err(err) => {
+                        if server.in_transaction() {
+                            server.mark_bad();
+                        }
+
+                        return Err(err);
+                    }
+                };
+
                 let original = message.clone(); // To be forwarded to the server
                 let code = message.get_u8() as char;
                 let _len = message.get_i32() as usize;
@@ -94,10 +114,17 @@ impl Client {
                     'Q' => {
                         server.send(original).await?;
                         let response = server.recv().await?;
-                        write_all_half(&mut self.write, response).await?;
+                        match write_all_half(&mut self.write, response).await {
+                            Ok(_) => (),
+                            Err(err) => {
+                                server.mark_bad();
+                                return Err(err);
+                            }
+                        };
 
-                         // Release server
+                        // Release server
                         if !server.in_transaction() {
+                            drop(server);
                             break;
                         }
                     }
@@ -131,10 +158,17 @@ impl Client {
                         self.buffer.clear();
 
                         let response = server.recv().await?;
-                        write_all_half(&mut self.write, response).await?;
+                        match write_all_half(&mut self.write, response).await {
+                            Ok(_) => (),
+                            Err(err) => {
+                                server.mark_bad();
+                                return Err(err);
+                            }
+                        };
 
                         // Release server
                         if !server.in_transaction() {
+                            drop(server);
                             break;
                         }
                     }

src/main.rs

@@ -1,17 +1,17 @@
+extern crate async_trait;
+extern crate bb8;
 extern crate bytes;
 extern crate md5;
 extern crate tokio;
-extern crate async_trait;
-extern crate bb8;
 
-use tokio::net::TcpListener;
 use bb8::Pool;
+use tokio::net::TcpListener;
 
 mod client;
 mod errors;
 mod messages;
-mod server;
 mod pool;
+mod server;
 
 #[tokio::main]
 async fn main() {
@@ -42,7 +42,7 @@ async fn main() {
         // Client goes to another thread, bye.
         tokio::task::spawn(async move {
             println!(">> Client {:?} connected.", addr);
-            
+
             let pool = pool.clone();
 
             match client::Client::startup(socket).await {

src/pool.rs

@@ -1,9 +1,8 @@
-
 use async_trait::async_trait;
 use bb8::{ManageConnection, PooledConnection};
 
-use crate::server::Server;
 use crate::errors::Error;
+use crate::server::Server;
 
 pub struct ServerPool {
     host: String,
@@ -30,10 +29,17 @@ impl ManageConnection for ServerPool {
     type Connection = Server;
     type Error = Error;
 
-   /// Attempts to create a new connection.
+    /// Attempts to create a new connection.
     async fn connect(&self) -> Result<Self::Connection, Self::Error> {
         println!(">> Getting connetion from pool");
-        Ok(Server::startup(&self.host, &self.port, &self.user, &self.password, &self.database).await?)
+        Ok(Server::startup(
+            &self.host,
+            &self.port,
+            &self.user,
+            &self.password,
+            &self.database,
+        )
+        .await?)
     }
 
     /// Determines if the connection is still connected to the database.
@@ -42,7 +48,7 @@ impl ManageConnection for ServerPool {
     }
 
     /// Synchronously determine if the connection is no longer usable, if possible.
-    fn has_broken(&self, _conn: &mut Self::Connection) -> bool {
-        false
+    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
+        conn.is_bad()
     }
-}
+}

src/server.rs

@@ -11,6 +11,7 @@ pub struct Server {
     write: OwnedWriteHalf,
     buffer: BytesMut,
     in_transaction: bool,
+    bad: bool,
 }
 
 impl Server {
@@ -114,6 +115,7 @@ impl Server {
                         write: write,
                         buffer: BytesMut::with_capacity(8196),
                         in_transaction: false,
+                        bad: false,
                     });
                 }
 
@@ -126,12 +128,26 @@ impl Server {
     }
 
     pub async fn send(&mut self, messages: BytesMut) -> Result<(), Error> {
-        Ok(write_all_half(&mut self.write, messages).await?)
+        match write_all_half(&mut self.write, messages).await {
+            Ok(_) => Ok(()),
+            Err(err) => {
+                println!(">> Terminating server because of: {:?}", err);
+                self.bad = true;
+                Err(err)
+            }
+        }
     }
 
     pub async fn recv(&mut self) -> Result<BytesMut, Error> {
         loop {
-            let mut message = read_message(&mut self.read).await?;
+            let mut message = match read_message(&mut self.read).await {
+                Ok(message) => message,
+                Err(err) => {
+                    println!(">> Terminating server because of: {:?}", err);
+                    self.bad = true;
+                    return Err(err);
+                }
+            };
 
             // Buffer the message we'll forward to the client in a bit.
             self.buffer.put(&message[..]);
@@ -143,23 +159,30 @@ impl Server {
                 'Z' => {
                     // Ready for query, time to forward buffer to client.
                     let transaction_state = message.get_u8() as char;
-                    
+
                     match transaction_state {
                         'T' => {
                             self.in_transaction = true;
-                        },
+                        }
 
                         'I' => {
                             self.in_transaction = false;
-                        },
+                        }
+
+                        // Error client didn't clean up!
+                        // We shuold drop this server
+                        'E' => {
+                            self.bad = true;
+                        }
 
                         _ => {
-                            self.in_transaction = false;
-                        },
+                            self.bad = true;
+                            return Err(Error::ProtocolSyncError);
+                        }
                     };
 
                     break;
-                },
+                }
 
                 _ => {
                     // Keep buffering,
@@ -176,4 +199,31 @@ impl Server {
     pub fn in_transaction(&self) -> bool {
         self.in_transaction
     }
+
+    pub fn is_bad(&self) -> bool {
+        self.bad
+    }
+
+    pub fn mark_bad(&mut self) {
+        println!(">> Server marked bad");
+        self.bad = true;
+    }
+
+    pub async fn set_name(&mut self, name: &str) -> Result<(), Error> {
+        let mut query = BytesMut::from(&format!("SET application_name = {}", name).as_bytes()[..]);
+        query.put_u8(0);
+
+        let len = query.len() as i32 + 4;
+
+        let mut msg = BytesMut::with_capacity(len as usize + 1);
+
+        msg.put_u8(b'Q');
+        msg.put_i32(len);
+        msg.put_slice(&query[..]);
+
+        self.send(msg).await?;
+        let _ = self.recv().await?;
+
+        Ok(())
+    }
 }