test(prepared_statements): avoid corrupting PS after setting number of PS to 0

pgdog/src/backend/server.rs

@@ -2809,6 +2809,35 @@ pub mod test {
         }
     }
 
+    /// Names Postgres itself thinks are prepared on this connection.
+    pub(crate) async fn prepared_in_postgres(server: &mut Server) -> Vec<String> {
+        server
+            .fetch_all::<String>("SELECT name FROM pg_prepared_statements")
+            .await
+            .unwrap()
+    }
+
+    /// Run a named statement the way a client would, leaving it to the
+    /// connection to prepare it first if it isn't already.
+    pub(crate) async fn execute_prepared(
+        server: &mut Server,
+        name: &str,
+        param: &[u8],
+    ) -> Vec<i64> {
+        use crate::net::bind::Parameter;
+
+        let request = ServerRequest {
+            messages: vec![
+                Bind::new_params(name, &[Parameter::new(param)]).into(),
+                Execute::new().into(),
+                Sync::new().into(),
+            ],
+            expected: 1,
+        };
+
+        server.fetch_all::<i64>(request).await.unwrap()
+    }
+
     #[tokio::test]
     async fn test_deallocate_all_clears_cache() {
         let mut server = test_server().await;

pgdog/src/frontend/prepared_statements/global_cache.rs

@@ -238,17 +238,6 @@ impl GlobalCache {
         }
     }
 
-    /// Clear the global cache. Test-only: rolling the name counter back
-    /// would reuse global statement names.
-    #[cfg(test)]
-    pub fn reset(&mut self) {
-        self.statements.clear();
-        self.names.clear();
-        self.unused.clear();
-        self.counter = 0;
-        self.versions = 0;
-    }
-
     /// Get the query string stored in the global cache
     /// for the given globally unique prepared statement name.
     #[inline]

pgdog/src/frontend/prepared_statements/mod.rs

@@ -194,10 +194,62 @@ pub fn run_maintenance() {
 
 #[cfg(test)]
 mod test {
+    use crate::backend::server::test::{execute_prepared, prepared_in_postgres, test_server};
     use crate::net::messages::Bind;
 
     use super::*;
 
+    #[tokio::test]
+    async fn test_close_unused_does_not_reuse_names() {
+        let mut client = PreparedStatements::new();
+
+        let mut first = Parse::named("client_a", "SELECT $1::bigint");
+        client.insert(&mut first);
+        let first_name = first.name();
+
+        let mut server = test_server().await;
+        assert_eq!(execute_prepared(&mut server, first_name, b"1").await, [1]);
+        assert_eq!(prepared_in_postgres(&mut server).await, [first_name]);
+
+        client.close("client_a");
+        PreparedStatements::global().write().close_unused(0);
+        assert!(PreparedStatements::global().read().is_empty());
+
+        assert_eq!(prepared_in_postgres(&mut server).await, [first_name]);
+
+        let mut second = Parse::named("client_b", "SELECT $1::bigint + 100");
+        client.insert(&mut second);
+        let second_name = second.name();
+        assert_ne!(second_name, first_name);
+
+        assert_eq!(
+            execute_prepared(&mut server, second_name, b"1").await,
+            [101]
+        );
+    }
+
+    #[tokio::test]
+    async fn test_close_unused_keeps_statements_clients_still_hold() {
+        let mut client = PreparedStatements::new();
+
+        let mut parse = Parse::named("client_a", "SELECT $1::bigint");
+        client.insert(&mut parse);
+        let name = parse.name();
+
+        let mut warm = test_server().await;
+        assert_eq!(execute_prepared(&mut warm, name, b"1").await, [1]);
+        assert_eq!(prepared_in_postgres(&mut warm).await, [name]);
+
+        PreparedStatements::global().write().close_unused(0);
+
+        assert_eq!(execute_prepared(&mut warm, name, b"1").await, [1]);
+
+        let mut cold = test_server().await;
+        assert!(prepared_in_postgres(&mut cold).await.is_empty());
+        assert_eq!(execute_prepared(&mut cold, name, b"1").await, [1]);
+        assert_eq!(prepared_in_postgres(&mut cold).await, [name]);
+    }
+
     #[test]
     fn test_maybe_rewrite() {
         let mut statements = PreparedStatements::default();