Merge pull request #5 from levkk/levkk-replica-primary

#4 Primary/replica selection

Author: levkk

README.md

@@ -34,8 +34,9 @@ See [sharding README](./tests/sharding/README.md) for sharding logic testing.
 3. `COPY` protocol support.
 4. Query cancellation.
 5. Round-robin load balancing of replicas.
-6. Banlist & failover
+6. Banlist & failover.
 7. Sharding!
+8. Explicit query routing to primary or replicas.
 
 ### Session mode
 Each client owns its own server for the duration of the session. Commands like `SET` are allowed.
@@ -56,7 +57,8 @@ this might be relevant given than this is a transactional pooler but if you're n
 ### Round-robin load balancing
 This is the novel part. PgBouncer doesn't support it and suggests we use DNS or a TCP proxy instead.
 We prefer to have everything as part of one package; arguably, it's easier to understand and optimize.
-This pooler will round-robin between multiple replicas keeping load reasonably even.
+This pooler will round-robin between multiple replicas keeping load reasonably even. If the primary is in
+the pool as well, it'll be treated as a replica for read-only queries.
 
 ### Banlist & failover
 This is where it gets even more interesting. If we fail to connect to one of the replicas or it fails a health check,
@@ -82,6 +84,19 @@ SET SHARDING KEY TO '1234';
 
 This sharding key will be hashed and the pooler will select a shard to use for the next transaction. If the pooler is in session mode, this sharding key has to be set as the first query on startup & cannot be changed until the client re-connects.
 
+### Explicit read/write query routing
+
+If you want to have the primary and replicas in the same pooler, you'd probably want to
+route queries explicitely to the primary or replicas, depending if they are reads or writes (e.g `SELECT`s or `INSERT`/`UPDATE`, etc). To help with this, we introduce some more custom syntax:
+
+```sql
+SET SERVER ROLE TO 'primary';
+SET SERVER ROLE TO 'replica';
+```
+
+After executing this, the next transaction will be routed to the primary or replica respectively. By default, all queries will be load-balanced between all servers, so if the client wants to write or talk to the primary, they have to explicitely select it using the syntax above.
+
+
 
 ## Missing
 

pgcat.toml

@@ -43,26 +43,29 @@ password = "sharding_user"
 # Shard 0
 [shards.0]
 
-# [ host, port ]
+# [ host, port, role ]
 servers = [
-    [ "127.0.0.1", 5432 ],
-    [ "localhost", 5432 ],
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ],
+    # [ "127.0.1.1", 5432, "replica" ],
 ]
 # Database name (e.g. "postgres")
 database = "shard0"
 
 [shards.1]
-# [ host, port ]
+# [ host, port, role ]
 servers = [
-    [ "127.0.0.1", 5432 ],
-    [ "localhost", 5432 ],
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ],
+    # [ "127.0.1.1", 5432, "replica" ],
 ]
 database = "shard1"
 
 [shards.2]
-# [ host, port ]
+# [ host, port, role ]
 servers = [
-    [ "127.0.0.1", 5432 ],
-    [ "localhost", 5432 ],
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5432, "replica" ],
+    # [ "127.0.1.1", 5432, "replica" ],
 ]
-database = "shard2"
+database = "shard2"

src/client.rs

@@ -7,13 +7,15 @@ use tokio::io::{AsyncReadExt, BufReader};
 use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
 use tokio::net::TcpStream;
 
+use crate::config::Role;
 use crate::errors::Error;
 use crate::messages::*;
 use crate::pool::{ClientServerMap, ConnectionPool};
 use crate::server::Server;
 use crate::sharding::Sharder;
 
 const SHARDING_REGEX: &str = r"SET SHARDING KEY TO '[0-9]+';";
