Deallocate client prepared statements at checkin

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

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"

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"

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<Vec<Query>> = Lazy::new(Vec::new);
 /// client modifications.
 #[allow(dead_code)]
 pub struct Cleanup {
-    queries: &'static Vec<Query>,
+    queries: Cow<'static, [Query]>,
     reset: bool,
     dirty: bool,
-    deallocate: bool,
     close: Vec<Close>,
 }
 
 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,16 +100,15 @@ impl Cleanup {
     /// Cleanup prepared statements.
     pub fn prepared_statements() -> Self {
         Self {
-            queries: &*PREPARED,
-            deallocate: true,
+            queries: Cow::Borrowed(&PREPARED),
             ..Default::default()
         }
     }
 
     /// 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
-    }
 }

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<i32> = server.fetch_all("SELECT 1").await.unwrap();
         assert_eq!(one[0], 1);
     }

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<String> = 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<String> = 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.
 }