Only flush if necessary

sfackler · sfackler · commit 9d5a36e1739e · 2018-06-18T22:49:53.000-04:00
If poll_write was blocked trying to write a request out to the socket,
it's because the write triggered a flush and the socket wasn't ready. We
don't want to try to flush again, since it's at best a waste of time and
at worst can cause a deadlock if the socket becomes available after the
poll_write and before the poll_flush. If that happens, we should be in
poll_write again writing more data but that wouldn't happen.
diff --git a/tokio-postgres/src/proto/connection.rs b/tokio-postgres/src/proto/connection.rs
@@ -75,10 +75,10 @@ impl Connection {
         self.stream.poll()
     }
 
-    fn poll_read(&mut self) -> Poll<(), Error> {
+    fn poll_read(&mut self) -> Result<(), Error> {
         if self.state != State::Active {
             trace!("poll_read: done");
-            return Ok(Async::Ready(()));
+            return Ok(());
         }
 
         loop {
@@ -89,7 +89,7 @@ impl Connection {
                 }
                 Async::NotReady => {
                     trace!("poll_read: waiting on response");
-                    return Ok(Async::NotReady);
+                    return Ok(());
                 }
             };
 
@@ -131,7 +131,7 @@ impl Connection {
                     self.responses.push_front(sender);
                     self.pending_response = Some(message);
                     trace!("poll_read: waiting on socket");
-                    return Ok(Async::NotReady);
+                    return Ok(());
                 }
             }
         }
@@ -153,11 +153,11 @@ impl Connection {
         }
     }
 
-    fn poll_write(&mut self) -> Poll<(), Error> {
+    fn poll_write(&mut self) -> Result<bool, Error> {
         loop {
             if self.state == State::Closing {
                 trace!("poll_write: done");
-                return Ok(Async::Ready(()));
+                return Ok(false);
             }
 
             let request = match self.poll_request()? {
@@ -174,12 +174,12 @@ impl Connection {
                         "poll_write: at eof, pending responses {}",
                         self.responses.len(),
                     );
-                    return Ok(Async::Ready(()));
+                    return Ok(true);
                 }
                 Async::NotReady => {
                     trace!("poll_write: waiting on request");
                     self.taker.want();
-                    return Ok(Async::NotReady);
+                    return Ok(true);
                 }
             };
 
@@ -193,15 +193,24 @@ impl Connection {
                 AsyncSink::NotReady(request) => {
                     trace!("poll_write: waiting on socket");
                     self.pending_request = Some(request);
-                    return Ok(Async::NotReady);
+                    return Ok(false);
                 }
             }
         }
     }
 