+const ROLE_REGEX: &str = r"SET SERVER ROLE TO '(PRIMARY|REPLICA)';";
 
 /// The client state. One of these is created per client.
 pub struct Client {
@@ -45,6 +47,9 @@ pub struct Client {
 
     // sharding regex
     sharding_regex: Regex,
+
+    // role detection regex
+    role_regex: Regex,
 }
 
 impl Client {
@@ -57,6 +62,7 @@ impl Client {
         transaction_mode: bool,
     ) -> Result<Client, Error> {
         let sharding_regex = Regex::new(SHARDING_REGEX).unwrap();
+        let role_regex = Regex::new(ROLE_REGEX).unwrap();
 
         loop {
             // Could be StartupMessage or SSLRequest
@@ -114,6 +120,7 @@ impl Client {
                         secret_key: secret_key,
                         client_server_map: client_server_map,
                         sharding_regex: sharding_regex,
+                        role_regex: role_regex,
                     });
                 }
 
@@ -134,6 +141,7 @@ impl Client {
                         secret_key: secret_key,
                         client_server_map: client_server_map,
                         sharding_regex: sharding_regex,
+                        role_regex: role_regex,
                     });
                 }
 
@@ -145,7 +153,7 @@ impl Client {
     }
 
     /// Client loop. We handle all messages between the client and the database here.
