From a1fc412f9e5036149981ea3400159c47602d280a Mon Sep 17 00:00:00 2001 From: Steven Fackler Date: Sun, 23 Dec 2018 15:02:42 -0800 Subject: [PATCH] Simplify sync copy_in --- postgres/src/client.rs | 99 ++++++------------------------------------ 1 file changed, 13 insertions(+), 86 deletions(-) diff --git a/postgres/src/client.rs b/postgres/src/client.rs index 6ce11650..99d64842 100644 --- a/postgres/src/client.rs +++ b/postgres/src/client.rs @@ -1,7 +1,6 @@ use bytes::{Buf, Bytes}; use futures::stream; -use futures::sync::mpsc; -use futures::{Async, AsyncSink, Future, Poll, Sink, Stream}; +use futures::{Async, Future, Poll, Stream}; use std::io::{self, BufRead, Cursor, Read}; use std::marker::PhantomData; use tokio_postgres::types::{ToSql, Type}; @@ -68,17 +67,9 @@ impl Client { R: Read, { let statement = query.__statement(self)?; - let (sender, receiver) = mpsc::channel(1); - let future = self.0.copy_in(&statement.0, params, CopyInStream(receiver)); - - CopyInFuture { - future, - sender, - reader, - pending: None, - done: false, - } - .wait() + self.0 + .copy_in(&statement.0, params, CopyInStream(reader)) + .wait() } pub fn copy_out( @@ -125,84 +116,20 @@ impl From for Client { } } -enum CopyData { - Data(Vec), - Error(io::Error), - Done, -} +struct CopyInStream(R); -struct CopyInStream(mpsc::Receiver); - -impl Stream for CopyInStream { +impl Stream for CopyInStream +where + R: Read, +{ type Item = Vec; type Error = io::Error; fn poll(&mut self) -> Poll>, io::Error> { - match self.0.poll().expect("mpsc::Receiver can't error") { - Async::Ready(Some(CopyData::Data(buf))) => Ok(Async::Ready(Some(buf))), - Async::Ready(Some(CopyData::Error(e))) => Err(e), - Async::Ready(Some(CopyData::Done)) => Ok(Async::Ready(None)), - Async::Ready(None) => Err(io::Error::new(io::ErrorKind::Other, "writer disconnected")), - Async::NotReady => Ok(Async::NotReady), - } - } -} - -struct CopyInFuture { - future: tokio_postgres::CopyIn, - sender: mpsc::Sender, - reader: R, - pending: Option, - done: bool, -} - -impl Future for CopyInFuture -where - R: Read, -{ - type Item = u64; - type Error = Error; - - fn poll(&mut self) -> Poll { - loop { - if let Async::Ready(n) = self.future.poll()? { - return Ok(Async::Ready(n)); - } - - let data = match self.pending.take() { - Some(pending) => pending, - None => { - if self.done { - continue; - } - - let mut buf = vec![]; - match self.reader.by_ref().take(4096).read_to_end(&mut buf) { - Ok(0) => { - self.done = true; - CopyData::Done - } - Ok(_) => CopyData::Data(buf), - Err(e) => { - self.done = true; - CopyData::Error(e) - } - } - } - }; - - match self.sender.start_send(data) { - Ok(AsyncSink::Ready) => {} - Ok(AsyncSink::NotReady(pending)) => { - self.pending = Some(pending); - return Ok(Async::NotReady); - } - // the future's hung up on its end of the channel, so we'll wait for it to error - Err(_) => { - self.done = true; - return Ok(Async::NotReady); - } - } + let mut buf = vec![]; + match self.0.by_ref().take(4096).read_to_end(&mut buf)? { + 0 => Ok(Async::Ready(None)), + _ => Ok(Async::Ready(Some(buf))), } } }