-    fn poll_flush(&mut self) -> Poll<(), Error> {
-        trace!("flushing");
-        self.stream.poll_complete().map_err(Into::into)
+    fn poll_flush(&mut self) -> Result<(), Error> {
+        match self.stream.poll_complete() {
+            Ok(Async::Ready(())) => {
+                trace!("poll_flush: flushed");
+                Ok(())
+            }
+            Ok(Async::NotReady) => {
+                trace!("poll_flush: waiting on socket");
+                Ok(())
+            }
+            Err(e) => Err(Error::from(e)),
+        }
     }
 
     fn poll_shutdown(&mut self) -> Poll<(), Error> {
@@ -229,8 +238,10 @@ impl Future for Connection {
 
     fn poll(&mut self) -> Poll<(), Error> {
         self.poll_read()?;
-        self.poll_write()?;
-        self.poll_flush()?;
+        let want_flush = self.poll_write()?;
+        if want_flush {
+            self.poll_flush()?;
+        }
         self.poll_shutdown()
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -75,10 +75,10 @@ impl Connection {`
`75`	`75`	`self.stream.poll()`
`76`	`76`	`}`
`77`	`77`
`78`		`- fn poll_read(&mut self) -> Poll<(), Error> {`
	`78`	`+ fn poll_read(&mut self) -> Result<(), Error> {`
`79`	`79`	`if self.state != State::Active {`
`80`	`80`	`trace!("poll_read: done");`
`81`		`- return Ok(Async::Ready(()));`
	`81`	`+ return Ok(());`
`82`	`82`	`}`
`83`	`83`
`84`	`84`	`loop {`
`@@ -89,7 +89,7 @@ impl Connection {`
`89`	`89`	`}`
`90`	`90`	`Async::NotReady => {`
`91`	`91`	`trace!("poll_read: waiting on response");`
`92`		`- return Ok(Async::NotReady);`
	`92`	`+ return Ok(());`
`93`	`93`	`}`
`94`	`94`	`};`
`95`	`95`
`@@ -131,7 +131,7 @@ impl Connection {`
`131`	`131`	`self.responses.push_front(sender);`
`132`	`132`	`self.pending_response = Some(message);`
`133`	`133`	`trace!("poll_read: waiting on socket");`
`134`		`- return Ok(Async::NotReady);`
	`134`	`+ return Ok(());`
`135`	`135`	`}`
`136`	`136`	`}`
`137`	`137`	`}`
`@@ -153,11 +153,11 @@ impl Connection {`
`153`	`153`	`}`
`154`	`154`	`}`
`155`	`155`
`156`		`- fn poll_write(&mut self) -> Poll<(), Error> {`
	`156`	`+ fn poll_write(&mut self) -> Result<bool, Error> {`
`157`	`157`	`loop {`
`158`	`158`	`if self.state == State::Closing {`
`159`	`159`	`trace!("poll_write: done");`
`160`		`- return Ok(Async::Ready(()));`
	`160`	`+ return Ok(false);`
`161`	`161`	`}`
`162`	`162`
`163`	`163`	`let request = match self.poll_request()? {`
`@@ -174,12 +174,12 @@ impl Connection {`
`174`	`174`	`"poll_write: at eof, pending responses {}",`
`175`	`175`	`self.responses.len(),`
`176`	`176`	`);`
`177`		`- return Ok(Async::Ready(()));`
	`177`	`+ return Ok(true);`
`178`	`178`	`}`
`179`	`179`	`Async::NotReady => {`
`180`	`180`	`trace!("poll_write: waiting on request");`
`181`	`181`	`self.taker.want();`
`182`		`- return Ok(Async::NotReady);`
	`182`	`+ return Ok(true);`
`183`	`183`	`}`
`184`	`184`	`};`
`185`	`185`
`@@ -193,15 +193,24 @@ impl Connection {`
`193`	`193`	`AsyncSink::NotReady(request) => {`
`194`	`194`	`trace!("poll_write: waiting on socket");`
`195`	`195`	`self.pending_request = Some(request);`
`196`		`- return Ok(Async::NotReady);`
	`196`	`+ return Ok(false);`
`197`	`197`	`}`
`198`	`198`	`}`
`199`	`199`	`}`
`200`	`200`	`}`
`201`	`201`
`202`		`- fn poll_flush(&mut self) -> Poll<(), Error> {`
`203`		`- trace!("flushing");`
`204`		`- self.stream.poll_complete().map_err(Into::into)`
	`202`	`+ fn poll_flush(&mut self) -> Result<(), Error> {`
	`203`	`+ match self.stream.poll_complete() {`
	`204`	`+ Ok(Async::Ready(())) => {`
	`205`	`+ trace!("poll_flush: flushed");`
	`206`	`+ Ok(())`
	`207`	`+ }`
	`208`	`+ Ok(Async::NotReady) => {`
	`209`	`+ trace!("poll_flush: waiting on socket");`
	`210`	`+ Ok(())`
	`211`	`+ }`
	`212`	`+ Err(e) => Err(Error::from(e)),`
	`213`	`+ }`
`205`	`214`	`}`
`206`	`215`
`207`	`216`	`fn poll_shutdown(&mut self) -> Poll<(), Error> {`
`@@ -229,8 +238,10 @@ impl Future for Connection {`
`229`	`238`
`230`	`239`	`fn poll(&mut self) -> Poll<(), Error> {`
`231`	`240`	`self.poll_read()?;`
`232`		`- self.poll_write()?;`
`233`		`- self.poll_flush()?;`
	`241`	`+ let want_flush = self.poll_write()?;`
	`242`	`+ if want_flush {`
	`243`	`+ self.poll_flush()?;`
	`244`	`+ }`
`234`	`245`	`self.poll_shutdown()`
`235`	`246`	`}`
`236`	`247`	`}`