Skip to content

Commit 265ad67

Browse files
authored
fix(client): more reliably detect closed pooled connections (#1434)
1 parent 8fb84d2 commit 265ad67

File tree

5 files changed

+264
-60
lines changed

5 files changed

+264
-60
lines changed

src/client/cancel.rs

+148
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
use std::sync::Arc;
2+
use std::sync::atomic::{AtomicBool, Ordering};
3+
4+
use futures::{Async, Future, Poll};
5+
use futures::task::{self, Task};
6+
7+
use common::Never;
8+
9+
use self::lock::Lock;
10+
11+
#[derive(Clone)]
12+
pub struct Cancel {
13+
inner: Arc<Inner>,
14+
}
15+
16+
pub struct Canceled {
17+
inner: Arc<Inner>,
18+
}
19+
20+
struct Inner {
21+
is_canceled: AtomicBool,
22+
task: Lock<Option<Task>>,
23+
}
24+
25+
impl Cancel {
26+
pub fn new() -> (Cancel, Canceled) {
27+
let inner = Arc::new(Inner {
28+
is_canceled: AtomicBool::new(false),
29+
task: Lock::new(None),
30+
});
31+
let inner2 = inner.clone();
32+
(
33+
Cancel {
34+
inner: inner,
35+
},
36+
Canceled {
37+
inner: inner2,
38+
},
39+
)
40+
}
41+
42+
pub fn cancel(&self) {
43+
if !self.inner.is_canceled.swap(true, Ordering::SeqCst) {
44+
if let Some(mut locked) = self.inner.task.try_lock() {
45+
if let Some(task) = locked.take() {
46+
task.notify();
47+
}
48+
}
49+
// if we couldn't take the lock, Canceled was trying to park.
50+
// After parking, it will check is_canceled one last time,
51+
// so we can just stop here.
52+
}
53+
}
54+
55+
pub fn is_canceled(&self) -> bool {
56+
self.inner.is_canceled.load(Ordering::SeqCst)
57+
}
58+
}
59+
60+
impl Future for Canceled {
61+
type Item = ();
62+
type Error = Never;
63+
64+
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
65+
if self.inner.is_canceled.load(Ordering::SeqCst) {
66+
Ok(Async::Ready(()))
67+
} else {
68+
if let Some(mut locked) = self.inner.task.try_lock() {
69+
if locked.is_none() {
70+
// it's possible a Cancel just tried to cancel on another thread,
71+
// and we just missed it. Once we have the lock, we should check
72+
// one more time before parking this task and going away.
73+
if self.inner.is_canceled.load(Ordering::SeqCst) {
74+
return Ok(Async::Ready(()));
75+
}
76+
*locked = Some(task::current());
77+
}
78+
Ok(Async::NotReady)
79+
} else {
80+
// if we couldn't take the lock, then a Cancel taken has it.
81+
// The *ONLY* reason is because it is in the process of canceling.
82+
Ok(Async::Ready(()))
83+
}
84+
}
85+
}
86+
}
87+
88+
impl Drop for Canceled {
89+
fn drop(&mut self) {
90+
self.inner.is_canceled.store(true, Ordering::SeqCst);
91+
}
92+
}
93+
94+
95+
// a sub module just to protect unsafety
96+
mod lock {
97+
use std::cell::UnsafeCell;
98+
use std::ops::{Deref, DerefMut};
99+
use std::sync::atomic::{AtomicBool, Ordering};
100+
101+
pub struct Lock<T> {
102+
is_locked: AtomicBool,
103+
value: UnsafeCell<T>,
104+
}
105+
106+
impl<T> Lock<T> {
107+
pub fn new(val: T) -> Lock<T> {
108+
Lock {
109+
is_locked: AtomicBool::new(false),
110+
value: UnsafeCell::new(val),
111+
}
112+
}
113+
114+
pub fn try_lock(&self) -> Option<Locked<T>> {
115+
if !self.is_locked.swap(true, Ordering::SeqCst) {
116+
Some(Locked { lock: self })
117+
} else {
118+
None
119+
}
120+
}
121+
}
122+
123+
unsafe impl<T: Send> Send for Lock<T> {}
124+
unsafe impl<T: Send> Sync for Lock<T> {}
125+
126+
pub struct Locked<'a, T: 'a> {
127+
lock: &'a Lock<T>,
128+
}
129+
130+
impl<'a, T> Deref for Locked<'a, T> {
131+
type Target = T;
132+
fn deref(&self) -> &T {
133+
unsafe { &*self.lock.value.get() }
134+
}
135+
}
136+
137+
impl<'a, T> DerefMut for Locked<'a, T> {
138+
fn deref_mut(&mut self) -> &mut T {
139+
unsafe { &mut *self.lock.value.get() }
140+
}
141+
}
142+
143+
impl<'a, T> Drop for Locked<'a, T> {
144+
fn drop(&mut self) {
145+
self.lock.is_locked.store(false, Ordering::SeqCst);
146+
}
147+
}
148+
}

