Clean connection state up after protocol named prepared statement (#163)

* Clean connection state up after protocol named prepared statement

* Avoid cloning + add test

* fmt

Author: drdrsh

src/client.rs

@@ -879,6 +879,23 @@ where
 
                         self.buffer.put(&original[..]);
 
+                        // Clone after freeze does not allocate
+                        let first_message_code = (*self.buffer.get(0).unwrap_or(&0)) as char;
+
+                        // Almost certainly true
+                        if first_message_code == 'P' {
+                            // Message layout
+                            // P followed by 32 int followed by null-terminated statement name
+                            // So message code should be in offset 0 of the buffer, first character
+                            // in prepared statement name would be index 5
+                            let first_char_in_name = *self.buffer.get(5).unwrap_or(&0);
+                            if first_char_in_name != 0 {
+                                // This is a named prepared statement
+                                // Server connection state will need to be cleared at checkin
+                                server.mark_dirty();
+                            }
+                        }
+
                         self.send_and_receive_loop(
                             code,
                             self.buffer.clone(),

src/server.rs

@@ -638,6 +638,11 @@ impl Server {
     pub fn last_activity(&self) -> SystemTime {
         self.last_activity
     }
+
+    // Marks a connection as needing DISCARD ALL at checkin
+    pub fn mark_dirty(&mut self) {
+        self.needs_cleanup = true;
+    }
 }
 
 impl Drop for Server {

tests/ruby/misc_spec.rb

@@ -159,13 +159,20 @@
           conn.async_exec("PREPARE prepared_q (int) AS SELECT $1")
           conn.close
         end
+
+        15.times do
+          conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+          conn.prepare("prepared_q", "SELECT $1")
+          conn.close
+        end
       end
 
       it "Does not send DISCARD ALL unless necessary" do
         10.times do
           conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
           conn.async_exec("SET SERVER ROLE to 'primary'")
           conn.async_exec("SELECT 1")
+          conn.exec_params("SELECT $1", [1])
           conn.close
         end