From 92e428c947b8c76b40f5cd226fa4c8fd5ed7983b Mon Sep 17 00:00:00 2001 From: Andrew Tanner Date: Fri, 5 May 2023 13:50:21 -0700 Subject: [PATCH 1/5] First try with role rest --- src/server.rs | 25 ++++++++----------------- tests/ruby/misc_spec.rb | 12 ++++++++++++ 2 files changed, 20 insertions(+), 17 deletions(-) diff --git a/src/server.rs b/src/server.rs index 152f0524..e5a32d88 100644 --- a/src/server.rs +++ b/src/server.rs @@ -705,10 +705,7 @@ impl Server { Ok(()) } Err(err) => { - error!( - "Terminating server {:?} because of: {:?}", - self.address, err - ); + error!("Terminating server because of: {:?}", err); self.bad = true; Err(err) } @@ -723,10 +720,7 @@ impl Server { let mut message = match read_message(&mut self.stream).await { Ok(message) => message, Err(err) => { - error!( - "Terminating server {:?} because of: {:?}", - self.address, err - ); + error!("Terminating server because of: {:?}", err); self.bad = true; return Err(err); } @@ -962,7 +956,8 @@ impl Server { // it before each checkin. if self.needs_cleanup { warn!("Server returned with session state altered, discarding state"); - self.query("DISCARD ALL").await?; + self.query("DISCARD ALL;").await?; + self.query("RESET ROLE;").await?; self.needs_cleanup = false; } @@ -1141,18 +1136,14 @@ impl Drop for Server { _ => debug!("Dirty shutdown"), }; + // Should not matter. + self.bad = true; + let now = chrono::offset::Utc::now().naive_utc(); let duration = now - self.connected_at; - let message = if self.bad { - "Server connection terminated" - } else { - "Server connection closed" - }; - info!( - "{} {:?}, session duration: {}", - message, + "Server connection closed {:?}, session duration: {}", self.address, crate::format_duration(&duration) ); diff --git a/tests/ruby/misc_spec.rb b/tests/ruby/misc_spec.rb index 1a04c5d6..c2c3c041 100644 --- a/tests/ruby/misc_spec.rb +++ b/tests/ruby/misc_spec.rb @@ -241,6 +241,18 @@ expect(processes.primary.count_query("DISCARD ALL")).to eq(10) end + + it "Resets server roles correctly" 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.async_exec("SET statement_timeout to 5000") + conn.close + end + + expect(processes.primary.count_query("RESET ROLE;")).to eq(10) + end end context "transaction mode" do From 401f355be36e29e271d91b2eddbed617739848eb Mon Sep 17 00:00:00 2001 From: Andrew Tanner Date: Fri, 5 May 2023 13:58:42 -0700 Subject: [PATCH 2/5] update --- src/server.rs | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/src/server.rs b/src/server.rs index e5a32d88..adad2290 100644 --- a/src/server.rs +++ b/src/server.rs @@ -705,7 +705,10 @@ impl Server { Ok(()) } Err(err) => { - error!("Terminating server because of: {:?}", err); + error!( + "Terminating server {:?} because of: {:?}", + self.address, err + ); self.bad = true; Err(err) } @@ -720,7 +723,10 @@ impl Server { let mut message = match read_message(&mut self.stream).await { Ok(message) => message, Err(err) => { - error!("Terminating server because of: {:?}", err); + error!( + "Terminating server {:?} because of: {:?}", + self.address, err + ); self.bad = true; return Err(err); } @@ -956,7 +962,7 @@ impl Server { // it before each checkin. if self.needs_cleanup { warn!("Server returned with session state altered, discarding state"); - self.query("DISCARD ALL;").await?; + self.query("DISCARD ALL").await?; self.query("RESET ROLE;").await?; self.needs_cleanup = false; } @@ -1136,16 +1142,20 @@ impl Drop for Server { _ => debug!("Dirty shutdown"), }; - // Should not matter. - self.bad = true; - let now = chrono::offset::Utc::now().naive_utc(); let duration = now - self.connected_at; + let message = if self.bad { + "Server connection terminated" + } else { + "Server connection closed" + }; + info!( - "Server connection closed {:?}, session duration: {}", + "{} {:?}, session duration: {}", + message, self.address, crate::format_duration(&duration) ); } -} +} \ No newline at end of file From 089b8fd3497d02849294c6c524e40a1eb537965c Mon Sep 17 00:00:00 2001 From: Andrew Tanner Date: Fri, 5 May 2023 14:16:56 -0700 Subject: [PATCH 3/5] extra line --- src/server.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/server.rs b/src/server.rs index adad2290..e9ac2819 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1158,4 +1158,4 @@ impl Drop for Server { crate::format_duration(&duration) ); } -} \ No newline at end of file +} From 3692d3d1a815575b03f196557143843cf3fd700a Mon Sep 17 00:00:00 2001 From: Andrew Tanner Date: Fri, 5 May 2023 17:48:42 -0400 Subject: [PATCH 4/5] Update src/server.rs Co-authored-by: Lev Kokotov --- src/server.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/server.rs b/src/server.rs index e9ac2819..8f0ece99 100644 --- a/src/server.rs +++ b/src/server.rs @@ -963,7 +963,7 @@ impl Server { if self.needs_cleanup { warn!("Server returned with session state altered, discarding state"); self.query("DISCARD ALL").await?; - self.query("RESET ROLE;").await?; + self.query("RESET ROLE").await?; self.needs_cleanup = false; } From 41acb0b53924f4a1b45ffdf1b97bc19518f8b255 Mon Sep 17 00:00:00 2001 From: Andrew Tanner Date: Fri, 5 May 2023 17:48:51 -0400 Subject: [PATCH 5/5] Update tests/ruby/misc_spec.rb Co-authored-by: Lev Kokotov --- tests/ruby/misc_spec.rb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/ruby/misc_spec.rb b/tests/ruby/misc_spec.rb index c2c3c041..e4d6f6fb 100644 --- a/tests/ruby/misc_spec.rb +++ b/tests/ruby/misc_spec.rb @@ -251,7 +251,7 @@ conn.close end - expect(processes.primary.count_query("RESET ROLE;")).to eq(10) + expect(processes.primary.count_query("RESET ROLE")).to eq(10) end end