diff --git a/src/io/aio.rs b/src/io/aio.rs index 9f5d943..9262f0e 100644 --- a/src/io/aio.rs +++ b/src/io/aio.rs @@ -126,6 +126,26 @@ impl Config { pub struct TcpConnection { con: C, buf: Vec, + /// Set when a query or pipeline future was dropped before it consumed the whole + /// reply. The stream position is then unknown; every later operation on + /// this connection fails instead of reading another request's reply. + poisoned: bool, +} + +/// Sets `poisoned` on drop unless the exchange completed. This covers cancellation +/// at any await point (including a partial write) and I/O/protocol errors, after +/// which the next byte boundary is unknown. +struct PoisonOnDrop<'a> { + poisoned: &'a mut bool, + armed: bool, +} + +impl Drop for PoisonOnDrop<'_> { + fn drop(&mut self) { + if self.armed { + *self.poisoned = true; + } + } } impl TcpConnection { @@ -133,8 +153,21 @@ impl TcpConnection { Self { con, buf: Vec::with_capacity(crate::BUFSIZE), + poisoned: false, } } + /// Whether a previous operation was cancelled or failed before its reply + /// was fully consumed. Such a connection must be discarded: the pool does + /// this in `has_broken`, while a direct user must reconnect. + pub(crate) fn is_poisoned(&self) -> bool { + self.poisoned + } + fn poisoned_error() -> Error { + Error::IoError(std::io::Error::new( + std::io::ErrorKind::ConnectionAborted, + "skytable: connection poisoned by an incomplete operation", + )) + } async fn _handshake(mut self, cfg: &Config) -> ClientResult { let handshake = ClientHandshake::new(cfg); self.con.write_all(handshake.inner()).await?; @@ -146,31 +179,47 @@ impl TcpConnection { } } /// Execute a pipeline. The server returns the queries in the order they were sent (unless otherwise set). + /// + /// If the returned future is dropped after it has been polled, the + /// connection is poisoned: the pipeline may have been written but its + /// reply was not consumed. Later operations return `ConnectionAborted`. + /// Drop the connection and reconnect, or use a pool, which discards + /// poisoned connections. pub async fn execute_pipeline(&mut self, pipeline: &Pipeline) -> ClientResult> { - self.buf.clear(); - self.buf.push(b'P'); + if self.poisoned { + return Err(Self::poisoned_error()); + } + let Self { con, buf, poisoned } = self; + let mut guard = PoisonOnDrop { + poisoned, + armed: true, + }; + buf.clear(); + buf.push(b'P'); // packet size - self.buf - .extend(itoa::Buffer::new().format(pipeline.buf().len()).as_bytes()); - self.buf.push(b'\n'); + buf.extend(itoa::Buffer::new().format(pipeline.buf().len()).as_bytes()); + buf.push(b'\n'); // write - self.con.write_all(&self.buf).await?; - self.con.write_all(pipeline.buf()).await?; - self.buf.clear(); + con.write_all(buf).await?; + con.write_all(pipeline.buf()).await?; + buf.clear(); // read let mut cursor = 0; let mut state = MRespState::default(); loop { - let mut buf = [0u8; crate::BUFSIZE]; - let n = self.con.read(&mut buf).await?; + let mut chunk = [0u8; crate::BUFSIZE]; + let n = con.read(&mut chunk).await?; if n == 0 { return Err(Error::IoError(std::io::ErrorKind::ConnectionReset.into())); } - self.buf.extend_from_slice(&buf[..n]); + buf.extend_from_slice(&chunk[..n]); let (_state, _position) = - Decoder::new(&self.buf, cursor).validate_pipe(pipeline.query_count(), state); + Decoder::new(buf, cursor).validate_pipe(pipeline.query_count(), state); match _state { - PipelineResult::Completed(r) => return Ok(r), + PipelineResult::Completed(r) => { + guard.armed = false; + return Ok(r); + } PipelineResult::Pending(_state) => { cursor = _position; state = _state; @@ -180,27 +229,46 @@ impl TcpConnection { } } /// Run a query and return a raw [`Response`] + /// + /// If the returned future is dropped after it has been polled, the + /// connection is poisoned: the query may have been written but its reply + /// was not consumed. Later operations return `ConnectionAborted`. Drop the + /// connection and reconnect, or use a pool, which discards poisoned + /// connections. pub async fn query(&mut self, q: &Query) -> ClientResult { - self.buf.clear(); - q.write_packet(&mut self.buf).unwrap(); - self.con.write_all(&self.buf).await?; - self.buf.clear(); + if self.poisoned { + return Err(Self::poisoned_error()); + } + // Split the borrow so the guard can watch `poisoned` while the body + // uses the socket and the response buffer. + let Self { con, buf, poisoned } = self; + let mut guard = PoisonOnDrop { + poisoned, + armed: true, + }; + buf.clear(); + q.write_packet(buf).unwrap(); + con.write_all(buf).await?; + buf.clear(); let mut state = RState::default(); let mut cursor = 0; let mut expected = Decoder::MIN_READBACK; loop { - let mut buf = [0u8; crate::BUFSIZE]; - let n = self.con.read(&mut buf).await?; + let mut chunk = [0u8; crate::BUFSIZE]; + let n = con.read(&mut chunk).await?; if n == 0 { return Err(Error::IoError(std::io::ErrorKind::ConnectionReset.into())); } if n < expected { continue; } - self.buf.extend_from_slice(&buf[..n]); - let (_state, _position) = Decoder::new(&self.buf, cursor).validate_response(state); + buf.extend_from_slice(&chunk[..n]); + let (_state, _position) = Decoder::new(buf, cursor).validate_response(state); match _state { - DecodeState::Completed(resp) => return Ok(resp), + DecodeState::Completed(resp) => { + guard.armed = false; + return Ok(resp); + } DecodeState::ChangeState(_state) => { expected = 1; state = _state; @@ -220,3 +288,224 @@ impl TcpConnection { self.buf.shrink_to_fit() } } + +#[cfg(test)] +mod cancellation_tests { + use super::*; + use std::{ + future::{poll_fn, Future}, + pin::Pin, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + task::{Context, Poll}, + }; + use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + + struct PendingReply(Arc); + + /// Accepts at most `limit` bytes and then stays pending, so the query + /// future is cancelled after a partial write. + struct PartialWrite { + writes: Arc, + accepted: usize, + limit: usize, + } + impl AsyncRead for PartialWrite { + fn poll_read( + self: Pin<&mut Self>, + _: &mut Context<'_>, + _: &mut ReadBuf<'_>, + ) -> Poll> { + Poll::Pending + } + } + impl AsyncWrite for PartialWrite { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + if self.accepted >= self.limit { + return Poll::Pending; + } + let n = bytes.len().min(self.limit - self.accepted); + self.accepted += n; + self.writes.fetch_add(n, Ordering::SeqCst); + Poll::Ready(Ok(n)) + } + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + impl AsyncRead for PendingReply { + fn poll_read( + self: Pin<&mut Self>, + _: &mut Context<'_>, + _: &mut ReadBuf<'_>, + ) -> Poll> { + Poll::Pending + } + } + impl AsyncWrite for PendingReply { + fn poll_write( + self: Pin<&mut Self>, + _: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + self.0.fetch_add(bytes.len(), Ordering::SeqCst); + Poll::Ready(Ok(bytes.len())) + } + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + async fn drop_at_pending(future: impl Future) { + let mut future = Box::pin(future); + poll_fn(|cx| { + assert!(future.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + } + #[tokio::test] + async fn dropped_query_blocks_both_entry_points_without_writing() { + let writes = Arc::new(AtomicUsize::new(0)); + let mut conn = TcpConnection::new(PendingReply(writes.clone())); + drop_at_pending(conn.query(&Query::new("sysctl report status"))).await; + assert!(conn.is_poisoned()); + let before = writes.load(Ordering::SeqCst); + assert!(conn + .query(&Query::new("sysctl report status")) + .await + .is_err()); + assert!(conn + .query_parse::<()>(&Query::new("sysctl report status")) + .await + .is_err()); + assert!(conn.execute_pipeline(&Pipeline::new()).await.is_err()); + assert_eq!(before, writes.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn partially_written_query_poisons_the_connection() { + let writes = Arc::new(AtomicUsize::new(0)); + let mut conn = TcpConnection::new(PartialWrite { + writes: writes.clone(), + accepted: 0, + limit: 1, + }); + drop_at_pending(conn.query(&Query::new("sysctl report status"))).await; + assert!(conn.is_poisoned()); + assert_eq!(writes.load(Ordering::SeqCst), 1); + assert!(conn + .query(&Query::new("sysctl report status")) + .await + .is_err()); + } + + #[tokio::test] + async fn partially_read_query_poisons_the_connection() { + let (socket, mut peer) = tokio::io::duplex(256); + // A response type byte without the rest of the reply: the decoder is + // still pending when the future is dropped. + peer.write_all(b"\x10").await.unwrap(); + let mut conn = TcpConnection::new(socket); + drop_at_pending(conn.query(&Query::new("sysctl report status"))).await; + assert!(conn.is_poisoned()); + assert!(conn + .query(&Query::new("sysctl report status")) + .await + .is_err()); + } + + #[tokio::test] + async fn protocol_error_poisons_the_connection() { + let (socket, mut peer) = tokio::io::duplex(256); + peer.write_all(b"\xff").await.unwrap(); + let mut conn = TcpConnection::new(socket); + let err = conn + .query(&Query::new("sysctl report status")) + .await + .unwrap_err(); + assert!(matches!(err, Error::ProtocolError(_)), "{:?}", err); + assert!(conn.is_poisoned()); + assert!(conn + .query(&Query::new("sysctl report status")) + .await + .is_err()); + } + + #[tokio::test] + async fn eof_poisons_the_connection() { + let (socket, peer) = tokio::io::duplex(256); + drop(peer); + let mut conn = TcpConnection::new(socket); + let err = conn + .query(&Query::new("sysctl report status")) + .await + .unwrap_err(); + assert!(matches!(err, Error::IoError(_)), "{:?}", err); + assert!(conn.is_poisoned()); + assert!(conn + .query(&Query::new("sysctl report status")) + .await + .is_err()); + } + #[tokio::test] + async fn dropped_pipeline_also_poisons_the_connection() { + let writes = Arc::new(AtomicUsize::new(0)); + let mut conn = TcpConnection::new(PendingReply(writes.clone())); + let mut pipeline = Pipeline::new(); + pipeline.push(&Query::new("sysctl report status")); + drop_at_pending(conn.execute_pipeline(&pipeline)).await; + assert!(conn.is_poisoned()); + let before = writes.load(Ordering::SeqCst); + assert!(conn.execute_pipeline(&pipeline).await.is_err()); + assert!(conn + .query(&Query::new("sysctl report status")) + .await + .is_err()); + assert_eq!(before, writes.load(Ordering::SeqCst)); + } + #[tokio::test] + async fn completed_replies_do_not_poison_query_or_pipeline() { + for reply in [&b"\x12"[..], &b"\x10\x01\x00"[..]] { + let (socket, mut peer) = tokio::io::duplex(256); + peer.write_all(reply).await.unwrap(); + let mut conn = TcpConnection::new(socket); + conn.query(&Query::new("sysctl report status")) + .await + .unwrap(); + assert!(!conn.is_poisoned()); + peer.write_all(reply).await.unwrap(); + let mut pipeline = Pipeline::new(); + pipeline.push(&Query::new("sysctl report status")); + assert_eq!(conn.execute_pipeline(&pipeline).await.unwrap().len(), 1); + assert!(!conn.is_poisoned()); + } + } + #[test] + fn unpolled_query_does_not_poison() { + let mut conn = TcpConnection::new(PendingReply(Arc::new(AtomicUsize::new(0)))); + let query = Query::new("sysctl report status"); + drop(conn.query(&query)); + assert!(!conn.is_poisoned()); + } + + #[test] + fn unpolled_pipeline_does_not_poison() { + let mut conn = TcpConnection::new(PendingReply(Arc::new(AtomicUsize::new(0)))); + let mut pipeline = Pipeline::new(); + pipeline.push(&Query::new("sysctl report status")); + drop(conn.execute_pipeline(&pipeline)); + assert!(!conn.is_poisoned()); + } +} diff --git a/src/pool.rs b/src/pool.rs index 35292da..8e061e8 100644 --- a/src/pool.rs +++ b/src/pool.rs @@ -44,6 +44,10 @@ //! To create a pool of TLS connections you can use the [`get_tls`] and [`get_tls_async`] methods, passing a PEM certificate //! as a string. //! +//! Cancelling a query or pipeline poisons its connection. The async managers +//! report the connection as broken, so the pool drops it instead of returning +//! it to the pool and the next checkout opens a fresh connection. +//! use crate::{error::Error, Config, Connection, ConnectionAsync, ConnectionTls, ConnectionTlsAsync}; @@ -115,10 +119,16 @@ impl bb8::ManageConnection for ConnectionMgrTcp { self.config.connect_async().await } async fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> { + if conn.is_poisoned() { + return Err(Error::IoError(std::io::Error::new( + std::io::ErrorKind::ConnectionAborted, + "skytable: connection poisoned by an incomplete operation", + ))); + } conn.query_parse::<()>(&query!(QUERY_SYSCTL_STATUS)).await } - fn has_broken(&self, _: &mut Self::Connection) -> bool { - false + fn has_broken(&self, conn: &mut Self::Connection) -> bool { + conn.is_poisoned() } } @@ -160,9 +170,15 @@ impl bb8::ManageConnection for ConnectionMgrTls { self.config.connect_tls_async(&self.pem_cert).await } async fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> { + if conn.is_poisoned() { + return Err(Error::IoError(std::io::Error::new( + std::io::ErrorKind::ConnectionAborted, + "skytable: connection poisoned by an incomplete operation", + ))); + } conn.query_parse::<()>(&query!(QUERY_SYSCTL_STATUS)).await } - fn has_broken(&self, _: &mut Self::Connection) -> bool { - false + fn has_broken(&self, conn: &mut Self::Connection) -> bool { + conn.is_poisoned() } }