refactor: use noop waker for no blocking Connection io

Author: kezhuw

Cargo.toml

@@ -33,6 +33,7 @@ rustls = "0.23.2"
 rustls-pemfile = "2"
 webpki-roots = "0.26.1"
 derive-where = "1.2.7"
+tokio-rustls = "0.26.0"
 
 [dev-dependencies]
 test-log = "0.2.12"

src/session/connection.rs

@@ -1,128 +1,129 @@
-use std::io::{self, Read, Write};
-use std::sync::Arc;
+use std::io::{ErrorKind, IoSlice, Result};
+use std::pin::Pin;
+use std::ptr;
+use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
 
-use rustls::pki_types::ServerName;
-use rustls::{ClientConfig, ClientConnection};
+use bytes::buf::BufMut;
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio::net::TcpStream;
+use tokio_rustls::client::TlsStream;
 
-use crate::error::Error;
+const NOOP_VTABLE: RawWakerVTable =
+    RawWakerVTable::new(|_| RawWaker::new(ptr::null(), &NOOP_VTABLE), |_| {}, |_| {}, |_| {});
+const NOOP_WAKER: RawWaker = RawWaker::new(ptr::null(), &NOOP_VTABLE);
 
-pub struct Connection {
-    tls: Option<ClientConnection>,
-    stream: TcpStream,
+pub enum Connection {
+    Tls(TlsStream<TcpStream>),
+    Raw(TcpStream),
 }
 
-struct WrappingStream<'a> {
-    stream: &'a TcpStream,
-}
-
-impl io::Read for WrappingStream<'_> {
-    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
-        self.stream.try_read_buf(&mut buf)
+impl AsyncRead for Connection {
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<Result<()>> {
+        match self.get_mut() {
+            Self::Raw(stream) => Pin::new(stream).poll_read(cx, buf),
+            Self::Tls(stream) => Pin::new(stream).poll_read(cx, buf),
+        }
     }
 }
 
-impl io::Write for WrappingStream<'_> {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        self.stream.try_write(buf)
+impl AsyncWrite for Connection {
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
+        match self.get_mut() {
+            Self::Raw(stream) => Pin::new(stream).poll_write(cx, buf),
+            Self::Tls(stream) => Pin::new(stream).poll_write(cx, buf),
+        }
     }
 
-    fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
-        self.stream.try_write_vectored(bufs)
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        match self.get_mut() {
+            Self::Raw(stream) => Pin::new(stream).poll_flush(cx),
+            Self::Tls(stream) => Pin::new(stream).poll_flush(cx),
+        }
     }
 
