docker/pgcat
1
0
Fork 0
mirror of https://github.com/postgresml/pgcat.git synced 2026-03-23 01:16:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
Mostafa Abdelraouf
2023-09-11 13:47:28 -05:00
committed by GitHub
parent 33db0dffa8
commit 0b01d70b55
16 changed files with 449 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
src/client.rs
View File

@@ -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);
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 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?;
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,

81
src/config.rs
View File

@@ -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,

1
src/errors.rs
View File

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

83
src/pool.rs
View File

@@ -10,6 +10,7 @@ use rand::thread_rng;
use regex::Regex;
use std::collections::HashMap;
use std::fmt::{Display, Formatter};
use std::sync::atomic::AtomicU64;
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
@@ -18,7 +19,7 @@ use std::time::Instant;
use tokio::sync::Notify;
use crate::config::{
get_config, Address, General, LoadBalancingMode, Plugins, PoolMode, Role, User,
get_config, Address, DefaultShard, General, LoadBalancingMode, Plugins, PoolMode, Role, User,
};
use crate::errors::Error;
@@ -140,6 +141,9 @@ pub struct PoolSettings {
// Regex for searching for the shard id in SQL statements
pub shard_id_regex: Option<Regex>,
// What to do when no shard is selected in a sharded system
pub default_shard: DefaultShard,
// Limit how much of each query is searched for a potential shard regex match
pub regex_search_limit: usize,
@@ -173,6 +177,7 @@ impl Default for PoolSettings {
sharding_key_regex: None,
shard_id_regex: None,
regex_search_limit: 1000,
default_shard: DefaultShard::Shard(0),
auth_query: None,
auth_query_user: None,
auth_query_password: None,
@@ -299,6 +304,7 @@ impl ConnectionPool {
pool_name: pool_name.clone(),
mirrors: vec![],
stats: Arc::new(AddressStats::default()),
error_count: Arc::new(AtomicU64::new(0)),
});
address_id += 1;
}
@@ -317,6 +323,7 @@ impl ConnectionPool {
pool_name: pool_name.clone(),
mirrors: mirror_addresses,
stats: Arc::new(AddressStats::default()),
error_count: Arc::new(AtomicU64::new(0)),
};
address_id += 1;
@@ -482,6 +489,7 @@ impl ConnectionPool {
.clone()
.map(|regex| Regex::new(regex.as_str()).unwrap()),
regex_search_limit: pool_config.regex_search_limit.unwrap_or(1000),
default_shard: pool_config.default_shard.clone(),
auth_query: pool_config.auth_query.clone(),
auth_query_user: pool_config.auth_query_user.clone(),
auth_query_password: pool_config.auth_query_password.clone(),
@@ -603,19 +611,51 @@ impl ConnectionPool {
/// Get a connection from the pool.
pub async fn get(
&self,
shard: usize, // shard number
shard: Option<usize>, // shard number
role: Option<Role>, // primary or replica
client_stats: &ClientStats, // client id
) -> Result<(PooledConnection<'_, ServerPool>, Address), Error> {
let mut candidates: Vec<&Address> = self.addresses[shard]
.iter()
.filter(|address| address.role == role)
.collect();
let effective_shard_id = if self.shards() == 1 {
// The base, unsharded case
Some(0)
} else {
if !self.valid_shard_id(shard) {
// None is valid shard ID so it is safe to unwrap here
return Err(Error::InvalidShardId(shard.unwrap()));
}
shard
};
// We shuffle even if least_outstanding_queries is used to avoid imbalance
// in cases where all candidates have more or less the same number of outstanding
// queries
let mut candidates = self
.addresses
.iter()
.flatten()
.filter(|address| address.role == role)
.collect::<Vec<&Address>>();
// We start with a shuffled list of addresses even if we end up resorting
// this is meant to avoid hitting instance 0 everytime if the sorting metric
// ends up being the same for all instances
candidates.shuffle(&mut thread_rng());
match effective_shard_id {
Some(shard_id) => candidates.retain(|address| address.shard == shard_id),
None => match self.settings.default_shard {
DefaultShard::Shard(shard_id) => {
candidates.retain(|address| address.shard == shard_id)
}
DefaultShard::Random => (),
DefaultShard::RandomHealthy => {
candidates.sort_by(|a, b| {
b.error_count
.load(Ordering::Relaxed)
.partial_cmp(&a.error_count.load(Ordering::Relaxed))
.unwrap()
});
}
},
};
if self.settings.load_balancing_mode == LoadBalancingMode::LeastOutstandingConnections {
candidates.sort_by(|a, b| {
self.busy_connection_count(b)
@@ -651,7 +691,10 @@ impl ConnectionPool {
.get()
.await
{
Ok(conn) => conn,
Ok(conn) => {
address.reset_error_count();
conn
}
Err(err) => {
error!(
"Connection checkout error for instance {:?}, error: {:?}",
@@ -766,6 +809,18 @@ impl ConnectionPool {
/// traffic for any new transactions. Existing transactions on that replica
/// will finish successfully or error out to the clients.
pub fn ban(&self, address: &Address, reason: BanReason, client_info: Option<&ClientStats>) {
// Count the number of errors since the last successful checkout
// This is used to determine if the shard is down
match reason {
BanReason::FailedHealthCheck
| BanReason::FailedCheckout
| BanReason::MessageSendFailed
| BanReason::MessageReceiveFailed => {
address.increment_error_count();
}
_ => (),
};
// Primary can never be banned
if address.role == Role::Primary {
return;
@@ -920,6 +975,7 @@ impl ConnectionPool {
self.original_server_parameters.read().clone()
}
/// Get the number of checked out connection for an address
fn busy_connection_count(&self, address: &Address) -> u32 {
let state = self.pool_state(address.shard, address.address_index);
let idle = state.idle_connections;
@@ -933,6 +989,13 @@ impl ConnectionPool {
debug!("{:?} has {:?} busy connections", address, busy);
return busy;
}
fn valid_shard_id(&self, shard: Option<usize>) -> bool {
match shard {
None => true,
Some(shard) => shard < self.shards(),
}
}
}
/// Wrapper for the bb8 connection pool.

57
src/query_router.rs
View File

@@ -143,13 +143,14 @@ impl QueryRouter {
let code = message_cursor.get_u8() as char;
let len = message_cursor.get_i32() as usize;
let comment_shard_routing_enabled = self.pool_settings.shard_id_regex.is_some()
|| self.pool_settings.sharding_key_regex.is_some();
// Check for any sharding regex matches in any queries
match code as char {
// For Parse and Query messages peek to see if they specify a shard_id as a comment early in the statement
'P' | 'Q' => {
if self.pool_settings.shard_id_regex.is_some()
|| self.pool_settings.sharding_key_regex.is_some()
{
if comment_shard_routing_enabled {
match code as char {
// For Parse and Query messages peek to see if they specify a shard_id as a comment early in the statement
'P' | 'Q' => {
// Check only the first block of bytes configured by the pool settings
let seg = cmp::min(len - 5, self.pool_settings.regex_search_limit);
@@ -166,7 +167,7 @@ impl QueryRouter {
});
if let Some(shard_id) = shard_id {
debug!("Setting shard to {:?}", shard_id);
self.set_shard(shard_id);
self.set_shard(Some(shard_id));
// Skip other command processing since a sharding command was found
return None;
}
@@ -188,8 +189,8 @@ impl QueryRouter {
}
}
}
_ => {}
}
_ => {}
}
// Only simple protocol supported for commands processed below
@@ -248,7 +249,9 @@ impl QueryRouter {
}
}
Command::ShowShard => self.shard().to_string(),
Command::ShowShard => self
.shard()
.map_or_else(|| "unset".to_string(), |x| x.to_string()),
Command::ShowServerRole => match self.active_role {
Some(Role::Primary) => Role::Primary.to_string(),
Some(Role::Replica) => Role::Replica.to_string(),
@@ -581,7 +584,7 @@ impl QueryRouter {
// TODO: Support multi-shard queries some day.
if shards.len() == 1 {
debug!("Found one sharding key");
self.set_shard(*shards.first().unwrap());
self.set_shard(Some(*shards.first().unwrap()));
true
} else {
debug!("Found no sharding keys");
@@ -865,7 +868,7 @@ impl QueryRouter {
self.pool_settings.sharding_function,
);
let shard = sharder.shard(sharding_key);
self.set_shard(shard);
self.set_shard(Some(shard));
self.active_shard
}
@@ -875,12 +878,12 @@ impl QueryRouter {
}
/// Get desired shard we should be talking to.
pub fn shard(&self) -> usize {
self.active_shard.unwrap_or(0)
pub fn shard(&self) -> Option<usize> {
self.active_shard
}
pub fn set_shard(&mut self, shard: usize) {
self.active_shard = Some(shard);
pub fn set_shard(&mut self, shard: Option<usize>) {
self.active_shard = shard;
}
/// Should we attempt to parse queries?
@@ -1090,7 +1093,7 @@ mod test {
qr.try_execute_command(&query),
Some((Command::SetShardingKey, String::from("0")))
);
assert_eq!(qr.shard(), 0);
assert_eq!(qr.shard().unwrap(), 0);
// SetShard
let query = simple_query("SET SHARD TO '1'");
@@ -1098,7 +1101,7 @@ mod test {
qr.try_execute_command(&query),
Some((Command::SetShard, String::from("1")))
);
assert_eq!(qr.shard(), 1);
assert_eq!(qr.shard().unwrap(), 1);
// ShowShard
let query = simple_query("SHOW SHARD");
@@ -1204,6 +1207,7 @@ mod test {
ban_time: PoolSettings::default().ban_time,
sharding_key_regex: None,
shard_id_regex: None,
default_shard: crate::config::DefaultShard::Shard(0),
regex_search_limit: 1000,
auth_query: None,
auth_query_password: None,
@@ -1281,6 +1285,7 @@ mod test {
ban_time: PoolSettings::default().ban_time,
sharding_key_regex: Some(Regex::new(r"/\* sharding_key: (\d+) \*/").unwrap()),
shard_id_regex: Some(Regex::new(r"/\* shard_id: (\d+) \*/").unwrap()),
default_shard: crate::config::DefaultShard::Shard(0),
regex_search_limit: 1000,
auth_query: None,
auth_query_password: None,
@@ -1331,7 +1336,7 @@ mod test {
.unwrap(),
)
.is_ok());
assert_eq!(qr.shard(), 2);
assert_eq!(qr.shard().unwrap(), 2);
assert!(qr
.infer(
@@ -1341,7 +1346,7 @@ mod test {
.unwrap()
)
.is_ok());
assert_eq!(qr.shard(), 0);
assert_eq!(qr.shard().unwrap(), 0);
assert!(qr
.infer(
@@ -1354,7 +1359,7 @@ mod test {
.unwrap()
)
.is_ok());
assert_eq!(qr.shard(), 2);
assert_eq!(qr.shard().unwrap(), 2);
// Shard did not move because we couldn't determine the sharding key since it could be ambiguous
// in the query.
@@ -1366,7 +1371,7 @@ mod test {
.unwrap()
)
.is_ok());
assert_eq!(qr.shard(), 2);
assert_eq!(qr.shard().unwrap(), 2);
assert!(qr
.infer(
@@ -1376,7 +1381,7 @@ mod test {
.unwrap()
)
.is_ok());
assert_eq!(qr.shard(), 0);
assert_eq!(qr.shard().unwrap(), 0);
assert!(qr
.infer(
@@ -1386,7 +1391,7 @@ mod test {
.unwrap()
)
.is_ok());
assert_eq!(qr.shard(), 2);
assert_eq!(qr.shard().unwrap(), 2);
// Super unique sharding key
qr.pool_settings.automatic_sharding_key = Some("*.unique_enough_column_name".to_string());
@@ -1398,7 +1403,7 @@ mod test {
.unwrap()
)
.is_ok());
assert_eq!(qr.shard(), 0);
assert_eq!(qr.shard().unwrap(), 0);
assert!(qr
.infer(
@@ -1406,7 +1411,7 @@ mod test {
.unwrap()
)
.is_ok());
assert_eq!(qr.shard(), 0);
assert_eq!(qr.shard().unwrap(), 0);
}
#[test]
@@ -1434,7 +1439,7 @@ mod test {
assert_eq!(qr.placeholders.len(), 1);
assert!(qr.infer_shard_from_bind(&bind));
assert_eq!(qr.shard(), 2);
assert_eq!(qr.shard().unwrap(), 2);
assert!(qr.placeholders.is_empty());
}