Enforce prepared statements cache limits at runtime, add memory limit

.schema/pgdog.schema.json

@@ -83,6 +83,7 @@
         "port": 6432,
         "prepared_statements": "extended",
         "prepared_statements_limit": 9223372036854775807,
+        "prepared_statements_memory_limit": 0,
         "pub_sub_channel_size": 0,
         "query_cache_limit": 1000,
         "query_log": null,
@@ -974,6 +975,13 @@
           "default": 9223372036854775807,
           "minimum": 0
         },
+        "prepared_statements_memory_limit": {
+          "description": "Approximate memory limit (bytes) for the global prepared statements cache. Statements no client is holding are evicted once the cache grows past it. `0` disables the limit.\n\n**Note:** A limit smaller than the working set causes constant eviction and re-preparation of statements; size it well above what the active workload keeps in flight.\n\n_Default:_ `0`\n\n<https://docs.pgdog.dev/configuration/pgdog.toml/general/#prepared_statements_memory_limit>",
+          "type": "integer",
+          "format": "uint",
+          "default": 0,
+          "minimum": 0
+        },
         "pub_sub_channel_size": {
           "description": "Enables support for pub/sub and configures the size of the background task queue.\n\n<https://docs.pgdog.dev/configuration/pgdog.toml/general/#pub_sub_channel_size>",
           "type": "integer",

pgdog-config/src/general.rs

@@ -382,6 +382,16 @@ pub struct General {
     #[serde(default = "General::prepared_statements_limit")]
     pub prepared_statements_limit: usize,
 
+    /// Approximate memory limit (bytes) for the global prepared statements cache. Statements no client is holding are evicted once the cache grows past it. `0` disables the limit.
+    ///
+    /// **Note:** A limit smaller than the working set causes constant eviction and re-preparation of statements; size it well above what the active workload keeps in flight.
+    ///
+    /// _Default:_ `0`
+    ///
+    /// <https://docs.pgdog.dev/configuration/pgdog.toml/general/#prepared_statements_memory_limit>
+    #[serde(default = "General::prepared_statements_memory_limit")]
+    pub prepared_statements_memory_limit: usize,
+
     /// Limit on the number of statements saved in the statement cache used to accelerate query parsing.
     ///
     /// _Default:_ `50000`
@@ -887,6 +897,7 @@ impl Default for General {
             regex_parser_limit: Self::regex_parser_limit(),
             query_parser_engine: QueryParserEngine::default(),
             prepared_statements_limit: Self::prepared_statements_limit(),
+            prepared_statements_memory_limit: Self::prepared_statements_memory_limit(),
             query_cache_limit: Self::query_cache_limit(),
             passthrough_auth: Self::default_passthrough_auth(),
             connect_timeout: Self::default_connect_timeout(),
@@ -1381,6 +1392,10 @@ impl General {
         Self::env_or_default("PGDOG_PREPARED_STATEMENTS_LIMIT", i64::MAX as usize)
     }
 
+    pub fn prepared_statements_memory_limit() -> usize {
+        Self::env_or_default("PGDOG_PREPARED_STATEMENTS_MEMORY_LIMIT", 0)
+    }
+
     pub fn query_cache_limit() -> usize {
         Self::env_or_default("PGDOG_QUERY_CACHE_LIMIT", 1_000)
     }
@@ -1826,12 +1841,14 @@ mod tests {
         let _guard = set_env_var("PGDOG_MIRROR_EXPOSURE", "0.5");
         let _guard = set_env_var("PGDOG_DNS_TTL", "60000");
         let _guard = set_env_var("PGDOG_PUB_SUB_CHANNEL_SIZE", "100");
+        let _guard = set_env_var("PGDOG_PREPARED_STATEMENTS_MEMORY_LIMIT", "4294967296");
         let _guard = set_env_var("PGDOG_LOG_MIN_DURATION_PARSE", "5");
         let _guard = set_env_var("PGDOG_LOG_QUERY_SAMPLE_LENGTH", "200");
 
         assert_eq!(General::broadcast_port(), 7432);
         assert_eq!(General::openmetrics_port(), Some(9090));
         assert_eq!(General::prepared_statements_limit(), 1000);
+        assert_eq!(General::prepared_statements_memory_limit(), 4294967296);
         assert_eq!(General::query_cache_limit(), 500);
         assert_eq!(General::connect_attempts(), 3);
         assert_eq!(General::mirror_queue(), 256);
@@ -1844,6 +1861,7 @@ mod tests {
         let _guard = remove_env_var("PGDOG_BROADCAST_PORT");
         let _guard = remove_env_var("PGDOG_OPENMETRICS_PORT");
         let _guard = remove_env_var("PGDOG_PREPARED_STATEMENTS_LIMIT");
+        let _guard = remove_env_var("PGDOG_PREPARED_STATEMENTS_MEMORY_LIMIT");
         let _guard = remove_env_var("PGDOG_QUERY_CACHE_LIMIT");
         let _guard = remove_env_var("PGDOG_CONNECT_ATTEMPTS");
         let _guard = remove_env_var("PGDOG_MIRROR_QUEUE");
@@ -1856,6 +1874,7 @@ mod tests {
         assert_eq!(General::broadcast_port(), General::port() + 1);
         assert_eq!(General::openmetrics_port(), None);
         assert_eq!(General::prepared_statements_limit(), i64::MAX as usize);
+        assert_eq!(General::prepared_statements_memory_limit(), 0);
         assert_eq!(General::query_cache_limit(), 1_000);
         assert_eq!(General::connect_attempts(), 1);
         assert_eq!(General::mirror_queue(), 128);

pgdog/src/admin/parser.rs

@@ -255,6 +255,14 @@ mod tests {
         assert!(matches!(result, Ok(ParseResult::ResetQueryCache(_))));
     }
 
+    #[test]
+    fn parses_reset_prepared_command() {
+        // The exact string the prepared_statements_limit=0 warning advises:
+        // if this stops parsing, the advice is broken.
+        let result = Parser::parse(super::super::reset_prepared::RESET_PREPARED);
+        assert!(matches!(result, Ok(ParseResult::ResetPrepared(_))));
+    }
+
     #[test]
     fn rejects_unknown_admin_command() {
         let result = Parser::parse("FOO BAR");

pgdog/src/admin/reset_prepared.rs

@@ -1,26 +1,29 @@
 //! RESET PREPARED.
-use crate::config::config;
 use crate::frontend::prepared_statements::PreparedStatements;
 
 use super::prelude::*;
 
+/// The admin console spelling of this command. The limit-0 warning quotes
+/// it, and a parser test keeps the advice parseable.
+pub(super) const RESET_PREPARED: &str = "RESET PREPARED";
+
 pub struct ResetPrepared;
 
 #[async_trait]
 impl Command for ResetPrepared {
     fn name(&self) -> String {
-        "RESET PREPARED".into()
+        RESET_PREPARED.into()
     }
 
     fn parse(_: &str) -> Result<Self, Error> {
         Ok(Self)
     }
 
     async fn execute(&self) -> Result<Vec<Message>, Error> {
-        let config = config();
-        PreparedStatements::global()
-            .write()
-            .close_unused(config.config.general.prepared_statements_limit);
+        // Explicit 0: drop everything not in use, whatever the configured
+        // limit is. With the default (unlimited) limit this would otherwise
+        // be a no-op.
+        PreparedStatements::global().write().close_unused(0);
         Ok(vec![])
     }
 }

pgdog/src/admin/set.rs

@@ -117,9 +117,25 @@ impl Command for Set {
 
             "prepared_statements_limit" => {
                 config.config.general.prepared_statements_limit = self.value.parse()?;
-                PreparedStatements::global()
-                    .write()
-                    .close_unused(config.config.general.prepared_statements_limit);
+                if config.config.general.prepared_statements_limit == 0 {
+                    tracing::warn!(
+                        "prepared_statements_limit set to 0, which now means unlimited; \
+                        to clear the cache, use {}",
+                        super::reset_prepared::RESET_PREPARED,
+                    );
+                }
+                PreparedStatements::global().write().configure(
+                    config.config.general.prepared_statements_limit,
+                    config.config.general.prepared_statements_memory_limit,
+                );
+            }
+
+            "prepared_statements_memory_limit" => {
+                config.config.general.prepared_statements_memory_limit = self.value.parse()?;
+                PreparedStatements::global().write().configure(
+                    config.config.general.prepared_statements_limit,
+                    config.config.general.prepared_statements_memory_limit,
+                );
             }
 
             "prepared_statements" => {

pgdog/src/backend/databases.rs

@@ -107,6 +107,12 @@ pub fn init() -> Result<(), Error> {
     // Resize query cache
     Cache::resize(config.config.general.query_cache_limit);
 
+    // Apply prepared statements cache limits.
+    PreparedStatements::global().write().configure(
+        config.config.general.prepared_statements_limit,
+        config.config.general.prepared_statements_memory_limit,
+    );
+
     // Start two-pc manager.
     let _monitor = Manager::get();
 
@@ -147,10 +153,12 @@ pub fn reload() -> Result<(), Error> {
     // Reload TLS connectors.
     tls::reload()?;
 
-    // Remove any unused prepared statements.
-    PreparedStatements::global()
-        .write()
-        .close_unused(new_config.config.general.prepared_statements_limit);
+    // Apply prepared statements cache limits, dropping anything
+    // unused over the new caps.
+    PreparedStatements::global().write().configure(
+        new_config.config.general.prepared_statements_limit,
+        new_config.config.general.prepared_statements_memory_limit,
+    );
 
     // Resize query cache.
     Cache::resize(new_config.config.general.query_cache_limit);