diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index 745cb39d2..00344aa37 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -31,6 +31,7 @@ "ban_timeout": 300000, "broadcast_address": null, "broadcast_port": 6433, + "canonicalize_type_information": false, "checkout_timeout": 5000, "client_connection_recovery": "drop", "client_idle_in_transaction_timeout": 9223372036854775807, @@ -634,6 +635,11 @@ "maximum": 65535, "minimum": 0 }, + "canonicalize_type_information": { + "description": "Controls whether PgDog maps each shard's type information to a canonical source.\n\n_Default:_ `false`\n\n", + "type": "boolean", + "default": false + }, "checkout_timeout": { "description": "Maximum amount of time a client is allowed to wait for a connection from the pool.\n\n_Default:_ `5000`\n\n", "type": "integer", diff --git a/integration/rust/Cargo.toml b/integration/rust/Cargo.toml index 3cbc0ecb7..427b729b1 100644 --- a/integration/rust/Cargo.toml +++ b/integration/rust/Cargo.toml @@ -6,6 +6,9 @@ edition = "2024" [lib] test = true +[features] +new_parser = [] + [dependencies] tokio-postgres = {version = "0.7.13", features = ["with-uuid-1"]} sqlx = { version = "0.8.6", features = ["postgres", "runtime-tokio", "tls-native-tls", "bigdecimal", "chrono", "json", "rust_decimal"]} diff --git a/integration/rust/tests/integration/cross_shard_oid_drift.rs b/integration/rust/tests/integration/cross_shard_oid_drift.rs new file mode 100644 index 000000000..0c3aeb76a --- /dev/null +++ b/integration/rust/tests/integration/cross_shard_oid_drift.rs @@ -0,0 +1,77 @@ +#![cfg(feature = "new_parser")] +use crate::setup::{admin_sqlx, connections_sqlx}; +use sqlx::postgres::types::Oid; +use sqlx::{Column, Executor, Row}; + +#[derive(sqlx::Type, Debug, Clone, PartialEq)] +#[sqlx(type_name = "test_oid_drift_composite")] +struct Composite { + a: String, + b: String, +} + +#[tokio::test] +async fn test_oid_drift() { + let conn = connections_sqlx().await.pop().unwrap(); + let admin = admin_sqlx().await; + + // Intentionally cause the OID of the type to differ between shards + conn.execute("/* pgdog_shard: 0 */ CREATE SEQUENCE foo; DROP SEQUENCE foo;") + .await + .unwrap(); + conn.execute("DROP TYPE IF EXISTS test_oid_drift_composite CASCADE") + .await + .unwrap(); + conn.execute("CREATE TYPE test_oid_drift_composite AS (a text, b text)") + .await + .unwrap(); + conn.execute("DROP TABLE IF EXISTS test_oid_drift") + .await + .unwrap(); + conn.execute( + "CREATE TABLE test_oid_drift (customer_id BIGINT, composite test_oid_drift_composite)", + ) + .await + .unwrap(); + admin + .execute("SET canonicalize_type_information TO true") + .await + .unwrap(); + + let composite = Composite { + a: String::from("a"), + b: String::from("b"), + }; + for i in 1..=20 { + sqlx::query("INSERT INTO test_oid_drift VALUES ($1, $2)") + .bind(i) + .bind(&composite) + .execute(&conn) + .await + .unwrap(); + } + + let rows: Vec = sqlx::query_scalar("SELECT composite FROM test_oid_drift") + .fetch_all(&conn) + .await + .unwrap(); + assert_eq!(rows, vec![composite.clone(); 20]); + + let simple_rows = conn + .fetch_all("SELECT composite FROM test_oid_drift") + .await + .unwrap(); + + let expected_oid: Oid = + sqlx::query_scalar("SELECT oid FROM pg_type WHERE typname = 'test_oid_drift_composite'") + .fetch_one(&conn) + .await + .unwrap(); + let given_oid = simple_rows.first().unwrap().column(0).type_info().oid(); + assert_eq!(given_oid, Some(expected_oid)); + + let simple_data: Vec = simple_rows.into_iter().map(|row| row.get(0)).collect(); + assert_eq!(simple_data, vec![composite; 20]); + + admin.execute("RELOAD").await.unwrap(); +} diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index 55e70c4cb..3983863c7 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/integration/rust/tests/integration/mod.rs @@ -10,6 +10,7 @@ pub mod client_ids; pub mod connection_recovery; pub mod copy; pub mod cross_shard_disabled; +mod cross_shard_oid_drift; pub mod distinct; pub mod explain; pub mod fake_transactions; diff --git a/pgdog-config/src/general.rs b/pgdog-config/src/general.rs index f82d21b4b..947839cd8 100644 --- a/pgdog-config/src/general.rs +++ b/pgdog-config/src/general.rs @@ -767,6 +767,14 @@ pub struct General { #[serde(default = "General::load_schema")] pub load_schema: LoadSchema, + /// Controls whether PgDog maps each shard's type information to a canonical source. + /// + /// _Default:_ `false` + /// + /// + #[serde(default = "General::canonicalize_type_information")] + pub canonicalize_type_information: bool, + /// Replication lag threshold (in bytes) at which PgDog will pause traffic automatically during a traffic cutover. /// /// _Default:_ `1000000` @@ -915,6 +923,7 @@ impl Default for General { resharding_replication_retry_min_delay: Self::resharding_replication_retry_min_delay(), reload_schema_on_ddl: Self::reload_schema_on_ddl(), load_schema: Self::load_schema(), + canonicalize_type_information: Self::canonicalize_type_information(), cutover_replication_lag_threshold: Self::cutover_replication_lag_threshold(), cutover_traffic_stop_threshold: Self::cutover_traffic_stop_threshold(), cutover_last_transaction_delay: Self::cutover_last_transaction_delay(), @@ -1327,6 +1336,10 @@ impl General { Self::env_enum_or_default("PGDOG_LOAD_SCHEMA") } + fn canonicalize_type_information() -> bool { + Self::env_or_default("PGDOG_CANONICALIZE_TYPE_INFORMATION", false) + } + pub fn mirror_queue() -> usize { Self::env_or_default("PGDOG_MIRROR_QUEUE", 128) } diff --git a/pgdog/benches/comment_parser.rs b/pgdog/benches/comment_parser.rs index d5086ca5a..22afa9ed7 100644 --- a/pgdog/benches/comment_parser.rs +++ b/pgdog/benches/comment_parser.rs @@ -1,7 +1,6 @@ use brunch::{Bench, benches}; #[cfg(not(feature = "new_parser"))] use pg_query::scan_raw; -use pgdog::backend::ShardingSchema; use pgdog::frontend::router::parser::comment::parse_edge_comment; const QUERY_WITH_LEADING: &str = @@ -13,20 +12,20 @@ const QUERY_NO_COMMENT: &str = "SELECT * FROM users WHERE id = $1 AND name = $2" #[cfg(feature = "new_parser")] benches!( Bench::new("parse_edge_comment(leading)") - .run(|| parse_edge_comment(QUERY_WITH_LEADING, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_WITH_LEADING, &Default::default())), Bench::new("parse_edge_comment(trailing)") - .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &Default::default())), Bench::new("parse_edge_comment(no comment)") - .run(|| parse_edge_comment(QUERY_NO_COMMENT, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_NO_COMMENT, &Default::default())), ); #[cfg(not(feature = "new_parser"))] benches!( Bench::new("parse_edge_comment(leading)") - .run(|| parse_edge_comment(QUERY_WITH_LEADING, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_WITH_LEADING, &Default::default())), Bench::new("parse_edge_comment(trailing)") - .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_WITH_TRAILING, &Default::default())), Bench::new("parse_edge_comment(no comment)") - .run(|| parse_edge_comment(QUERY_NO_COMMENT, &ShardingSchema::default())), + .run(|| parse_edge_comment(QUERY_NO_COMMENT, &Default::default())), Bench::new("scan_raw(leading)").run(|| scan_raw(QUERY_WITH_LEADING)), Bench::new("scan_raw(trailing)").run(|| scan_raw(QUERY_WITH_TRAILING)), Bench::new("scan_raw(no comment)").run(|| scan_raw(QUERY_NO_COMMENT)), diff --git a/pgdog/src/admin/probe.rs b/pgdog/src/admin/probe.rs index 1fe397552..4cc55bf8d 100644 --- a/pgdog/src/admin/probe.rs +++ b/pgdog/src/admin/probe.rs @@ -32,6 +32,7 @@ impl Command for Probe { &Address::try_from(self.url.clone()).map_err(|_| Error::InvalidAddress)?, ServerOptions::default(), ConnectReason::Probe, + Default::default(), ), ) .await? diff --git a/pgdog/src/admin/set.rs b/pgdog/src/admin/set.rs index 28859c079..cae78e626 100644 --- a/pgdog/src/admin/set.rs +++ b/pgdog/src/admin/set.rs @@ -211,6 +211,10 @@ impl Command for Set { config.config.general.connect_timeout = self.value.parse()?; } + "canonicalize_type_information" => { + config.config.general.canonicalize_type_information = Self::from_json(&self.value)?; + } + _ => return Ok(vec![]), } diff --git a/pgdog/src/backend/error.rs b/pgdog/src/backend/error.rs index 2514ca6c1..6b344c820 100644 --- a/pgdog/src/backend/error.rs +++ b/pgdog/src/backend/error.rs @@ -143,6 +143,9 @@ pub enum Error { #[error("cannot ignore response for message type: {0}")] UnsupportedHandleIgnore(char), + + #[error("missing canonical oid for type {0}")] + MissingCanonicalOid(String), } impl From for Error { diff --git a/pgdog/src/backend/mod.rs b/pgdog/src/backend/mod.rs index a4f051759..608fceb81 100644 --- a/pgdog/src/backend/mod.rs +++ b/pgdog/src/backend/mod.rs @@ -21,7 +21,9 @@ pub mod validation; pub use connect_reason::ConnectReason; pub use disconnect_reason::DisconnectReason; pub use error::Error; -pub use pool::{Cluster, ClusterShardConfig, LoadBalancer, Pool, Shard, ShardingSchema}; +pub(crate) use pool::{ + CanonicalOids, Cluster, ClusterShardConfig, Oids, Pool, Shard, ShardingSchema, +}; pub use prepared_statements::PreparedStatements; pub use protocol::*; pub use pub_sub::{PubSubClient, PubSubListener}; diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index f66654c95..4839433c8 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -23,7 +23,9 @@ use crate::{ net::{Query, messages::FrontendPid}, }; -use super::{Address, Config, Error, Guard, MirrorStats, Request, Shard, ShardConfig}; +use super::{ + Address, CanonicalOids, Config, Error, Guard, MirrorStats, Request, Shard, ShardConfig, +}; use crate::config::LoadBalancingStrategy; use launch::Readiness; @@ -85,6 +87,7 @@ pub struct Cluster { tls_client_certificate_required: bool, #[debug(skip)] schema_loader: Box, + canonical_oids: Option>, } /// Sharding configuration from the cluster. @@ -174,6 +177,7 @@ pub struct ClusterConfig<'a> { identity: &'a Option, tls_client_certificate_required: bool, schema_cache: SchemaCache, + canonicalize_oids: bool, } impl<'a> ClusterConfig<'a> { @@ -245,6 +249,7 @@ impl<'a> ClusterConfig<'a> { identity: &user.identity, tls_client_certificate_required: user.tls_client_certificate_required.unwrap_or(true), schema_cache, + canonicalize_oids: general.canonicalize_type_information, } } } @@ -293,12 +298,14 @@ impl Cluster { identity, tls_client_certificate_required, schema_cache, + canonicalize_oids, } = config; let identifier = Arc::new(DatabaseUser { user: user.to_owned(), database: name.to_owned(), }); + let canonical_oids = canonicalize_oids.then(|| schema_cache.canonical_oids(name)); Self { identifier: identifier.clone(), @@ -357,6 +364,7 @@ impl Cluster { identity: identity.clone(), tls_client_certificate_required, schema_loader: Box::new(schema_loader::FromServer), + canonical_oids, } } @@ -411,7 +419,7 @@ impl Cluster { } /// Get all shards. - pub fn shards(&self) -> &[Shard] { + pub(crate) fn shards(&self) -> &[Shard] { &self.shards } @@ -678,6 +686,11 @@ impl Cluster { Ok(()) } + + #[cfg(feature = "new_parser")] + pub(crate) fn is_canonicalizing_oids(&self) -> bool { + self.canonical_oids.is_some() + } } #[cfg(test)] @@ -813,6 +826,11 @@ mod test { rewrite: config.config.rewrite.clone(), two_phase_commit: config.config.general.two_phase_commit, two_phase_commit_auto: config.config.general.two_phase_commit_auto.unwrap_or(false), + canonical_oids: config + .config + .general + .canonicalize_type_information + .then(Default::default), ..Default::default() } } diff --git a/pgdog/src/backend/pool/cluster/schema_loader.rs b/pgdog/src/backend/pool/cluster/schema_loader.rs index 21c7f244f..79f4dd254 100644 --- a/pgdog/src/backend/pool/cluster/schema_loader.rs +++ b/pgdog/src/backend/pool/cluster/schema_loader.rs @@ -3,6 +3,7 @@ use crate::backend::pool::ee::schema_changed_hook; use crate::tasks; use crate::util::safe_sleep; use dyn_clone::DynClone; +use std::sync::Arc; use std::time::Duration; use tokio::select; use tracing::error; @@ -25,6 +26,42 @@ impl SchemaLoader for FromServer { return; } + // For now we treat shard 0 as the canonical OID source + if let Some(shard) = cluster.shards().first() + && let Some(canonical_oids) = &cluster.canonical_oids + { + let canonical_oids = Arc::clone(canonical_oids); + let shard = shard.clone(); + tasks::spawn("load canonical oids", async move { + loop { + let result = tasks::shutdown_signal() + .run_until_cancelled(async { + canonical_oids + .load(&mut *shard.primary_or_replica(&Default::default()).await?) + .await + }) + .await; + + match result { + Some(Ok(_)) | None => break, + Some(Err(err)) => { + if shard.online() { + error!("error loading canonical type information: {err}"); + safe_sleep(Duration::from_millis(100)).await; + } else { + // Cluster is shutting down + break; + } + } + } + } + }); + } else { + for shard in cluster.shards() { + shard.skip_loading_oids() + } + } + for shard in cluster.shards() { let identifier = cluster.identifier(); let shard = shard.clone(); diff --git a/pgdog/src/backend/pool/inner.rs b/pgdog/src/backend/pool/inner.rs index a44004f8b..0d3bdc4ed 100644 --- a/pgdog/src/backend/pool/inner.rs +++ b/pgdog/src/backend/pool/inner.rs @@ -324,6 +324,7 @@ impl Inner { for conn in idle.iter_mut() { conn.stats_mut().set_pool_id(destination.id()); + conn.replace_oids(&destination.inner().oids) } (idle, taken) @@ -352,6 +353,7 @@ impl Inner { if moved.id() != self.id { server.stats_mut().set_pool_id(moved.id()); server.stats().update(); + server.replace_oids(&moved.inner().oids); moved.lock().maybe_check_in(server, now, stats, true)?; return Ok(result); } diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index 91e39a198..7c37068f2 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -18,7 +18,7 @@ use crate::{ net::Parameters, }; -use super::{Error, Guard, Pool, PoolConfig, Request}; +use super::{Error, Guard, Oids, Pool, PoolConfig, Request}; use crate::util::safe_timeout; pub mod ban; @@ -91,11 +91,12 @@ pub struct LoadBalancer { impl LoadBalancer { /// Create new replicas pools. - pub fn new( + pub(crate) fn new( primary: &Option, addrs: &[PoolConfig], lb_strategy: LoadBalancingStrategy, rw_split: ReadWriteSplit, + oids: Arc, ) -> LoadBalancer { let checkout_timeout = primary .as_ref() @@ -111,7 +112,12 @@ impl LoadBalancer { let mut targets: Vec<_> = addrs .iter() - .map(|config| Target::new(Pool::new(config), config.address.configured_role)) + .map(|config| { + Target::new( + Pool::with_oid_mapping(config, Arc::clone(&oids)), + config.address.configured_role, + ) + }) .collect(); let primary_target = primary diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index f0d0dfb1b..35bbf99dd 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -41,6 +41,7 @@ fn setup_test_replicas() -> LoadBalancer { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); replicas @@ -69,6 +70,7 @@ async fn test_include_primary_if_replica_banned_only_primary() { &[], LoadBalancingStrategy::default(), ReadWriteSplit::IncludePrimaryIfReplicaBanned, + Default::default(), ); lb.launch(); @@ -212,6 +214,7 @@ async fn test_primary_pool_banning() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -368,6 +371,7 @@ async fn test_read_write_split_exclude_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::ExcludePrimary, + Default::default(), ); replicas.launch(); @@ -404,6 +408,7 @@ async fn test_read_write_split_include_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -457,6 +462,7 @@ async fn test_prefer_primary_optin_read_honors_read_write_split() { &replica_configs, LoadBalancingStrategy::RoundRobin, split, + Default::default(), ); lb.launch(); @@ -504,6 +510,7 @@ async fn test_read_write_split_exclude_primary_no_replicas() { &replica_configs, LoadBalancingStrategy::RoundRobin, ReadWriteSplit::ExcludePrimary, + Default::default(), ); replicas.launch(); @@ -540,6 +547,7 @@ async fn test_read_write_split_exclude_primary_no_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::ExcludePrimary, + Default::default(), ); replicas.launch(); @@ -570,6 +578,7 @@ async fn test_read_write_split_include_primary_no_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -601,6 +610,7 @@ async fn test_read_write_split_with_banned_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -641,6 +651,7 @@ async fn test_read_write_split_with_banned_replicas() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -681,6 +692,7 @@ async fn test_prefer_primary_with_banned_replicas_falls_back_to_primary() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::PreferPrimary, + Default::default(), ); replicas.launch(); @@ -719,6 +731,7 @@ async fn test_read_write_split_exclude_primary_with_round_robin() { &replica_configs, LoadBalancingStrategy::RoundRobin, ReadWriteSplit::ExcludePrimary, + Default::default(), ); replicas.launch(); @@ -764,6 +777,7 @@ async fn test_monitor_shuts_down_on_notify() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas @@ -826,6 +840,7 @@ async fn test_monitor_does_not_ban_single_target() { &[pool_config], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -898,6 +913,7 @@ async fn test_monitor_does_not_ban_with_zero_ban_timeout() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -954,6 +970,7 @@ async fn test_include_primary_if_replica_banned_no_bans() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimaryIfReplicaBanned, + Default::default(), ); replicas.launch(); @@ -990,6 +1007,7 @@ async fn test_include_primary_if_replica_banned_with_ban() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimaryIfReplicaBanned, + Default::default(), ); replicas.launch(); @@ -1039,6 +1057,7 @@ async fn test_has_replicas_with_primary_and_replicas() { &replica_configs, LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1058,6 +1077,7 @@ async fn test_has_replicas_primary_only() { &[], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1073,6 +1093,7 @@ async fn test_has_replicas_empty() { &[], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(!lb.has_replicas()); @@ -1119,6 +1140,7 @@ async fn test_can_move_conns_to_same_config() { &[pool_config1.clone(), pool_config2.clone()], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let lb2 = LoadBalancer::new( @@ -1126,6 +1148,7 @@ async fn test_can_move_conns_to_same_config() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(lb1.can_move_conns_to(&lb2)); @@ -1143,6 +1166,7 @@ async fn test_can_move_conns_to_with_removed_replica() { &[pool_config1.clone(), pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let lb2 = LoadBalancer::new( @@ -1150,6 +1174,7 @@ async fn test_can_move_conns_to_with_removed_replica() { &[pool_config1], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(!lb1.can_move_conns_to(&lb2)); @@ -1167,6 +1192,7 @@ async fn test_can_move_conns_to_with_added_replica() { std::slice::from_ref(&pool_config1), LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let lb_new = LoadBalancer::new( @@ -1174,6 +1200,7 @@ async fn test_can_move_conns_to_with_added_replica() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(lb_old.can_move_conns_to(&lb_new)); @@ -1193,6 +1220,7 @@ async fn test_move_conns_to_with_added_replica_matches_by_address() { std::slice::from_ref(&pool_config1), LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb_old.launch(); @@ -1201,6 +1229,7 @@ async fn test_move_conns_to_with_added_replica_matches_by_address() { &[pool_config1.clone(), pool_config2.clone()], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb_new.launch(); @@ -1243,6 +1272,7 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan std::slice::from_ref(&existing_replica_config), LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); set_lsn_stats(&lb_old.targets[0], true, 100); @@ -1259,6 +1289,7 @@ async fn test_redetect_roles_marks_added_auto_target_replica_when_primary_unchan &[existing_replica_config, added_replica_config], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb_old.move_conns_to(&lb_new).unwrap(); @@ -1295,6 +1326,7 @@ async fn test_redetect_roles_leaves_auto_targets_pending_when_stats_are_invalid( &[config1, config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(lb.targets.iter().all(|target| target.role() == Role::Auto)); @@ -1323,6 +1355,7 @@ async fn test_redetect_roles_marks_auto_targets_replicas_when_all_valid_targets_ &[config1, config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); set_lsn_stats(&lb.targets[0], true, 100); @@ -1356,6 +1389,7 @@ async fn test_can_move_conns_to_different_addresses() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let lb2 = LoadBalancer::new( @@ -1363,6 +1397,7 @@ async fn test_can_move_conns_to_different_addresses() { &[pool_config3.clone(), pool_config3], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); assert!(!lb1.can_move_conns_to(&lb2)); @@ -1439,6 +1474,7 @@ async fn test_weighted_round_robin_smooth_distribution() { &[pool_config1, pool_config2], LoadBalancingStrategy::WeightedRoundRobin, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1473,6 +1509,7 @@ async fn test_weighted_round_robin_equal_weights() { &[pool_config1, pool_config2], LoadBalancingStrategy::WeightedRoundRobin, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1504,6 +1541,7 @@ async fn test_weighted_round_robin_zero_weight_never_selected() { &[pool_config1, pool_config2], LoadBalancingStrategy::WeightedRoundRobin, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1532,6 +1570,7 @@ async fn test_weighted_round_robin_proportional_distribution() { &[pool_config1, pool_config2], LoadBalancingStrategy::WeightedRoundRobin, ReadWriteSplit::IncludePrimary, + Default::default(), ); lb.launch(); @@ -1563,6 +1602,7 @@ async fn test_least_active_connections_prefers_pool_with_fewer_checked_out() { &[pool_config1, pool_config2], LoadBalancingStrategy::LeastActiveConnections, ReadWriteSplit::IncludePrimary, + Default::default(), ); replicas.launch(); @@ -1599,6 +1639,7 @@ fn setup_test_replicas_no_launch() -> LoadBalancer { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ) } @@ -1799,6 +1840,7 @@ fn test_ban_check_does_not_ban_single_target() { &[pool_config], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); // Don't launch - we're unit testing ban_check @@ -1864,6 +1906,7 @@ fn test_ban_check_does_not_ban_with_zero_ban_timeout() { &[pool_config1, pool_config2], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); // Set target as unhealthy @@ -2171,6 +2214,7 @@ async fn test_params_returns_all_replicas_down_when_empty() { &[], LoadBalancingStrategy::Random, ReadWriteSplit::IncludePrimary, + Default::default(), ); let request = Request::default(); diff --git a/pgdog/src/backend/pool/mod.rs b/pgdog/src/backend/pool/mod.rs index 5550a0e68..ada2d6e7e 100644 --- a/pgdog/src/backend/pool/mod.rs +++ b/pgdog/src/backend/pool/mod.rs @@ -40,7 +40,7 @@ pub use monitor::Monitor; pub use password::Password; pub use pool_impl::Pool; pub use request::Request; -pub use shard::Shard; +pub(crate) use shard::{CanonicalOids, Oids, Shard}; pub use state::State; pub use stats::Stats; diff --git a/pgdog/src/backend/pool/monitor.rs b/pgdog/src/backend/pool/monitor.rs index c6d36bb04..38e318726 100644 --- a/pgdog/src/backend/pool/monitor.rs +++ b/pgdog/src/backend/pool/monitor.rs @@ -43,6 +43,7 @@ //! The loop exits when the pool shuts down (e.g. on config reload), preventing //! refresh tasks from leaking across reloads. +use std::sync::Arc; use std::time::Duration; use super::{Error, Guard, Healtcheck, Pool, Request}; @@ -444,7 +445,12 @@ impl Monitor { for attempt in 0..connect_attempts { match safe_timeout( connect_timeout, - Server::connect(pool.addr(), options.clone(), reason), + Server::connect( + pool.addr(), + options.clone(), + reason, + Arc::clone(&pool.inner().oids), + ), ) .await { diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 1b161e22d..62fd13310 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -20,8 +20,8 @@ use crate::net::{Parameter, Parameters}; use super::inner::CheckInResult; use super::{ - Address, Comms, Config, Error, Guard, Healtcheck, Inner, Monitor, PoolConfig, Request, State, - Waiting, + Address, Comms, Config, Error, Guard, Healtcheck, Inner, Monitor, Oids, PoolConfig, Request, + State, Waiting, lb::TargetHealth, lsn_monitor::{LsnMonitor, ReplicaLag}, }; @@ -48,6 +48,7 @@ pub(crate) struct InnerSync { pub(super) params: OnceCell, pub(super) lsn_stats: RwLock, pub(super) lsn_role_change: Notify, + pub(super) oids: Arc, } impl std::fmt::Debug for Pool { @@ -59,8 +60,13 @@ impl std::fmt::Debug for Pool { } impl Pool { + #[cfg(test)] + pub(crate) fn new(config: &PoolConfig) -> Self { + Self::with_oid_mapping(config, Default::default()) + } + /// Create new connection pool. - pub fn new(config: &PoolConfig) -> Self { + pub(crate) fn with_oid_mapping(config: &PoolConfig, oids: Arc) -> Self { let id = next_pool_id(); Self { inner: Arc::new(InnerSync { @@ -73,6 +79,7 @@ impl Pool { params: OnceCell::new(), lsn_stats: RwLock::new(LsnStats::default()), lsn_role_change: Notify::new(), + oids, }), } } diff --git a/pgdog/src/backend/pool/shard/mod.rs b/pgdog/src/backend/pool/shard/mod.rs index 0373d0ba5..79c812e62 100644 --- a/pgdog/src/backend/pool/shard/mod.rs +++ b/pgdog/src/backend/pool/shard/mod.rs @@ -1,6 +1,7 @@ //! A shard is a collection of replicas and an optional primary. use arc_swap::ArcSwap; +use futures::try_join; use std::ops::Deref; use std::sync::Arc; use std::time::Duration; @@ -22,9 +23,11 @@ use crate::net::messages::FrontendPid; use super::{Error, Guard, LoadBalancer, Pool, PoolConfig, Request}; pub mod monitor; +mod oids; pub mod role_detector; use monitor::*; +pub(crate) use oids::{CanonicalOids, Oids}; use role_detector::*; #[cfg_attr(test, derive(Default))] @@ -53,7 +56,7 @@ pub(super) struct ShardConfig<'a> { /// /// Includes a primary and replicas. #[derive(Clone, Debug)] -pub struct Shard { +pub(crate) struct Shard { inner: Arc, } @@ -137,7 +140,7 @@ impl Shard { // This is syncrhonized by database/shard number, so this prevents // a thundering herd with 100s of users, for example, all fetching // the same schema. - let schema = self.schema_cache.get(self).await?; + let (_, schema) = try_join!(self.oids.load(self), self.schema_cache.get(self))?; self.schema.set(schema).expect("schema was not initialized"); Ok(true) @@ -164,12 +167,19 @@ impl Shard { /// We don't need it for this shard. pub(super) fn schema_not_needed(&self) { let _ = self.schema.set(Schema::default()); + self.skip_loading_oids() + } + + /// Skip loading this shard's type information + pub(super) fn skip_loading_oids(&self) { + self.oids.skip_load(); } /// Wait for the shard to load the schema. /// If the schema is loaded already, this returns immediately. pub(super) async fn wait_schema_loaded(&self) { self.schema.wait().await; + self.oids.wait().await; } /// Check that the shard LB targets are all launched. @@ -321,8 +331,9 @@ impl Deref for Shard { /// Shard connection pools /// and internal state. -#[derive(Default, Debug)] -pub struct ShardInner { +#[cfg_attr(test, derive(Default))] +#[derive(Debug)] +pub(crate) struct ShardInner { number: usize, lb: LoadBalancer, comms: Arc, @@ -331,6 +342,7 @@ pub struct ShardInner { schema: SetOnce, pub_sub_enabled: bool, schema_cache: SchemaCache, + oids: Arc, } impl ShardInner { @@ -346,8 +358,9 @@ impl ShardInner { pub_sub_enabled, schema_cache, } = shard; - let primary = primary.map(Pool::new); - let lb = LoadBalancer::new(&primary, replicas, lb_strategy, rw_split); + let oids = schema_cache.oids(&identifier.database, number); + let primary = primary.map(|config| Pool::with_oid_mapping(config, Arc::clone(&oids))); + let lb = LoadBalancer::new(&primary, replicas, lb_strategy, rw_split, Arc::clone(&oids)); let comms = Arc::new(ShardComms { shutdown: CancellationToken::new(), lsn_check_interval, @@ -362,6 +375,7 @@ impl ShardInner { schema: SetOnce::new(), pub_sub_enabled, schema_cache, + oids, } } } diff --git a/pgdog/src/backend/pool/shard/oids.rs b/pgdog/src/backend/pool/shard/oids.rs new file mode 100644 index 000000000..4805716ca --- /dev/null +++ b/pgdog/src/backend/pool/shard/oids.rs @@ -0,0 +1,137 @@ +use super::{Request, Shard}; +use crate::{ + backend::{Error, Server}, + net::DataRow, + sync::SetOnceCell, +}; +use std::collections::HashMap; +use std::sync::Arc; +use tracing::info; + +#[derive(Debug)] +/// The mapping from a shards type OID to a canonical one +pub(crate) struct Oids { + canonical_oids: Arc, + mappings: SetOnceCell, +} + +impl Oids { + pub(crate) fn new(canonical_oids: &Arc) -> Arc { + Arc::new(Self { + canonical_oids: Arc::clone(canonical_oids), + mappings: Default::default(), + }) + } + + pub(crate) async fn load(&self, shard: &Shard) -> Result<&OidMappings, Error> { + self.mappings + .get_or_try_init(|| async { + let mut server = shard.primary_or_replica(&Request::default()).await?; + let oids = load_oids(&mut server).await?; + let server_addr = server.addr().clone(); + drop(server); + + let canonical = self.canonical_oids.oids.wait().await; + let mut canonical_to_shard = HashMap::new(); + let mut shard_to_canonical = HashMap::new(); + for (type_name, oid) in oids { + let canonical = canonical + .get(&type_name) + .copied() + .ok_or(Error::MissingCanonicalOid(type_name))?; + if canonical == oid { + continue; + } + canonical_to_shard.insert(canonical, oid); + shard_to_canonical.insert(oid, canonical); + } + + debug_assert_eq!(canonical_to_shard.len(), shard_to_canonical.len()); + info!( + "loaded type info for {} types on shard {} [{}]", + canonical_to_shard.len(), + shard.number(), + server_addr, + ); + + Ok(OidMappings { + canonical_to_shard, + shard_to_canonical, + }) + }) + .await + } + + pub(crate) async fn wait(&self) { + self.mappings.wait().await; + } + + /// Sets this to an empty mapping if it has not already been loaded + pub(crate) fn skip_load(&self) { + let _ = self.mappings.set(Default::default()); + } + + /// Get the mappings. Returns `None` if no mappings have been loaded + pub(crate) fn get(&self) -> Option<&OidMappings> { + self.mappings.get() + } + + #[cfg(test)] + pub(crate) fn from_canonical(canonical_to_shard: HashMap) -> Arc { + let shard_to_canonical = canonical_to_shard.iter().map(|(&k, &v)| (v, k)).collect(); + Arc::new(Self { + canonical_oids: Default::default(), + mappings: SetOnceCell::from(OidMappings { + canonical_to_shard, + shard_to_canonical, + }), + }) + } +} + +impl Default for Oids { + fn default() -> Self { + Self { + canonical_oids: Default::default(), + mappings: SetOnceCell::from(OidMappings::default()), + } + } +} + +#[derive(Debug, Default)] +pub(crate) struct OidMappings { + pub(crate) canonical_to_shard: HashMap, + pub(crate) shard_to_canonical: HashMap, +} + +#[derive(Debug, Default)] +pub(crate) struct CanonicalOids { + oids: SetOnceCell>, +} + +impl CanonicalOids { + pub(crate) async fn load(&self, server: &mut Server) -> Result<(), Error> { + self.oids + .get_or_try_init(|| async { Ok(load_oids(server).await?.collect()) }) + .await + .map(|_| ()) + } +} + +async fn load_oids( + server: &mut Server, +) -> Result + use<>, Error> { + // OIDs < 10,000 are reserved for PG's internal use and are assumed to be stable + Ok(server + .fetch_all::( + "SELECT nspname || '.' || typname, pg_type.oid FROM pg_type INNER JOIN pg_namespace ON typnamespace = pg_namespace.oid WHERE pg_type.oid >= 10000", + ) + .await? + .into_iter() + .map(|row| { + ( + row.get_text(0).expect("selected 2 columns"), + row.get_int(1, true).expect("selected 2 columns") as u32, + ) + })) +} diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 40f720af9..dbb264f0e 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -6,13 +6,13 @@ use crate::{ net::{ Close, CloseComplete, FromBytes, Message, ParseComplete, Protocol, ProtocolMessage, ToBytes, - messages::{RowDescription, parse::Parse}, + messages::{ParameterDescription, RowDescription, parse::Parse}, }, }; use parking_lot::RwLock; use pgdog_config::PreparedStatements as PreparedStatementsLevel; -use super::Error; +use super::{Error, Oids}; use super::{ protocol::{ProtocolState, state::Action}, state::ExecutionCode, @@ -24,7 +24,7 @@ fn str_mem(s: &str) -> usize { s.len() + std::mem::size_of::() } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub enum HandleResult { Forward, Drop, @@ -53,17 +53,19 @@ pub struct PreparedStatements { capacity: usize, memory_used: usize, level: PreparedStatementsLevel, + oids: Arc, } +#[cfg(test)] impl Default for PreparedStatements { fn default() -> Self { - Self::new() + Self::new(Default::default()) } } impl PreparedStatements { /// New server prepared statements. - pub fn new() -> Self { + pub(crate) fn new(oids: Arc) -> Self { Self { global_cache: frontend::PreparedStatements::global(), local_cache: LruCache::unbounded(), @@ -73,6 +75,7 @@ impl PreparedStatements { capacity: usize::MAX, memory_used: 0, level: PreparedStatementsLevel::default(), + oids, } } @@ -208,24 +211,28 @@ impl PreparedStatements { } ProtocolMessage::Parse(parse) => { + let mut parse = parse.clone(); + let mut rewritten = self.rewrite_parse_data_types(&mut parse); + if !parse.anonymous() { if self.contains(parse.name()) { self.state.add_simulated(ParseComplete.message()?); return Ok(HandleResult::Drop); } else { - self.state.add('1'); self.parses.push_back(parse.name().to_string()); } // The client is sending named prepared statements, // but we're in ExtendedAnonymous mode so we rewrite // them to anonymous to avoid storing them in Postgres. if self.level.rewrite_anonymous() { - let mut parse = parse.clone(); parse.anonymize(); - return Ok(HandleResult::Rewrite(ProtocolMessage::Parse(parse))); + rewritten = true; } - } else { - self.state.add('1'); + } + + self.state.add('1'); + if rewritten { + return Ok(HandleResult::Rewrite(ProtocolMessage::Parse(parse))); } } @@ -272,7 +279,7 @@ impl PreparedStatements { } /// Should we forward the message to the client. - pub fn forward(&mut self, message: &Message) -> Result { + pub(crate) fn forward(&mut self, message: &mut Message) -> Result { let code = message.code(); let action = self.state.action(code)?; @@ -287,11 +294,12 @@ impl PreparedStatements { } 'T' => { + let maybe_row_description = self.parse_and_rewrite_row_description(message)?; if let Some(describe) = self.describes.pop_front() { - self.add_row_description( - &describe, - &RowDescription::from_bytes(message.to_bytes())?, - ); + let row_description = maybe_row_description + .map(Ok) + .unwrap_or_else(|| RowDescription::from_bytes(message.payload()))?; + self.add_row_description(&describe, row_description); }; } @@ -315,6 +323,10 @@ impl PreparedStatements { self.state.action(code)?; } + 't' => { + self.rewrite_parameter_description_data_types(message)?; + } + _ => (), } @@ -383,18 +395,18 @@ impl PreparedStatements { /// Get the Parse message stored in the global prepared statements /// cache for this statement. pub(crate) fn parse(&self, name: &str) -> Option { - self.global_cache.read().rewritten_parse(name) - } - - /// Get the globally stored RowDescription for this prepared statement, - /// if any. - pub fn row_description(&self, name: &str) -> Option { - self.global_cache.read().row_description(name) + self.global_cache + .read() + .rewritten_parse(name) + .map(|mut parse| { + self.rewrite_parse_data_types(&mut parse); + parse + }) } /// Handle a Describe message, storing the RowDescription for the /// statement in the global cache. - fn add_row_description(&self, name: &str, row_description: &RowDescription) { + fn add_row_description(&self, name: &str, row_description: RowDescription) { self.global_cache .write() .insert_row_description(name, row_description); @@ -459,6 +471,54 @@ impl PreparedStatements { close } + + pub(crate) fn replace_oids(&mut self, oids: &Arc) { + self.oids = Arc::clone(oids) + } + + fn rewrite_parse_data_types(&self, parse: &mut Parse) -> bool { + let Some(mappings) = self.oids.get() else { + return false; + }; + parse.rewrite_data_types(&mappings.canonical_to_shard) + } + + /// Rewrite the given RowDescription Message to have the canonical set of + /// OIDs. Returns the parsed RowDescription if parsing occurred + fn parse_and_rewrite_row_description( + &self, + message: &mut Message, + ) -> Result, Error> { + // RowDescription is emitted during cluster startup, so we can't + // require OIDs to be loaded. + let empty_mapping = Default::default(); + let mappings = &self.oids.get().unwrap_or(&empty_mapping).shard_to_canonical; + + if !mappings.is_empty() { + let mut row_description = RowDescription::from_bytes(message.payload())?; + if row_description.rewrite_data_types(mappings) { + message.replace_payload(row_description.to_bytes()); + } + Ok(Some(row_description)) + } else { + Ok(None) + } + } + + fn rewrite_parameter_description_data_types(&self, message: &mut Message) -> Result<(), Error> { + let Some(mappings) = self.oids.get() else { + return Ok(()); + }; + let mappings = &mappings.shard_to_canonical; + if mappings.is_empty() { + return Ok(()); + } + + let mut parameter_description = ParameterDescription::from_bytes(message.payload())?; + parameter_description.rewrite_data_types(mappings); + message.replace_payload(parameter_description.to_bytes()); + Ok(()) + } } #[cfg(test)] @@ -473,14 +533,14 @@ mod test { /// Build a PreparedStatements instance configured for ExtendedAnonymous mode. fn new_extended_anonymous() -> PreparedStatements { - let mut ps = PreparedStatements::new(); + let mut ps = PreparedStatements::default(); ps.set_prepared_statements_level(PreparedStatementsLevel::ExtendedAnonymous); ps } /// Build a PreparedStatements instance configured for Extended (default) mode. fn new_extended() -> PreparedStatements { - let mut ps = PreparedStatements::new(); + let mut ps = PreparedStatements::default(); ps.set_prepared_statements_level(PreparedStatementsLevel::Extended); ps } @@ -709,8 +769,8 @@ mod test { // Simulate a ReadyForQuery message. // First we need to add a 'Z' to the state so we can action it. ps.state.add('Z'); - let rfq = Message::new(ReadyForQuery::idle().to_bytes()); - ps.forward(&rfq).unwrap(); + let mut rfq = Message::new(ReadyForQuery::idle().to_bytes()); + ps.forward(&mut rfq).unwrap(); // In extended_anonymous mode, cache should be cleared after done. assert_eq!(ps.len(), 0, "local cache should be cleared after RFQ"); @@ -724,8 +784,8 @@ mod test { assert_eq!(ps.len(), 2); ps.state.add('Z'); - let rfq = Message::new(ReadyForQuery::idle().to_bytes()); - ps.forward(&rfq).unwrap(); + let mut rfq = Message::new(ReadyForQuery::idle().to_bytes()); + ps.forward(&mut rfq).unwrap(); // In extended mode, cache should be preserved. assert_eq!( @@ -807,6 +867,21 @@ mod test { assert!(matches!(result, HandleResult::Forward)); } + #[test] + fn parse_rewrites_if_oids_change() { + let mut ps = new_extended(); + ps.oids = Oids::from_canonical([(10000, 10001)].into_iter().collect()); + + let parse = Parse::named("stmt1", "SELECT $1, $2"); + let client_parse = parse.with_data_types(&[10000, 10002]); + let result = ps.handle(&ProtocolMessage::Parse(client_parse)).unwrap(); + let expected = parse.with_data_types(&[10001, 10002]); + assert_eq!( + result, + HandleResult::Rewrite(ProtocolMessage::Parse(expected)) + ); + } + // ------------------------------------------------------- // Simple query is unaffected by mode // ------------------------------------------------------- diff --git a/pgdog/src/backend/replication/logical/publisher/slot.rs b/pgdog/src/backend/replication/logical/publisher/slot.rs index d34e11a12..b9fb0bc20 100644 --- a/pgdog/src/backend/replication/logical/publisher/slot.rs +++ b/pgdog/src/backend/replication/logical/publisher/slot.rs @@ -110,6 +110,7 @@ impl ReplicationSlot { &self.address, ServerOptions::new_replication(), ConnectReason::Replication, + Default::default(), ) .await?, ); @@ -127,6 +128,7 @@ impl ReplicationSlot { &self.address, ServerOptions::default(), ConnectReason::Replication, + Default::default(), ) .await?, ); diff --git a/pgdog/src/backend/schema/cache/mod.rs b/pgdog/src/backend/schema/cache/mod.rs index 3bfc62295..795beab59 100644 --- a/pgdog/src/backend/schema/cache/mod.rs +++ b/pgdog/src/backend/schema/cache/mod.rs @@ -5,7 +5,7 @@ use dashmap::DashMap; use std::sync::Arc; use tokio::sync::Mutex; -use crate::backend::{Schema, Shard}; +use crate::backend::{CanonicalOids, Oids, Schema, Shard}; type Entry = Arc>; @@ -13,7 +13,12 @@ type Entry = Arc>; #[derive(Debug, Default, Clone)] pub(crate) struct SchemaCache { // Database => shard => Schema - cache: Arc>>, + cache: Arc>, + /// The canonical mapping of type names to OID + /// A cluster's canonical mapping of type names to OID + canonical_oids: Arc>>, + /// Each database, shard pair's mappings to the canonical type OIDs + shard_oids: Arc>>, } impl SchemaCache { @@ -26,9 +31,7 @@ impl SchemaCache { // This is synchronized. let entry = self .cache - .entry(shard.identifier().database.clone()) - .or_default() - .entry(shard.number()) + .entry((shard.identifier().database.clone(), shard.number())) .or_default() .clone(); @@ -46,4 +49,17 @@ impl SchemaCache { Ok(schema) } + + pub(crate) fn canonical_oids(&self, database: &str) -> Arc { + Arc::clone(&self.canonical_oids.entry(database.to_owned()).or_default()) + } + + pub(crate) fn oids(&self, database: &str, shard_number: usize) -> Arc { + Arc::clone( + &self + .shard_oids + .entry((database.to_owned(), shard_number)) + .or_insert_with(|| Oids::new(&self.canonical_oids(database))), + ) + } } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index aa715aa81..e98e108ee 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -1,6 +1,6 @@ //! PostgreSQL server connection. -use std::{ops::Deref, time::Duration}; +use std::{ops::Deref, sync::Arc, time::Duration}; use bytes::{BufMut, BytesMut}; use rustls_pki_types::ServerName; @@ -13,7 +13,7 @@ use tokio::{ use tracing::{debug, error, info, trace, warn}; use super::{ - ConnectReason, DisconnectReason, Error, PreparedStatements, ServerOptions, Stats, + ConnectReason, DisconnectReason, Error, Oids, PreparedStatements, ServerOptions, Stats, pool::Address, prepared_statements::HandleResult, }; use crate::{ @@ -96,10 +96,11 @@ impl MemoryUsage for Server { impl Server { /// Create new PostgreSQL server connection. - pub async fn connect( + pub(crate) async fn connect( addr: &Address, options: ServerOptions, connect_reason: ConnectReason, + oids: Arc, ) -> Result { let (user, auth_secrets) = addr.auth_credentials().await?; let total = auth_secrets.len(); @@ -110,6 +111,7 @@ impl Server { options.clone(), connect_reason, &auth_secret, + Arc::clone(&oids), ) .await { @@ -151,6 +153,7 @@ impl Server { options: ServerOptions, connect_reason: ConnectReason, auth_secret: &super::pool::Password, + oids: Arc, ) -> Result { debug!("=> {}", addr); let stream = TcpStream::connect(addr.addr().await?).await?; @@ -341,7 +344,7 @@ impl Server { params, changed_params: Parameters::default(), client_params: Parameters::default(), - prepared_statements: PreparedStatements::new(), + prepared_statements: PreparedStatements::new(oids), dirty: false, streaming: false, schema_changed: false, @@ -477,8 +480,8 @@ impl Server { Ok(message) => { // INVARIANT: omni dedup in multi_shard relies on this being process-unique; // never substitute a non-unique value here. - let message = message.stream(self.streaming).backend(self.id); - match self.prepared_statements.forward(&message) { + let mut message = message.stream(self.streaming).backend(self.id); + match self.prepared_statements.forward(&mut message) { Ok(forward) => { if forward { break message; @@ -1174,6 +1177,10 @@ impl Server { }, } } + + pub(crate) fn replace_oids(&mut self, oids: &Arc) { + self.prepared_statements.replace_oids(oids); + } } impl Drop for Server { @@ -1247,7 +1254,7 @@ pub mod test { &ServerOptions::default(), &Memory::default(), ), - prepared_statements: super::PreparedStatements::new(), + prepared_statements: super::PreparedStatements::default(), addr, dirty: false, streaming: false, @@ -1277,11 +1284,12 @@ pub mod test { } } - pub async fn test_server() -> Server { + pub(crate) async fn test_server() -> Server { Server::connect( &Address::new_test(), ServerOptions::default(), ConnectReason::Other, + Default::default(), ) .await .unwrap() @@ -1298,6 +1306,7 @@ pub mod test { }, ServerOptions::default(), ConnectReason::Other, + Default::default(), ) .await .unwrap() @@ -1308,6 +1317,7 @@ pub mod test { &Address::new_test(), ServerOptions::new_replication(), ConnectReason::Replication, + Default::default(), ) .await .unwrap() @@ -1366,7 +1376,13 @@ pub mod test { expected_secret, SystemTime::now() + Duration::from_secs(3600), ); - let result = Server::connect(&addr, ServerOptions::default(), ConnectReason::Other).await; + let result = Server::connect( + &addr, + ServerOptions::default(), + ConnectReason::Other, + Default::default(), + ) + .await; TokenCache::global().evict(&addr); let server = result.unwrap(); @@ -1426,7 +1442,13 @@ pub mod test { expected_secret, SystemTime::now() + Duration::from_secs(3600), ); - let result = Server::connect(&addr, ServerOptions::default(), ConnectReason::Other).await; + let result = Server::connect( + &addr, + ServerOptions::default(), + ConnectReason::Other, + Default::default(), + ) + .await; TokenCache::global().evict(&addr); let server = result.unwrap(); @@ -1600,13 +1622,6 @@ pub mod test { assert_eq!(c, msg.code()); } - // RowDescription saved. - let global = server.prepared_statements.parse(&name).unwrap(); - server - .prepared_statements - .row_description(global.name()) - .unwrap(); - server .send( &vec![ diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index f61d388a7..50eed73ba 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -224,11 +224,11 @@ impl GlobalCache { /// Client sent a Describe for a prepared statement and received a RowDescription. /// We record the RowDescription for later use by the results decoder. - pub fn insert_row_description(&mut self, name: &str, row_description: &RowDescription) { + pub fn insert_row_description(&mut self, name: &str, row_description: RowDescription) { if let Some(ref mut entry) = self.names.get_mut(name) && entry.row_description.is_none() { - entry.row_description = Some(row_description.clone()); + entry.row_description = Some(row_description); } } diff --git a/pgdog/src/frontend/router/parser/context.rs b/pgdog/src/frontend/router/parser/context.rs index 650c891d6..250dbb166 100644 --- a/pgdog/src/frontend/router/parser/context.rs +++ b/pgdog/src/frontend/router/parser/context.rs @@ -113,4 +113,9 @@ impl<'a> QueryParserContext<'a> { pub(super) fn is_session_mode(&self) -> bool { self.router_context.cluster.pooler_mode() == crate::config::PoolerMode::Session } + + #[cfg(feature = "new_parser")] + pub(super) fn is_canonicalizing_oids(&self) -> bool { + self.router_context.cluster.is_canonicalizing_oids() + } } diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 6f76db565..71e9f12d1 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -332,7 +332,16 @@ impl QueryParser { return Ok(Command::Deallocate); } - Node::SelectStmt(stmt) => self.select(&statement, stmt, context), + Node::SelectStmt(stmt) => { + if context.is_canonicalizing_oids() && references_pg_type(stmt) { + // Shard 0 is considered the canonical source for now + return Ok(Command::Query(Route::read( + ShardWithPriority::new_override_canonical_schema_info(Shard::Direct(0)), + ))); + } else { + self.select(&statement, stmt, context) + } + } Node::CopyStmt(stmt) => Self::copy(stmt, context), @@ -975,5 +984,34 @@ cfg_select! { _ => {} } +#[cfg(feature = "new_parser")] +fn references_pg_type(stmt: &nodes::SelectStmt) -> bool { + use pg_raw_parse::walk::{self, Recurse}; + use std::ops::ControlFlow; + + walk::walk_manual(stmt.into(), |node| match node { + Node::RangeVar(rv) if rv.relname() == Some("pg_type") => ControlFlow::Break(true), + // atttypid references pg_type.oid + Node::RangeVar(rv) if rv.relname() == Some("pg_attribute") => ControlFlow::Break(true), + Node::TypeCast(tc) + if let Some(tn) = tc.type_name() + && tn.names().iter().map(|s| s.sval()).eq([Some("regtype")]) => + { + ControlFlow::Break(true) + } + Node::FuncCall(fc) + if fc + .funcname() + .iter() + .map(Node::as_str) + .eq([Some("to_regtype")]) => + { + ControlFlow::Break(true) + } + _ => Recurse::yes(), + }) + .unwrap_or_default() +} + #[cfg(test)] mod test; diff --git a/pgdog/src/frontend/router/parser/query/test/test_select.rs b/pgdog/src/frontend/router/parser/query/test/test_select.rs index 9eea2eb8e..d06c1c1ea 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_select.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_select.rs @@ -247,9 +247,9 @@ fn test_system_catalog_sharded() { let mut updated = config().deref().clone(); updated.config.general.system_catalogs = SystemCatalogsBehavior::Sharded; - config::set(updated).unwrap(); + updated.config.general.canonicalize_type_information = true; - let mut test = QueryParserTest::new_with_config(&config()); + let mut test = QueryParserTest::new_with_config(&updated); let command = test.execute(vec![Query::new("SELECT * FROM pg_class").into()]); assert_eq!( @@ -259,7 +259,7 @@ fn test_system_catalog_sharded() { ); let command = test.execute(vec![ - Query::new("SELECT * FROM pg_type WHERE typname = 'int4'").into(), + Query::new("SELECT * FROM pg_class WHERE relname = $1").into(), ]); assert_eq!( command.route().shard(), @@ -268,10 +268,29 @@ fn test_system_catalog_sharded() { ); assert!(!command.route().is_omnisharded()); - // Reset to default - let mut updated = config().deref().clone(); - updated.config.general.system_catalogs = SystemCatalogsBehavior::default(); - config::set(updated).unwrap(); + #[cfg(feature = "new_parser")] + { + let command = test.execute(vec![Query::new("SELECT * FROM pg_type").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "pg_type queries should go to shard 0", + ); + + let command = test.execute(vec![Query::new("SELECT $1::regtype").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "regtype casts should go to shard 0", + ); + + let command = test.execute(vec![Query::new("SELECT to_regtype($1)").into()]); + assert_eq!( + command.route().shard(), + &Shard::Direct(0), + "to_regtype should go to shard 0", + ); + } } #[test] diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index 180c8547a..f1344ab59 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -420,6 +420,7 @@ pub enum OverrideReason { OnlyOneShard, RewriteUpdate, CrossShardFunction, + CanonicalSchemaInfo, } #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] @@ -508,6 +509,14 @@ impl ShardWithPriority { } } + #[cfg(feature = "new_parser")] + pub(crate) fn new_override_canonical_schema_info(shard: Shard) -> Self { + Self { + shard, + source: ShardSource::Override(OverrideReason::CanonicalSchemaInfo), + } + } + pub fn new_default_unset(shard: Shard) -> Self { Self { shard, diff --git a/pgdog/src/lib.rs b/pgdog/src/lib.rs index c1011ac13..dbba68f63 100644 --- a/pgdog/src/lib.rs +++ b/pgdog/src/lib.rs @@ -18,6 +18,7 @@ pub mod plugin; pub mod sighup; pub mod state; pub mod stats; +pub(crate) mod sync; pub mod tasks; #[cfg(test)] pub mod test_utils; diff --git a/pgdog/src/net/messages/mod.rs b/pgdog/src/net/messages/mod.rs index e4d83b878..733fd3e75 100644 --- a/pgdog/src/net/messages/mod.rs +++ b/pgdog/src/net/messages/mod.rs @@ -288,6 +288,10 @@ impl Message { pub fn transaction_error(&self) -> bool { self.code() == 'Z' && self.payload[5] as char == 'E' } + + pub fn replace_payload(&mut self, bytes: Bytes) { + self.payload = bytes; + } } /// Check that the message we received is what we expected. diff --git a/pgdog/src/net/messages/parameter_description.rs b/pgdog/src/net/messages/parameter_description.rs index cedc4ebc2..8229cadd6 100644 --- a/pgdog/src/net/messages/parameter_description.rs +++ b/pgdog/src/net/messages/parameter_description.rs @@ -1,5 +1,6 @@ use super::code; use super::prelude::*; +use std::collections::HashMap; #[derive(Debug, Clone, Default)] pub struct ParameterDescription { @@ -47,6 +48,14 @@ impl ParameterDescription { pub fn from_params(params: Vec) -> Self { Self { params } } + + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) { + for param in &mut self.params { + if let Some(&canonical) = mapping.get(&(*param as u32)) { + *param = canonical as i32; + } + } + } } #[cfg(test)] diff --git a/pgdog/src/net/messages/parse.rs b/pgdog/src/net/messages/parse.rs index e02ea9789..0bb8fb36d 100644 --- a/pgdog/src/net/messages/parse.rs +++ b/pgdog/src/net/messages/parse.rs @@ -1,8 +1,10 @@ //! Parse (F) message. use crate::net::c_string_buf_len; +use bytes::BytesMut; +use std::collections::HashMap; use std::fmt::Debug; use std::io::Cursor; -use std::mem::size_of; +use std::mem; use std::str::from_utf8; use std::str::from_utf8_unchecked; @@ -101,13 +103,6 @@ impl Parse { } } - pub fn data_types(&self) -> DataTypesIter<'_> { - DataTypesIter { - data_types: &self.data_types, - offset: 0, - } - } - pub fn data_types_ref(&self) -> Bytes { self.data_types.clone() } @@ -117,34 +112,45 @@ impl Parse { self.query = Bytes::from(query.to_string() + "\0"); self.original = None; } -} -#[derive(Debug)] -pub struct DataTypesIter<'a> { - data_types: &'a Bytes, - offset: usize, -} + /// Rewrite the data types of this message using the given mapping. + /// Returns whether any data types were changed. + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) -> bool { + if mapping.is_empty() { + return false; + } -impl DataTypesIter<'_> { - pub fn len(&self) -> usize { - (self.data_types.len() - size_of::()) / size_of::() + let mut rewritten = false; + // FIXME: This will always copy the bytes, we should only copy on write + let mut bytes = Cursor::new(BytesMut::from(mem::take(&mut self.data_types))); + for _ in 0..bytes.get_u16() { + let canonical_oid = bytes.get_u32(); + if let Some(&shard_oid) = mapping.get(&canonical_oid) { + // Note: We can't use BufMut here, as that writes to the end of the + // buffer, which will allocate. We just want to write to the + // existing bytes. + let pos = bytes.position() as usize; + let oid_bytes = &mut bytes.get_mut()[(pos - 4)..pos]; + oid_bytes.copy_from_slice(&shard_oid.to_be_bytes()); + rewritten = true; + self.original = None; + } + } + self.data_types = bytes.into_inner().freeze(); + rewritten } -} -impl Iterator for DataTypesIter<'_> { - type Item = i32; - - fn next(&mut self) -> Option { - let pos = self.offset * size_of::() + size_of::(); - self.offset += 1; - let mut cursor = Cursor::new(self.data_types); - cursor.advance(pos); - - if cursor.remaining() >= size_of::() { - Some(cursor.get_i32()) - } else { - None + #[cfg(test)] + pub(crate) fn with_data_types(&self, data_types: &[u32]) -> Self { + let mut bytes = BytesMut::new(); + bytes.put_u16(data_types.len() as _); + for &oid in data_types { + bytes.put_u32(oid); } + + let mut this = self.clone(); + this.data_types = bytes.freeze(); + this } } @@ -210,8 +216,6 @@ impl Protocol for Parse { #[cfg(test)] mod test { - use bytes::BytesMut; - use super::*; #[test] @@ -223,19 +227,7 @@ mod test { #[test] fn test_parse_from_bytes() { - let mut parse = Parse::named("__pgdog_1", "SELECT * FROM users"); - let mut data_types = BytesMut::new(); - data_types.put_i16(3); - data_types.put_i32(1); - data_types.put_i32(2); - data_types.put_i32(3); - parse.data_types = data_types.freeze(); - - let iter = parse.data_types(); - assert_eq!(iter.len(), 3); - for (i, v) in iter.enumerate() { - assert_eq!(i as i32 + 1, v); - } + let parse = Parse::named("__pgdog_1", "SELECT * FROM users"); assert_eq!(parse.name(), "__pgdog_1"); assert_eq!(parse.query(), "SELECT * FROM users"); @@ -252,8 +244,33 @@ mod test { let parse = Parse::from_bytes(b.freeze()).unwrap(); assert_eq!(parse.name(), "__pgdog_1"); assert_eq!(parse.query(), "SELECT * FROM users"); - assert_eq!(parse.data_types().len(), 0); + assert_eq!(parse.data_types_ref().get_i16(), 0); assert!(Parse::new_anonymous("SELECT 1").anonymous()); } + + #[test] + fn test_parse_rewrite_oids() { + let mut parse = Parse::named("", "").with_data_types(&[10_001, 10_011, 10_091]); + + let mapping = [(10_001, 10_001), (10_011, 10_002), (10_091, 10_003)] + .into_iter() + .collect(); + assert!(parse.rewrite_data_types(&mapping)); + + let expected = Parse::named("", "").with_data_types(&[10_001, 10_002, 10_003]); + assert_eq!(parse, expected); + + let mut parse_without_custom_types = Parse::named("", "").with_data_types(&[1]); + assert!(!parse_without_custom_types.rewrite_data_types(&mapping)); + } + + #[test] + fn test_rewriting_oids_clears_original_bytes() { + let mut parse = Parse::named("", "").with_data_types(&[10_001]); + parse.original = Some(Bytes::new()); + let mapping = [(10_001, 10_002)].into_iter().collect(); + assert!(parse.rewrite_data_types(&mapping)); + assert!(!parse.to_bytes().is_empty()); + } } diff --git a/pgdog/src/net/messages/row_description.rs b/pgdog/src/net/messages/row_description.rs index 9e370a4e4..4b8c119f0 100644 --- a/pgdog/src/net/messages/row_description.rs +++ b/pgdog/src/net/messages/row_description.rs @@ -1,6 +1,6 @@ //! RowDescription (B) message. -use std::collections::BTreeSet; +use std::collections::{BTreeSet, HashMap}; use std::ops::Deref; use std::sync::Arc; @@ -266,6 +266,20 @@ impl RowDescription { true } + + /// Replaces the data types of each field using the given mapping. + /// Returns whether any changes actually occurred. + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) -> bool { + let mut changed = false; + for field in Arc::make_mut(&mut self.fields) { + if let Some(&canonical) = mapping.get(&(field.type_oid as u32)) { + changed = true; + field.type_oid = canonical as i32; + } + } + + changed + } } impl Deref for RowDescription { diff --git a/pgdog/src/sync.rs b/pgdog/src/sync.rs new file mode 100644 index 000000000..ffd9fe563 --- /dev/null +++ b/pgdog/src/sync.rs @@ -0,0 +1,96 @@ +use std::fmt; +use tokio::sync::{Semaphore, SemaphorePermit, SetError, SetOnce, TryAcquireError}; + +/// A combination of `SetOnce` and `OnceCell`. +/// It allows only a single writer to run, with no contention on future reads, +/// while also allowing callers to wait on the value. +pub(crate) struct SetOnceCell { + value: SetOnce, + semaphore: Semaphore, +} + +impl SetOnceCell { + /// Equivalent to [`tokio::sync::OnceCell::get_or_try_init`] + pub(crate) async fn get_or_try_init(&self, f: F) -> Result<&T, E> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + if let Some(val) = self.get() { + Ok(val) + } else { + // Ensure only a single writer attempts to write. + // Err indicates that another task successfully initialized this + if let Ok(permit) = self.semaphore.acquire().await { + self.set_value(f().await?, permit); + } + + Ok(self.get().expect("always initialized")) + } + } + + /// Equivalent to [`SetOnce::wait`] + pub(crate) async fn wait(&self) -> &T { + self.value.wait().await + } + + /// Equivalent to [`SetOnce::get`] + pub(crate) fn get(&self) -> Option<&T> { + self.value.get() + } + + /// Equivalent to [`tokio::sync::OnceCell::set`] + pub(crate) fn set(&self, value: T) -> Result<(), SetError> { + if self.value.initialized() { + return Err(SetError::AlreadyInitializedError(value)); + } + + // Ensure no task is currently attempting to initialize this value + match self.semaphore.try_acquire() { + Ok(permit) => { + self.set_value(value, permit); + Ok(()) + } + Err(TryAcquireError::NoPermits) => Err(SetError::InitializingError(value)), + Err(TryAcquireError::Closed) => Err(SetError::AlreadyInitializedError(value)), + } + } + + /// Set the value and close the semaphore + fn set_value(&self, value: T, permit: SemaphorePermit) { + let Ok(_) = self.value.set(value) else { + panic!("set_value called when already initialized"); + }; + self.semaphore.close(); + permit.forget(); + } +} + +impl Default for SetOnceCell { + fn default() -> Self { + Self { + value: Default::default(), + semaphore: Semaphore::new(1), + } + } +} + +impl fmt::Debug for SetOnceCell +where + SetOnce: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("SetOnceCell").field(&self.value).finish() + } +} + +impl From for SetOnceCell { + fn from(value: T) -> Self { + let semaphore = Semaphore::new(0); + semaphore.close(); + Self { + value: SetOnce::from(value), + semaphore, + } + } +}