src/client/dispatch.rs

+73
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
use futures::{Async, Future, Poll, Stream};
2+
use futures::sync::{mpsc, oneshot};
3+
4+
use common::Never;
5+
use super::cancel::{Cancel, Canceled};
6+
7+
pub type Callback<U> = oneshot::Sender<::Result<U>>;
8+
pub type Promise<U> = oneshot::Receiver<::Result<U>>;
9+
10+
pub fn channel<T, U>() -> (Sender<T, U>, Receiver<T, U>) {
11+
let (tx, rx) = mpsc::unbounded();
12+
let (cancel, canceled) = Cancel::new();
13+
let tx = Sender {
14+
cancel: cancel,
15+
inner: tx,
16+
};
17+
let rx = Receiver {
18+
canceled: canceled,
19+
inner: rx,
20+
};
21+
(tx, rx)
22+
}
23+
24+
pub struct Sender<T, U> {
25+
cancel: Cancel,
26+
inner: mpsc::UnboundedSender<(T, Callback<U>)>,
27+
}
28+
29+
impl<T, U> Sender<T, U> {
30+
pub fn is_closed(&self) -> bool {
31+
self.cancel.is_canceled()
32+
}
33+
34+
pub fn cancel(&self) {
35+
self.cancel.cancel();
36+
}
37+
38+
pub fn send(&self, val: T) -> Result<Promise<U>, T> {
39+
let (tx, rx) = oneshot::channel();
40+
self.inner.unbounded_send((val, tx))
41+
.map(move |_| rx)
42+
.map_err(|e| e.into_inner().0)
43+
}
44+
}
45+
46+
impl<T, U> Clone for Sender<T, U> {
47+
fn clone(&self) -> Sender<T, U> {
48+
Sender {
49+
cancel: self.cancel.clone(),
50+
inner: self.inner.clone(),
51+
}
52+
}
53+
}
54+
55+
pub struct Receiver<T, U> {
56+
canceled: Canceled,
57+
inner: mpsc::UnboundedReceiver<(T, Callback<U>)>,
58+
}
59+
60+
impl<T, U> Stream for Receiver<T, U> {
61+
type Item = (T, Callback<U>);
62+
type Error = Never;
63+
64+
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
65+
if let Async::Ready(()) = self.canceled.poll()? {
66+
return Ok(Async::Ready(None));
67+
}
68+
self.inner.poll()
69+
.map_err(|()| unreachable!("mpsc never errors"))
70+
}
71+
}
72+
73+
//TODO: Drop for Receiver should consume inner

src/client/mod.rs

+33-45
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
//! HTTP Client
22
3-
use std::cell::{Cell, RefCell};
3+
use std::cell::Cell;
44
use std::fmt;
55
use std::io;
66
use std::marker::PhantomData;
77
use std::rc::Rc;
88
use std::time::Duration;
99

10-
use futures::{Future, Poll, Stream};
11-
use futures::future::{self, Executor};
10+
use futures::{Async, Future, Poll, Stream};
11+
use futures::future::{self, Either, Executor};
1212
#[cfg(feature = "compat")]
1313
use http;
1414
use tokio::reactor::Handle;
@@ -28,7 +28,10 @@ pub use self::connect::{HttpConnector, Connect};
2828

2929
use self::background::{bg, Background};
3030

31+
mod cancel;
3132
mod connect;
33+
//TODO(easy): move cancel and dispatch into common instead
34+
pub(crate) mod dispatch;
3235
mod dns;
3336
mod pool;
3437
#[cfg(feature = "compat")]
@@ -189,20 +192,16 @@ where C: Connect,
189192
head.headers.set_pos(0, host);
190193
}
191194

