Minor stats fixed and cleanup

Author: zainkabani

src/admin.rs

@@ -744,6 +744,7 @@ where
         ("age_seconds", DataType::Numeric),
         ("prepare_cache_hit", DataType::Numeric),
         ("prepare_cache_miss", DataType::Numeric),
+        ("prepare_cache_eviction", DataType::Numeric),
         ("prepare_cache_size", DataType::Numeric),
     ];
 
@@ -776,6 +777,10 @@ where
                 .prepared_miss_count
                 .load(Ordering::Relaxed)
                 .to_string(),
+            server
+                .prepared_eviction_count
+                .load(Ordering::Relaxed)
+                .to_string(),
             server
                 .prepared_cache_size
                 .load(Ordering::Relaxed)

src/client.rs

@@ -109,7 +109,8 @@ pub struct Client<S, T> {
     /// Mapping of client named prepared statement to rewritten parse messages
     prepared_statements: HashMap<String, (Arc<Parse>, u64)>,
 
-    extended_protocol_data: VecDeque<ExtendedProtocolData>,
+    /// Buffered extended protocol data
+    extended_protocol_data_buffer: VecDeque<ExtendedProtocolData>,
 }
 
 /// Client entrypoint.
@@ -726,7 +727,7 @@ where
             shutdown,
             prepared_statements_enabled,
             prepared_statements: HashMap::new(),
-            extended_protocol_data: VecDeque::new(),
+            extended_protocol_data_buffer: VecDeque::new(),
         })
     }
 
@@ -764,7 +765,7 @@ where
             shutdown,
             prepared_statements_enabled: false,
             prepared_statements: HashMap::new(),
-            extended_protocol_data: VecDeque::new(),
+            extended_protocol_data_buffer: VecDeque::new(),
         })
     }
 
@@ -970,7 +971,7 @@ where
                 }
 
                 'E' => {
-                    self.extended_protocol_data
+                    self.extended_protocol_data_buffer
                         .push_back(ExtendedProtocolData::create_new_execute(message));
                     continue;
                 }
@@ -1259,7 +1260,7 @@ where
                     // Execute
                     // Execute a prepared statement prepared in `P` and bound in `B`.
                     'E' => {
-                        self.extended_protocol_data
+                        self.extended_protocol_data_buffer
                             .push_back(ExtendedProtocolData::create_new_execute(message));
                     }
 
@@ -1310,7 +1311,9 @@ where
                         //              ReadyForQuery
 
                         // Iterate over our extended protocol data that we've buffered
