Allow configuring routing decision when no shard is selected (#578)

The TL;DR for the change is that we allow QueryRouter to set the active shard to None. This signals to the Pool::get method that we have no shard selected. The get method follows a no_shard_specified_behavior config to know how to route the query.

Original PR description
Ruby-pg library makes a startup query to SET client_encoding to ... if Encoding.default_internal value is set (Code). This query is troublesome because we cannot possibly attach a routing comment to it. PgCat, by default, will route that query to the default shard.

Everything is fine until shard 0 has issues, Clients will all be attempting to send this query to shard0 which increases the connection latency significantly for all clients, even those not interested in shard0

This PR introduces no_shard_specified_behavior that defines the behavior in case we have routing-by-comment enabled but we get a query without a comment. The allowed behaviors are

random: Picks a shard at random
random_healthy: Picks a shard at random favoring shards with the least number of recent connection/checkout errors
shard_<number>: e.g. shard_0, shard_4, etc. picks a specific shard, everytime
In order to achieve this, this PR introduces an error_count on the Address Object that tracks the number of errors since the last checkout and uses that metric to sort shards by error count before making a routing decision.
I didn't want to use address stats to avoid introducing a routing dependency on internal stats (We might do that in the future but I prefer to avoid this for the time being.

I also made changes to the test environment to replace Ruby's TOML reader library, It appears to be abandoned and does not support mixed arrays (which we use in the config toml), and it also does not play nicely with single-quoted regular expressions. I opted for using yj which is a CLI tool that can convert from toml to JSON and back. So I refactor the tests to use that library.

Author: drdrsh

.circleci/config.yml

@@ -9,7 +9,7 @@ jobs:
     # Specify the execution environment. You can specify an image from Dockerhub or use one of our Convenience Images from CircleCI's Developer Hub.
     # See: https://circleci.com/docs/2.0/configuration-reference/#docker-machine-macos-windows-executor
     docker:
-      - image: ghcr.io/levkk/pgcat-ci:1.67
+      - image: ghcr.io/postgresml/pgcat-ci:latest
         environment:
           RUST_LOG: info
           LLVM_PROFILE_FILE: /tmp/pgcat-%m-%p.profraw

pgcat.toml

@@ -177,6 +177,12 @@ primary_reads_enabled = true
 # shard_id_regex = '/\* shard_id: (\d+) \*/'
 # regex_search_limit = 1000 # only look at the first 1000 characters of SQL statements
 
+# Defines the behavior when no shard is selected in a sharded system.
+# `random`: picks a shard at random
+# `random_healthy`: picks a shard at random favoring shards with the least number of recent errors
+# `shard_<number>`: e.g. shard_0, shard_4, etc. picks a specific shard, everytime
+# no_shard_specified_behavior = "shard_0"
+
 # So what if you wanted to implement a different hashing function,
 # or you've already built one and you want this pooler to use it?
 # Current options:

src/client.rs

@@ -1009,23 +1009,27 @@ where
 
                 // SET SHARD TO
                 Some((Command::SetShard, _)) => {
-                    // Selected shard is not configured.
-                    if query_router.shard() >= pool.shards() {
-                        // Set the shard back to what it was.
-                        query_router.set_shard(current_shard);
-
-                        error_response(
-                            &mut self.write,
-                            &format!(
-                                "shard {} is more than configured {}, staying on shard {} (shard numbers start at 0)",
-                                query_router.shard(),
-                                pool.shards(),
-                                current_shard,
-                            ),
-                        )
-                            .await?;
-                    } else {
-                        custom_protocol_response_ok(&mut self.write, "SET SHARD").await?;
+                    match query_router.shard() {
+                        None => (),
+                        Some(selected_shard) => {
+                            if selected_shard >= pool.shards() {
+                                // Bad shard number, send error message to client.
+                                query_router.set_shard(current_shard);
+
+                                error_response(
+                                    &mut self.write,
+                                    &format!(
+                                        "shard {} is not configured {}, staying on shard {:?} (shard numbers start at 0)",
+                                        selected_shard,
+                                        pool.shards(),
+                                        current_shard,
+                                    ),
+                                )
+                                    .await?;
+                            } else {
+                                custom_protocol_response_ok(&mut self.write, "SET SHARD").await?;
+                            }
+                        }
                     }
                     continue;
                 }
@@ -1093,8 +1097,11 @@ where
                         self.buffer.clear();
                     }
 
-                    error_response(&mut self.write, "could not get connection from the pool")
-                        .await?;
+                    error_response(
+                        &mut self.write,
+                        format!("could not get connection from the pool - {}", err).as_str(),
+                    )
+                    .await?;
 
                     error!(
                         "Could not get connection from pool: \
@@ -1234,7 +1241,7 @@ where
                                     {{ \
                                         pool_name: {}, \
                                         username: {}, \
-                                        shard: {}, \
+                                        shard: {:?}, \
                                         role: \"{:?}\" \
                                     }}",
                                     self.pool_name,

src/config.rs

@@ -3,11 +3,14 @@ use arc_swap::ArcSwap;
 use log::{error, info};
 use once_cell::sync::Lazy;
 use regex::Regex;
+use serde::{Deserializer, Serializer};
 use serde_derive::{Deserialize, Serialize};
+
 use std::collections::hash_map::DefaultHasher;
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::hash::{Hash, Hasher};
 use std::path::Path;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use tokio::fs::File;
 use tokio::io::AsyncReadExt;
@@ -101,6 +104,9 @@ pub struct Address {
 
     /// Address stats
     pub stats: Arc<AddressStats>,
+
+    /// Number of errors encountered since last successful checkout
+    pub error_count: Arc<AtomicU64>,
 }
 
 impl Default for Address {
@@ -118,6 +124,7 @@ impl Default for Address {
             pool_name: String::from("pool_name"),
             mirrors: Vec::new(),
             stats: Arc::new(AddressStats::default()),
+            error_count: Arc::new(AtomicU64::new(0)),
         }
     }
 }
@@ -182,6 +189,18 @@ impl Address {
             ),
         }
     }
+
+    pub fn error_count(&self) -> u64 {
+        self.error_count.load(Ordering::Relaxed)
+    }
+
+    pub fn increment_error_count(&self) {
+        self.error_count.fetch_add(1, Ordering::Relaxed);
+    }
+
+    pub fn reset_error_count(&self) {
+        self.error_count.store(0, Ordering::Relaxed);
+    }
 }
 
 /// PostgreSQL user.