-    fn flush(&mut self) -> io::Result<()> {
-        Ok(())
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        match self.get_mut() {
+            Self::Raw(stream) => Pin::new(stream).poll_shutdown(cx),
+            Self::Tls(stream) => Pin::new(stream).poll_shutdown(cx),
+        }
     }
 }
 
 impl Connection {
     pub fn new_raw(stream: TcpStream) -> Self {
-        Self { tls: None, stream }
+        Self::Raw(stream)
     }
 
-    pub fn new_tls(host: &str, config: Arc<ClientConfig>, stream: TcpStream) -> Result<Self, Error> {
-        let name = match ServerName::try_from(host) {
-            Err(_) => return Err(Error::BadArguments(&"invalid server dns name")),
-            Ok(name) => name.to_owned(),
-        };
-        let client = match ClientConnection::new(config, name) {
-            Err(err) => return Err(Error::other(format!("fail to create tls client for host({host}): {err}"), err)),
-            Ok(client) => client,
-        };
-        Ok(Self { tls: Some(client), stream })
+    pub fn new_tls(stream: TlsStream<TcpStream>) -> Self {
+        Self::Tls(stream)
     }
 
-    pub fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
-        let Some(client) = self.tls.as_mut() else {
-            return self.stream.try_write_vectored(bufs);
-        };
-        let n = client.writer().write_vectored(bufs)?;
-        let mut stream = WrappingStream { stream: &self.stream };
-        client.write_tls(&mut stream)?;
-        Ok(n)
+    pub fn try_write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> Result<usize> {
+        let waker = unsafe { Waker::from_raw(NOOP_WAKER) };
+        let mut context = Context::from_waker(&waker);
+        match Pin::new(self).poll_write_vectored(&mut context, bufs) {
+            Poll::Pending => Err(ErrorKind::WouldBlock.into()),
+            Poll::Ready(result) => result,
+        }
     }
 
-    pub fn read_buf(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
-        let Some(client) = self.tls.as_mut() else {
-            return self.stream.try_read_buf(buf);
-        };
-        let mut stream = WrappingStream { stream: &self.stream };
-        let mut read_bytes = 0;
-        loop {
-            match client.read_tls(&mut stream) {
-                // We may have plaintext to return though tcp stream has been closed.
-                // If not, read_bytes should be zero.
-                Ok(0) => break,
-                Ok(_) => {},
-                Err(err) => match err.kind() {
-                    // backpressure: tls buffer is full, let's process_new_packets.
-                    io::ErrorKind::Other => {},
-                    io::ErrorKind::WouldBlock if read_bytes == 0 => {
-                        return Err(err);
-                    },
-                    _ => break,
-                },
-            }
-            let state = client.process_new_packets().map_err(io::Error::other)?;
-            let n = state.plaintext_bytes_to_read();
-            buf.reserve(n);
-            let slice = unsafe { &mut std::slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.len() + n)[buf.len()..] };
-            client.reader().read_exact(slice).unwrap();
-            unsafe { buf.set_len(buf.len() + n) };
-            read_bytes += n;
+    pub fn try_read_buf(&mut self, buf: &mut impl BufMut) -> Result<usize> {
+        let waker = unsafe { Waker::from_raw(NOOP_WAKER) };
+        let mut context = Context::from_waker(&waker);
+        let chunk = buf.chunk_mut();
+        let mut read_buf = unsafe { ReadBuf::uninit(chunk.as_uninit_slice_mut()) };
+        match Pin::new(self).poll_read(&mut context, &mut read_buf) {
+            Poll::Pending => Err(ErrorKind::WouldBlock.into()),
+            Poll::Ready(Err(err)) => Err(err),
+            Poll::Ready(Ok(())) => {
+                let n = read_buf.filled().len();
+                unsafe { buf.advance_mut(n) };
+                Ok(n)
+            },
         }
-        Ok(read_bytes)
     }
 
-    pub async fn readable(&self) -> io::Result<()> {
-        let Some(client) = self.tls.as_ref() else {
-            return self.stream.readable().await;
-        };
-        if client.wants_read() {
-            self.stream.readable().await
-        } else {
-            // plaintext data are available for read
-            std::future::ready(Ok(())).await
+    pub async fn readable(&self) -> Result<()> {
+        match self {
+            Self::Raw(stream) => stream.readable().await,
+            Self::Tls(stream) => {
+                let (stream, session) = stream.get_ref();
+                if session.wants_read() {
+                    stream.readable().await
+                } else {
+                    // plaintext data are available for read
+                    std::future::ready(Ok(())).await
+                }
+            },
         }
     }
 
-    pub async fn writable(&self) -> io::Result<()> {
-        self.stream.writable().await
+    pub async fn writable(&self) -> Result<()> {
+        match self {
+            Self::Raw(stream) => stream.writable().await,
+            Self::Tls(stream) => {
+                let (stream, _session) = stream.get_ref();
+                stream.writable().await
+            },
+        }
     }
 
     pub fn wants_write(&self) -> bool {
-        self.tls.as_ref().map(|tls| tls.wants_write()).unwrap_or(false)
+        match self {
+            Self::Raw(_) => false,
+            Self::Tls(stream) => {
+                let (_stream, session) = stream.get_ref();
+                session.wants_write()
+            },
+        }
     }
 
-    pub fn flush(&mut self) -> io::Result<()> {
-        let Some(client) = self.tls.as_mut() else {
-            return Ok(());
-        };
-        let mut stream = WrappingStream { stream: &self.stream };
-        while client.wants_write() {
-            client.write_tls(&mut stream)?;
+    pub fn try_flush(&mut self) -> Result<()> {
+        let waker = unsafe { Waker::from_raw(NOOP_WAKER) };
+        let mut context = Context::from_waker(&waker);
+        match Pin::new(self).poll_flush(&mut context) {
+            Poll::Pending => Err(ErrorKind::WouldBlock.into()),
+            Poll::Ready(result) => result,
         }
-        Ok(())
     }
 }

src/session/depot.rs

@@ -206,15 +206,15 @@ impl Depot {
 
     pub fn write_operations(&mut self, conn: &mut Connection) -> Result<(), Error> {
         if !self.has_pending_writes() {
-            if let Err(err) = conn.flush() {
+            if let Err(err) = conn.try_flush() {
                 if err.kind() == io::ErrorKind::WouldBlock {
                     return Ok(());
                 }
                 return Err(Error::other_from(err));
             }
             return Ok(());
         }
-        let result = conn.write_vectored(self.writing_slices.as_slice());
+        let result = conn.try_write_vectored(self.writing_slices.as_slice());
         let mut written_bytes = match result {
             Err(err) => {
                 if err.kind() == io::ErrorKind::WouldBlock {

src/session/mod.rs

@@ -11,11 +11,13 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use ignore_result::Ignore;
+use rustls::pki_types::ServerName;
 use rustls::ClientConfig;
 use tokio::net::TcpStream;
 use tokio::select;
 use tokio::sync::mpsc;
 use tokio::time::{self, Instant, Sleep};
+use tokio_rustls::TlsConnector;
 
 use self::connection::Connection;
 pub use self::depot::Depot;
@@ -60,7 +62,7 @@ pub struct Session {
     readonly: bool,
     detached: bool,
 
-    tls_config: Arc<ClientConfig>,
+    tls_connector: TlsConnector,
 
     configured_connection_timeout: Duration,
 
@@ -104,7 +106,7 @@ impl Session {
         let mut session = Session {
             readonly,
             detached,
-            tls_config: Arc::new(tls_config),
+            tls_connector: TlsConnector::from(Arc::new(tls_config)),
 
             configured_connection_timeout: connection_timeout,
 
@@ -367,7 +369,7 @@ impl Session {
     }
 
     fn read_connection(&mut self, conn: &mut Connection, buf: &mut Vec<u8>) -> Result<(), Error> {
-        match conn.read_buf(buf) {
+        match conn.try_read_buf(buf) {
             Ok(0) => {
                 return Err(Error::ConnectionLoss);
             },
@@ -505,7 +507,15 @@ impl Session {
                     },
                     Ok(sock) => {
                         let connection = if tls {
-                            Connection::new_tls(host, self.tls_config.clone(), sock)?
+                            let domain = ServerName::try_from(host).unwrap().to_owned();
+                            let stream = match self.tls_connector.connect(domain, sock).await {
+                                Err(err) => {
+                                    log::debug!("ZooKeeper fails to complete tls session to {}:{} due to {}", host, port, err);
+                                    return Err(Error::ConnectionLoss);
+                                },
+                                Ok(stream) => stream,
+                            };
+                            Connection::new_tls(stream)
                         } else {
                             Connection::new_raw(sock)
                         };