diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index f574ec87b..60d88e501 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -2310,7 +2310,7 @@ pub mod test { let mut prep = PreparedStatements::new(); let mut parse = Parse::named("test", "SELECT 1::bigint"); - prep.insert_anyway(&mut parse); + prep.insert_prepare(&mut parse); assert_eq!(parse.name(), "__pgdog_1"); server diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs index e427cf26f..08f87f829 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_simple_prepared.rs @@ -51,3 +51,56 @@ async fn test_rewrite_prepare() { "expected rewritten prepared statement" ); } + +fn rewritten_query(messages: &[ProtocolMessage]) -> String { + match &messages[0] { + ProtocolMessage::Query(query) => query.query().to_string(), + other => panic!("expected Query, got {other:#?}"), + } +} + +#[tokio::test] +async fn test_reprepare_releases_previous_statement() { + load_test(); + + change_config(|general| { + general.prepared_statements = PreparedStatements::Full; + }); + + let global = crate::frontend::PreparedStatements::global(); + let mut client = Client::new_test(Stream::dev_null(), Parameters::default()); + + assert_eq!(global.read().len(), 0); + + let first = rewritten_query( + &run_test( + &mut client, + &[Query::new("PREPARE reused AS SELECT $1::bigint").into()], + ) + .await, + ); + + // The client name is replaced, the statement is not. + assert!(first.starts_with("PREPARE __pgdog_")); + assert!(first.ends_with(" AS SELECT $1::bigint")); + + assert_eq!(global.read().len(), 1); + + let second = rewritten_query( + &run_test( + &mut client, + &[Query::new("PREPARE reused AS SELECT $1::bigint + 1").into()], + ) + .await, + ); + + assert!(second.starts_with("PREPARE __pgdog_")); + assert!(second.ends_with(" AS SELECT $1::bigint + 1")); + + // Re-PREPARE mints a new global name. + assert_eq!(global.read().len(), 2); + + // Only the statement the client replaced is evictable. + assert_eq!(global.write().close_unused(0), 1); + assert_eq!(global.read().len(), 1); +} diff --git a/pgdog/src/frontend/prepared_statements/global_cache.rs b/pgdog/src/frontend/prepared_statements/global_cache.rs index 9d50e84fc..4ebdb3908 100644 --- a/pgdog/src/frontend/prepared_statements/global_cache.rs +++ b/pgdog/src/frontend/prepared_statements/global_cache.rs @@ -8,19 +8,20 @@ use std::{collections::hash_map::HashMap, str::from_utf8}; use fnv::FnvHashSet as HashSet; +/// Identity of a prepared statement inside the global cache. +pub type Counter = usize; + // Format the globally unique prepared statement // name based on the counter. -fn global_name(counter: usize) -> String { +fn global_name(counter: Counter) -> String { format!("__pgdog_{}", counter) } -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone)] pub struct Statement { parse: Parse, rewrite: Option, row_description: Option, - #[allow(dead_code)] - version: usize, cache_key: CacheKey, } @@ -28,7 +29,7 @@ impl MemoryUsage for Statement { #[inline] fn memory_usage(&self) -> usize { self.parse.len() - + if let Some(ref row_description) = self.row_description { + + if let Some(row_description) = &self.row_description { row_description.memory_usage() } else { 0 @@ -49,16 +50,19 @@ impl Statement { /// Prepared statements cache key. /// -/// If these match, it's effectively the same statement. +/// If two `Extended` keys match, it's effectively the same statement. /// If they don't, e.g. client sent the same query but /// with different data types, we can't re-use it and /// need to plan a new one. /// -#[derive(Debug, Clone, PartialEq, Hash, Eq, Default)] -pub struct CacheKey { - pub query: Bytes, - pub data_types: Bytes, - pub version: usize, +/// A `Simple` key comes from SQL `PREPARE` and matches nothing but itself. +/// Its declared argument types are not captured, so two of those +/// statements are never known to be the same. +/// +#[derive(Debug, Clone, PartialEq, Hash, Eq)] +pub enum CacheKey { + Extended { query: Bytes, data_types: Bytes }, + Simple { query: Bytes, unique: Counter }, } impl MemoryUsage for CacheKey { @@ -70,15 +74,24 @@ impl MemoryUsage for CacheKey { } impl CacheKey { + fn query_ref(&self) -> &Bytes { + match self { + Self::Extended { query, .. } => query, + Self::Simple { query, .. } => query, + } + } + pub fn query(&self) -> Result<&str, crate::net::Error> { + let query = self.query_ref(); + // Postgres string. - Ok(from_utf8(&self.query[0..self.query.len() - 1])?) + Ok(from_utf8(&query[0..query.len() - 1])?) } } #[derive(Debug, Copy, Clone)] pub struct CachedStmt { - pub counter: usize, + pub counter: Counter, pub used: usize, } @@ -101,6 +114,8 @@ impl CachedStmt { /// /// 1. Mapping between unique prepared statement identifiers (query and result data types), /// and the global unique prepared statement name used in all server connections. +/// Statements created by SQL `PREPARE` carry a key of their own, so they are +/// never handed to a second client. /// /// 2. Mapping between the global unique names and Parse & RowDescription messages /// used to prepare the statement on server connections and to decode @@ -110,9 +125,8 @@ impl CachedStmt { pub struct GlobalCache { statements: HashMap, names: HashMap, - unused: HashSet, - counter: usize, - versions: usize, + unused: HashSet, + counter: Counter, } impl MemoryUsage for GlobalCache { @@ -121,7 +135,6 @@ impl MemoryUsage for GlobalCache { self.statements.memory_usage() + self.names.memory_usage() + self.counter.memory_usage() - + self.versions.memory_usage() + self.unused.capacity() * 1usize.memory_usage() } } @@ -133,10 +146,9 @@ impl GlobalCache { /// If the statement exists, no entry is created /// and the global name is returned instead. pub fn insert(&mut self, parse: &Parse) -> (bool, String) { - let parse_key = CacheKey { + let parse_key = CacheKey::Extended { query: parse.query_ref(), data_types: parse.data_types_ref(), - version: 0, }; if let Some(entry) = self.statements.get_mut(&parse_key) { @@ -144,6 +156,7 @@ impl GlobalCache { self.unused.remove(&entry.counter); } entry.used += 1; + (false, global_name(entry.counter)) } else { self.counter += 1; @@ -158,10 +171,9 @@ impl GlobalCache { // so the buffer grows monotonically. let parse = parse.renamed(&name); - let cache_key = CacheKey { + let cache_key = CacheKey::Extended { query: parse.query_ref(), data_types: parse.data_types_ref(), - version: 0, }; self.statements.insert( @@ -177,7 +189,8 @@ impl GlobalCache { Statement { parse, cache_key, - ..Default::default() + rewrite: None, + row_description: None, }, ); @@ -187,21 +200,29 @@ impl GlobalCache { /// Insert a prepared statement into the global cache ignoring /// duplicate check. - pub fn insert_anyway(&mut self, parse: &Parse) -> String { + /// + /// SQL `PREPARE` gets a key of its own, so it is never handed to + /// a second client. It is tracked and evicted like any other statement. + pub fn insert_prepare(&mut self, parse: &Parse) -> String { self.counter += 1; - self.versions += 1; let name = global_name(self.counter); let parse = parse.renamed(&name); - - let key = CacheKey { + // insert_anyway is used for the simple query PREPARE call + // and here the `unique` field based on counter defines + // that this statement won't be reused with other clients + // i.e. it'll always have `used <= 1` and will be closed + // only by the specific client that created it. + // The close happens when the client re-uses the same PREPARE + // name, or on client disconnect in the close_all call. + // TODO: a direct DEALLOCATE won't close it yet. + let cache_key = CacheKey::Simple { query: parse.query_ref(), - data_types: parse.data_types_ref(), - version: self.versions, + unique: self.counter, }; self.statements.insert( - key.clone(), + cache_key.clone(), CachedStmt { counter: self.counter, used: 1, @@ -212,9 +233,9 @@ impl GlobalCache { name.clone(), Statement { parse, - version: self.versions, - cache_key: key, - ..Default::default() + cache_key, + rewrite: None, + row_description: None, }, ); @@ -231,7 +252,7 @@ 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) { - if let Some(ref mut entry) = self.names.get_mut(name) + if let Some(entry) = self.names.get_mut(name) && entry.row_description.is_none() { entry.row_description = Some(row_description); @@ -342,18 +363,6 @@ impl GlobalCache { } } - /// Decrement usage of prepared statement without removing it. - pub fn decrement(&mut self, name: &str) { - if let Some(stmt) = self.names.get(name) - && let Some(stmt) = self.statements.get_mut(stmt.cache_key()) - { - stmt.used = stmt.used.saturating_sub(1); - if stmt.used == 0 { - self.unused.insert(stmt.counter); - } - } - } - /// Get all prepared statements by name. pub fn names(&self) -> &HashMap { &self.names @@ -395,8 +404,8 @@ mod test { let map_key = cache.statements.keys().next().unwrap(); let owned = stored.parse.query_ref(); - assert_eq!(owned.as_ptr(), stored.cache_key.query.as_ptr()); - assert_eq!(owned.as_ptr(), map_key.query.as_ptr()); + assert_eq!(owned.as_ptr(), stored.cache_key.query_ref().as_ptr()); + assert_eq!(owned.as_ptr(), map_key.query_ref().as_ptr()); assert_ne!(owned.as_ptr(), source.query_ref().as_ptr()); } @@ -431,11 +440,44 @@ mod test { assert_eq!(entry.used, 0); assert!(cache.unused.contains(&1)); // __pgdog_1 - let name = cache.insert_anyway(&parse); + let name = cache.insert_prepare(&parse); cache.close(&name); assert!(cache.unused.contains(&2)); // __pgdog_2 } + fn used(cache: &GlobalCache, name: &str) -> usize { + let statement = cache.names.get(name).unwrap(); + cache.statements.get(statement.cache_key()).unwrap().used + } + + #[test] + fn test_simple_prepared_is_never_shared() { + let mut cache = GlobalCache::default(); + let parse = Parse::named("client_stmt", "SELECT $1"); + + let first = cache.insert_prepare(&parse); + let second = cache.insert_prepare(&parse); + + assert_ne!(first, second); + assert_eq!(cache.len(), 2); + assert_eq!(used(&cache, &first), 1); + assert_eq!(used(&cache, &second), 1); + + // A Parse never re-uses a SQL PREPARE statement. + let (new, extended) = cache.insert(&parse); + assert!(new); + assert_ne!(extended, first); + assert_ne!(extended, second); + assert_eq!(cache.len(), 3); + + // A Parse re-uses another Parse. + let (new_again, shared) = cache.insert(&parse); + assert!(!new_again); + assert_eq!(shared, extended); + assert_eq!(cache.len(), 3); + assert_eq!(used(&cache, &extended), 2); + } + #[test] fn test_remove_unused() { let mut cache = GlobalCache::default(); @@ -533,7 +575,7 @@ mod test { } #[test] - fn test_decrement_marks_as_unused() { + fn test_close_marks_as_unused() { let mut cache = GlobalCache::default(); let parse = Parse::named("test", "SELECT 1"); @@ -545,18 +587,18 @@ mod test { let entry = cache.statements.get(stmt.cache_key()).unwrap(); assert_eq!(entry.used, 3); - cache.decrement(&name); + cache.close(&name); let entry = cache.statements.get(stmt.cache_key()).unwrap(); assert_eq!(entry.used, 2); assert!(cache.unused.is_empty()); - cache.decrement(&name); - cache.decrement(&name); + cache.close(&name); + cache.close(&name); let entry = cache.statements.get(stmt.cache_key()).unwrap(); assert_eq!(entry.used, 0); assert!(cache.unused.contains(&1)); - cache.decrement(&name); + cache.close(&name); let entry = cache.statements.get(stmt.cache_key()).unwrap(); assert_eq!(entry.used, 0); } diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs index f41dacc83..c623f8bc2 100644 --- a/pgdog/src/frontend/prepared_statements/mod.rs +++ b/pgdog/src/frontend/prepared_statements/mod.rs @@ -79,7 +79,7 @@ impl PreparedStatements { // Key already existed, only value changed. self.memory_used = self.memory_used.saturating_sub(str_mem(&old_value)); self.memory_used += str_mem(&name); - self.global.write().decrement(&old_value); + self.global.write().close(&old_value); } else { // New entry. self.memory_used += str_mem(key) + str_mem(&name); @@ -89,8 +89,8 @@ impl PreparedStatements { } /// Insert statement into the cache bypassing duplicate checks. - pub fn insert_anyway(&mut self, parse: &mut Parse) { - let name = { self.global.write().insert_anyway(parse) }; + pub fn insert_prepare(&mut self, parse: &mut Parse) { + let name = { self.global.write().insert_prepare(parse) }; let key = parse.name(); let existed = self.local.insert(key.to_owned(), name.clone()); @@ -98,6 +98,7 @@ impl PreparedStatements { // Key already existed, only value changed. self.memory_used = self.memory_used.saturating_sub(str_mem(&old_value)); self.memory_used += str_mem(&name); + self.global.write().close(&old_value); } else { // New entry. self.memory_used += str_mem(key) + str_mem(&name); @@ -382,5 +383,21 @@ mod test { // The other two should have used=0. let unused = global.statements().values().filter(|s| s.used == 0).count(); assert_eq!(unused, 2, "old statements should be unused"); + + let live = global + .statements() + .values() + .find(|stmt| stmt.used == 1) + .unwrap() + .name(); + drop(global); + + // Both replaced statements are evicted. + assert_eq!(statements.global.write().close_unused(0), 2); + + // The statement the client still holds survives. + let global = statements.global.read(); + assert_eq!(global.len(), 1); + assert!(global.names().contains_key(&live)); } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs index 30a73288a..34a854e7a 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -78,7 +78,7 @@ fn rewrite_single_prepared<'a>( stmt.name().expect("PREPARE always has a name"), query.as_str(), ); - prepared_statements.insert_anyway(&mut parse); + prepared_statements.insert_prepare(&mut parse); stmt.set_name(Some(mem.copy_string(parse.name()))); Ok(SimplePreparedRewrite::Prepared)