forked/hyper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(client): close connections when Response Future or Body is dropped

Browse Source 
Closes #1397
This commit is contained in:
Sean McArthur
2017-12-14 13:57:31 -08:00
parent 8f6931b349
commit ef4008121e
3 changed files with 227 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
src/proto/body.rs
View File

@@ -1,6 +1,6 @@
use bytes::Bytes;
use futures::{Poll, Stream};
use futures::sync::mpsc;
use futures::{Async, AsyncSink, Future, Poll, Sink, StartSend, Stream};
use futures::sync::{mpsc, oneshot};
use tokio_proto;
use std::borrow::Cow;
@@ -12,20 +12,36 @@ pub type BodySender = mpsc::Sender<Result<Chunk, ::Error>>;
/// A `Stream` for `Chunk`s used in requests and responses.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Body(TokioBody);
pub struct Body(Inner);
#[derive(Debug)]
enum Inner {
Tokio(TokioBody),
Hyper {
close_tx: oneshot::Sender<()>,
rx: mpsc::Receiver<Result<Chunk, ::Error>>,
}
}
//pub(crate)
#[derive(Debug)]
pub struct ChunkSender {
close_rx: oneshot::Receiver<()>,
tx: BodySender,
}
impl Body {
/// Return an empty body stream
#[inline]
pub fn empty() -> Body {
Body(TokioBody::empty())
Body(Inner::Tokio(TokioBody::empty()))
}
/// Return a body stream with an associated sender half
#[inline]
pub fn pair() -> (mpsc::Sender<Result<Chunk, ::Error>>, Body) {
let (tx, rx) = TokioBody::pair();
let rx = Body(rx);
let rx = Body(Inner::Tokio(rx));
(tx, rx)
}
}
@@ -43,7 +59,51 @@ impl Stream for Body {
#[inline]
fn poll(&mut self) -> Poll<Option<Chunk>, ::Error> {
self.0.poll()
match self.0 {
Inner::Tokio(ref mut rx) => rx.poll(),
Inner::Hyper { ref mut rx, .. } => match rx.poll().expect("mpsc cannot error") {
Async::Ready(Some(Ok(chunk))) => Ok(Async::Ready(Some(chunk))),
Async::Ready(Some(Err(err))) => Err(err),
Async::Ready(None) => Ok(Async::Ready(None)),
Async::NotReady => Ok(Async::NotReady),
},
}
}
}
//pub(crate)
pub fn channel() -> (ChunkSender, Body) {
let (tx, rx) = mpsc::channel(0);
let (close_tx, close_rx) = oneshot::channel();
let tx = ChunkSender {
close_rx: close_rx,
tx: tx,
};
let rx = Body(Inner::Hyper {
close_tx: close_tx,
rx: rx,
});
(tx, rx)
}
impl ChunkSender {
pub fn poll_ready(&mut self) -> Poll<(), ()> {
match self.close_rx.poll() {
Ok(Async::Ready(())) | Err(_) => return Err(()),
Ok(Async::NotReady) => (),
}
self.tx.poll_ready().map_err(|_| ())
}
pub fn start_send(&mut self, msg: Result<Chunk, ::Error>) -> StartSend<(), ()> {
match self.tx.start_send(msg) {
Ok(AsyncSink::Ready) => Ok(AsyncSink::Ready),
Ok(AsyncSink::NotReady(_)) => Ok(AsyncSink::NotReady(())),
Err(_) => Err(()),
}
}
}
@@ -52,7 +112,14 @@ impl Stream for Body {
impl From<Body> for tokio_proto::streaming::Body<Chunk, ::Error> {
#[inline]
fn from(b: Body) -> tokio_proto::streaming::Body<Chunk, ::Error> {
b.0
match b.0 {
Inner::Tokio(b) => b,
Inner::Hyper { close_tx, rx } => {
warn!("converting hyper::Body into a tokio_proto Body is deprecated");
::std::mem::forget(close_tx);
rx.into()
}
}
}
}
@@ -61,42 +128,42 @@ impl From<Body> for tokio_proto::streaming::Body<Chunk, ::Error> {
impl From<tokio_proto::streaming::Body<Chunk, ::Error>> for Body {
#[inline]
fn from(tokio_body: tokio_proto::streaming::Body<Chunk, ::Error>) -> Body {
Body(tokio_body)
Body(Inner::Tokio(tokio_body))
}
}
impl From<mpsc::Receiver<Result<Chunk, ::Error>>> for Body {
#[inline]
fn from(src: mpsc::Receiver<Result<Chunk, ::Error>>) -> Body {
Body(src.into())
TokioBody::from(src).into()
}
}
impl From<Chunk> for Body {
#[inline]
fn from (chunk: Chunk) -> Body {
Body(TokioBody::from(chunk))
TokioBody::from(chunk).into()
}
}
impl From<Bytes> for Body {
#[inline]
fn from (bytes: Bytes) -> Body {
Body(TokioBody::from(Chunk::from(bytes)))
Body::from(TokioBody::from(Chunk::from(bytes)))
}
}
impl From<Vec<u8>> for Body {
#[inline]
fn from (vec: Vec<u8>) -> Body {
Body(TokioBody::from(Chunk::from(vec)))
Body::from(TokioBody::from(Chunk::from(vec)))
}
}
impl From<&'static [u8]> for Body {
#[inline]
fn from (slice: &'static [u8]) -> Body {
Body(TokioBody::from(Chunk::from(slice)))
Body::from(TokioBody::from(Chunk::from(slice)))
}
}
@@ -113,14 +180,14 @@ impl From<Cow<'static, [u8]>> for Body {
impl From<String> for Body {
#[inline]
fn from (s: String) -> Body {
Body(TokioBody::from(Chunk::from(s.into_bytes())))
Body::from(TokioBody::from(Chunk::from(s.into_bytes())))
}
}
impl From<&'static str> for Body {
#[inline]
fn from(slice: &'static str) -> Body {
Body(TokioBody::from(Chunk::from(slice.as_bytes())))
Body::from(TokioBody::from(Chunk::from(slice.as_bytes())))
}
}

