connection pool

src/client.rs

@@ -1,4 +1,4 @@
-use tokio::io::{AsyncReadExt, BufReader};
+use tokio::io::{AsyncReadExt, BufReader, Interest};
 use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
 /// PostgreSQL client (frontend).
 /// We are pretending to be the backend.
@@ -8,7 +8,9 @@ use bytes::{Buf, BufMut, BytesMut};
 
 use crate::errors::Error;
 use crate::messages::*;
-use crate::server::Server;
+
+use bb8::Pool;
+use crate::pool::ServerPool;
 
 pub struct Client {
     read: BufReader<OwnedReadHalf>,
@@ -69,54 +71,77 @@ impl Client {
         }
     }
 
-    pub async fn handle(&mut self, mut server: Server) -> Result<(), Error> {
+    pub async fn handle(&mut self, pool: Pool<ServerPool>) -> Result<(), Error> {
+        
         loop {
-            let mut message = read_message(&mut self.read).await?;
-            let original = message.clone(); // To be forwarded to the server
-            let code = message.get_u8() as char;
-            let _len = message.get_i32() as usize;
+            // Only grab a connection once we have some traffic on the socket
+            // TODO: this is not the most optimal way to share servers.
+            match self.read.get_ref().ready(Interest::READABLE).await {
+                Ok(_) => (),
+                Err(_) => return Err(Error::ClientDisconnected),
+            };
 
-            match code {
-                'Q' => {
-                    server.send(original).await?;
-                    let response = server.recv().await?;
-                    write_all_half(&mut self.write, response).await?;
-                }
+            let mut proxy = pool.get().await.unwrap();
+            let server = &mut *proxy;
 
-                'X' => {
-                    // Client closing
-                    return Ok(());
-                }
+            loop {
+                let mut message = read_message(&mut self.read).await?;
+                let original = message.clone(); // To be forwarded to the server
+                let code = message.get_u8() as char;
+                let _len = message.get_i32() as usize;
 
-                'P' => {
-                    // Extended protocol, let's buffer most of it
-                    self.buffer.put(&original[..]);
-                }
+                match code {
+                    'Q' => {
+                        server.send(original).await?;
+                        let response = server.recv().await?;
+                        write_all_half(&mut self.write, response).await?;
 
-                'B' => {
-                    self.buffer.put(&original[..]);
-                }
+                         // Release server
+                        if !server.in_transaction() {
+                            break;
+                        }
+                    }
 
-                'D' => {
-                    self.buffer.put(&original[..]);
-                }
+                    'X' => {
+                        // Client closing
+                        return Ok(());
+                    }
 
-                'E' => {
-                    self.buffer.put(&original[..]);
-                }
+                    'P' => {
+                        // Extended protocol, let's buffer most of it
+                        self.buffer.put(&original[..]);
+                    }
 
-                'S' => {
-                    // Extended protocol, client requests sync
-                    self.buffer.put(&original[..]);
-                    server.send(self.buffer.clone()).await?;
-                    self.buffer.clear();
+                    'B' => {
+                        self.buffer.put(&original[..]);
+                    }
 
-                    let response = server.recv().await?;
-                    write_all_half(&mut self.write, response).await?;
-                }
+                    'D' => {
+                        self.buffer.put(&original[..]);
+                    }
 
-                _ => {
-                    println!(">>> Unexpected code: {}", code);
+                    'E' => {
+                        self.buffer.put(&original[..]);
+                    }
+
+                    'S' => {
+                        // Extended protocol, client requests sync
+                        self.buffer.put(&original[..]);
+                        server.send(self.buffer.clone()).await?;
+                        self.buffer.clear();
+
+                        let response = server.recv().await?;
+                        write_all_half(&mut self.write, response).await?;
+
+                        // Release server
+                        if !server.in_transaction() {
+                            break;
+                        }
+                    }
+
+                    _ => {
+                        println!(">>> Unexpected code: {}", code);
+                    }
                 }
             }
         }

src/errors.rs

@@ -1,7 +1,7 @@
 #[derive(Debug, PartialEq)]
 pub enum Error {
     SocketError,
-    // ClientDisconnected,
+    ClientDisconnected,
     ClientBadStartup,
     ProtocolSyncError,
     ServerError,

src/main.rs

@@ -1,13 +1,17 @@
 extern crate bytes;
 extern crate md5;
 extern crate tokio;
+extern crate async_trait;
+extern crate bb8;
 
 use tokio::net::TcpListener;
+use bb8::Pool;
 
 mod client;
 mod errors;
 mod messages;
 mod server;
+mod pool;
 
 #[tokio::main]
 async fn main() {
@@ -21,7 +25,12 @@ async fn main() {
         }
     };
 
+    let manager = pool::ServerPool::new("127.0.0.1", "5432", "lev", "lev", "lev");
+    let pool = Pool::builder().max_size(15).build(manager).await.unwrap();
+
     loop {
+        let pool = pool.clone();
+
         let (socket, addr) = match listener.accept().await {
             Ok((socket, addr)) => (socket, addr),
             Err(err) => {
@@ -33,19 +42,14 @@ 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 {
                 Ok(mut client) => {
                     println!(">> Client {:?} authenticated successfully!", addr);
-                    let server =
-                        match server::Server::startup("127.0.0.1", "5432", "lev", "lev", "lev")
-                            .await
-                        {
-                            Ok(server) => server,
-                            Err(_) => return,
-                        };
 
-                    match client.handle(server).await {
+                    match client.handle(pool).await {
                         Ok(()) => {
                             println!(">> Client {:?} disconnected.", addr);
                         }

src/pool.rs

@@ -0,0 +1,48 @@
+
+use async_trait::async_trait;
+use bb8::{ManageConnection, PooledConnection};
+
+use crate::server::Server;
+use crate::errors::Error;
+
+pub struct ServerPool {
+    host: String,
+    port: String,
+    user: String,
+    password: String,
+    database: String,
+}
+
+impl ServerPool {
+    pub fn new(host: &str, port: &str, user: &str, password: &str, database: &str) -> ServerPool {
+        ServerPool {
+            host: host.to_string(),
+            port: port.to_string(),
+            user: user.to_string(),
+            password: password.to_string(),
+            database: database.to_string(),
+        }
+    }
+}
+
+#[async_trait]
+impl ManageConnection for ServerPool {
+    type Connection = Server;
+    type Error = Error;
+
+   /// 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?)
+    }
+
+    /// Determines if the connection is still connected to the database.
+    async fn is_valid(&self, _conn: &mut PooledConnection<'_, Self>) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
+    /// Synchronously determine if the connection is no longer usable, if possible.
+    fn has_broken(&self, _conn: &mut Self::Connection) -> bool {
+        false
+    }
+}

src/server.rs

@@ -10,6 +10,7 @@ pub struct Server {
     read: BufReader<OwnedReadHalf>,
     write: OwnedWriteHalf,
     buffer: BytesMut,
+    in_transaction: bool,
 }
 
 impl Server {
@@ -112,6 +113,7 @@ impl Server {
                         read: BufReader::new(read),
                         write: write,
                         buffer: BytesMut::with_capacity(8196),
+                        in_transaction: false,
                     });
                 }
 
@@ -135,12 +137,29 @@ impl Server {
             self.buffer.put(&message[..]);
 
             let code = message.get_u8() as char;
+            let _len = message.get_i32();
 
             match code {
                 '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;
+                        },
+
+                        _ => {
+                            self.in_transaction = false;
+                        },
+                    };
+
                     break;
-                }
+                },
 
                 _ => {
                     // Keep buffering,
@@ -153,4 +172,8 @@ impl Server {
 
         Ok(bytes)
     }
+
+    pub fn in_transaction(&self) -> bool {
+        self.in_transaction
+    }
 }