docker/pgcat
1
0
Fork 0
mirror of https://github.com/postgresml/pgcat.git synced 2026-03-26 18:36:28 +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

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());
}