diff --git a/integration/pgdog.toml b/integration/pgdog.toml index d5fd64d75..eed85ee83 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -53,6 +53,18 @@ host = "127.0.0.1" role = "replica" read_only = true +# ------------------------------------------------------------------------------ +# ----- Database :: pgdog_leak ------------------------------------------------- +# Exactly one server connection, kept around: tests for session state leaking +# between clients need the next client to get the same connection back. + +[[databases]] +name = "pgdog_leak" +host = "127.0.0.1" +database_name = "pgdog" +pool_size = 1 +min_pool_size = 1 + # ------------------------------------------------------------------------------ # ----- Database :: pgdog_sharded ---------------------------------------------- diff --git a/integration/python/test_pg_dump.py b/integration/python/test_pg_dump.py new file mode 100644 index 000000000..0b06361ff --- /dev/null +++ b/integration/python/test_pg_dump.py @@ -0,0 +1,70 @@ +"""pg_dump must work against a pooled connection that already served one. + +pg_dump creates a SQL-level prepared statement ("dumpfunc"). In transaction +pooling the server connection goes back into the pool at the end of the +transaction, so the next dump that lands on it fails with "prepared statement +already exists" unless the pooler cleans that state up. + +Runs against a database with a single server connection, so every dump lands +on the one another dump just used. +""" + +import os +import subprocess + +import psycopg + +DUMPS = 3 + + +def connect(): + conn = psycopg.connect( + user="pgdog", + password="pgdog", + dbname="pgdog_leak", + host="127.0.0.1", + port=6432, + ) + conn.autocommit = True + return conn + + +def pg_dump(): + return subprocess.run( + ["pg_dump", "-h", "127.0.0.1", "-p", "6432", "-U", "pgdog", "-d", "pgdog_leak"], + env=dict(os.environ, PGPASSWORD="pgdog"), + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + ) + + +def test_repeated_dumps(): + # pg_dump only prepares its statement when there are functions to dump, + # so don't rely on whatever else happens to live in the database. + conn = connect() + conn.execute("CREATE OR REPLACE FUNCTION public.pg_dump_probe() RETURNS int AS $$ SELECT 1 $$ LANGUAGE SQL") + conn.close() + + for attempt in range(DUMPS): + result = pg_dump() + assert result.returncode == 0, ( + f"dump {attempt + 1} of {DUMPS} failed: {result.stderr.strip()}" + ) + + +def test_prepared_statement_with_dirty_connection(): + """A connection can need both a parameter reset and a deallocate.""" + conn = connect() + conn.execute("SET pgdog.pin TO true") + conn.execute("PREPARE pg_dump_probe_stmt AS SELECT 1") + conn.close() + + conn = connect() + left = conn.execute( + "SELECT count(*) FROM pg_catalog.pg_prepared_statements " + "WHERE name = 'pg_dump_probe_stmt'" + ).fetchone()[0] + conn.close() + + assert left == 0, "prepared statement outlived its client's checkin" diff --git a/integration/users.toml b/integration/users.toml index bba115a85..360246b5b 100644 --- a/integration/users.toml +++ b/integration/users.toml @@ -3,6 +3,11 @@ name = "pgdog" database = "pgdog" password = "pgdog" +[[users]] +name = "pgdog" +database = "pgdog_leak" +password = "pgdog" + [[users]] name = "pgdog_migrator" database = "pgdog" diff --git a/pgdog/src/backend/pool/cleanup.rs b/pgdog/src/backend/pool/cleanup.rs index 712b127f3..db10c5adf 100644 --- a/pgdog/src/backend/pool/cleanup.rs +++ b/pgdog/src/backend/pool/cleanup.rs @@ -1,4 +1,6 @@ //! Cleanup queries for servers altered by client behavior. +use std::borrow::Cow; + use once_cell::sync::Lazy; use crate::net::{Close, Query}; @@ -27,20 +29,18 @@ static NONE: Lazy> = Lazy::new(Vec::new); /// client modifications. #[allow(dead_code)] pub struct Cleanup { - queries: &'static Vec, + queries: Cow<'static, [Query]>, reset: bool, dirty: bool, - deallocate: bool, close: Vec, } impl Default for Cleanup { fn default() -> Self { Self { - queries: &*NONE, + queries: Cow::Borrowed(&NONE), reset: false, dirty: false, - deallocate: false, close: vec![], } } @@ -63,11 +63,20 @@ impl std::fmt::Display for Cleanup { impl Cleanup { /// New cleanup operation. pub fn new(guard: &Guard, server: &mut Server) -> Self { + // A client that prepared statements with SQL leaves them on the + // connection. They belong to its session, so drop them before another + // client gets the connection and collides with their names. + let deallocate = server.schema_changed() || server.sync_prepared(); + let mut clean = if guard.reset { Self::all() } else if server.dirty() { - Self::parameters() - } else if server.schema_changed() { + let mut clean = Self::parameters(); + if deallocate { + clean.add(&PREPARED); + } + clean + } else if deallocate { Self::prepared_statements() } else { Self::none() @@ -78,6 +87,11 @@ impl Cleanup { clean } + /// Append more queries to run during the same cleanup. + fn add(&mut self, queries: &'static [Query]) { + self.queries.to_mut().extend_from_slice(queries); + } + /// Number of queries to run for cleanup. pub fn len(&self) -> usize { self.queries.len() @@ -86,8 +100,7 @@ impl Cleanup { /// Cleanup prepared statements. pub fn prepared_statements() -> Self { Self { - queries: &*PREPARED, - deallocate: true, + queries: Cow::Borrowed(&PREPARED), ..Default::default() } } @@ -95,7 +108,7 @@ impl Cleanup { /// Cleanup parameters. pub fn parameters() -> Self { Self { - queries: &*DIRTY, + queries: Cow::Borrowed(&DIRTY), dirty: true, ..Default::default() } @@ -106,8 +119,7 @@ impl Cleanup { Self { reset: true, dirty: true, - deallocate: true, - queries: &*ALL, + queries: Cow::Borrowed(&ALL), close: vec![], } } @@ -124,7 +136,7 @@ impl Cleanup { /// Get queries to execute on the server to perform cleanup. pub fn queries(&self) -> &[Query] { - self.queries + &self.queries } /// Prepared statemens to close. @@ -135,8 +147,4 @@ impl Cleanup { pub fn is_reset_params(&self) -> bool { self.dirty } - - pub fn is_deallocate(&self) -> bool { - self.deallocate - } } diff --git a/pgdog/src/backend/pool/guard.rs b/pgdog/src/backend/pool/guard.rs index 03d27cf76..7cfb591ae 100644 --- a/pgdog/src/backend/pool/guard.rs +++ b/pgdog/src/backend/pool/guard.rs @@ -137,7 +137,6 @@ impl Guard { conn_recovery: ConnectionRecovery, ) -> Result<(), Error> { let schema_changed = server.schema_changed(); - let sync_prepared = server.sync_prepared(); let needs_drain = server.needs_drain(); if needs_drain { @@ -180,11 +179,9 @@ impl Guard { server.stats().get_state(), server.addr() ); + // The cache is dropped by the DEALLOCATE ALL / DISCARD ALL + // response, so there is nothing to clear here. server.execute_batch(cleanup.queries()).await?; - - if cleanup.is_deallocate() { - server.prepared_statements_mut().clear(); - } server.cleaned(); debug!( @@ -202,15 +199,6 @@ impl Guard { server.reset_params(); } - if sync_prepared { - debug!( - "[cleanup] syncing prepared statements, server in \"{}\" state [{}]", - server.stats().get_state(), - server.addr() - ); - server.sync_prepared_statements().await?; - } - Ok(()) } } @@ -762,7 +750,7 @@ mod test { } #[tokio::test] - async fn test_cleanup_syncs_prepared_statements() { + async fn test_cleanup_deallocates_client_prepared_statements() { crate::logger(); let mut server = Guard::new( @@ -802,10 +790,16 @@ mod test { ); assert!( - server.prepared_statements_mut().contains("test_stmt"), - "Statement should be in local cache after sync" + !server.prepared_statements_mut().contains("test_stmt"), + "statement prepared by a client must not outlive its checkin" ); + // The next client can use the same name, which is what pg_dump does. + server + .execute("PREPARE test_stmt AS SELECT $1::bigint") + .await + .unwrap(); + let one: Vec = server.fetch_all("SELECT 1").await.unwrap(); assert_eq!(one[0], 1); } diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index 35af9706d..2bab4d84f 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/pgdog/src/backend/pool/test/mod.rs @@ -437,14 +437,14 @@ async fn test_prepared_statements_limit() { assert_eq!(guard.prepared_statements_mut().len(), 2); // Let's make sure Postgres agreees. - guard.sync_prepared_statements().await.unwrap(); + let named: Vec = guard + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); // It's random! - assert!( - guard.prepared_statements_mut().contains("__pgdog_99") - || guard.prepared_statements_mut().contains("__pgdog_98") - ); - assert_eq!(guard.prepared_statements_mut().len(), 2); + assert!(named.contains(&"__pgdog_99".to_string()) || named.contains(&"__pgdog_98".to_string())); + assert_eq!(named.len(), 2); assert_eq!(guard.stats().total().prepared_statements, 2); // stats are accurate. let pool = pool_with_prepared_capacity(100); @@ -476,10 +476,13 @@ async fn test_prepared_statements_limit() { assert_eq!(guard.stats().total().prepared_statements, 100); // stats are accurate. // Let's make sure Postgres agreees. - guard.sync_prepared_statements().await.unwrap(); + let named: Vec = guard + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); - assert!(guard.prepared_statements_mut().contains("__pgdog_99")); - assert_eq!(guard.prepared_statements_mut().len(), 100); + assert!(named.contains(&"__pgdog_99".to_string())); + assert_eq!(named.len(), 100); assert_eq!(guard.stats().total().prepared_statements, 100); // stats are accurate. } diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index aa715aa81..4b7650d8c 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -564,9 +564,9 @@ impl Server { let cmd = CommandComplete::from_bytes(message.to_bytes())?; match cmd.command() { "PREPARE" | "DEALLOCATE" => self.sync_prepared = true, - "DEALLOCATE ALL" => self.prepared_statements.clear(), + "DEALLOCATE ALL" => self.clear_prepared_statements(), "DISCARD ALL" => { - self.prepared_statements.clear(); + self.clear_prepared_statements(); self.client_params.clear(); } "RESET" => self.client_params.clear(), // Someone reset params, we're gonna need to re-sync. @@ -902,25 +902,6 @@ impl Server { Ok(()) } - /// Synchronize prepared statements from Postgres. - pub(super) async fn sync_prepared_statements(&mut self) -> Result<(), Error> { - let names = self - .fetch_all::("SELECT name FROM pg_prepared_statements") - .await?; - - for name in names { - self.prepared_statements.prepared(&name); - } - - debug!("prepared statements synchronized [{}]", self.addr()); - - let count = self.prepared_statements.len(); - self.stats.set_prepared_statements(count); - self.sync_prepared = false; - - Ok(()) - } - /// Close any prepared statements that exceed cache capacity. pub(super) fn ensure_prepared_capacity(&mut self) -> Vec { let close = self.prepared_statements.ensure_capacity(); @@ -977,7 +958,14 @@ impl Server { #[inline] pub fn reset_schema_changed(&mut self) { self.schema_changed = false; + self.clear_prepared_statements(); + } + + /// Drop the prepared statements cache, and the stat that counts them. + #[inline] + fn clear_prepared_statements(&mut self) { self.prepared_statements.clear(); + self.stats.clear_prepared_statements(); } #[inline] @@ -1116,6 +1104,7 @@ impl Server { #[inline] pub(super) fn cleaned(&mut self) { self.dirty = false; + self.sync_prepared = false; self.stats.cleaned(); } @@ -2115,7 +2104,7 @@ pub mod test { assert_eq!(msg.code(), c); } assert!(server.sync_prepared()); - server.sync_prepared_statements().await.unwrap(); + server.prepared_statements.prepared("__pgdog_1"); assert!(server.prepared_statements.contains("__pgdog_1")); let describe = Describe::new_statement("__pgdog_1"); @@ -2634,6 +2623,25 @@ pub mod test { assert!(server.done()); } + #[tokio::test] + async fn test_reset_schema_changed_clears_cache() { + let mut server = test_server().await; + + server + .send(&vec![Parse::named("__pgdog_1", "SELECT 1").into(), Flush.into()].into()) + .await + .unwrap(); + let msg = server.read().await.unwrap(); + assert_eq!(msg.code(), '1'); + assert!(!server.prepared_statements.is_empty()); + + // A schema change invalidates everything we cached for this connection. + server.reset_schema_changed(); + + assert!(server.prepared_statements.is_empty()); + assert_eq!(server.stats().total().prepared_statements, 0); + } + #[tokio::test] async fn test_discard_all_clears_cache() { let mut server = test_server().await; @@ -2841,11 +2849,11 @@ pub mod test { "sync_prepared flag should be set after PREPARE command" ); - server.sync_prepared_statements().await.unwrap(); + server.cleaned(); assert!( !server.sync_prepared(), - "sync_prepared flag should be cleared after sync_prepared_statements()" + "sync_prepared flag should be cleared once the connection is cleaned" ); server.execute("SELECT 1").await.unwrap(); @@ -3779,10 +3787,12 @@ pub mod test { "cache should be cleared after RFQ in extended_anonymous mode" ); // Verify Postgres has no named prepared statements stored. - server.sync_prepared_statements().await.unwrap(); - assert_eq!( - server.prepared_statements.len(), - 0, + let named: Vec = server + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); + assert!( + named.is_empty(), "Postgres should have no prepared statements in extended_anonymous mode" ); } @@ -3820,8 +3830,11 @@ pub mod test { assert!(server.done()); assert_eq!(server.prepared_statements.len(), 0); // Verify Postgres has no named prepared statements stored. - server.sync_prepared_statements().await.unwrap(); - assert_eq!(server.prepared_statements.len(), 0); + let named: Vec = server + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); + assert!(named.is_empty()); } #[tokio::test] @@ -3871,10 +3884,12 @@ pub mod test { assert!(server.done()); } // Verify Postgres has no named prepared statements stored. - server.sync_prepared_statements().await.unwrap(); - assert_eq!( - server.prepared_statements.len(), - 0, + let named: Vec = server + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); + assert!( + named.is_empty(), "Postgres should have no prepared statements after repeated anonymous usage" ); } @@ -3908,8 +3923,11 @@ pub mod test { assert!(server.done()); // Verify Postgres has no named prepared statements stored. - server.sync_prepared_statements().await.unwrap(); - assert_eq!(server.prepared_statements.len(), 0); + let named: Vec = server + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); + assert!(named.is_empty()); } #[tokio::test] @@ -3942,8 +3960,11 @@ pub mod test { // Server should still be usable. verify_server_usable(&mut server).await; // Verify Postgres has no named prepared statements stored. - server.sync_prepared_statements().await.unwrap(); - assert_eq!(server.prepared_statements.len(), 0); + let named: Vec = server + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); + assert!(named.is_empty()); } #[tokio::test] @@ -3993,8 +4014,11 @@ pub mod test { assert!(server.done()); // Verify Postgres has no named prepared statements stored. - server.sync_prepared_statements().await.unwrap(); - assert_eq!(server.prepared_statements.len(), 0); + let named: Vec = server + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); + assert!(named.is_empty()); } #[tokio::test] @@ -4041,10 +4065,12 @@ pub mod test { assert!(server.prepared_statements.ensure_capacity().is_empty()); } // Verify Postgres has no named prepared statements stored. - server.sync_prepared_statements().await.unwrap(); - assert_eq!( - server.prepared_statements.len(), - 0, + let named: Vec = server + .fetch_all("SELECT name FROM pg_prepared_statements") + .await + .unwrap(); + assert!( + named.is_empty(), "Postgres should have no prepared statements despite many named parses" ); } diff --git a/pgdog/src/backend/stats.rs b/pgdog/src/backend/stats.rs index 25b7fab03..43df2ad90 100644 --- a/pgdog/src/backend/stats.rs +++ b/pgdog/src/backend/stats.rs @@ -170,11 +170,9 @@ impl Stats { self.local.last_checkout.prepared_statements += 1; } - /// Overwrite how many prepared statements we have in the cache for stats. - pub fn set_prepared_statements(&mut self, size: usize) { - self.local.total.prepared_statements = size; - self.local.total.prepared_sync += 1; - self.local.last_checkout.prepared_sync += 1; + /// Prepared statements are gone from the server, so the cache is empty. + pub fn clear_prepared_statements(&mut self) { + self.local.total.prepared_statements = 0; self.sync_to_shared(); }