mirror of
https://github.com/postgresml/pgcat.git
synced 2026-03-25 02:06:28 +00:00
Reimplement prepared statements with LRU cache and statement deduplication (#618)
* Initial commit * Cleanup and add stats * Use an arc instead of full clones to store the parse packets * Use mutex instead * fmt * clippy * fmt * fix? * fix? * fmt * typo * Update docs * Refactor custom protocol * fmt * move custom protocol handling to before parsing * Support describe * Add LRU for server side statement cache * rename variable * Refactoring * Move docs * Fix test * fix * Update tests * trigger build * Add more tests * Reorder handling sync * Support when a named describe is sent along with Parse (go pgx) and expecting results * don't talk to client if not needed when client sends Parse * fmt :( * refactor tests * nit * Reduce hashing * Reducing work done to decode describe and parse messages * minor refactor * Merge branch 'main' into zain/reimplment-prepared-statements-with-global-lru-cache * Rewrite extended and prepared protocol message handling to better support mocking response packets and close * An attempt to better handle if there are DDL changes that might break cached plans with ideas about how to further improve it * fix * Minor stats fixed and cleanup * Cosmetic fixes (#64) * Cosmetic fixes * fix test * Change server drop for statement cache error to a `deallocate all` * Updated comments and added new idea for handling DDL changes impacting cached plans * fix test? * Revert test change * trigger build, flakey test * Avoid potential race conditions by changing get_or_insert to promote for pool LRU * remove ps enabled variable on the server in favor of using an option * Add close to the Extended Protocol buffer --------- Co-authored-by: Lev Kokotov <levkk@users.noreply.github.com>
This commit is contained in:
Zain Kabani
194
src/server.rs
194
src/server.rs
@@ -3,12 +3,14 @@
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use fallible_iterator::FallibleIterator;
|
||||
use log::{debug, error, info, trace, warn};
|
||||
use lru::LruCache;
|
||||
use once_cell::sync::Lazy;
|
||||
use parking_lot::{Mutex, RwLock};
|
||||
use postgres_protocol::message;
|
||||
use std::collections::{BTreeSet, HashMap, HashSet};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::mem;
|
||||
use std::net::IpAddr;
|
||||
use std::num::NonZeroUsize;
|
||||
use std::sync::Arc;
|
||||
use std::time::SystemTime;
|
||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, BufStream};
|
||||
@@ -16,7 +18,7 @@ use tokio::net::TcpStream;
|
||||
use tokio_rustls::rustls::{OwnedTrustAnchor, RootCertStore};
|
||||
use tokio_rustls::{client::TlsStream, TlsConnector};
|
||||
|
||||
use crate::config::{get_config, get_prepared_statements_cache_size, Address, User};
|
||||
use crate::config::{get_config, Address, User};
|
||||
use crate::constants::*;
|
||||
use crate::dns_cache::{AddrSet, CACHED_RESOLVER};
|
||||
use crate::errors::{Error, ServerIdentifier};
|
||||
@@ -322,7 +324,7 @@ pub struct Server {
|
||||
log_client_parameter_status_changes: bool,
|
||||
|
||||
/// Prepared statements
|
||||
prepared_statements: BTreeSet<String>,
|
||||
prepared_statement_cache: Option<LruCache<String, ()>>,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
@@ -338,6 +340,7 @@ impl Server {
|
||||
auth_hash: Arc<RwLock<Option<String>>>,
|
||||
cleanup_connections: bool,
|
||||
log_client_parameter_status_changes: bool,
|
||||
prepared_statement_cache_size: usize,
|
||||
) -> Result<Server, Error> {
|
||||
let cached_resolver = CACHED_RESOLVER.load();
|
||||
let mut addr_set: Option<AddrSet> = None;
|
||||
@@ -713,7 +716,7 @@ impl Server {
|
||||
}
|
||||
};
|
||||
|
||||
let fields = match PgErrorMsg::parse(error) {
|
||||
let fields = match PgErrorMsg::parse(&error) {
|
||||
Ok(f) => f,
|
||||
Err(err) => {
|
||||
return Err(err);
|
||||
@@ -818,7 +821,12 @@ impl Server {
|
||||
},
|
||||
cleanup_connections,
|
||||
log_client_parameter_status_changes,
|
||||
prepared_statements: BTreeSet::new(),
|
||||
prepared_statement_cache: match prepared_statement_cache_size {
|
||||
0 => None,
|
||||
_ => Some(LruCache::new(
|
||||
NonZeroUsize::new(prepared_statement_cache_size).unwrap(),
|
||||
)),
|
||||
},
|
||||
};
|
||||
|
||||
return Ok(server);
|
||||
@@ -957,6 +965,20 @@ impl Server {
|
||||
if self.in_copy_mode {
|
||||
self.in_copy_mode = false;
|
||||
}
|
||||
|
||||
if self.prepared_statement_cache.is_some() {
|
||||
let error_message = PgErrorMsg::parse(&message)?;
|
||||
if error_message.message == "cached plan must not change result type" {
|
||||
warn!("Server {:?} changed schema, dropping connection to clean up prepared statements", self.address);
|
||||
// This will still result in an error to the client, but this server connection will drop all cached prepared statements
|
||||
// so that any new queries will be re-prepared
|
||||
// TODO: Other ideas to solve errors when there are DDL changes after a statement has been prepared
|
||||
// - Recreate entire connection pool to force recreation of all server connections
|
||||
// - Clear the ConnectionPool's statement cache so that new statement names are generated
|
||||
// - Implement a retry (re-prepare) so the client doesn't see an error
|
||||
self.cleanup_state.needs_cleanup_prepare = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// CommandComplete
|
||||
@@ -1067,115 +1089,92 @@ impl Server {
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
/// Add the prepared statement to being tracked by this server.
|
||||
/// The client is processing data that will create a prepared statement on this server.
|
||||
pub fn will_prepare(&mut self, name: &str) {
|
||||
debug!("Will prepare `{}`", name);
|
||||
// Determines if the server already has a prepared statement with the given name
|
||||
// Increments the prepared statement cache hit counter
|
||||
pub fn has_prepared_statement(&mut self, name: &str) -> bool {
|
||||
let cache = match &mut self.prepared_statement_cache {
|
||||
Some(cache) => cache,
|
||||
None => return false,
|
||||
};
|
||||
|
||||
self.prepared_statements.insert(name.to_string());
|
||||
self.stats.prepared_cache_add();
|
||||
}
|
||||
|
||||
/// Check if we should prepare a statement on the server.
|
||||
pub fn should_prepare(&self, name: &str) -> bool {
|
||||
let should_prepare = !self.prepared_statements.contains(name);
|
||||
|
||||
debug!("Should prepare `{}`: {}", name, should_prepare);
|
||||
|
||||
if should_prepare {
|
||||
self.stats.prepared_cache_miss();
|
||||
} else {
|
||||
let has_it = cache.get(name).is_some();
|
||||
if has_it {
|
||||
self.stats.prepared_cache_hit();
|
||||
} else {
|
||||
self.stats.prepared_cache_miss();
|
||||
}
|
||||
|
||||
should_prepare
|
||||
has_it
|
||||
}
|
||||
|
||||
/// Create a prepared statement on the server.
|
||||
pub async fn prepare(&mut self, parse: &Parse) -> Result<(), Error> {
|
||||
debug!("Preparing `{}`", parse.name);
|
||||
pub fn add_prepared_statement_to_cache(&mut self, name: &str) -> Option<String> {
|
||||
let cache = match &mut self.prepared_statement_cache {
|
||||
Some(cache) => cache,
|
||||
None => return None,
|
||||
};
|
||||
|
||||
let bytes: BytesMut = parse.try_into()?;
|
||||
self.send(&bytes).await?;
|
||||
self.send(&flush()).await?;
|
||||
|
||||
// Read and discard ParseComplete (B)
|
||||
match read_message(&mut self.stream).await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
self.bad = true;
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
|
||||
self.prepared_statements.insert(parse.name.to_string());
|
||||
self.stats.prepared_cache_add();
|
||||
|
||||
debug!("Prepared `{}`", parse.name);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Maintain adequate cache size on the server.
|
||||
pub async fn maintain_cache(&mut self) -> Result<(), Error> {
|
||||
debug!("Cache maintenance run");
|
||||
|
||||
let max_cache_size = get_prepared_statements_cache_size();
|
||||
let mut names = Vec::new();
|
||||
|
||||
while self.prepared_statements.len() >= max_cache_size {
|
||||
// The prepared statmeents are alphanumerically sorted by the BTree.
|
||||
// FIFO.
|
||||
if let Some(name) = self.prepared_statements.pop_last() {
|
||||
names.push(name);
|
||||
// If we evict something, we need to close it on the server
|
||||
if let Some((evicted_name, _)) = cache.push(name.to_string(), ()) {
|
||||
if evicted_name != name {
|
||||
debug!(
|
||||
"Evicted prepared statement {} from cache, replaced with {}",
|
||||
evicted_name, name
|
||||
);
|
||||
return Some(evicted_name);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if !names.is_empty() {
|
||||
self.deallocate(names).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
None
|
||||
}
|
||||
|
||||
/// Remove the prepared statement from being tracked by this server.
|
||||
/// The client is processing data that will cause the server to close the prepared statement.
|
||||
pub fn will_close(&mut self, name: &str) {
|
||||
debug!("Will close `{}`", name);
|
||||
pub fn remove_prepared_statement_from_cache(&mut self, name: &str) {
|
||||
let cache = match &mut self.prepared_statement_cache {
|
||||
Some(cache) => cache,
|
||||
None => return,
|
||||
};
|
||||
|
||||
self.prepared_statements.remove(name);
|
||||
self.stats.prepared_cache_remove();
|
||||
cache.pop(name);
|
||||
}
|
||||
|
||||
/// Close a prepared statement on the server.
|
||||
pub async fn deallocate(&mut self, names: Vec<String>) -> Result<(), Error> {
|
||||
for name in &names {
|
||||
debug!("Deallocating prepared statement `{}`", name);
|
||||
pub async fn register_prepared_statement(
|
||||
&mut self,
|
||||
parse: &Parse,
|
||||
should_send_parse_to_server: bool,
|
||||
) -> Result<(), Error> {
|
||||
if !self.has_prepared_statement(&parse.name) {
|
||||
let mut bytes = BytesMut::new();
|
||||
|
||||
let close = Close::new(name);
|
||||
let bytes: BytesMut = close.try_into()?;
|
||||
if should_send_parse_to_server {
|
||||
let parse_bytes: BytesMut = parse.try_into()?;
|
||||
bytes.extend_from_slice(&parse_bytes);
|
||||
}
|
||||
|
||||
self.send(&bytes).await?;
|
||||
}
|
||||
|
||||
if !names.is_empty() {
|
||||
self.send(&flush()).await?;
|
||||
}
|
||||
|
||||
// Read and discard CloseComplete (3)
|
||||
for name in &names {
|
||||
match read_message(&mut self.stream).await {
|
||||
Ok(_) => {
|
||||
self.prepared_statements.remove(name);
|
||||
self.stats.prepared_cache_remove();
|
||||
debug!("Closed `{}`", name);
|
||||
}
|
||||
|
||||
Err(err) => {
|
||||
self.bad = true;
|
||||
return Err(err);
|
||||
}
|
||||
// If we evict something, we need to close it on the server
|
||||
// We do this by adding it to the messages we're sending to the server before the sync
|
||||
if let Some(evicted_name) = self.add_prepared_statement_to_cache(&parse.name) {
|
||||
self.remove_prepared_statement_from_cache(&evicted_name);
|
||||
let close_bytes: BytesMut = Close::new(&evicted_name).try_into()?;
|
||||
bytes.extend_from_slice(&close_bytes);
|
||||
};
|
||||
}
|
||||
|
||||
// If we have a parse or close we need to send to the server, send them and sync
|
||||
if !bytes.is_empty() {
|
||||
bytes.extend_from_slice(&sync());
|
||||
|
||||
self.send(&bytes).await?;
|
||||
|
||||
loop {
|
||||
self.recv(None).await?;
|
||||
|
||||
if !self.is_data_available() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1312,6 +1311,10 @@ impl Server {
|
||||
|
||||
if self.cleanup_state.needs_cleanup_prepare {
|
||||
reset_string.push_str("DEALLOCATE ALL;");
|
||||
// Since we deallocated all prepared statements, we need to clear the cache
|
||||
if let Some(cache) = &mut self.prepared_statement_cache {
|
||||
cache.clear();
|
||||
}
|
||||
};
|
||||
|
||||
self.query(&reset_string).await?;
|
||||
@@ -1377,6 +1380,7 @@ impl Server {
|
||||
Arc::new(RwLock::new(None)),
|
||||
true,
|
||||
false,
|
||||
0,
|
||||
)
|
||||
.await?;
|
||||
debug!("Connected!, sending query.");
|
||||
|
||||
Reference in New Issue
Block a user