fix(client): fix polling dispatch channel after it has closed

seanmonstar · seanmonstar · commit 039281b89cf1 · 2019-10-30T14:34:19.000-07:00
diff --git a/src/proto/h1/dispatch.rs b/src/proto/h1/dispatch.rs
@@ -37,6 +37,7 @@ pub struct Server<S: HttpService<B>, B> {
 pub struct Client<B> {
     callback: Option<crate::client::dispatch::Callback<Request<B>, Response<Body>>>,
     rx: ClientRx<B>,
+    rx_closed: bool,
 }
 
 type ClientRx<B> = crate::client::dispatch::Receiver<Request<B>, Response<Body>>;
@@ -490,7 +491,8 @@ impl<B> Client<B> {
     pub fn new(rx: ClientRx<B>) -> Client<B> {
         Client {
             callback: None,
-            rx: rx,
+            rx,
+            rx_closed: false,
         }
     }
 }
@@ -505,6 +507,7 @@ where
     type RecvItem = ResponseHead;
 
     fn poll_msg(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Result<(Self::PollItem, Self::PollBody), Never>>> {
+        debug_assert!(!self.rx_closed);
         match self.rx.poll_next(cx) {
             Poll::Ready(Some((req, mut cb))) => {
                 // check that future hasn't been canceled already
@@ -526,8 +529,9 @@ where
                 }
             },
             Poll::Ready(None) => {
-                trace!("client tx closed");
                 // user has dropped sender handle
+                trace!("client tx closed");
+                self.rx_closed = true;
                 Poll::Ready(None)
             },
             Poll::Pending => Poll::Pending,
@@ -555,7 +559,7 @@ where
                 if let Some(cb) = self.callback.take() {
                     let _ = cb.send(Err((err, None)));
                     Ok(())
-                } else {
+                } else if !self.rx_closed {
                     self.rx.close();
                     if let Some((req, cb)) = self.rx.try_recv() {
                         trace!("canceling queued request with connection error: {}", err);
@@ -566,6 +570,8 @@ where
                     } else {
                         Err(err)
                     }
+                } else {
+                    Err(err)
                 }
             }
         }

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,7 @@ pub struct Server<S: HttpService<B>, B> {`
`37`	`37`	`pub struct Client<B> {`
`38`	`38`	`callback: Option<crate::client::dispatch::Callback<Request<B>, Response<Body>>>,`
`39`	`39`	`rx: ClientRx<B>,`
	`40`	`+ rx_closed: bool,`
`40`	`41`	`}`
`41`	`42`
`42`	`43`	`type ClientRx<B> = crate::client::dispatch::Receiver<Request<B>, Response<Body>>;`
`@@ -490,7 +491,8 @@ impl<B> Client<B> {`
`490`	`491`	`pub fn new(rx: ClientRx<B>) -> Client<B> {`
`491`	`492`	`Client {`
`492`	`493`	`callback: None,`
`493`		`- rx: rx,`
	`494`	`+ rx,`
	`495`	`+ rx_closed: false,`
`494`	`496`	`}`
`495`	`497`	`}`
`496`	`498`	`}`
`@@ -505,6 +507,7 @@ where`
`505`	`507`	`type RecvItem = ResponseHead;`
`506`	`508`
`507`	`509`	`fn poll_msg(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Result<(Self::PollItem, Self::PollBody), Never>>> {`
	`510`	`+ debug_assert!(!self.rx_closed);`
`508`	`511`	`match self.rx.poll_next(cx) {`
`509`	`512`	`Poll::Ready(Some((req, mut cb))) => {`
`510`	`513`	`// check that future hasn't been canceled already`
`@@ -526,8 +529,9 @@ where`
`526`	`529`	`}`
`527`	`530`	`},`
`528`	`531`	`Poll::Ready(None) => {`
`529`		`- trace!("client tx closed");`
`530`	`532`	`// user has dropped sender handle`
	`533`	`+ trace!("client tx closed");`
	`534`	`+ self.rx_closed = true;`
`531`	`535`	`Poll::Ready(None)`
`532`	`536`	`},`
`533`	`537`	`Poll::Pending => Poll::Pending,`
`@@ -555,7 +559,7 @@ where`
`555`	`559`	`if let Some(cb) = self.callback.take() {`
`556`	`560`	`let _ = cb.send(Err((err, None)));`
`557`	`561`	`Ok(())`
`558`		`- } else {`
	`562`	`+ } else if !self.rx_closed {`
`559`	`563`	`self.rx.close();`
`560`	`564`	`if let Some((req, cb)) = self.rx.try_recv() {`
`561`	`565`	`trace!("canceling queued request with connection error: {}", err);`
`@@ -566,6 +570,8 @@ where`
`566`	`570`	`} else {`
`567`	`571`	`Err(err)`
`568`	`572`	`}`
	`573`	`+ } else {`
	`574`	`+ Err(err)`
`569`	`575`	`}`
`570`	`576`	`}`
`571`	`577`	`}`