sqlx-core: Implement poll methods and Sink for BufferedSocket

Author: joeydewaal

sqlx-core/src/net/socket/buffered.rs

@@ -1,7 +1,10 @@
 use crate::error::Error;
-use crate::net::Socket;
+use crate::net::{Socket, SocketExt};
 use bytes::BytesMut;
+use futures_util::Sink;
 use std::ops::ControlFlow;
+use std::pin::Pin;
+use std::task::{ready, Context, Poll};
 use std::{cmp, io};
 
 use crate::io::{AsyncRead, AsyncReadExt, ProtocolDecode, ProtocolEncode};
@@ -13,6 +16,7 @@ pub struct BufferedSocket<S> {
     socket: S,
     write_buf: WriteBuffer,
     read_buf: ReadBuffer,
+    wants_bytes: usize,
 }
 
 pub struct WriteBuffer {
@@ -42,6 +46,7 @@ impl<S: Socket> BufferedSocket<S> {
                 read: BytesMut::new(),
                 available: BytesMut::with_capacity(DEFAULT_BUF_SIZE),
             },
+            wants_bytes: 0,
         }
     }
 
@@ -56,6 +61,25 @@ impl<S: Socket> BufferedSocket<S> {
         .await
     }
 
+    pub fn poll_try_read<F, R>(
+        &mut self,
+        cx: &mut Context<'_>,
+        mut try_read: F,
+    ) -> Poll<Result<R, Error>>
+    where
+        F: FnMut(&mut BytesMut) -> Result<ControlFlow<R, usize>, Error>,
+    {
+        loop {
+            // Read if we want bytes
+            ready!(self.poll_handle_read(cx)?);
+
+            match try_read(&mut self.read_buf.read)? {
+                ControlFlow::Continue(read_len) => self.wants_bytes = read_len,
+                ControlFlow::Break(ret) => return Poll::Ready(Ok(ret)),
+            };
+        }
+    }
+
     /// Retryable read operation.
     ///
     /// The callback should check the contents of the buffer passed to it and either:
@@ -125,8 +149,8 @@ impl<S: Socket> BufferedSocket<S> {
     pub async fn flush(&mut self) -> io::Result<()> {
         while !self.write_buf.is_empty() {
             let written = self.socket.write(self.write_buf.get()).await?;
+            // Consume does the sanity check
             self.write_buf.consume(written);
-            self.write_buf.sanity_check();
         }
 
         self.socket.flush().await?;
@@ -154,8 +178,37 @@ impl<S: Socket> BufferedSocket<S> {
             socket: Box::new(self.socket),
             write_buf: self.write_buf,
             read_buf: self.read_buf,
+            wants_bytes: self.wants_bytes,
         }
     }
+
+    fn poll_handle_read(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        // Because of how `BytesMut` works, we should only be shifting capacity back and forth
+        // between `read` and `available` unless we have to read an oversize message.
+        while self.read_buf.len() < self.wants_bytes {
+            self.read_buf
+                .reserve(self.wants_bytes - self.read_buf.len());
+
+            let read = ready!(self.socket.poll_read(cx, &mut self.read_buf.available)?);
+
+            if read == 0 {
+                return Poll::Ready(Err(io::Error::new(
+                    io::ErrorKind::UnexpectedEof,
+                    format!(
+                        "expected to read {} bytes, got {} bytes at EOF",
+                        self.wants_bytes,
+                        self.read_buf.len()
+                    ),
+                )));
+            }
+
+            self.read_buf.advance(read);
+        }
+
+        // we've read at least enough for `wants_bytes`, so we don't want more.
+        self.wants_bytes = 0;
+        Poll::Ready(Ok(()))
+    }
 }
 
 impl WriteBuffer {
@@ -326,4 +379,41 @@ impl ReadBuffer {
             self.available = BytesMut::with_capacity(DEFAULT_BUF_SIZE);
         }
     }
+
+    fn len(&self) -> usize {
+        self.read.len()
+    }
+}
+
+impl<S: Socket> Sink<&[u8]> for BufferedSocket<S> {
+    type Error = crate::Error;
+
+    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
+        if self.write_buf.bytes_written >= DEFAULT_BUF_SIZE {
+            self.poll_flush(cx)
+        } else {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    fn start_send(mut self: Pin<&mut Self>, item: &[u8]) -> crate::Result<()> {
+        self.write_buffer_mut().put_slice(item);
+        Ok(())
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
+        let this = &mut *self;
+
+        while !this.write_buf.is_empty() {
+            let written = ready!(this.socket.poll_write(cx, this.write_buf.get())?);
+            // Consume does the sanity check
+            this.write_buf.consume(written);
+        }
+        this.socket.poll_flush(cx).map_err(Into::into)
+    }
+
+    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
+        ready!(self.as_mut().poll_flush(cx))?;
+        self.socket.poll_shutdown(cx).map_err(Into::into)
+    }
 }

sqlx-postgres/src/connection/tls.rs

@@ -1,6 +1,6 @@
 use crate::error::Error;
 use crate::net::tls::{self, TlsConfig};
-use crate::net::{Socket, SocketIntoBox, WithSocket};
+use crate::net::{Socket, SocketExt, SocketIntoBox, WithSocket};
 
 use crate::message::SslRequest;
 use crate::{PgConnectOptions, PgSslMode};