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

* Clean connection state up after protocol named prepared statement

* Avoid cloning + add test

* fmt
This commit is contained in:
Mostafa Abdelraouf
2022-09-07 22:37:17 -05:00
committed by GitHub
parent 6d41640ea9
commit 9514b3b2d1
3 changed files with 29 additions and 0 deletions

View File

@@ -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(),

View File

@@ -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 {

View File

@@ -159,6 +159,12 @@ describe "Miscellaneous" do
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
@@ -166,6 +172,7 @@ describe "Miscellaneous" 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