Refactor query routing into its own module (#22)

* Refactor query routing into its own module

* commments; tests; dead code

* error message

* safer startup

* hm

* dont have to be public

* wow

* fix ci

* ok

* nl

* no more silent errors

Author: levkk

src/client.rs

@@ -2,8 +2,6 @@
 /// We are pretending to the server in this scenario,
 /// and this module implements that.
 use bytes::{Buf, BufMut, BytesMut};
-use once_cell::sync::OnceCell;
-use regex::Regex;
 use tokio::io::{AsyncReadExt, BufReader};
 use tokio::net::{
     tcp::{OwnedReadHalf, OwnedWriteHalf},
@@ -17,16 +15,10 @@ use crate::constants::*;
 use crate::errors::Error;
 use crate::messages::*;
 use crate::pool::{ClientServerMap, ConnectionPool};
+use crate::query_router::QueryRouter;
 use crate::server::Server;
-use crate::sharding::Sharder;
 use crate::stats::Reporter;
 
-pub const SHARDING_REGEX: &str = r"SET SHARDING KEY TO '[0-9]+';";
-pub const ROLE_REGEX: &str = r"SET SERVER ROLE TO '(PRIMARY|REPLICA)';";
-
-pub static SHARDING_REGEX_RE: OnceCell<Regex> = OnceCell::new();
-pub static ROLE_REGEX_RE: OnceCell<Regex> = OnceCell::new();
-
 /// The client state. One of these is created per client.
 pub struct Client {
     // The reads are buffered (8K by default).
@@ -199,15 +191,11 @@ impl Client {
             return Ok(Server::cancel(&address, &port, process_id, secret_key).await?);
         }
 
-        // Active shard we're talking to.
-        // The lifetime of this depends on the pool mode:
-        // - if in session mode, this lives until the client disconnects,
-        // - if in transaction mode, this lives for the duration of one transaction.
-        let mut shard: Option<usize> = None;
-
-        // Active database role we want to talk to, e.g. primary or replica.
-        let mut role: Option<Role> = self.default_server_role;
+        let mut query_router = QueryRouter::new(self.default_server_role, pool.shards());
 
+        // Our custom protocol loop.
+        // We expect the client to either start a transaction with regular queries
+        // or issue commands for our sharding and server selection protocols.
         loop {
             // Read a complete message from the client, which normally would be
             // either a `Q` (query) or `P` (prepare, extended protocol).
@@ -218,32 +206,31 @@ impl Client {
 
             // Parse for special select shard command.
             // SET SHARDING KEY TO 'bigint';
-            match self.select_shard(message.clone(), pool.shards()) {
-                Some(s) => {
-                    custom_protocol_response_ok(&mut self.write, "SET SHARDING KEY").await?;
-                    shard = Some(s);
-                    continue;
-                }
-                None => (),
-            };
+            if query_router.select_shard(message.clone()) {
+                custom_protocol_response_ok(
+                    &mut self.write,
+                    &format!("SET SHARD TO {}", query_router.shard()),
+                )
+                .await?;
+                continue;
+            }
 
             // Parse for special server role selection command.
             // SET SERVER ROLE TO '(primary|replica)';
-            match self.select_role(message.clone()) {
-                Some(r) => {
-                    custom_protocol_response_ok(&mut self.write, "SET SERVER ROLE").await?;
-                    role = Some(r);
-                    continue;
-                }
-                None => (),
-            };
+            if query_router.select_role(message.clone()) {
+                custom_protocol_response_ok(&mut self.write, "SET SERVER ROLE").await?;
+                continue;
+            }
 
-            // Grab a server from the pool.
-            let connection = match pool.get(shard, role).await {
+            // Grab a server from the pool: the client issued a regular query.
+            let connection = match pool.get(query_router.shard(), query_router.role()).await {
                 Ok(conn) => conn,
                 Err(err) => {
                     println!(">> Could not get connection from pool: {:?}", err);
-                    return Err(err);
+                    error_response(&mut self.write, "could not get connection from the pool")
+                        .await?;
+                    query_router.reset();
+                    continue;
                 }
             };
 
@@ -264,11 +251,8 @@ impl Client {
                         Err(err) => {
                             // Client disconnected without warning.
                             if server.in_transaction() {
-                                // TODO: this is what PgBouncer does
-                                // which leads to connection thrashing.
-                                //
-                                // I think we could issue a ROLLBACK here instead.
-                                // server.mark_bad();
+                                // Client left dirty server. Clean up and proceed
+                                // without thrashing this connection.
                                 server.query("ROLLBACK; DISCARD ALL;").await?;
                             }
 
@@ -328,8 +312,7 @@ impl Client {
                                 // Report this client as idle.
                                 self.stats.client_idle();
 
-                                shard = None;
-                                role = self.default_server_role;
+                                query_router.reset();
 
                                 break;
                             }
@@ -414,8 +397,7 @@ impl Client {
                             if self.transaction_mode {
                                 self.stats.client_idle();
 
-                                shard = None;
-                                role = self.default_server_role;
+                                query_router.reset();
 
                                 break;
                             }
@@ -450,8 +432,7 @@ impl Client {
                             self.stats.transaction();
 
                             if self.transaction_mode {
-                                shard = None;
-                                role = self.default_server_role;
+                                query_router.reset();
 
                                 break;
                             }
@@ -476,77 +457,4 @@ impl Client {
         let mut guard = self.client_server_map.lock().unwrap();
         guard.remove(&(self.process_id, self.secret_key));
     }
-
-    /// Determine if the query is part of our special syntax, extract
-    /// the shard key, and return the shard to query based on Postgres'
-    /// PARTITION BY HASH function.
-    fn select_shard(&self, mut buf: BytesMut, shards: usize) -> Option<usize> {
-        let code = buf.get_u8() as char;
-
-        // Only supporting simpe protocol here, so
-        // one would have to execute something like this:
-        // psql -c "SET SHARDING KEY TO '1234'"
-        // after sanitizing the value manually, which can be just done with an
-        // int parser, e.g. `let key = "1234".parse::<i64>().unwrap()`.
-        match code {
-            'Q' => (),
-            _ => return None,
-        };
-
-        let len = buf.get_i32();
-        let query = String::from_utf8_lossy(&buf[..len as usize - 4 - 1]).to_ascii_uppercase(); // Don't read the ternminating null
-
-        let rgx = match SHARDING_REGEX_RE.get() {
-            Some(r) => r,
-            None => return None,
-        };
-
-        if rgx.is_match(&query) {
-            let shard = query.split("'").collect::<Vec<&str>>()[1];
-
-            match shard.parse::<i64>() {
-                Ok(shard) => {
-                    let sharder = Sharder::new(shards);
-                    Some(sharder.pg_bigint_hash(shard))
-                }
-
-                Err(_) => None,
-            }
-        } else {
-            None
-        }
-    }
-
-    // Pick a primary or a replica from the pool.
-    fn select_role(&self, mut buf: BytesMut) -> Option<Role> {
-        let code = buf.get_u8() as char;
-
-        // Same story as select_shard() above.
-        match code {
-            'Q' => (),
-            _ => return None,
-        };
-
-        let len = buf.get_i32();
-        let query = String::from_utf8_lossy(&buf[..len as usize - 4 - 1]).to_ascii_uppercase();
-
-        let rgx = match ROLE_REGEX_RE.get() {
-            Some(r) => r,
-            None => return None,
-        };
-
-        // Copy / paste from above. If we get one more of these use cases,
-        // it'll be time to abstract :).
-        if rgx.is_match(&query) {
-            let role = query.split("'").collect::<Vec<&str>>()[1];
-
-            match role {
-                "PRIMARY" => Some(Role::Primary),
-                "REPLICA" => Some(Role::Replica),
-                _ => return None,
-            }
-        } else {
-            None
-        }
-    }
 }

src/main.rs

@@ -25,7 +25,6 @@ extern crate statsd;
 extern crate tokio;
 extern crate toml;
 
-use regex::Regex;
 use tokio::net::TcpListener;
 use tokio::signal;
 
@@ -39,6 +38,7 @@ mod constants;
 mod errors;
 mod messages;
 mod pool;
+mod query_router;
 mod server;
 mod sharding;
 mod stats;
@@ -54,12 +54,11 @@ use stats::{Collector, Reporter};
 async fn main() {
     println!("> Welcome to PgCat! Meow.");
 
-    client::SHARDING_REGEX_RE
-        .set(Regex::new(client::SHARDING_REGEX).unwrap())
-        .unwrap();
-    client::ROLE_REGEX_RE
-        .set(Regex::new(client::ROLE_REGEX).unwrap())
-        .unwrap();
+    // Prepare regexes
+    if !query_router::QueryRouter::setup() {
+        println!("> Could not setup query router.");
+        return;
+    }
 
     let config = match config::parse("pgcat.toml").await {
         Ok(config) => config,

src/messages.rs

@@ -185,6 +185,50 @@ pub async fn custom_protocol_response_ok(
     write_all_half(stream, res).await
 }
 
+/// Send a custom error message to the client.
+/// Tell the client we are ready for the next query and no rollback is necessary.
+/// Docs on error codes: https://www.postgresql.org/docs/12/errcodes-appendix.html
+pub async fn error_response(stream: &mut OwnedWriteHalf, message: &str) -> Result<(), Error> {
+    let mut error = BytesMut::new();
+
+    // Error level
+    error.put_u8(b'S');
+    error.put_slice(&b"FATAL\0"[..]);
+
+    // Error level (non-translatable)
+    error.put_u8(b'V');
+    error.put_slice(&b"FATAL\0"[..]);
+
+    // Error code: not sure how much this matters.
+    error.put_u8(b'C');
+    error.put_slice(&b"58000\0"[..]); // system_error, see Appendix A.
+
+    // The short error message.
+    error.put_u8(b'M');
+    error.put_slice(&format!("{}\0", message).as_bytes());
+
+    // No more fields follow.
+    error.put_u8(0);
+
+    // Ready for query, no rollback needed (I = idle).
+    let mut ready_for_query = BytesMut::new();
+
+    ready_for_query.put_u8(b'Z');
+    ready_for_query.put_i32(5);
+    ready_for_query.put_u8(b'I');
+
+    // Compose the two message reply.
+    let mut res = BytesMut::with_capacity(error.len() + ready_for_query.len() + 5);
+
+    res.put_u8(b'E');
+    res.put_i32(error.len() as i32 + 4);
+
+    res.put(error);
+    res.put(ready_for_query);
+
+    Ok(write_all_half(stream, res).await?)
+}
+
 /// Write all data in the buffer to the TcpStream.
 pub async fn write_all(stream: &mut TcpStream, buf: BytesMut) -> Result<(), Error> {
     match stream.write_all(&buf).await {

src/pool.rs

@@ -121,7 +121,7 @@ impl ConnectionPool {
 
         for shard in 0..self.shards() {
             for _ in 0..self.replicas(shard) {
-                let connection = match self.get(Some(shard), None).await {
+                let connection = match self.get(shard, None).await {
                     Ok(conn) => conn,
                     Err(err) => {
                         println!("> Shard {} down or misconfigured.", shard);
@@ -149,18 +149,13 @@ impl ConnectionPool {
     /// Get a connection from the pool.
     pub async fn get(
         &mut self,
-        shard: Option<usize>,
+        shard: usize,
         role: Option<Role>,
     ) -> Result<(PooledConnection<'_, ServerPool>, Address), Error> {
         // Set this to false to gain ~3-4% speed.
         let with_health_check = true;
         let now = Instant::now();
 
-        let shard = match shard {
-            Some(shard) => shard,
-            None => 0, // TODO: pick a shard at random
-        };
-
         // We are waiting for a server now.
         self.stats.client_waiting();
 
@@ -208,11 +203,8 @@ impl ConnectionPool {
             // as per request.
             match role {
                 Some(role) => {
-                    // If the client wants a specific role,
-                    // we'll do our best to pick it, but if we only
-                    // have one server in the cluster, it's probably only a primary
-                    // (or only a replica), so the client will just get what we have.
-                    if address.role != role && addresses.len() > 1 {
+                    // Find the specific role the client wants in the pool.
+                    if address.role != role {
                         continue;
                     }
                 }