@@ -540,6 +559,9 @@ pub struct Pool {
     pub shard_id_regex: Option<String>,
     pub regex_search_limit: Option<usize>,
 
+    #[serde(default = "Pool::default_default_shard")]
+    pub default_shard: DefaultShard,
+
     pub auth_query: Option<String>,
     pub auth_query_user: Option<String>,
     pub auth_query_password: Option<String>,
@@ -575,6 +597,10 @@ impl Pool {
         PoolMode::Transaction
     }
 
+    pub fn default_default_shard() -> DefaultShard {
+        DefaultShard::default()
+    }
+
     pub fn default_load_balancing_mode() -> LoadBalancingMode {
         LoadBalancingMode::Random
     }
@@ -666,6 +692,16 @@ impl Pool {
             None => None,
         };
 
+        match self.default_shard {
+            DefaultShard::Shard(shard_number) => {
+                if shard_number >= self.shards.len() {
+                    error!("Invalid shard {:?}", shard_number);
+                    return Err(Error::BadConfig);
+                }
+            }
+            _ => (),
+        }
+
         for (_, user) in &self.users {
             user.validate()?;
         }
@@ -693,6 +729,7 @@ impl Default for Pool {
             sharding_key_regex: None,
             shard_id_regex: None,
             regex_search_limit: Some(1000),
+            default_shard: Self::default_default_shard(),
             auth_query: None,
             auth_query_user: None,
             auth_query_password: None,
@@ -711,6 +748,50 @@ pub struct ServerConfig {
     pub role: Role,
 }
 
+// No Shard Specified handling.
+#[derive(Debug, PartialEq, Clone, Eq, Hash, Copy)]
+pub enum DefaultShard {
+    Shard(usize),
+    Random,
+    RandomHealthy,
+}
+impl Default for DefaultShard {
+    fn default() -> Self {
+        DefaultShard::Shard(0)
+    }
+}
+impl serde::Serialize for DefaultShard {
+    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+        match self {
+            DefaultShard::Shard(shard) => {
+                serializer.serialize_str(&format!("shard_{}", &shard.to_string()))
+            }
+            DefaultShard::Random => serializer.serialize_str("random"),
+            DefaultShard::RandomHealthy => serializer.serialize_str("random_healthy"),
+        }
+    }
+}
+impl<'de> serde::Deserialize<'de> for DefaultShard {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let s = String::deserialize(deserializer)?;
+        if s.starts_with("shard_") {
+            let shard = s[6..].parse::<usize>().map_err(serde::de::Error::custom)?;
+            return Ok(DefaultShard::Shard(shard));
+        }
+
+        match s.as_str() {
+            "random" => Ok(DefaultShard::Random),
+            "random_healthy" => Ok(DefaultShard::RandomHealthy),
+            _ => Err(serde::de::Error::custom(
+                "invalid value for no_shard_specified_behavior",
+            )),
+        }
+    }
+}
+
 #[derive(Clone, PartialEq, Serialize, Deserialize, Debug, Hash, Eq)]
 pub struct MirrorServerConfig {
     pub host: String,

src/errors.rs

@@ -28,6 +28,7 @@ pub enum Error {
     UnsupportedStatement,
     QueryRouterParserError(String),
     QueryRouterError(String),
+    InvalidShardId(usize),
 }
 
 #[derive(Clone, PartialEq, Debug)]