@@ -5,16 +5,18 @@ use std::{
5
5
} ;
6
6
7
7
use derive_where:: derive_where;
8
- use tokio:: sync:: { mpsc, Mutex } ;
8
+ use tokio:: sync:: { broadcast , mpsc, Mutex } ;
9
9
10
10
use super :: {
11
11
CmapEventEmitter ,
12
12
Connection ,
13
13
ConnectionGeneration ,
14
14
ConnectionInfo ,
15
+ Message ,
15
16
PendingConnection ,
16
17
PinnedConnectionHandle ,
17
18
PoolManager ,
19
+ RawCommandResponse ,
18
20
} ;
19
21
use crate :: {
20
22
bson:: oid:: ObjectId ,
@@ -50,7 +52,7 @@ pub(crate) struct PooledConnection {
50
52
}
51
53
52
54
/// The state of a pooled connection.
53
- #[ derive( Clone , Debug ) ]
55
+ #[ derive( Debug ) ]
54
56
enum PooledConnectionState {
55
57
/// The state associated with a connection checked into the connection pool.
56
58
CheckedIn { available_time : Instant } ,
@@ -59,6 +61,10 @@ enum PooledConnectionState {
59
61
CheckedOut {
60
62
/// The manager used to check this connection back into the pool.
61
63
pool_manager : PoolManager ,
64
+
65
+ /// The receiver to receive a cancellation notice. Only present on non-load-balanced
66
+ /// connections.
67
+ cancellation_receiver : Option < broadcast:: Receiver < ( ) > > ,
62
68
} ,
63
69
64
70
/// The state associated with a pinned connection.
@@ -140,6 +146,24 @@ impl PooledConnection {
140
146
. and_then ( |sd| sd. service_id )
141
147
}
142
148
149
+ /// Sends a message on this connection.
150
+ pub ( crate ) async fn send_message (
151
+ & mut self ,
152
+ message : impl TryInto < Message , Error = impl Into < Error > > ,
153
+ ) -> Result < RawCommandResponse > {
154
+ match self . state {
155
+ PooledConnectionState :: CheckedOut {
156
+ cancellation_receiver : Some ( ref mut cancellation_receiver) ,
157
+ ..
158
+ } => {
159
+ self . connection
160
+ . send_message_with_cancellation ( message, cancellation_receiver)
161
+ . await
162
+ }
163
+ _ => self . connection . send_message ( message) . await ,
164
+ }
165
+ }
166
+
143
167
/// Updates the state of the connection to indicate that it is checked into the pool.
144
168
pub ( crate ) fn mark_checked_in ( & mut self ) {
145
169
if !matches ! ( self . state, PooledConnectionState :: CheckedIn { .. } ) {
@@ -155,8 +179,15 @@ impl PooledConnection {
155
179
}
156
180
157
181
/// Updates the state of the connection to indicate that it is checked out of the pool.
158
- pub ( crate ) fn mark_checked_out ( & mut self , pool_manager : PoolManager ) {
159
- self . state = PooledConnectionState :: CheckedOut { pool_manager } ;
182
+ pub ( crate ) fn mark_checked_out (
183
+ & mut self ,
184
+ pool_manager : PoolManager ,
185
+ cancellation_receiver : Option < broadcast:: Receiver < ( ) > > ,
186
+ ) {
187
+ self . state = PooledConnectionState :: CheckedOut {
188
+ pool_manager,
189
+ cancellation_receiver,
190
+ } ;
160
191
}
161
192
162
193
/// Whether this connection is idle.
@@ -175,15 +206,14 @@ impl PooledConnection {
175
206
Instant :: now ( ) . duration_since ( available_time) >= max_idle_time
176
207
}
177
208
178
- /// Nullifies the internal state of this connection and returns it in a new [PooledConnection].
179
- /// If a state is provided, then the new connection will contain that state; otherwise, this
180
- /// connection's state will be cloned.
181
- fn take ( & mut self , state : impl Into < Option < PooledConnectionState > > ) -> Self {
209
+ /// Nullifies the internal state of this connection and returns it in a new [PooledConnection]
210
+ /// with the given state.
211
+ fn take ( & mut self , new_state : PooledConnectionState ) -> Self {
182
212
Self {
183
213
connection : self . connection . take ( ) ,
184
214
generation : self . generation ,
185
215
event_emitter : self . event_emitter . clone ( ) ,
186
- state : state . into ( ) . unwrap_or_else ( || self . state . clone ( ) ) ,
216
+ state : new_state ,
187
217
}
188
218
}
189
219
@@ -196,7 +226,9 @@ impl PooledConnection {
196
226
self . id
197
227
) ) )
198
228
}
199
- PooledConnectionState :: CheckedOut { ref pool_manager } => {
229
+ PooledConnectionState :: CheckedOut {
230
+ ref pool_manager, ..
231
+ } => {
200
232
let ( tx, rx) = mpsc:: channel ( 1 ) ;
201
233
self . state = PooledConnectionState :: Pinned {
202
234
// Mark the connection as in-use while the operation currently using the
@@ -286,10 +318,11 @@ impl Drop for PooledConnection {
286
318
// Nothing needs to be done when a checked-in connection is dropped.
287
319
PooledConnectionState :: CheckedIn { .. } => Ok ( ( ) ) ,
288
320
// A checked-out connection should be sent back to the connection pool.
289
- PooledConnectionState :: CheckedOut { pool_manager } => {
321
+ PooledConnectionState :: CheckedOut { pool_manager, .. } => {
290
322
let pool_manager = pool_manager. clone ( ) ;
291
- let mut dropped_connection = self . take ( None ) ;
292
- dropped_connection. mark_checked_in ( ) ;
323
+ let dropped_connection = self . take ( PooledConnectionState :: CheckedIn {
324
+ available_time : Instant :: now ( ) ,
325
+ } ) ;
293
326
pool_manager. check_in ( dropped_connection)
294
327
}
295
328
// A pinned connection should be returned to its pinner or to the connection pool.
@@ -339,7 +372,11 @@ impl Drop for PooledConnection {
339
372
}
340
373
// The pinner of this connection has been dropped while the connection was
341
374
// sitting in its channel, so the connection should be returned to the pool.
342
- PinnedState :: Returned { .. } => pool_manager. check_in ( self . take ( None ) ) ,
375
+ PinnedState :: Returned { .. } => {
376
+ pool_manager. check_in ( self . take ( PooledConnectionState :: CheckedIn {
377
+ available_time : Instant :: now ( ) ,
378
+ } ) )
379
+ }
343
380
}
344
381
}
345
382
} ;
0 commit comments