diff --git a/tarpc/src/client.rs b/tarpc/src/client.rs index a9ea9549..0c58979a 100644 --- a/tarpc/src/client.rs +++ b/tarpc/src/client.rs @@ -19,7 +19,6 @@ use futures::{prelude::*, ready, stream::Fuse, task::*}; use in_flight_requests::InFlightRequests; use pin_project::pin_project; use std::{ - any::Any, convert::TryFrom, fmt, pin::Pin, @@ -279,7 +278,10 @@ where #[must_use] #[pin_project()] #[derive(Debug)] -pub struct RequestDispatch { +pub struct RequestDispatch +where + C: Sink>, +{ /// Writes requests to the wire and reads responses off the wire. #[pin] transport: Fuse, @@ -293,9 +295,8 @@ pub struct RequestDispatch { config: Config, /// Produces errors that can be sent in response to any unprocessed requests at the time /// RequestDispatch is dropped. Correctness note: this field should only be populated by - /// RequestDispatch::poll, which relies on downcasting the Any to a concrete error type - /// determined within the poll function. - terminal_error: Option>, + /// RequestDispatch::poll. + terminal_error: Option>, } impl RequestDispatch @@ -355,7 +356,7 @@ where fn terminal_error_mut<'a>( self: &'a mut Pin<&mut Self>, - ) -> &'a mut Option> { + ) -> &'a mut Option> { self.as_mut().project().terminal_error } @@ -662,17 +663,13 @@ where loop { if let Some(e) = self.terminal_error_mut() { tracing::debug!("RpcError::Channel"); - let e: ChannelError = e - .clone() - .downcast() - .expect("Invariant: ChannelError must store a C::Error"); ready!(self.shut_down_with_terminal_error(cx, e.clone().upcast_error())); - return Poll::Ready(Err(e)); + return Poll::Ready(Err(e.clone())); } let result = ready!(self.run(cx)); match result { Ok(()) => return Poll::Ready(Ok(())), - Err(e) => *self.terminal_error_mut() = Some(e.upcast_any()), + Err(e) => *self.terminal_error_mut() = Some(e), } } } @@ -692,7 +689,8 @@ struct DispatchRequest { #[cfg(test)] mod tests { use super::{ - Channel, DispatchRequest, RequestDispatch, ResponseGuard, RpcError, cancellations, + Channel, DispatchRequest, NewClient, RequestDispatch, ResponseGuard, RpcError, + cancellations, }; use crate::{ ChannelError, ClientMessage, Response, @@ -888,6 +886,51 @@ mod tests { assert_matches!(resp.response().await, Err(RpcError::Channel(_))); } + /// `shut_down_with_terminal_error` can return `Pending` while an mpsc permit is held. + /// The next poll must resume shutdown, not call `run()` again. + #[tokio::test] + async fn test_terminal_error_shutdown_resumes_after_pending() { + let _ = tracing_subscriber::fmt().with_test_writer().try_init(); + let writes = Arc::new(AtomicUsize::new(0)); + let (mut dispatch, channel, mut cx) = + set_up_one_shot_err(TransportError::Read, writes.clone()); + + // Hold a permit so draining pending_requests after close() stays Pending. + let permit = channel.to_dispatch.reserve().await.unwrap(); + assert_eq!(dispatch.as_mut().poll(&mut cx), Poll::Pending); + + drop(permit); + drop(channel); + + assert_eq!( + dispatch.as_mut().poll(&mut cx), + Poll::Ready(Err(ChannelError::Read(Arc::new(TransportError::Read)))) + ); + assert_eq!(writes.load(Ordering::SeqCst), 0); + } + + /// A request sent on a permit that outlived the terminal error must be failed + /// with `RpcError::Channel`, not written to the transport by a second `run()`. + #[tokio::test] + async fn test_terminal_error_shutdown_drains_late_permit_send() { + let _ = tracing_subscriber::fmt().with_test_writer().try_init(); + let writes = Arc::new(AtomicUsize::new(0)); + let (mut dispatch, mut channel, mut cx) = + set_up_one_shot_err(TransportError::Read, writes.clone()); + let (tx, mut rx) = oneshot::channel(); + let permit = reserve_for_send(&mut channel, tx, &mut rx).await; + + assert_eq!(dispatch.as_mut().poll(&mut cx), Poll::Pending); + + let resp = permit("late"); + assert_eq!( + dispatch.as_mut().poll(&mut cx), + Poll::Ready(Err(ChannelError::Read(Arc::new(TransportError::Read)))) + ); + assert_matches!(resp.response().await, Err(RpcError::Channel(_))); + assert_eq!(writes.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn test_shutdown() { let _ = tracing_subscriber::fmt().with_test_writer().try_init(); @@ -997,6 +1040,36 @@ mod tests { (dispatch, channel, cx) } + /// Transport that yields `cause` once, then stays healthy/pending. + /// + /// A transport that *keeps* failing masks the shutdown-resume bug: a second + /// `run()` would just error again and accidentally finish shutdown. + fn set_up_one_shot_err( + cause: TransportError, + writes: Arc, + ) -> ( + Pin>>, + Channel, + Context<'static>, + ) { + let NewClient { + client: channel, + dispatch, + } = super::new( + Config { + max_in_flight_requests: 1_000, + pending_request_buffer: 1, + }, + OneShotErrorTransport { + cause, + errored: false, + writes, + }, + ); + let cx = Context::from_waker(noop_waker_ref()); + (Box::pin(dispatch), channel, cx) + } + struct AlwaysErrorTransport(TransportError, PhantomData); #[derive(Debug, Error, PartialEq, Eq, Clone, Copy)] @@ -1057,6 +1130,82 @@ mod tests { } } + struct OneShotErrorTransport { + cause: TransportError, + errored: bool, + writes: Arc, + } + + impl OneShotErrorTransport { + fn take_error(&mut self) -> Option { + if !self.errored { + self.errored = true; + Some(self.cause) + } else { + None + } + } + } + + impl Sink for OneShotErrorTransport { + type Error = TransportError; + fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + match this.cause { + TransportError::Ready => { + if let Some(e) = this.take_error() { + Poll::Ready(Err(e)) + } else { + Poll::Ready(Ok(())) + } + } + TransportError::Flush if !this.errored => Poll::Pending, + _ => Poll::Ready(Ok(())), + } + } + fn start_send(self: Pin<&mut Self>, _: S) -> Result<(), Self::Error> { + let this = self.get_mut(); + this.writes.fetch_add(1, Ordering::SeqCst); + if matches!(this.cause, TransportError::Write) { + if let Some(e) = this.take_error() { + return Err(e); + } + } + Ok(()) + } + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if matches!(this.cause, TransportError::Flush) { + if let Some(e) = this.take_error() { + return Poll::Ready(Err(e)); + } + } + Poll::Ready(Ok(())) + } + fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if matches!(this.cause, TransportError::Close) { + if let Some(e) = this.take_error() { + return Poll::Ready(Err(e)); + } + } + Poll::Ready(Ok(())) + } + } + + impl Stream for OneShotErrorTransport { + type Item = Result, TransportError>; + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if matches!(this.cause, TransportError::Read) { + if let Some(e) = this.take_error() { + return Poll::Ready(Some(Err(e))); + } + } + Poll::Pending + } + } + fn set_up() -> ( Pin< Box< diff --git a/tarpc/src/lib.rs b/tarpc/src/lib.rs index 7e194430..d393807a 100644 --- a/tarpc/src/lib.rs +++ b/tarpc/src/lib.rs @@ -250,7 +250,7 @@ pub(crate) mod util; pub use crate::transport::sealed::Transport; -use std::{any::Any, error::Error, io, sync::Arc, time::Instant}; +use std::{error::Error, io, sync::Arc, time::Instant}; /// A message from a client to a server. #[derive(Debug)] @@ -445,44 +445,6 @@ where } } -impl ChannelError -where - E: Send + Sync + 'static, -{ - /// Converts the ChannelError's source error type to a dyn Any. This is useful in type-erased - /// contexts, for example, storing a ChannelError in a non-generic type like - /// [`client::RpcError`]. - fn upcast_any(self) -> ChannelError { - use ChannelError::*; - match self { - Read(e) => Read(e), - Ready(e) => Ready(e), - Write(e) => Write(e), - Flush(e) => Flush(e), - Close(e) => Close(e), - } - } -} - -impl ChannelError { - /// Converts the ChannelError's source error type to a concrete type. This is useful in - /// type-erased contexts, for example, storing a ChannelError in a non-generic type like - /// [`Client::RpcError`]. - fn downcast(self) -> Result, Self> - where - E: Any + Send + Sync, - { - use ChannelError::*; - match self { - Read(e) => e.downcast::().map(Read).map_err(Read), - Ready(e) => e.downcast::().map(Ready).map_err(Ready), - Write(e) => e.downcast::().map(Write).map_err(Write), - Flush(e) => e.downcast::().map(Flush).map_err(Flush), - Close(e) => e.downcast::().map(Close).map_err(Close), - } - } -} - impl ServerError { /// Returns a new server error with `kind` and `detail`. pub fn new(kind: io::ErrorKind, detail: String) -> ServerError { @@ -497,30 +459,6 @@ impl Request { } } -#[test] -fn test_channel_any_casts() { - use assert_matches::assert_matches; - let any = ChannelError::Read(Arc::new("")).upcast_any(); - assert_matches!(any, ChannelError::Read(_)); - assert_matches!(any.downcast::<&'static str>(), Ok(ChannelError::Read(_))); - - let any = ChannelError::Ready(Arc::new("")).upcast_any(); - assert_matches!(any, ChannelError::Ready(_)); - assert_matches!(any.downcast::<&'static str>(), Ok(ChannelError::Ready(_))); - - let any = ChannelError::Write(Arc::new("")).upcast_any(); - assert_matches!(any, ChannelError::Write(_)); - assert_matches!(any.downcast::<&'static str>(), Ok(ChannelError::Write(_))); - - let any = ChannelError::Flush(Arc::new("")).upcast_any(); - assert_matches!(any, ChannelError::Flush(_)); - assert_matches!(any.downcast::<&'static str>(), Ok(ChannelError::Flush(_))); - - let any = ChannelError::Close(Arc::new("")).upcast_any(); - assert_matches!(any, ChannelError::Close(_)); - assert_matches!(any.downcast::<&'static str>(), Ok(ChannelError::Close(_))); -} - #[test] fn test_channel_error_upcast() { use assert_matches::assert_matches;