Send DISCARD ALL even if client is not in transaction (#152)

* Send DISCARD ALL even if client is not in transaction

* fmt

* Added tests + avoided sending extra discard all

* Adds set name logic to beginning of handle client

* fmt

* refactor dead code handling

* Refactor reading command tag

* remove unnecessary trim

* Removing debugging statement

* typo

* typo{

* documentation

* edit text

* un-unwrap

* run ci

* run ci

Co-authored-by: Zain Kabani <zain.kabani@instacart.com>

src/client.rs

@@ -59,6 +59,7 @@ pub struct Client<S, T> {
     client_server_map: ClientServerMap,
 
     /// Client parameters, e.g. user, client_encoding, etc.
+    #[allow(dead_code)]
     parameters: HashMap<String, String>,
 
     /// Statistics
@@ -82,6 +83,9 @@ pub struct Client<S, T> {
     /// Postgres user for this client (This comes from the user in the connection string)
     username: String,
 
+    /// Application name for this client (defaults to pgcat)
+    application_name: String,
+
     /// Used to notify clients about an impending shutdown
     shutdown: Receiver<()>,
 }
@@ -365,6 +369,11 @@ where
             None => return Err(Error::ClientError),
         };
 
+        let application_name = match parameters.get("application_name") {
+            Some(application_name) => application_name,
+            None => "pgcat",
+        };
+
         let admin = ["pgcat", "pgbouncer"]
             .iter()
             .filter(|db| *db == &pool_name)
@@ -493,6 +502,7 @@ where
             last_server_id: None,
             pool_name: pool_name.clone(),
             username: username.clone(),
+            application_name: application_name.to_string(),
             shutdown,
             connected_to_server: false,
         });
@@ -526,6 +536,7 @@ where
             last_server_id: None,
             pool_name: String::from("undefined"),
             username: String::from("undefined"),
+            application_name: String::from("undefined"),
             shutdown,
             connected_to_server: false,
         });
@@ -767,13 +778,10 @@ where
                 server.address()
             );
 
-            // Set application_name if any.
             // TODO: investigate other parameters and set them too.
-            if self.parameters.contains_key("application_name") {
-                server
-                    .set_name(&self.parameters["application_name"])
-                    .await?;
-            }
+
+            // Set application_name.
+            server.set_name(&self.application_name).await?;
 
             // Transaction loop. Multiple queries can be issued by the client here.
             // The connection belongs to the client until the transaction is over,
@@ -790,12 +798,7 @@ where
                         Err(err) => {
                             // Client disconnected inside a transaction.
                             // Clean up the server and re-use it.
-                            // This prevents connection thrashing by bad clients.
-                            if server.in_transaction() {
-                                server.query("ROLLBACK").await?;
-                                server.query("DISCARD ALL").await?;
-                                server.set_name("pgcat").await?;
-                            }
+                            server.checkin_cleanup().await?;
 
                             return Err(err);
                         }
@@ -837,16 +840,7 @@ where
 
                     // Terminate
                     'X' => {
-                        // Client closing. Rollback and clean up
-                        // connection before releasing into the pool.
-                        // Pgbouncer closes the connection which leads to
-                        // connection thrashing when clients misbehave.
-                        if server.in_transaction() {
-                            server.query("ROLLBACK").await?;
-                            server.query("DISCARD ALL").await?;
-                            server.set_name("pgcat").await?;
-                        }
-
+                        server.checkin_cleanup().await?;
                         self.release();
 
                         return Ok(());
@@ -950,8 +944,10 @@ where
 
             // The server is no longer bound to us, we can't cancel it's queries anymore.
             debug!("Releasing server back into the pool");
+            server.checkin_cleanup().await?;
             self.stats.server_idle(server.process_id(), address.id);
             self.connected_to_server = false;
+
             self.release();
             self.stats.client_idle(self.process_id, address.id);
         }

src/server.rs

@@ -1,7 +1,8 @@
 /// Implementation of the PostgreSQL server (database) protocol.
 /// Here we are pretending to the a Postgres client.
 use bytes::{Buf, BufMut, BytesMut};
-use log::{debug, error, info, trace};
+use log::{debug, error, info, trace, warn};
+use std::io::Read;
 use std::time::SystemTime;
 use tokio::io::{AsyncReadExt, BufReader};
 use tokio::net::{
@@ -48,6 +49,9 @@ pub struct Server {
     /// Is the server broken? We'll remote it from the pool if so.
     bad: bool,
 
+    /// If server connection requires a DISCARD ALL before checkin
+    needs_cleanup: bool,
+
     /// Mapping of clients and servers used for query cancellation.
     client_server_map: ClientServerMap,
 
@@ -316,6 +320,7 @@ impl Server {
                         in_transaction: false,
                         data_available: false,
                         bad: false,
+                        needs_cleanup: false,
                         client_server_map: client_server_map,
                         connected_at: chrono::offset::Utc::now().naive_utc(),
                         stats: stats,
@@ -440,6 +445,29 @@ impl Server {
                     break;
                 }
 
+                // CommandComplete
+                'C' => {
+                    let mut command_tag = String::new();
+                    match message.reader().read_to_string(&mut command_tag) {
+                        Ok(_) => {
+                            // Non-exhaustive list of commands that are likely to change session variables/resources
+                            // which can leak between clients. This is a best effort to block bad clients
+                            // from poisoning a transaction-mode pool by setting inappropriate session variables
+                            match command_tag.as_str() {
+                                "SET\0" | "PREPARE\0" => {
+                                    debug!("Server connection marked for clean up");
+                                    self.needs_cleanup = true;
+                                }
+                                _ => (),
+                            }
+                        }
+
+                        Err(err) => {
+                            warn!("Encountered an error while parsing CommandTag {}", err);
+                        }
+                    }
+                }
+
                 // DataRow
                 'D' => {
                     // More data is available after this message, this is not the end of the reply.
@@ -553,14 +581,43 @@ impl Server {
         Ok(())
     }
 
+    /// Perform any necessary cleanup before putting the server
+    /// connection back in the pool
+    pub async fn checkin_cleanup(&mut self) -> Result<(), Error> {
+        // Client disconnected with an open transaction on the server connection.
+        // Pgbouncer behavior is to close the server connection but that can cause
+        // server connection thrashing if clients repeatedly do this.
+        // Instead, we ROLLBACK that transaction before putting the connection back in the pool
+        if self.in_transaction() {
+            self.query("ROLLBACK").await?;
+        }
+
+        // Client disconnected but it perfromed session-altering operations such as
+        // SET statement_timeout to 1 or create a prepared statement. We clear that
+        // to avoid leaking state between clients. For performance reasons we only
+        // send `DISCARD ALL` if we think the session is altered instead of just sending
+        // it before each checkin.
+        if self.needs_cleanup {
+            self.query("DISCARD ALL").await?;
+            self.needs_cleanup = false;
+        }
+
+        return Ok(());
+    }
+
     /// A shorthand for `SET application_name = $1`.
-    #[allow(dead_code)]
     pub async fn set_name(&mut self, name: &str) -> Result<(), Error> {
         if self.application_name != name {
             self.application_name = name.to_string();
-            Ok(self
+            // We don't want `SET application_name` to mark the server connection
+            // as needing cleanup
+            let needs_cleanup_before = self.needs_cleanup;
+
+            let result = Ok(self
                 .query(&format!("SET application_name = '{}'", name))
-                .await?)
+                .await?);
+            self.needs_cleanup = needs_cleanup_before;
+            return result;
         } else {
             Ok(())
         }