diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs index 010a4518e..2584ef2ac 100644 --- a/pgdog-config/src/core.rs +++ b/pgdog-config/src/core.rs @@ -167,10 +167,28 @@ impl ConfigAndUsers { Ok(()) } - /// Prepared statements are enabled. + /// Prepared statements level for clients using the `[general]` pooler mode. + /// + /// Prefer [`Self::prepared_statements_for`] with the pooler mode the client's + /// pool actually runs in: users and databases can override `pooler_mode`. pub fn prepared_statements(&self) -> PreparedStatementsLevel { - // Disable prepared statements automatically in session mode - if self.config.general.pooler_mode == PoolerMode::Session { + self.prepared_statements_for(self.config.general.pooler_mode) + } + + /// Prepared statements level for a client whose pool runs in `pooler_mode`. + /// + /// A pool that multiplexes server connections (transaction and statement + /// mode, and with read/write split, primary and replica servers) must track + /// prepared statements and re-prepare them per server, even when the + /// `[general]` default is session mode and would otherwise disable tracking. + /// + /// Session-pooled clients keep the `[general]`-derived behaviour: tracking is + /// not forced off for them, because the router still needs the statement + /// text to route `Bind` and `Execute` on sharded clusters. + pub fn prepared_statements_for(&self, pooler_mode: PoolerMode) -> PreparedStatementsLevel { + if pooler_mode == PoolerMode::Session + && self.config.general.pooler_mode == PoolerMode::Session + { PreparedStatementsLevel::Disabled } else { self.config.general.prepared_statements @@ -1028,6 +1046,53 @@ tls_server_private_key = "/certs/replica-client.key" ); } + /// A `[general]` session default must not disable tracking for a user or + /// database configured for transaction pooling. Session-pooled clients + /// keep whatever the `[general]` default resolves to. + #[test] + fn test_prepared_statements_for_pooler_mode() { + let mut config = ConfigAndUsers::default(); + config.config.general.pooler_mode = PoolerMode::Session; + config.config.general.prepared_statements = PreparedStatementsLevel::Extended; + + assert_eq!( + config.prepared_statements(), + PreparedStatementsLevel::Disabled + ); + assert_eq!( + config.prepared_statements_for(PoolerMode::Transaction), + PreparedStatementsLevel::Extended, + "transaction-pooled client must track statements despite the session default" + ); + assert_eq!( + config.prepared_statements_for(PoolerMode::Statement), + PreparedStatementsLevel::Extended + ); + assert_eq!( + config.prepared_statements_for(PoolerMode::Session), + PreparedStatementsLevel::Disabled, + "session-pooled client follows the session default" + ); + + config.config.general.pooler_mode = PoolerMode::Transaction; + assert_eq!( + config.prepared_statements(), + PreparedStatementsLevel::Extended + ); + assert_eq!( + config.prepared_statements_for(PoolerMode::Session), + PreparedStatementsLevel::Extended, + "session-pooled client on a transaction default keeps tracking (needed for sharded routing)" + ); + + config.config.general.prepared_statements = PreparedStatementsLevel::Disabled; + assert_eq!( + config.prepared_statements_for(PoolerMode::Transaction), + PreparedStatementsLevel::Disabled, + "explicitly disabled stays disabled" + ); + } + #[test] fn test_mirroring_config() { let source = r#" diff --git a/pgdog/src/backend/pool/connection/mirror/mod.rs b/pgdog/src/backend/pool/connection/mirror/mod.rs index b196cc5b2..6b99f37fd 100644 --- a/pgdog/src/backend/pool/connection/mirror/mod.rs +++ b/pgdog/src/backend/pool/connection/mirror/mod.rs @@ -51,9 +51,11 @@ pub(crate) struct Mirror { } impl Mirror { - fn new(params: &Parameters, config: &ConfigAndUsers) -> Self { + fn new(params: &Parameters, config: &ConfigAndUsers, cluster: &Cluster) -> Self { let mut prepared_statements = PreparedStatements::new(); - prepared_statements.set_level(config.prepared_statements()); + // Follow the destination cluster's pooler mode, like a regular client + // of that cluster would, rather than the `[general]` default. + prepared_statements.set_level(config.prepared_statements_for(cluster.pooler_mode())); Self { id: FrontendPid::new(), @@ -96,7 +98,7 @@ impl Mirror { ]); // Mirror traffic handler. - let mut mirror = Self::new(¶ms, &config); + let mut mirror = Self::new(¶ms, &config, cluster); // Same query engine as the client, except with a potentially different database config. // Use mirror.id so pool checkout (Request) and comms keying share one identity. @@ -376,10 +378,22 @@ mod test { }, ]); - let mirror = Mirror::new(¶ms, &config); + // Destination cluster in session mode follows the `[general]` default. + let cluster = Cluster::new_test_session_mode(&config); + let mirror = Mirror::new(¶ms, &config, &cluster); assert_eq!( mirror.prepared_statements.level(), PreparedStatementsLevel::Disabled ); + + // Destination cluster in transaction mode (the test cluster default) + // must track statements even though `[general]` is session mode. + let cluster = Cluster::new_test(&config); + assert_eq!(cluster.pooler_mode(), PoolerMode::Transaction); + let mirror = Mirror::new(¶ms, &config, &cluster); + assert_eq!( + mirror.prepared_statements.level(), + PreparedStatementsLevel::Extended + ); } } diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 96eb23e9a..e0e66baf8 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::time::Duration; use chrono::{DateTime, Utc}; +use pgdog_config::PoolerMode; use pgdog_config::users::PasswordKind; use timeouts::Timeouts; use tokio::{select, spawn}; @@ -532,6 +533,7 @@ impl Client { } let client_state = query_engine.client_state(); + let pooler_mode = query_engine.pooler_mode(); let cancellation_token = query_engine.cancellation_token(); @@ -546,7 +548,7 @@ impl Client { self.server_message(&mut query_engine, message).await?; } - buffer = self.buffer(client_state, &cancellation_token) => { + buffer = self.buffer(client_state, pooler_mode, &cancellation_token) => { let event = buffer?; // Only send requests to the backend if they are complete. @@ -649,14 +651,21 @@ impl Client { async fn buffer( &mut self, state: State, + pooler_mode: Option, cancellation_token: &CancellationToken, ) -> Result { self.client_request.clear(); // Check config once per request. let config = config::config(); - // Configure prepared statements cache. - self.prepared_statements.level = config.prepared_statements(); + // Configure prepared statements cache using the pooler mode this + // client's pool actually runs in. The `[general]` default alone is + // wrong when it is session mode but this user is transaction pooled: + // untracked Parse and Bind could land on different servers. + self.prepared_statements.level = match pooler_mode { + Some(pooler_mode) => config.prepared_statements_for(pooler_mode), + None => config.prepared_statements(), + }; self.timeouts = Timeouts::from_config(&config.config.general); self.query_log_stdout = config.config.general.query_log_stdout; self.query_size_limit = config.config.general.query_size_limit; diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index 9196bcf8b..92b5897db 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -10,6 +10,7 @@ use crate::{ net::{ErrorResponse, Message, Parameters}, state::State, }; +use pgdog_config::PoolerMode; use tokio_util::sync::CancellationToken; use tracing::debug; @@ -130,6 +131,16 @@ impl QueryEngine { self.stats.state } + /// Pooler mode of the cluster this client is connected to, resolved from + /// the user, database and `[general]` config. `None` when the client has + /// no cluster (e.g. admin database). + pub(crate) fn pooler_mode(&self) -> Option { + self.backend + .cluster() + .ok() + .map(|cluster| cluster.pooler_mode()) + } + /// Handle client request. pub(crate) async fn handle( &mut self, diff --git a/pgdog/src/frontend/client/test/mod.rs b/pgdog/src/frontend/client/test/mod.rs index febb77a5e..72e75117b 100644 --- a/pgdog/src/frontend/client/test/mod.rs +++ b/pgdog/src/frontend/client/test/mod.rs @@ -255,7 +255,7 @@ async fn test_abrupt_disconnect() { drop(conn); let event = client - .buffer(State::Idle, &CancellationToken::new()) + .buffer(State::Idle, None, &CancellationToken::new()) .await .unwrap(); assert_eq!(event, BufferEvent::DisconnectAbrupt); @@ -277,7 +277,7 @@ async fn test_client_idle_timeout() { let start = Instant::now(); let res = client - .buffer(State::Idle, &CancellationToken::new()) + .buffer(State::Idle, None, &CancellationToken::new()) .await .unwrap(); assert_eq!(res, BufferEvent::DisconnectAbrupt); @@ -290,7 +290,7 @@ async fn test_client_idle_timeout() { assert!( timeout( Duration::from_millis(50), - client.buffer(State::IdleInTransaction, &CancellationToken::new()) + client.buffer(State::IdleInTransaction, None, &CancellationToken::new()) ) .await .is_err() @@ -643,7 +643,7 @@ async fn test_query_timeout() { conn.write_all(&buf).await.unwrap(); client - .buffer(State::Idle, &CancellationToken::new()) + .buffer(State::Idle, None, &CancellationToken::new()) .await .unwrap(); let result = client.client_messages(&mut engine).await; diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs index 6adadff89..c535d0d4e 100644 --- a/pgdog/src/frontend/client/test/test_client.rs +++ b/pgdog/src/frontend/client/test/test_client.rs @@ -238,8 +238,13 @@ impl TestClient { /// Process a request. pub(crate) async fn try_process(&mut self) -> Result<(), Box> { + let pooler_mode = self.engine.pooler_mode(); self.client - .buffer(self.engine.stats().state, &CancellationToken::new()) + .buffer( + self.engine.stats().state, + pooler_mode, + &CancellationToken::new(), + ) .await?; self.client.client_messages(&mut self.engine).await?;