62
src/proto/dispatch.rs
View File

@@ -1,6 +1,6 @@
use std::io;
use futures::{Async, AsyncSink, Future, Poll, Sink, Stream};
use futures::{Async, AsyncSink, Future, Poll, Stream};
use futures::sync::{mpsc, oneshot};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_service::Service;
@@ -11,7 +11,7 @@ use ::StatusCode;
pub struct Dispatcher<D, Bs, I, B, T, K> {
conn: Conn<I, B, T, K>,
dispatch: D,
body_tx: Option<super::body::BodySender>,
body_tx: Option<super::body::ChunkSender>,
body_rx: Option<Bs>,
is_closing: bool,
}
@@ -22,6 +22,7 @@ pub trait Dispatch {
type RecvItem;
fn poll_msg(&mut self) -> Poll<Option<(Self::PollItem, Option<Self::PollBody>)>, ::Error>;
fn recv_msg(&mut self, msg: ::Result<(Self::RecvItem, Option<Body>)>) -> ::Result<()>;
fn poll_ready(&mut self) -> Poll<(), ()>;
fn should_poll(&self) -> bool;
}
@@ -70,10 +71,22 @@ where
if self.is_closing {
return Ok(Async::Ready(()));
} else if self.conn.can_read_head() {
// can dispatch receive, or does it still care about, an incoming message?
match self.dispatch.poll_ready() {
Ok(Async::Ready(())) => (),
Ok(Async::NotReady) => unreachable!("dispatch not ready when conn is"),
Err(()) => {
trace!("dispatch no longer receiving messages");
self.is_closing = true;
return Ok(Async::Ready(()));
}
}
// dispatch is ready for a message, try to read one
match self.conn.read_head() {
Ok(Async::Ready(Some((head, has_body)))) => {
let body = if has_body {
let (tx, rx) = super::Body::pair();
let (mut tx, rx) = super::body::channel();
let _ = tx.poll_ready(); // register this task if rx is dropped
self.body_tx = Some(tx);
Some(rx)
} else {
@@ -111,6 +124,8 @@ where
self.conn.close_read();
return Ok(Async::Ready(()));
}
// else the conn body is done, and user dropped,
// so everything is fine!
}
}
if can_read_body {
@@ -133,7 +148,7 @@ where
}
},
Ok(Async::Ready(None)) => {
let _ = body.close();
// just drop, the body will close automatically
},
Ok(Async::NotReady) => {
self.body_tx = Some(body);
@@ -144,7 +159,7 @@ where
}
}
} else {
let _ = body.close();
// just drop, the body will close automatically
}
} else if !T::should_read_first() {
self.conn.try_empty_read()?;
@@ -305,6 +320,14 @@ where
Ok(())
}
fn poll_ready(&mut self) -> Poll<(), ()> {
if self.in_flight.is_some() {
Ok(Async::NotReady)
} else {
Ok(Async::Ready(()))
}
}
fn should_poll(&self) -> bool {
self.in_flight.is_some()
}
@@ -333,9 +356,18 @@ where
fn poll_msg(&mut self) -> Poll<Option<(Self::PollItem, Option<Self::PollBody>)>, ::Error> {
match self.rx.poll() {
Ok(Async::Ready(Some(ClientMsg::Request(head, body, cb)))) => {
self.callback = Some(cb);
Ok(Async::Ready(Some((head, body))))
Ok(Async::Ready(Some(ClientMsg::Request(head, body, mut cb)))) => {
// check that future hasn't been canceled already
match cb.poll_cancel().expect("poll_cancel cannot error") {
Async::Ready(()) => {
trace!("request canceled");
Ok(Async::Ready(None))
},
Async::NotReady => {
self.callback = Some(cb);
Ok(Async::Ready(Some((head, body))))
}
}
},
Ok(Async::Ready(Some(ClientMsg::Close))) |
Ok(Async::Ready(None)) => {
@@ -370,6 +402,20 @@ where
}
}
fn poll_ready(&mut self) -> Poll<(), ()> {
match self.callback {
Some(ref mut cb) => match cb.poll_cancel() {
Ok(Async::Ready(())) => {
trace!("callback receiver has dropped");
Err(())
},
Ok(Async::NotReady) => Ok(Async::Ready(())),
Err(_) => unreachable!("oneshot poll_cancel cannot error"),
},
None => Err(()),
}
}
fn should_poll(&self) -> bool {
self.callback.is_none()
}