192-
use futures::Sink;
193-
use futures::sync::{mpsc, oneshot};
194-
195195
let checkout = self.pool.checkout(domain.as_ref());
196196
let connect = {
197197
let executor = self.executor.clone();
198198
let pool = self.pool.clone();
199199
let pool_key = Rc::new(domain.to_string());
200200
self.connector.connect(url)
201201
.and_then(move |io| {
202-
// 1 extra slot for possible Close message
203-
let (tx, rx) = mpsc::channel(1);
202+
let (tx, rx) = dispatch::channel();
204203
let tx = HyperClient {
205-
tx: RefCell::new(tx),
204+
tx: tx,
206205
should_close: Cell::new(true),
207206
};
208207
let pooled = pool.pooled(pool_key, tx);
@@ -225,33 +224,26 @@ where C: Connect,
225224
});
226225

227226
let resp = race.and_then(move |client| {
228-
use proto::dispatch::ClientMsg;
229-
230-
let (callback, rx) = oneshot::channel();
231-
client.should_close.set(false);
232-
233-
match client.tx.borrow_mut().start_send(ClientMsg::Request(head, body, callback)) {
234-
Ok(_) => (),
235-
Err(e) => match e.into_inner() {
236-
ClientMsg::Request(_, _, callback) => {
237-
error!("pooled connection was not ready, this is a hyper bug");
238-
let err = io::Error::new(
239-
io::ErrorKind::BrokenPipe,
240-
"pool selected dead connection",
241-
);
242-
let _ = callback.send(Err(::Error::Io(err)));
243-
},
244-
_ => unreachable!("ClientMsg::Request was just sent"),
227+
match client.tx.send((head, body)) {
228+
Ok(rx) => {
229+
client.should_close.set(false);
230+
Either::A(rx.then(|res| {
231+
match res {
232+
Ok(Ok(res)) => Ok(res),
233+
Ok(Err(err)) => Err(err),
234+
Err(_) => panic!("dispatch dropped without returning error"),
235+
}
236+
}))
237+
},
238+
Err(_) => {
239+
error!("pooled connection was not ready, this is a hyper bug");
240+
let err = io::Error::new(
241+
io::ErrorKind::BrokenPipe,
242+
"pool selected dead connection",
243+
);
244+
Either::B(future::err(::Error::Io(err)))
245245
}
246246
}
247-
248-
rx.then(|res| {
249-
match res {
250-
Ok(Ok(res)) => Ok(res),
251-
Ok(Err(err)) => Err(err),
252-
Err(_) => panic!("dispatch dropped without returning error"),
253-
}
254-
})
255247
});
256248

257249
FutureResponse(Box::new(resp))
@@ -276,13 +268,8 @@ impl<C, B> fmt::Debug for Client<C, B> {
276268
}
277269

278270
struct HyperClient<B> {
279-
// A sentinel that is usually always true. If this is dropped
280-
// while true, this will try to shutdown the dispatcher task.
281-
//
282-
// This should be set to false whenever it is checked out of the
283-
// pool and successfully used to send a request.
284271
should_close: Cell<bool>,
285-
tx: RefCell<::futures::sync::mpsc::Sender<proto::dispatch::ClientMsg<B>>>,
272+
tx: dispatch::Sender<proto::dispatch::ClientMsg<B>, ::Response>,
286273
}
287274

288275
impl<B> Clone for HyperClient<B> {
@@ -296,18 +283,19 @@ impl<B> Clone for HyperClient<B> {
296283

297284
impl<B> self::pool::Ready for HyperClient<B> {
298285
fn poll_ready(&mut self) -> Poll<(), ()> {
299-
self.tx
300-
.borrow_mut()
301-
.poll_ready()
302-
.map_err(|_| ())
286+
if self.tx.is_closed() {
287+
Err(())
288+
} else {
289+
Ok(Async::Ready(()))
290+
}
303291
}
304292
}
305293

306294
impl<B> Drop for HyperClient<B> {
307295
fn drop(&mut self) {
308296
if self.should_close.get() {
309297
self.should_close.set(false);
310-
let _ = self.tx.borrow_mut().try_send(proto::dispatch::ClientMsg::Close);
298+
self.tx.cancel();
311299
}
312300
}
313301
}

src/common/mod.rs

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
11
pub use self::str::ByteStr;
22

33
mod str;
4+
5+
pub enum Never {}

0 commit comments

Comments
 (0)