From aa1348c11ac843436236735110e2f74f6773fb13 Mon Sep 17 00:00:00 2001 From: Tianyi Chen Date: Thu, 22 Oct 2026 02:24:19 +0900 Subject: [PATCH] Fail closed after an interrupted async exchange A cancelled query or pipeline can leave an unread reply on the connection. Reusing it associates that reply with the next request, including when checkout validation consumes the stale reply first. Poison incomplete exchanges and reject later operations without writing. Have both async pool managers discard poisoned connections, and document that direct users must reconnect. Complete responses, including server errors, remain reusable. Add regression coverage for cancellation, partial reads/writes, EOF, protocol errors, unpolled futures and completed replies. Assisted-by: LLM Signed-off-by: Tianyi Chen --- src/io/aio.rs | 333 ++++++++++++++++++++++++++++++++++++++++++++++---- src/pool.rs | 24 +++- 2 files changed, 331 insertions(+), 26 deletions(-) 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() } }