-                        while let Some(protocol_data) = self.extended_protocol_data.pop_front() {
+                        while let Some(protocol_data) =
+                            self.extended_protocol_data_buffer.pop_front()
+                        {
                             match protocol_data {
                                 ExtendedProtocolData::Parse { data, metadata } => {
                                     let (parse, hash) = match metadata {
@@ -1712,7 +1715,7 @@ where
 
         if client_given_name.is_empty() {
             debug!("Anonymous parse message");
-            self.extended_protocol_data
+            self.extended_protocol_data_buffer
                 .push_back(ExtendedProtocolData::create_new_parse(message, None));
             return Ok(());
         }
@@ -1741,7 +1744,7 @@ where
         self.prepared_statements
             .insert(client_given_name, (new_parse.clone(), hash));
 
-        self.extended_protocol_data
+        self.extended_protocol_data_buffer
             .push_back(ExtendedProtocolData::create_new_parse(
                 new_parse.as_ref().try_into()?,
                 Some((new_parse.clone(), hash)),
@@ -1762,7 +1765,7 @@ where
 
         if client_given_name.is_empty() {
             debug!("Anonymous bind message");
-            self.extended_protocol_data
+            self.extended_protocol_data_buffer
                 .push_back(ExtendedProtocolData::create_new_bind(message, None));
             return Ok(());
         }
@@ -1776,11 +1779,9 @@ where
                     client_given_name, rewritten_parse.name
                 );
 
-                self.extended_protocol_data
-                    .push_back(ExtendedProtocolData::create_new_bind(
-                        message,
-                        Some(client_given_name),
-                    ));
+                self.extended_protocol_data_buffer.push_back(
+                    ExtendedProtocolData::create_new_bind(message, Some(client_given_name)),
+                );
 
                 Ok(())
             }
@@ -1818,7 +1819,7 @@ where
 
         if describe.anonymous() {
             debug!("Anonymous describe message");
-            self.extended_protocol_data
+            self.extended_protocol_data_buffer
                 .push_back(ExtendedProtocolData::create_new_describe(message, None));
 
             return Ok(());
@@ -1835,11 +1836,12 @@ where
                     client_given_name, describe.statement_name
                 );
 
-                self.extended_protocol_data
-                    .push_back(ExtendedProtocolData::create_new_describe(
+                self.extended_protocol_data_buffer.push_back(
+                    ExtendedProtocolData::create_new_describe(
                         describe.try_into()?,
                         Some(client_given_name),
-                    ));
+                    ),
+                );
 
                 Ok(())
             }
@@ -1866,7 +1868,7 @@ where
 
     fn reset_buffered_state(&mut self) {
         self.buffer.clear();
-        self.extended_protocol_data.clear();
+        self.extended_protocol_data_buffer.clear();
         self.response_message_queue_buffer.clear();
     }
 

src/messages.rs

@@ -925,20 +925,6 @@ impl Parse {
     pub fn anonymous(&self) -> bool {
         self.name.is_empty()
     }
-
-    /// Reads the parse message from the buffer and throws them away
-    pub fn remove_from_buffer_start(buf: &mut BytesMut) -> Result<(), Error> {
-        buf.get_u8();
-        buf.get_i32();
-        buf.read_string()?;
-        buf.read_string()?;
-
-        for _ in 0..buf.get_i16() {
-            buf.get_i32();
-        }
-
-        Ok(())
-    }
 }
 
 /// Bind (B) message.
@@ -1104,11 +1090,6 @@ impl Bind {
         Ok(response_buf)
     }
 
-    pub fn reassign(mut self, new_prepared_statement_name: &str) -> Self {
-        self.prepared_statement = new_prepared_statement_name.to_string();
-        self
-    }
-
     pub fn anonymous(&self) -> bool {
         self.prepared_statement.is_empty()
     }

src/pool.rs

@@ -60,21 +60,18 @@ pub enum BanReason {
 pub type PreparedStatementCacheType = Arc<Mutex<PreparedStatementCache>>;
 
 // TODO: Add stats the this cache
+// TODO: Add application name to the cache value to help identify which application is using the cache
+// TODO: Create admin command to show which statements are in the cache
 #[derive(Debug)]
 pub struct PreparedStatementCache {
-    cache: LruCache<u64, Arc<Parse>>, // TODO: Store a tuple of BytesMut and String instead Parse to remove need to convert to BytesMut again
-}
-
-impl Default for PreparedStatementCache {
-    fn default() -> Self {
-        Self::new(1)
-    }
+    cache: LruCache<u64, Arc<Parse>>,
 }
 
 impl PreparedStatementCache {
-    pub fn new(size: usize) -> Self {
+    pub fn new(mut size: usize) -> Self {
+        // Cannot be zeros
         if size == 0 {
-            return Self::default();
+            size = 1;
         }
 
         PreparedStatementCache {

src/server.rs

@@ -1116,14 +1116,21 @@ impl Server {
         Ok(bytes)
     }
 
+    // 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 {
-        self.stats.prepared_cache_hit();
-        self.prepared_statement_cache.get(name).is_some()
+        let has_it = self.prepared_statement_cache.get(name).is_some();
+        if has_it {
+            self.stats.prepared_cache_hit();
+        } else {
+            self.stats.prepared_cache_miss();
+        }
+
+        has_it
     }
 
     pub fn add_prepared_statement_to_cache(&mut self, name: &str) -> Option<String> {
         self.stats.prepared_cache_add();
-        self.stats.prepared_cache_miss();
 
         // If we evict something, we need to close it on the server
         if let Some((evicted_name, _)) = self.prepared_statement_cache.push(name.to_string(), ()) {
@@ -1149,35 +1156,33 @@ impl Server {
         parse: &Parse,
         should_send_parse_to_server: bool,
     ) -> Result<(), Error> {
-        match self.prepared_statement_cache.get(&parse.name) {
-            Some(_) => self.stats.prepared_cache_hit(),
-            None => {
-                let mut bytes = BytesMut::new();
-
-                if should_send_parse_to_server {
-                    let parse_bytes: BytesMut = parse.try_into()?;
-                    bytes.extend_from_slice(&parse_bytes);
-                }
+        if !self.has_prepared_statement(&parse.name) {
+            let mut bytes = BytesMut::new();
 
-                // 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 should_send_parse_to_server {
+                let parse_bytes: BytesMut = parse.try_into()?;
+                bytes.extend_from_slice(&parse_bytes);
+            }
 
-                if !bytes.is_empty() {
-                    bytes.extend_from_slice(&sync());
+            // 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);
+            };
 
-                    self.send(&bytes).await?;
+            // 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());
 
-                    loop {
-                        self.recv(None).await?;
+                self.send(&bytes).await?;
 
-                        if !self.is_data_available() {
-                            break;
-                        }
+                loop {
+                    self.recv(None).await?;
+
+                    if !self.is_data_available() {
+                        break;
                     }
                 }
             }

src/stats/server.rs

@@ -49,6 +49,7 @@ pub struct ServerStats {
     pub error_count: Arc<AtomicU64>,
     pub prepared_hit_count: Arc<AtomicU64>,
     pub prepared_miss_count: Arc<AtomicU64>,
+    pub prepared_eviction_count: Arc<AtomicU64>,
     pub prepared_cache_size: Arc<AtomicU64>,
 }
 
@@ -68,6 +69,7 @@ impl Default for ServerStats {
             reporter: get_reporter(),
             prepared_hit_count: Arc::new(AtomicU64::new(0)),
             prepared_miss_count: Arc::new(AtomicU64::new(0)),
+            prepared_eviction_count: Arc::new(AtomicU64::new(0)),
             prepared_cache_size: Arc::new(AtomicU64::new(0)),
         }
     }
@@ -221,6 +223,7 @@ impl ServerStats {
     }
 
     pub fn prepared_cache_remove(&self) {
+        self.prepared_eviction_count.fetch_add(1, Ordering::Relaxed);
         self.prepared_cache_size.fetch_sub(1, Ordering::Relaxed);
     }
 }