Merge pull request #448 from sfackler/bufer-copy-in

Buffer copy_in messages

Author: sfackler

.circleci/config.yml

@@ -22,7 +22,7 @@ version: 2
 jobs:
   build:
     docker:
-      - image: rust:1.31.0
+      - image: rust:1.35.0
         environment:
           RUSTFLAGS: -D warnings
       - image: sfackler/rust-postgres-test:5

codegen/src/main.rs

@@ -1,4 +1,5 @@
 #![warn(clippy::all)]
+#![allow(clippy::write_with_newline)]
 
 extern crate linked_hash_map;
 extern crate marksman_escape;

tokio-postgres/src/proto/client.rs

@@ -212,8 +212,8 @@ impl Client {
         S::Error: Into<Box<dyn StdError + Sync + Send>>,
     {
         let (mut sender, receiver) = mpsc::channel(1);
-        let pending = PendingRequest(self.excecute_message(statement, params).map(|buf| {
-            match sender.start_send(CopyMessage::Data(buf)) {
+        let pending = PendingRequest(self.excecute_message(statement, params).map(|data| {
+            match sender.start_send(CopyMessage { data, done: false }) {
                 Ok(AsyncSink::Ready) => {}
                 _ => unreachable!("channel should have capacity"),
             }

tokio-postgres/src/proto/copy_in.rs

@@ -6,14 +6,15 @@ use postgres_protocol::message::backend::Message;
 use postgres_protocol::message::frontend;
 use state_machine_future::{transition, RentToOwn, StateMachineFuture};
 use std::error::Error as StdError;
+use std::mem;
 
 use crate::proto::client::{Client, PendingRequest};
 use crate::proto::statement::Statement;
 use crate::Error;
 
-pub enum CopyMessage {
-    Data(Vec<u8>),
-    Done,
+pub struct CopyMessage {
+    pub data: Vec<u8>,
+    pub done: bool,
 }
 
 pub struct CopyInReceiver {
@@ -40,13 +41,14 @@ impl Stream for CopyInReceiver {
         }
 
         match self.receiver.poll()? {
-            Async::Ready(Some(CopyMessage::Data(buf))) => Ok(Async::Ready(Some(buf))),
-            Async::Ready(Some(CopyMessage::Done)) => {
-                self.done = true;
-                let mut buf = vec![];
-                frontend::copy_done(&mut buf);
-                frontend::sync(&mut buf);
-                Ok(Async::Ready(Some(buf)))
+            Async::Ready(Some(mut data)) => {
+                if data.done {
+                    self.done = true;
+                    frontend::copy_done(&mut data.data);
+                    frontend::sync(&mut data.data);
+                }
+
+                Ok(Async::Ready(Some(data.data)))
             }
             Async::Ready(None) => {
                 self.done = true;
@@ -85,6 +87,7 @@ where
     #[state_machine_future(transitions(WriteCopyDone))]
     WriteCopyData {
         stream: S,
+        buf: Vec<u8>,
         pending_message: Option<CopyMessage>,
         sender: mpsc::Sender<CopyMessage>,
         receiver: mpsc::Receiver<Message>,
@@ -133,6 +136,7 @@ where
                     let state = state.take();
                     transition!(WriteCopyData {
                         stream: state.stream,
+                        buf: vec![],
                         pending_message: None,
                         sender: state.sender,
                         receiver: state.receiver
@@ -148,34 +152,58 @@ where
     fn poll_write_copy_data<'a>(
         state: &'a mut RentToOwn<'a, WriteCopyData<S>>,
     ) -> Poll<AfterWriteCopyData, Error> {
+        if let Some(message) = state.pending_message.take() {
+            match state
+                .sender
+                .start_send(message)
+                .map_err(|_| Error::closed())?
+            {
+                AsyncSink::Ready => {}
+                AsyncSink::NotReady(message) => {
+                    state.pending_message = Some(message);
+                    return Ok(Async::NotReady);
+                }
+            }
+        }
+
         loop {
-            let message = match state.pending_message.take() {
-                Some(message) => message,
-                None => match try_ready!(state.stream.poll().map_err(Error::copy_in_stream)) {
+            let done = loop {
+                match try_ready!(state.stream.poll().map_err(Error::copy_in_stream)) {
                     Some(data) => {
-                        let mut buf = vec![];
                         // FIXME avoid collect
-                        frontend::copy_data(&data.into_buf().collect::<Vec<_>>(), &mut buf)
+                        frontend::copy_data(&data.into_buf().collect::<Vec<_>>(), &mut state.buf)
                             .map_err(Error::encode)?;
-                        CopyMessage::Data(buf)
+                        if state.buf.len() > 4096 {
+                            break false;
+                        }
                     }
-                    None => {
-                        let state = state.take();
-                        transition!(WriteCopyDone {
-                            future: state.sender.send(CopyMessage::Done),
-                            receiver: state.receiver
-                        })
-                    }
-                },
+                    None => break true,
+                }
             };
 
-            match state.sender.start_send(message) {
-                Ok(AsyncSink::Ready) => {}
-                Ok(AsyncSink::NotReady(message)) => {
+            let message = CopyMessage {
+                data: mem::replace(&mut state.buf, vec![]),
+                done,
+            };
+
+            if done {
+                let state = state.take();
+                transition!(WriteCopyDone {
+                    future: state.sender.send(message),
+                    receiver: state.receiver,
+                });
+            }
+
+            match state
+                .sender
+                .start_send(message)
+                .map_err(|_| Error::closed())?
+            {
+                AsyncSink::Ready => {}
+                AsyncSink::NotReady(message) => {
                     state.pending_message = Some(message);
                     return Ok(Async::NotReady);
                 }
-                Err(_) => return Err(Error::closed()),
             }
         }
     }