@@ -103,6 +103,48 @@ impl StreamInner {
103
103
}
104
104
}
105
105
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
+
106
148
/// Server state.
107
149
pub struct Server {
108
150
/// Server host, e.g. localhost,
@@ -131,8 +173,8 @@ pub struct Server {
131
173
/// Is the server broken? We'll remote it from the pool if so.
132
174
bad : bool ,
133
175
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 ,
136
178
137
179
/// Mapping of clients and servers used for query cancellation.
138
180
client_server_map : ClientServerMap ,
@@ -630,7 +672,7 @@ impl Server {
630
672
in_transaction : false ,
631
673
data_available : false ,
632
674
bad : false ,
633
- needs_cleanup : false ,
675
+ cleanup_state : CleanupState :: new ( ) ,
634
676
client_server_map,
635
677
addr_set,
636
678
connected_at : chrono:: offset:: Utc :: now ( ) . naive_utc ( ) ,
@@ -793,12 +835,12 @@ impl Server {
793
835
// This will reduce amount of discard statements sent
794
836
if !self . in_transaction {
795
837
debug ! ( "Server connection marked for clean up" ) ;
796
- self . needs_cleanup = true ;
838
+ self . cleanup_state . needs_cleanup_set = true ;
797
839
}
798
840
}
799
841
"PREPARE\0 " => {
800
842
debug ! ( "Server connection marked for clean up" ) ;
801
- self . needs_cleanup = true ;
843
+ self . cleanup_state . needs_cleanup_prepare = true ;
802
844
}
803
845
_ => ( ) ,
804
846
}
@@ -960,11 +1002,11 @@ impl Server {
960
1002
// to avoid leaking state between clients. For performance reasons we only
961
1003
// send `DISCARD ALL` if we think the session is altered instead of just sending
962
1004
// 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 ) ;
965
1007
self . query ( "DISCARD ALL" ) . await ?;
966
1008
self . query ( "RESET ROLE" ) . await ?;
967
- self . needs_cleanup = false ;
1009
+ self . cleanup_state . reset ( ) ;
968
1010
}
969
1011
970
1012
Ok ( ( ) )
@@ -976,12 +1018,12 @@ impl Server {
976
1018
self . application_name = name. to_string ( ) ;
977
1019
// We don't want `SET application_name` to mark the server connection
978
1020
// as needing cleanup
979
- let needs_cleanup_before = self . needs_cleanup ;
1021
+ let needs_cleanup_before = self . cleanup_state ;
980
1022
981
1023
let result = Ok ( self
982
1024
. query ( & format ! ( "SET application_name = '{}'" , name) )
983
1025
. await ?) ;
984
- self . needs_cleanup = needs_cleanup_before;
1026
+ self . cleanup_state = needs_cleanup_before;
985
1027
result
986
1028
} else {
987
1029
Ok ( ( ) )
@@ -1006,7 +1048,7 @@ impl Server {
1006
1048
1007
1049
// Marks a connection as needing DISCARD ALL at checkin
1008
1050
pub fn mark_dirty ( & mut self ) {
1009
- self . needs_cleanup = true ;
1051
+ self . cleanup_state . set_true ( ) ;
1010
1052
}
1011
1053
1012
1054
pub fn mirror_send ( & mut self , bytes : & BytesMut ) {
0 commit comments