-    pub async fn handle(&mut self, pool: ConnectionPool) -> Result<(), Error> {
+    pub async fn handle(&mut self, mut pool: ConnectionPool) -> Result<(), Error> {
         // Special: cancelling existing running query
         if self.cancel_mode {
             let (process_id, secret_key, address, port) = {
@@ -172,6 +180,9 @@ impl Client {
         // - 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> = None;
+
         loop {
             // Read a complete message from the client, which normally would be
             // either a `Q` (query) or `P` (prepare, extended protocol).
@@ -182,18 +193,36 @@ impl Client {
 
             // Parse for special select shard command.
             // SET SHARDING KEY TO 'bigint';
-            match self.select_shard(message.clone(), pool.shards()).await {
+            match self.select_shard(message.clone(), pool.shards()) {
                 Some(s) => {
-                    set_sharding_key(&mut self.write).await?;
+                    custom_protocol_response_ok(&mut self.write, "SET SHARDING KEY").await?;
                     shard = Some(s);
                     continue;
                 }
                 None => (),
             };
 
+            // Parse for special server role selection command.
+            //
+            match self.select_role(message.clone()) {
+                Some(r) => {
+                    custom_protocol_response_ok(&mut self.write, "SET SERVER ROLE").await?;
+                    role = Some(r);
+                    continue;
+                }
+                None => (),
+            };
+
             // Grab a server from the pool.
             // None = any shard
-            let connection = pool.get(shard).await.unwrap();
+            let connection = match pool.get(shard, role).await {
+                Ok(conn) => conn,
+                Err(err) => {
+                    println!(">> Could not get connection from pool: {:?}", err);
+                    return Err(err);
+                }
+            };
+
             let mut proxy = connection.0;
             let _address = connection.1;
             let server = &mut *proxy;
@@ -232,10 +261,13 @@ impl Client {
 
                 match code {
                     'Q' => {
+                        // TODO: implement retries here for read-only transactions.
                         server.send(original).await?;
 
                         loop {
+                            // TODO: implement retries here for read-only transactions.
                             let response = server.recv().await?;
+
                             match write_all_half(&mut self.write, response).await {
                                 Ok(_) => (),
                                 Err(err) => {
@@ -252,6 +284,7 @@ impl Client {
                         // Release server
                         if !server.in_transaction() && self.transaction_mode {
                             shard = None;
+                            role = None;
                             break;
                         }
                     }
@@ -290,10 +323,13 @@ impl Client {
                     'S' => {
                         // Extended protocol, client requests sync
                         self.buffer.put(&original[..]);
+
+                        // TODO: retries for read-only transactions
                         server.send(self.buffer.clone()).await?;
                         self.buffer.clear();
 
                         loop {
+                            // TODO: retries for read-only transactions
                             let response = server.recv().await?;
                             match write_all_half(&mut self.write, response).await {
                                 Ok(_) => (),
@@ -311,6 +347,7 @@ impl Client {
                         // Release server
                         if !server.in_transaction() && self.transaction_mode {
                             shard = None;
+                            role = None;
                             break;
                         }
                     }
@@ -338,6 +375,7 @@ impl Client {
                         if !server.in_transaction() && self.transaction_mode {
                             println!("Releasing after copy done");
                             shard = None;
+                            role = None;
                             break;
                         }
                     }
@@ -361,7 +399,7 @@ impl Client {
     /// 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.
-    async fn select_shard(&mut self, mut buf: BytesMut, shards: usize) -> Option<usize> {
+    fn select_shard(&mut self, mut buf: BytesMut, shards: usize) -> Option<usize> {
         let code = buf.get_u8() as char;
 
         // Only supporting simpe protocol here, so
@@ -390,4 +428,31 @@ impl Client {
             None
         }
     }
+
+    // Pick a primary or a replica from the pool.
+    fn select_role(&mut 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();
+
+        // Copy / paste from above. If we get one more of these use cases,
+        // it'll be time to abstract :).
+        if self.role_regex.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/config.rs

@@ -3,14 +3,21 @@ use tokio::fs::File;
 use tokio::io::AsyncReadExt;
 use toml;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use crate::errors::Error;
 
+#[derive(Clone, PartialEq, Deserialize, Hash, std::cmp::Eq, Debug, Copy)]
+pub enum Role {
+    Primary,
+    Replica,
+}
+
 #[derive(Clone, PartialEq, Hash, std::cmp::Eq, Debug)]
 pub struct Address {
     pub host: String,
     pub port: String,
+    pub role: Role,
 }
 
 #[derive(Clone, PartialEq, Hash, std::cmp::Eq, Deserialize, Debug)]
@@ -32,7 +39,7 @@ pub struct General {
 
 #[derive(Deserialize, Debug, Clone)]
 pub struct Shard {
-    pub servers: Vec<(String, u16)>,
+    pub servers: Vec<(String, u16, String)>,
     pub database: String,
 }
 
@@ -70,6 +77,47 @@ pub async fn parse(path: &str) -> Result<Config, Error> {
         }
     };
 
+    // Quick config sanity check.
+    for shard in &config.shards {
+        // We use addresses as unique identifiers,
+        // let's make sure they are unique in the config as well.
+        let mut dup_check = HashSet::new();
+        let mut primary_count = 0;
+
+        for server in &shard.1.servers {
+            dup_check.insert(server);
+
+            // Check that we define only zero or one primary.
+            match server.2.as_ref() {
+                "primary" => primary_count += 1,
+                _ => (),
+            };
+
+            // Check role spelling.
+            match server.2.as_ref() {
+                "primary" => (),
+                "replica" => (),
+                _ => {
+                    println!(
+                        "> Shard {} server role must be either 'primary' or 'replica', got: '{}'",
+                        shard.0, server.2
+                    );
+                    return Err(Error::BadConfig);
+                }
+            };
+        }
+
+        if primary_count > 1 {
+            println!("> Shard {} has more than on primary configured.", &shard.0);
+            return Err(Error::BadConfig);
+        }
+
+        if dup_check.len() != shard.1.servers.len() {
+            println!("> Shard {} contains duplicate server configs.", &shard.0);
+            return Err(Error::BadConfig);
+        }
+    }
+
     Ok(config)
 }
 
@@ -83,5 +131,6 @@ mod test {
         assert_eq!(config.general.pool_size, 15);
         assert_eq!(config.shards.len(), 3);
         assert_eq!(config.shards["1"].servers[0].0, "127.0.0.1");
+        assert_eq!(config.shards["0"].servers[0].2, "primary");
     }
 }

src/errors.rs

@@ -8,4 +8,5 @@ pub enum Error {
     // ServerTimeout,
     // DirtyServer,
     BadConfig,
+    AllServersDown,
 }

src/main.rs

@@ -73,6 +73,7 @@ async fn main() {
         "> Healthcheck timeout: {}ms",
         config.general.healthcheck_timeout
     );
+    println!("> Connection timeout: {}ms", config.general.connect_timeout);
 
     let pool = ConnectionPool::from_config(config.clone(), client_server_map.clone()).await;
     let transaction_mode = config.general.pool_mode == "transaction";