1
1
use super :: { ConnectionLike , Runtime } ;
2
2
use crate :: aio:: setup_connection;
3
+ use crate :: aio:: DisconnectNotifier ;
3
4
use crate :: cmd:: Cmd ;
4
5
#[ cfg( any( feature = "tokio-comp" , feature = "async-std-comp" ) ) ]
5
6
use crate :: parser:: ValueCodec ;
@@ -23,6 +24,7 @@ use std::fmt;
23
24
use std:: fmt:: Debug ;
24
25
use std:: io;
25
26
use std:: pin:: Pin ;
27
+ use std:: sync:: atomic:: { AtomicBool , Ordering } ;
26
28
use std:: sync:: Arc ;
27
29
use std:: task:: { self , Poll } ;
28
30
use std:: time:: Duration ;
@@ -73,19 +75,11 @@ struct PipelineMessage<S> {
73
75
/// items being output by the `Stream` (the number is specified at time of sending). With the
74
76
/// interface provided by `Pipeline` an easy interface of request to response, hiding the `Stream`
75
77
/// and `Sink`.
78
+ #[ derive( Clone ) ]
76
79
struct Pipeline < SinkItem > {
77
80
sender : mpsc:: Sender < PipelineMessage < SinkItem > > ,
78
-
79
81
push_manager : Arc < ArcSwap < PushManager > > ,
80
- }
81
-
82
- impl < SinkItem > Clone for Pipeline < SinkItem > {
83
- fn clone ( & self ) -> Self {
84
- Pipeline {
85
- sender : self . sender . clone ( ) ,
86
- push_manager : self . push_manager . clone ( ) ,
87
- }
88
- }
82
+ is_stream_closed : Arc < AtomicBool > ,
89
83
}
90
84
91
85
impl < SinkItem > Debug for Pipeline < SinkItem >
@@ -104,14 +98,21 @@ pin_project! {
104
98
in_flight: VecDeque <InFlight >,
105
99
error: Option <RedisError >,
106
100
push_manager: Arc <ArcSwap <PushManager >>,
101
+ disconnect_notifier: Option <Box <dyn DisconnectNotifier >>,
102
+ is_stream_closed: Arc <AtomicBool >,
107
103
}
108
104
}
109
105
110
106
impl < T > PipelineSink < T >
111
107
where
112
108
T : Stream < Item = RedisResult < Value > > + ' static ,
113
109
{
114
- fn new < SinkItem > ( sink_stream : T , push_manager : Arc < ArcSwap < PushManager > > ) -> Self
110
+ fn new < SinkItem > (
111
+ sink_stream : T ,
112
+ push_manager : Arc < ArcSwap < PushManager > > ,
113
+ disconnect_notifier : Option < Box < dyn DisconnectNotifier > > ,
114
+ is_stream_closed : Arc < AtomicBool > ,
115
+ ) -> Self
115
116
where
116
117
T : Sink < SinkItem , Error = RedisError > + Stream < Item = RedisResult < Value > > + ' static ,
117
118
{
@@ -120,6 +121,8 @@ where
120
121
in_flight : VecDeque :: new ( ) ,
121
122
error : None ,
122
123
push_manager,
124
+ disconnect_notifier,
125
+ is_stream_closed,
123
126
}
124
127
}
125
128
@@ -130,7 +133,15 @@ where
130
133
Some ( result) => result,
131
134
// The redis response stream is not going to produce any more items so we `Err`
132
135
// to break out of the `forward` combinator and stop handling requests
133
- None => return Poll :: Ready ( Err ( ( ) ) ) ,
136
+ None => {
137
+ // this is the right place to notify about the passive TCP disconnect
138
+ // In other places we cannot distinguish between the active destruction of MultiplexedConnection and passive disconnect
139
+ if let Some ( disconnect_notifier) = self . as_mut ( ) . project ( ) . disconnect_notifier {
140
+ disconnect_notifier. notify_disconnect ( ) ;
141
+ }
142
+ self . is_stream_closed . store ( true , Ordering :: Relaxed ) ;
143
+ return Poll :: Ready ( Err ( ( ) ) ) ;
144
+ }
134
145
} ;
135
146
self . as_mut ( ) . send_result ( item) ;
136
147
}
@@ -296,7 +307,10 @@ impl<SinkItem> Pipeline<SinkItem>
296
307
where
297
308
SinkItem : Send + ' static ,
298
309
{
299
- fn new < T > ( sink_stream : T ) -> ( Self , impl Future < Output = ( ) > )
310
+ fn new < T > (
311
+ sink_stream : T ,
312
+ disconnect_notifier : Option < Box < dyn DisconnectNotifier > > ,
313
+ ) -> ( Self , impl Future < Output = ( ) > )
300
314
where
301
315
T : Sink < SinkItem , Error = RedisError > + Stream < Item = RedisResult < Value > > + ' static ,
302
316
T : Send + ' static ,
@@ -308,7 +322,13 @@ where
308
322
let ( sender, mut receiver) = mpsc:: channel ( BUFFER_SIZE ) ;
309
323
let push_manager: Arc < ArcSwap < PushManager > > =
310
324
Arc :: new ( ArcSwap :: new ( Arc :: new ( PushManager :: default ( ) ) ) ) ;
311
- let sink = PipelineSink :: new :: < SinkItem > ( sink_stream, push_manager. clone ( ) ) ;
325
+ let is_stream_closed = Arc :: new ( AtomicBool :: new ( false ) ) ;
326
+ let sink = PipelineSink :: new :: < SinkItem > (
327
+ sink_stream,
328
+ push_manager. clone ( ) ,
329
+ disconnect_notifier,
330
+ is_stream_closed. clone ( ) ,
331
+ ) ;
312
332
let f = stream:: poll_fn ( move |cx| receiver. poll_recv ( cx) )
313
333
. map ( Ok )
314
334
. forward ( sink)
@@ -317,6 +337,7 @@ where
317
337
Pipeline {
318
338
sender,
319
339
push_manager,
340
+ is_stream_closed,
320
341
} ,
321
342
f,
322
343
)
@@ -363,6 +384,10 @@ where
363
384
async fn set_push_manager ( & mut self , push_manager : PushManager ) {
364
385
self . push_manager . store ( Arc :: new ( push_manager) ) ;
365
386
}
387
+
388
+ pub fn is_closed ( & self ) -> bool {
389
+ self . is_stream_closed . load ( Ordering :: Relaxed )
390
+ }
366
391
}
367
392
368
393
/// A connection object which can be cloned, allowing requests to be be sent concurrently
@@ -392,6 +417,7 @@ impl MultiplexedConnection {
392
417
connection_info : & ConnectionInfo ,
393
418
stream : C ,
394
419
push_sender : Option < mpsc:: UnboundedSender < PushInfo > > ,
420
+ disconnect_notifier : Option < Box < dyn DisconnectNotifier > > ,
395
421
) -> RedisResult < ( Self , impl Future < Output = ( ) > ) >
396
422
where
397
423
C : Unpin + AsyncRead + AsyncWrite + Send + ' static ,
@@ -401,6 +427,7 @@ impl MultiplexedConnection {
401
427
stream,
402
428
std:: time:: Duration :: MAX ,
403
429
push_sender,
430
+ disconnect_notifier,
404
431
)
405
432
. await
406
433
}
@@ -412,6 +439,7 @@ impl MultiplexedConnection {
412
439
stream : C ,
413
440
response_timeout : std:: time:: Duration ,
414
441
push_sender : Option < mpsc:: UnboundedSender < PushInfo > > ,
442
+ disconnect_notifier : Option < Box < dyn DisconnectNotifier > > ,
415
443
) -> RedisResult < ( Self , impl Future < Output = ( ) > ) >
416
444
where
417
445
C : Unpin + AsyncRead + AsyncWrite + Send + ' static ,
@@ -429,7 +457,7 @@ impl MultiplexedConnection {
429
457
let codec = ValueCodec :: default ( )
430
458
. framed ( stream)
431
459
. and_then ( |msg| async move { msg } ) ;
432
- let ( mut pipeline, driver) = Pipeline :: new ( codec) ;
460
+ let ( mut pipeline, driver) = Pipeline :: new ( codec, disconnect_notifier ) ;
433
461
let driver = boxed ( driver) ;
434
462
let pm = PushManager :: default ( ) ;
435
463
if let Some ( sender) = push_sender {
@@ -560,6 +588,10 @@ impl ConnectionLike for MultiplexedConnection {
560
588
fn get_db ( & self ) -> i64 {
561
589
self . db
562
590
}
591
+
592
+ fn is_closed ( & self ) -> bool {
593
+ self . pipeline . is_closed ( )
594
+ }
563
595
}
564
596
impl MultiplexedConnection {
565
597
/// Subscribes to a new channel.
0 commit comments