Merge pull request rust-postgres#451 from sfackler/less-copy-copies

sfackler · web-flow · commit eaef62c34085 · 2019-06-25T19:27:21.000-07:00
Avoid copies in copy_in
diff --git a/postgres-protocol/src/message/frontend.rs b/postgres-protocol/src/message/frontend.rs
@@ -2,6 +2,8 @@
 #![allow(missing_docs)]
 
 use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
+use bytes::{Buf, BufMut, BytesMut, IntoBuf};
+use std::convert::TryFrom;
 use std::error::Error;
 use std::io;
 use std::marker;
@@ -263,6 +265,40 @@ pub fn copy_data(data: &[u8], buf: &mut Vec<u8>) -> io::Result<()> {
     })
 }
 
+pub struct CopyData<T> {
+    buf: T,
+    len: i32,
+}
+
+impl<T> CopyData<T>
+where
+    T: Buf,
+{
+    pub fn new<U>(buf: U) -> io::Result<CopyData<T>>
+    where
+        U: IntoBuf<Buf = T>,
+    {
+        let buf = buf.into_buf();
+
+        let len = buf
+            .remaining()
+            .checked_add(4)
+            .and_then(|l| i32::try_from(l).ok())
+            .ok_or_else(|| {
+                io::Error::new(io::ErrorKind::InvalidInput, "message length overflow")
+            })?;
+
+        Ok(CopyData { buf, len })
+    }
+
+    pub fn write(self, out: &mut BytesMut) {
+        out.reserve(self.len as usize + 1);
+        out.put_u8(b'd');
+        out.put_i32_be(self.len);
+        out.put(self.buf);
+    }
+}
+
 #[inline]
 pub fn copy_done(buf: &mut Vec<u8>) {
     buf.push(b'c');
diff --git a/tokio-postgres/src/impls.rs b/tokio-postgres/src/impls.rs
@@ -170,7 +170,7 @@ pub struct CopyIn<S>(pub(crate) proto::CopyInFuture<S>)
 where
     S: Stream,
     S::Item: IntoBuf,
-    <S::Item as IntoBuf>::Buf: Send,
+    <S::Item as IntoBuf>::Buf: 'static + Send,
     S::Error: Into<Box<dyn error::Error + Sync + Send>>;
 
 impl<S> Future for CopyIn<S>
diff --git a/tokio-postgres/src/lib.rs b/tokio-postgres/src/lib.rs
@@ -242,7 +242,7 @@ impl Client {
     where
         S: Stream,
         S::Item: IntoBuf,
-        <S::Item as IntoBuf>::Buf: Send,
+        <S::Item as IntoBuf>::Buf: 'static + Send,
         // FIXME error type?
         S::Error: Into<Box<dyn StdError + Sync + Send>>,
     {
diff --git a/tokio-postgres/src/proto/client.rs b/tokio-postgres/src/proto/client.rs
@@ -11,6 +11,7 @@ use std::sync::{Arc, Weak};
 use tokio_io::{AsyncRead, AsyncWrite};
 
 use crate::proto::bind::BindFuture;
+use crate::proto::codec::FrontendMessage;
 use crate::proto::connection::{Request, RequestMessages};
 use crate::proto::copy_in::{CopyInFuture, CopyInReceiver, CopyMessage};
 use crate::proto::copy_out::CopyOutStream;
@@ -185,8 +186,12 @@ impl Client {
         if let Ok(ref mut buf) = buf {
             frontend::sync(buf);
         }
-        let pending =
-            PendingRequest(buf.map(|m| (RequestMessages::Single(m), self.0.idle.guard())));
+        let pending = PendingRequest(buf.map(|m| {
+            (
+                RequestMessages::Single(FrontendMessage::Raw(m)),
+                self.0.idle.guard(),
+            )
+        }));
         BindFuture::new(self.clone(), pending, name, statement.clone())
     }
 
@@ -208,12 +213,12 @@ impl Client {
     where
         S: Stream,
         S::Item: IntoBuf,
-        <S::Item as IntoBuf>::Buf: Send,
+        <S::Item as IntoBuf>::Buf: 'static + Send,
         S::Error: Into<Box<dyn StdError + Sync + Send>>,
     {
         let (mut sender, receiver) = mpsc::channel(1);
         let pending = PendingRequest(self.excecute_message(statement, params).map(|data| {
-            match sender.start_send(CopyMessage { data, done: false }) {
+            match sender.start_send(CopyMessage::Message(data)) {
                 Ok(AsyncSink::Ready) => {}
                 _ => unreachable!("channel should have capacity"),
             }
@@ -278,7 +283,7 @@ impl Client {
         frontend::sync(&mut buf);
         let (sender, _) = mpsc::channel(0);
         let _ = self.0.sender.unbounded_send(Request {
-            messages: RequestMessages::Single(buf),
+            messages: RequestMessages::Single(FrontendMessage::Raw(buf)),
             sender,
             idle: None,
         });
@@ -326,20 +331,23 @@ impl Client {
         &self,
         statement: &Statement,
         params: &[&dyn ToSql],
-    ) -> Result<Vec<u8>, Error> {
+    ) -> Result<FrontendMessage, Error> {
         let mut buf = self.bind_message(statement, "", params)?;
         frontend::execute("", 0, &mut buf).map_err(Error::parse)?;
         frontend::sync(&mut buf);
-        Ok(buf)
+        Ok(FrontendMessage::Raw(buf))
     }
 
     fn pending<F>(&self, messages: F) -> PendingRequest
     where
         F: FnOnce(&mut Vec<u8>) -> Result<(), Error>,
     {
         let mut buf = vec![];
-        PendingRequest(
-            messages(&mut buf).map(|()| (RequestMessages::Single(buf), self.0.idle.guard())),
-        )
+        PendingRequest(messages(&mut buf).map(|()| {
+            (
+                RequestMessages::Single(FrontendMessage::Raw(buf)),
+                self.0.idle.guard(),
+            )
+        }))
     }
 }
diff --git a/tokio-postgres/src/proto/codec.rs b/tokio-postgres/src/proto/codec.rs
@@ -1,16 +1,26 @@
-use bytes::BytesMut;
+use bytes::{Buf, BytesMut};
 use postgres_protocol::message::backend;
+use postgres_protocol::message::frontend::CopyData;
 use std::io;
 use tokio_codec::{Decoder, Encoder};
 
+pub enum FrontendMessage {
+    Raw(Vec<u8>),
+    CopyData(CopyData<Box<dyn Buf + Send>>),
+}
+
 pub struct PostgresCodec;
 
 impl Encoder for PostgresCodec {
-    type Item = Vec<u8>;
+    type Item = FrontendMessage;
     type Error = io::Error;
 
-    fn encode(&mut self, item: Vec<u8>, dst: &mut BytesMut) -> Result<(), io::Error> {
-        dst.extend_from_slice(&item);
+    fn encode(&mut self, item: FrontendMessage, dst: &mut BytesMut) -> Result<(), io::Error> {
+        match item {
+            FrontendMessage::Raw(buf) => dst.extend_from_slice(&buf),
+            FrontendMessage::CopyData(data) => data.write(dst),
+        }
+
         Ok(())
     }
 }
diff --git a/tokio-postgres/src/proto/connect_raw.rs b/tokio-postgres/src/proto/connect_raw.rs
@@ -11,7 +11,7 @@ use std::collections::HashMap;
 use tokio_codec::Framed;
 use tokio_io::{AsyncRead, AsyncWrite};
 
-use crate::proto::{Client, Connection, MaybeTlsStream, PostgresCodec, TlsFuture};
+use crate::proto::{Client, Connection, FrontendMessage, MaybeTlsStream, PostgresCodec, TlsFuture};
 use crate::tls::ChannelBinding;
 use crate::{Config, Error, TlsConnect};
 
@@ -111,7 +111,7 @@ where
         let stream = Framed::new(stream, PostgresCodec);
 
         transition!(SendingStartup {
-            future: stream.send(buf),
+            future: stream.send(FrontendMessage::Raw(buf)),
             config: state.config,
             idx: state.idx,
             channel_binding,
@@ -156,7 +156,7 @@ where
                 let mut buf = vec![];
                 frontend::password_message(pass, &mut buf).map_err(Error::encode)?;
                 transition!(SendingPassword {
-                    future: state.stream.send(buf),
+                    future: state.stream.send(FrontendMessage::Raw(buf)),
                     config: state.config,
                     idx: state.idx,
                 })
@@ -178,7 +178,7 @@ where
                 let mut buf = vec![];
                 frontend::password_message(output.as_bytes(), &mut buf).map_err(Error::encode)?;
                 transition!(SendingPassword {
-                    future: state.stream.send(buf),
+                    future: state.stream.send(FrontendMessage::Raw(buf)),
                     config: state.config,
                     idx: state.idx,
                 })
@@ -235,7 +235,7 @@ where
                     .map_err(Error::encode)?;
 
                 transition!(SendingSasl {
-                    future: state.stream.send(buf),
+                    future: state.stream.send(FrontendMessage::Raw(buf)),
                     scram,
                     config: state.config,
                     idx: state.idx,
@@ -293,7 +293,7 @@ where
                 let mut buf = vec![];
                 frontend::sasl_response(state.scram.message(), &mut buf).map_err(Error::encode)?;
                 transition!(SendingSasl {
-                    future: state.stream.send(buf),
+                    future: state.stream.send(FrontendMessage::Raw(buf)),
                     scram: state.scram,
                     config: state.config,
                     idx: state.idx,
diff --git a/tokio-postgres/src/proto/connection.rs b/tokio-postgres/src/proto/connection.rs
@@ -8,17 +8,17 @@ use std::io;
 use tokio_codec::Framed;
 use tokio_io::{AsyncRead, AsyncWrite};
 
-use crate::proto::codec::PostgresCodec;
+use crate::proto::codec::{FrontendMessage, PostgresCodec};
 use crate::proto::copy_in::CopyInReceiver;
 use crate::proto::idle::IdleGuard;
 use crate::{AsyncMessage, Notification};
 use crate::{DbError, Error};
 
 pub enum RequestMessages {
-    Single(Vec<u8>),
+    Single(FrontendMessage),
     CopyIn {
         receiver: CopyInReceiver,
-        pending_message: Option<Vec<u8>>,
+        pending_message: Option<FrontendMessage>,
     },
 }
 
@@ -188,7 +188,7 @@ where
                     self.state = State::Terminating;
                     let mut request = vec![];
                     frontend::terminate(&mut request);
-                    RequestMessages::Single(request)
+                    RequestMessages::Single(FrontendMessage::Raw(request))
                 }
                 Async::Ready(None) => {
                     trace!(
diff --git a/tokio-postgres/src/proto/copy_in.rs b/tokio-postgres/src/proto/copy_in.rs
diff --git a/tokio-postgres/src/proto/mod.rs b/tokio-postgres/src/proto/mod.rs
diff --git a/tokio-postgres/tests/test/main.rs b/tokio-postgres/tests/test/main.rs

Original file line number	Diff line number	Diff line change
`@@ -242,7 +242,7 @@ impl Client {`
`242`	`242`	`where`
`243`	`243`	`S: Stream,`
`244`	`244`	`S::Item: IntoBuf,`
`245`		`- <S::Item as IntoBuf>::Buf: Send,`
	`245`	`+ <S::Item as IntoBuf>::Buf: 'static + Send,`
`246`	`246`	`// FIXME error type?`
`247`	`247`	`S::Error: Into<Box<dyn StdError + Sync + Send>>,`
`248`	`248`	`{`