Implement Close for prepared statements (#482)

* Partial support for Close

* Close

* respect config value

* prepared spec

* Hmm

* Print cache size

src/server.rs

@@ -15,7 +15,7 @@ use tokio::net::TcpStream;
 use tokio_rustls::rustls::{OwnedTrustAnchor, RootCertStore};
 use tokio_rustls::{client::TlsStream, TlsConnector};
 
-use crate::config::{get_config, Address, User};
+use crate::config::{get_config, get_prepared_statements_cache_size, Address, User};
 use crate::constants::*;
 use crate::dns_cache::{AddrSet, CACHED_RESOLVER};
 use crate::errors::{Error, ServerIdentifier};
@@ -914,12 +914,16 @@ 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);
 
         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);
 
@@ -934,6 +938,7 @@ impl Server {
         should_prepare
     }
 
+    /// Create a prepared statement on the server.
     pub async fn prepare(&mut self, parse: &Parse) -> Result<(), Error> {
         debug!("Preparing `{}`", parse.name);
 
@@ -942,15 +947,82 @@ impl Server {
         self.send(&flush()).await?;
 
         // Read and discard ParseComplete (B)
-        let _ = read_message(&mut self.stream).await?;
+        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);
+            }
+        }
+
+        self.deallocate(names).await?;
+
+        Ok(())
+    }
+
+    /// 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);
+
+        self.prepared_statements.remove(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);
+
+            let close = Close::new(name);
+            let bytes: BytesMut = close.try_into()?;
+
+            self.send(&bytes).await?;
+        }
+
+        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);
+                }
+            };
+        }
+
+        Ok(())
+    }
+
     /// If the server is still inside a transaction.
     /// If the client disconnects while the server is in a transaction, we will clean it up.
     pub fn in_transaction(&self) -> bool {