diff --git a/zebra-network/src/peer/client.rs b/zebra-network/src/peer/client.rs index 16101bd4..faa6fcab 100644 --- a/zebra-network/src/peer/client.rs +++ b/zebra-network/src/peer/client.rs @@ -19,7 +19,7 @@ use tokio::{sync::broadcast, task::JoinHandle}; use tower::Service; use crate::{ - peer::error::AlreadyErrored, + peer::error::{AlreadyErrored, ErrorSlot, PeerError, SharedPeerError}, peer_set::InventoryChange, protocol::{ external::{types::Version, InventoryHash}, @@ -27,8 +27,6 @@ use crate::{ }, }; -use super::{ErrorSlot, PeerError, SharedPeerError}; - #[cfg(any(test, feature = "proptest-impl"))] pub mod tests; @@ -491,11 +489,12 @@ impl Client { /// Poll for space in the shared request sender channel. fn poll_request(&mut self, cx: &mut Context<'_>) -> Poll> { - if ready!(self.server_tx.poll_ready(cx)).is_err() { + let server_result = ready!(self.server_tx.poll_ready(cx)); + if server_result.is_err() { Poll::Ready(Err(self .error_slot .try_get_error() - .expect("failed servers must set their error slot"))) + .unwrap_or_else(|| PeerError::ConnectionTaskExited.into()))) } else if let Some(error) = self.error_slot.try_get_error() { Poll::Ready(Err(error)) } else { @@ -569,13 +568,15 @@ impl Service for Client { }) { Err(e) => { if e.is_disconnected() { - let ClientRequest { tx, .. } = e.into_inner(); - let _ = tx.send(Err(PeerError::ConnectionClosed.into())); - future::ready(Err(self + let peer_error = self .error_slot .try_get_error() - .expect("failed servers must set their error slot"))) - .boxed() + .unwrap_or_else(|| PeerError::ConnectionTaskExited.into()); + + let ClientRequest { tx, .. } = e.into_inner(); + let _ = tx.send(Err(peer_error.clone())); + + future::ready(Err(peer_error)).boxed() } else { // sending fails when there's not enough // channel space, but we called poll_ready