Skip to content

Commit 0907f1b

Browse files
authored
Improve logging for connection cleanup (#428)
* initial commit * fix * fmt
1 parent 7326069 commit 0907f1b

File tree

1 file changed

+53
-11
lines changed

1 file changed

+53
-11
lines changed

src/server.rs

+53-11
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,48 @@ impl StreamInner {
103103
}
104104
}
105105

106+
#[derive(Copy, Clone)]
107+
struct CleanupState {
108+
/// If server connection requires DISCARD ALL before checkin because of set statement
109+
needs_cleanup_set: bool,
110+
111+
/// If server connection requires DISCARD ALL before checkin because of prepare statement
112+
needs_cleanup_prepare: bool,
113+
}
114+
115+
impl CleanupState {
116+
fn new() -> Self {
117+
CleanupState {
118+
needs_cleanup_set: false,
119+
needs_cleanup_prepare: false,
120+
}
121+
}
122+
123+
fn needs_cleanup(&self) -> bool {
124+
self.needs_cleanup_set || self.needs_cleanup_prepare
125+
}
126+
127+
fn set_true(&mut self) {
128+
self.needs_cleanup_set = true;
129+
self.needs_cleanup_prepare = true;
130+
}
131+
132+
fn reset(&mut self) {
133+
self.needs_cleanup_set = false;
134+
self.needs_cleanup_prepare = false;
135+
}
136+
}
137+
138+
impl std::fmt::Display for CleanupState {
139+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140+
write!(
141+
f,
142+
"SET: {}, PREPARE: {}",
143+
self.needs_cleanup_set, self.needs_cleanup_prepare
144+
)
145+
}
146+
}
147+
106148
/// Server state.
107149
pub struct Server {
108150
/// Server host, e.g. localhost,
@@ -131,8 +173,8 @@ pub struct Server {
131173
/// Is the server broken? We'll remote it from the pool if so.
132174
bad: bool,
133175

134-
/// If server connection requires a DISCARD ALL before checkin
135-
needs_cleanup: bool,
176+
/// If server connection requires DISCARD ALL before checkin
177+
cleanup_state: CleanupState,
136178

137179
/// Mapping of clients and servers used for query cancellation.
138180
client_server_map: ClientServerMap,
@@ -630,7 +672,7 @@ impl Server {
630672
in_transaction: false,
631673
data_available: false,
632674
bad: false,
633-
needs_cleanup: false,
675+
cleanup_state: CleanupState::new(),
634676
client_server_map,
635677
addr_set,
636678
connected_at: chrono::offset::Utc::now().naive_utc(),
@@ -793,12 +835,12 @@ impl Server {
793835
// This will reduce amount of discard statements sent
794836
if !self.in_transaction {
795837
debug!("Server connection marked for clean up");
796-
self.needs_cleanup = true;
838+
self.cleanup_state.needs_cleanup_set = true;
797839
}
798840
}
799841
"PREPARE\0" => {
800842
debug!("Server connection marked for clean up");
801-
self.needs_cleanup = true;
843+
self.cleanup_state.needs_cleanup_prepare = true;
802844
}
803845
_ => (),
804846
}
@@ -960,11 +1002,11 @@ impl Server {
9601002
// to avoid leaking state between clients. For performance reasons we only
9611003
// send `DISCARD ALL` if we think the session is altered instead of just sending
9621004
// it before each checkin.
963-
if self.needs_cleanup {
964-
warn!("Server returned with session state altered, discarding state");
1005+
if self.cleanup_state.needs_cleanup() {
1006+
warn!("Server returned with session state altered, discarding state ({}) for application {}", self.cleanup_state, self.application_name);
9651007
self.query("DISCARD ALL").await?;
9661008
self.query("RESET ROLE").await?;
967-
self.needs_cleanup = false;
1009+
self.cleanup_state.reset();
9681010
}
9691011

9701012
Ok(())
@@ -976,12 +1018,12 @@ impl Server {
9761018
self.application_name = name.to_string();
9771019
// We don't want `SET application_name` to mark the server connection
9781020
// as needing cleanup
979-
let needs_cleanup_before = self.needs_cleanup;
1021+
let needs_cleanup_before = self.cleanup_state;
9801022

9811023
let result = Ok(self
9821024
.query(&format!("SET application_name = '{}'", name))
9831025
.await?);
984-
self.needs_cleanup = needs_cleanup_before;
1026+
self.cleanup_state = needs_cleanup_before;
9851027
result
9861028
} else {
9871029
Ok(())
@@ -1006,7 +1048,7 @@ impl Server {
10061048

10071049
// Marks a connection as needing DISCARD ALL at checkin
10081050
pub fn mark_dirty(&mut self) {
1009-
self.needs_cleanup = true;
1051+
self.cleanup_state.set_true();
10101052
}
10111053

10121054
pub fn mirror_send(&mut self, bytes: &BytesMut) {

0 commit comments

Comments
 (0)