use super::{ proto::{ self, AnyTypedEnvelope, EnvelopedMessage, MessageStream, PeerId, Receipt, RequestMessage, TypedEnvelope, }, Connection,};use anyhow::{anyhow, Context as _, Result};use collections::HashMap;use futures::{ channel::{mpsc, oneshot}, stream::BoxStream, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt,};use parking_lot::{Mutex, RwLock};use proto::{ErrorCode, ErrorCodeExt, ErrorExt, RpcError};use serde::{ser::SerializeStruct, Serialize};use std::{ fmt, future, future::Future, sync::atomic::Ordering::SeqCst, sync::{ atomic::{self, AtomicU32}, Arc, }, time::Duration, time::Instant,};use tracing::instrument; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Serialize)]pub struct ConnectionId { pub owner_id: u32, pub id: u32,} impl From<ConnectionId> for PeerId { fn from(id: ConnectionId) -> Self { PeerId { owner_id: id.owner_id, id: id.id, } }} impl From<PeerId> for ConnectionId { fn from(peer_id: PeerId) -> Self { Self { owner_id: peer_id.owner_id, id: peer_id.id, } }} impl fmt::Display for ConnectionId { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{}/{}", self.owner_id, self.id) }} pub struct Peer { epoch: AtomicU32, pub connections: RwLock<HashMap<ConnectionId, ConnectionState>>, next_connection_id: AtomicU32,} #[derive(Clone, Serialize)]pub struct ConnectionState { #[serde(skip)] outgoing_tx: mpsc::UnboundedSender<proto::Message>, next_message_id: Arc<AtomicU32>, #[allow(clippy::type_complexity)] #[serde(skip)] response_channels: Arc< Mutex< Option< HashMap< u32, oneshot::Sender<(proto::Envelope, std::time::Instant, oneshot::Sender<()>)>, >, >, >, >, #[allow(clippy::type_complexity)] #[serde(skip)] stream_response_channels: Arc< Mutex< Option< HashMap<u32, mpsc::UnboundedSender<(Result<proto::Envelope>, oneshot::Sender<()>)>>, >, >, >,} const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(1);const WRITE_TIMEOUT: Duration = Duration::from_secs(2);pub const RECEIVE_TIMEOUT: Duration = Duration::from_secs(10); impl Peer { pub fn new(epoch: u32) -> Arc<Self> { Arc::new(Self { epoch: AtomicU32::new(epoch), connections: Default::default(), next_connection_id: Default::default(), }) } pub fn epoch(&self) -> u32 { self.epoch.load(SeqCst) } #[instrument(skip_all)] pub fn add_connection<F, Fut, Out>( self: &Arc<Self>, connection: Connection, create_timer: F, ) -> ( ConnectionId, impl Future<Output = anyhow::Result<()>> + Send, BoxStream<'static, Box<dyn AnyTypedEnvelope>>, ) where F: Send + Fn(Duration) -> Fut, Fut: Send + Future<Output = Out>, Out: Send, { // For outgoing messages, use an unbounded channel so that application code // can always send messages without yielding. For incoming messages, use a // bounded channel so that other peers will receive backpressure if they send // messages faster than this peer can process them. #[cfg(any(test, feature = "test-support"))] const INCOMING_BUFFER_SIZE: usize = 1; #[cfg(not(any(test, feature = "test-support")))] const INCOMING_BUFFER_SIZE: usize = 256; let (mut incoming_tx, incoming_rx) = mpsc::channel(INCOMING_BUFFER_SIZE); let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded(); let connection_id = ConnectionId { owner_id: self.epoch.load(SeqCst), id: self.next_connection_id.fetch_add(1, SeqCst), }; let connection_state = ConnectionState { outgoing_tx, next_message_id: Default::default(), response_channels: Arc::new(Mutex::new(Some(Default::default()))), stream_response_channels: Arc::new(Mutex::new(Some(Default::default()))), }; let mut writer = MessageStream::new(connection.tx); let mut reader = MessageStream::new(connection.rx); let this = self.clone(); let response_channels = connection_state.response_channels.clone(); let stream_response_channels = connection_state.stream_response_channels.clone(); let handle_io = async move { tracing::trace!(%connection_id, "handle io future: start"); let _end_connection = util::defer(|| { response_channels.lock().take(); if let Some(channels) = stream_response_channels.lock().take() { for channel in channels.values() { let _ = channel.unbounded_send(( Err(anyhow!("connection closed")), oneshot::channel().0, )); } } this.connections.write().remove(&connection_id); tracing::trace!(%connection_id, "handle io future: end"); }); // Send messages on this frequency so the connection isn't closed. let keepalive_timer = create_timer(KEEPALIVE_INTERVAL).fuse(); futures::pin_mut!(keepalive_timer); // Disconnect if we don't receive messages at least this frequently. let receive_timeout = create_timer(RECEIVE_TIMEOUT).fuse(); futures::pin_mut!(receive_timeout); loop { tracing::trace!(%connection_id, "outer loop iteration start"); let read_message = reader.read().fuse(); futures::pin_mut!(read_message); loop { tracing::trace!(%connection_id, "inner loop iteration start"); futures::select_biased! { outgoing = outgoing_rx.next().fuse() => match outgoing { Some(outgoing) => { tracing::trace!(%connection_id, "outgoing rpc message: writing"); futures::select_biased! { result = writer.write(outgoing).fuse() => { tracing::trace!(%connection_id, "outgoing rpc message: done writing"); result.context("failed to write RPC message")?; tracing::trace!(%connection_id, "keepalive interval: resetting after sending message"); keepalive_timer.set(create_timer(KEEPALIVE_INTERVAL).fuse()); } _ = create_timer(WRITE_TIMEOUT).fuse() => { tracing::trace!(%connection_id, "outgoing rpc message: writing timed out"); Err(anyhow!("timed out writing message"))?; } } } None => { tracing::trace!(%connection_id, "outgoing rpc message: channel closed"); return Ok(()) }, }, _ = keepalive_timer => { tracing::trace!(%connection_id, "keepalive interval: pinging"); futures::select_biased! { result = writer.write(proto::Message::Ping).fuse() => { tracing::trace!(%connection_id, "keepalive interval: done pinging"); result.context("failed to send keepalive")?; tracing::trace!(%connection_id, "keepalive interval: resetting after pinging"); keepalive_timer.set(create_timer(KEEPALIVE_INTERVAL).fuse()); } _ = create_timer(WRITE_TIMEOUT).fuse() => { tracing::trace!(%connection_id, "keepalive interval: pinging timed out"); Err(anyhow!("timed out sending keepalive"))?; } } } incoming = read_message => { let incoming = incoming.context("error reading rpc message from socket")?; tracing::trace!(%connection_id, "incoming rpc message: received"); tracing::trace!(%connection_id, "receive timeout: resetting"); receive_timeout.set(create_timer(RECEIVE_TIMEOUT).fuse()); if let (proto::Message::Envelope(incoming), received_at) = incoming { tracing::trace!(%connection_id, "incoming rpc message: processing"); futures::select_biased! { result = incoming_tx.send((incoming, received_at)).fuse() => match result { Ok(_) => { tracing::trace!(%connection_id, "incoming rpc message: processed"); } Err(_) => { tracing::trace!(%connection_id, "incoming rpc message: channel closed"); return Ok(()) } }, _ = create_timer(WRITE_TIMEOUT).fuse() => { tracing::trace!(%connection_id, "incoming rpc message: processing timed out"); Err(anyhow!("timed out processing incoming message"))? } } } break; }, _ = receive_timeout => { tracing::trace!(%connection_id, "receive timeout: delay between messages too long"); Err(anyhow!("delay between messages too long"))? } } } } }; let response_channels = connection_state.response_channels.clone(); let stream_response_channels = connection_state.stream_response_channels.clone(); self.connections .write() .insert(connection_id, connection_state); let incoming_rx = incoming_rx.filter_map(move |(incoming, received_at)| { let response_channels = response_channels.clone(); let stream_response_channels = stream_response_channels.clone(); async move { let message_id = incoming.id; tracing::trace!(?incoming, "incoming message future: start"); let _end = util::defer(move || { tracing::trace!(%connection_id, message_id, "incoming message future: end"); }); if let Some(responding_to) = incoming.responding_to { tracing::trace!( %connection_id, message_id, responding_to, "incoming response: received" ); let response_channel = response_channels.lock().as_mut()?.remove(&responding_to); let stream_response_channel = stream_response_channels .lock() .as_ref()? .get(&responding_to) .cloned(); if let Some(tx) = response_channel { let requester_resumed = oneshot::channel(); if let Err(error) = tx.send((incoming, received_at, requester_resumed.0)) { tracing::trace!( %connection_id, message_id, responding_to = responding_to, ?error, "incoming response: request future dropped", ); } tracing::trace!( %connection_id, message_id, responding_to, "incoming response: waiting to resume requester" ); let _ = requester_resumed.1.await; tracing::trace!( %connection_id, message_id, responding_to, "incoming response: requester resumed" ); } else if let Some(tx) = stream_response_channel { let requester_resumed = oneshot::channel(); if let Err(error) = tx.unbounded_send((Ok(incoming), requester_resumed.0)) { tracing::debug!( %connection_id, message_id, responding_to = responding_to, ?error, "incoming stream response: request future dropped", ); } tracing::debug!( %connection_id, message_id, responding_to, "incoming stream response: waiting to resume requester" ); let _ = requester_resumed.1.await; tracing::debug!( %connection_id, message_id, responding_to, "incoming stream response: requester resumed" ); } else { let message_type = proto::build_typed_envelope( connection_id.into(), received_at, incoming, ) .map(|p| p.payload_type_name()); tracing::warn!( %connection_id, message_id, responding_to, message_type, "incoming response: unknown request" ); } None } else { tracing::trace!(%connection_id, message_id, "incoming message: received"); proto::build_typed_envelope(connection_id.into(), received_at, incoming) .or_else(|| { tracing::error!( %connection_id, message_id, "unable to construct a typed envelope" ); None }) } } }); (connection_id, handle_io, incoming_rx.boxed()) } #[cfg(any(test, feature = "test-support"))] pub fn add_test_connection( self: &Arc<Self>, connection: Connection, executor: gpui::BackgroundExecutor, ) -> ( ConnectionId, impl Future<Output = anyhow::Result<()>> + Send, BoxStream<'static, Box<dyn AnyTypedEnvelope>>, ) { let executor = executor.clone(); self.add_connection(connection, move |duration| executor.timer(duration)) } pub fn disconnect(&self, connection_id: ConnectionId) { self.connections.write().remove(&connection_id); } #[cfg(any(test, feature = "test-support"))] pub fn reset(&self, epoch: u32) { self.next_connection_id.store(0, SeqCst); self.epoch.store(epoch, SeqCst); } pub fn teardown(&self) { self.connections.write().clear(); } /// Make a request and wait for a response. pub fn request<T: RequestMessage>( &self, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<T::Response>> { self.request_internal(None, receiver_id, request) .map_ok(|envelope| envelope.payload) } pub fn request_envelope<T: RequestMessage>( &self, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> { self.request_internal(None, receiver_id, request) } pub fn forward_request<T: RequestMessage>( &self, sender_id: ConnectionId, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<T::Response>> { self.request_internal(Some(sender_id), receiver_id, request) .map_ok(|envelope| envelope.payload) } fn request_internal<T: RequestMessage>( &self, original_sender_id: Option<ConnectionId>, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> { let envelope = request.into_envelope(0, None, original_sender_id.map(Into::into)); let response = self.request_dynamic(receiver_id, envelope, T::NAME); async move { let (response, received_at) = response.await?; Ok(TypedEnvelope { message_id: response.id, sender_id: receiver_id.into(), original_sender_id: response.original_sender_id, payload: T::Response::from_envelope(response) .ok_or_else(|| anyhow!("received response of the wrong type"))?, received_at, }) } } /// Make a request and wait for a response. /// /// The caller must make sure to deserialize the response into the request's /// response type. This interface is only useful in trait objects, where /// generics can't be used. If you have a concrete type, use `request`. pub fn request_dynamic( &self, receiver_id: ConnectionId, mut envelope: proto::Envelope, type_name: &'static str, ) -> impl Future<Output = Result<(proto::Envelope, Instant)>> { let (tx, rx) = oneshot::channel(); let send = self.connection_state(receiver_id).and_then(|connection| { envelope.id = connection.next_message_id.fetch_add(1, SeqCst); connection .response_channels .lock() .as_mut() .ok_or_else(|| anyhow!("connection was closed"))? .insert(envelope.id, tx); connection .outgoing_tx .unbounded_send(proto::Message::Envelope(envelope)) .map_err(|_| anyhow!("connection was closed"))?; Ok(()) }); async move { send?; let (response, received_at, _barrier) = rx.await.map_err(|_| anyhow!("connection was closed"))?; if let Some(proto::envelope::Payload::Error(error)) = &response.payload { return Err(RpcError::from_proto(error, type_name)); } Ok((response, received_at)) } } pub fn request_stream<T: RequestMessage>( &self, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<impl Unpin + Stream<Item = Result<T::Response>>>> { let (tx, rx) = mpsc::unbounded(); let send = self.connection_state(receiver_id).and_then(|connection| { let message_id = connection.next_message_id.fetch_add(1, SeqCst); let stream_response_channels = connection.stream_response_channels.clone(); stream_response_channels .lock() .as_mut() .ok_or_else(|| anyhow!("connection was closed"))? .insert(message_id, tx);Lines 1–500 of 1061. Load more
use super::{ proto::{ self, AnyTypedEnvelope, EnvelopedMessage, MessageStream, PeerId, Receipt, RequestMessage, TypedEnvelope, }, Connection, }; use anyhow::{anyhow, Context as _, Result}; use collections::HashMap; use futures::{ channel::{mpsc, oneshot}, stream::BoxStream, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt, }; use parking_lot::{Mutex, RwLock}; use proto::{ErrorCode, ErrorCodeExt, ErrorExt, RpcError}; use serde::{ser::SerializeStruct, Serialize}; use std::{ fmt, future, future::Future, sync::atomic::Ordering::SeqCst, sync::{ atomic::{self, AtomicU32}, Arc, }, time::Duration, time::Instant, }; use tracing::instrument; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Serialize)] pub struct ConnectionId { pub owner_id: u32, pub id: u32, } impl From<ConnectionId> for PeerId { fn from(id: ConnectionId) -> Self { PeerId { owner_id: id.owner_id, id: id.id, } } } impl From<PeerId> for ConnectionId { fn from(peer_id: PeerId) -> Self { Self { owner_id: peer_id.owner_id, id: peer_id.id, } } } impl fmt::Display for ConnectionId { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{}/{}", self.owner_id, self.id) } } pub struct Peer { epoch: AtomicU32, pub connections: RwLock<HashMap<ConnectionId, ConnectionState>>, next_connection_id: AtomicU32, } #[derive(Clone, Serialize)] pub struct ConnectionState { #[serde(skip)] outgoing_tx: mpsc::UnboundedSender<proto::Message>, next_message_id: Arc<AtomicU32>, #[allow(clippy::type_complexity)] #[serde(skip)] response_channels: Arc< Mutex< Option< HashMap< u32, oneshot::Sender<(proto::Envelope, std::time::Instant, oneshot::Sender<()>)>, >, >, >, >, #[allow(clippy::type_complexity)] #[serde(skip)] stream_response_channels: Arc< Mutex< Option< HashMap<u32, mpsc::UnboundedSender<(Result<proto::Envelope>, oneshot::Sender<()>)>>, >, >, >, } const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(1); const WRITE_TIMEOUT: Duration = Duration::from_secs(2); pub const RECEIVE_TIMEOUT: Duration = Duration::from_secs(10); impl Peer { pub fn new(epoch: u32) -> Arc<Self> { Arc::new(Self { epoch: AtomicU32::new(epoch), connections: Default::default(), next_connection_id: Default::default(), }) } pub fn epoch(&self) -> u32 { self.epoch.load(SeqCst) } #[instrument(skip_all)] pub fn add_connection<F, Fut, Out>( self: &Arc<Self>, connection: Connection, create_timer: F, ) -> ( ConnectionId, impl Future<Output = anyhow::Result<()>> + Send, BoxStream<'static, Box<dyn AnyTypedEnvelope>>, ) where F: Send + Fn(Duration) -> Fut, Fut: Send + Future<Output = Out>, Out: Send, { // For outgoing messages, use an unbounded channel so that application code // can always send messages without yielding. For incoming messages, use a // bounded channel so that other peers will receive backpressure if they send // messages faster than this peer can process them. #[cfg(any(test, feature = "test-support"))] const INCOMING_BUFFER_SIZE: usize = 1; #[cfg(not(any(test, feature = "test-support")))] const INCOMING_BUFFER_SIZE: usize = 256; let (mut incoming_tx, incoming_rx) = mpsc::channel(INCOMING_BUFFER_SIZE); let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded(); let connection_id = ConnectionId { owner_id: self.epoch.load(SeqCst), id: self.next_connection_id.fetch_add(1, SeqCst), }; let connection_state = ConnectionState { outgoing_tx, next_message_id: Default::default(), response_channels: Arc::new(Mutex::new(Some(Default::default()))), stream_response_channels: Arc::new(Mutex::new(Some(Default::default()))), }; let mut writer = MessageStream::new(connection.tx); let mut reader = MessageStream::new(connection.rx); let this = self.clone(); let response_channels = connection_state.response_channels.clone(); let stream_response_channels = connection_state.stream_response_channels.clone(); let handle_io = async move { tracing::trace!(%connection_id, "handle io future: start"); let _end_connection = util::defer(|| { response_channels.lock().take(); if let Some(channels) = stream_response_channels.lock().take() { for channel in channels.values() { let _ = channel.unbounded_send(( Err(anyhow!("connection closed")), oneshot::channel().0, )); } } this.connections.write().remove(&connection_id); tracing::trace!(%connection_id, "handle io future: end"); }); // Send messages on this frequency so the connection isn't closed. let keepalive_timer = create_timer(KEEPALIVE_INTERVAL).fuse(); futures::pin_mut!(keepalive_timer); // Disconnect if we don't receive messages at least this frequently. let receive_timeout = create_timer(RECEIVE_TIMEOUT).fuse(); futures::pin_mut!(receive_timeout); loop { tracing::trace!(%connection_id, "outer loop iteration start"); let read_message = reader.read().fuse(); futures::pin_mut!(read_message); loop { tracing::trace!(%connection_id, "inner loop iteration start"); futures::select_biased! { outgoing = outgoing_rx.next().fuse() => match outgoing { Some(outgoing) => { tracing::trace!(%connection_id, "outgoing rpc message: writing"); futures::select_biased! { result = writer.write(outgoing).fuse() => { tracing::trace!(%connection_id, "outgoing rpc message: done writing"); result.context("failed to write RPC message")?; tracing::trace!(%connection_id, "keepalive interval: resetting after sending message"); keepalive_timer.set(create_timer(KEEPALIVE_INTERVAL).fuse()); } _ = create_timer(WRITE_TIMEOUT).fuse() => { tracing::trace!(%connection_id, "outgoing rpc message: writing timed out"); Err(anyhow!("timed out writing message"))?; } } } None => { tracing::trace!(%connection_id, "outgoing rpc message: channel closed"); return Ok(()) }, }, _ = keepalive_timer => { tracing::trace!(%connection_id, "keepalive interval: pinging"); futures::select_biased! { result = writer.write(proto::Message::Ping).fuse() => { tracing::trace!(%connection_id, "keepalive interval: done pinging"); result.context("failed to send keepalive")?; tracing::trace!(%connection_id, "keepalive interval: resetting after pinging"); keepalive_timer.set(create_timer(KEEPALIVE_INTERVAL).fuse()); } _ = create_timer(WRITE_TIMEOUT).fuse() => { tracing::trace!(%connection_id, "keepalive interval: pinging timed out"); Err(anyhow!("timed out sending keepalive"))?; } } } incoming = read_message => { let incoming = incoming.context("error reading rpc message from socket")?; tracing::trace!(%connection_id, "incoming rpc message: received"); tracing::trace!(%connection_id, "receive timeout: resetting"); receive_timeout.set(create_timer(RECEIVE_TIMEOUT).fuse()); if let (proto::Message::Envelope(incoming), received_at) = incoming { tracing::trace!(%connection_id, "incoming rpc message: processing"); futures::select_biased! { result = incoming_tx.send((incoming, received_at)).fuse() => match result { Ok(_) => { tracing::trace!(%connection_id, "incoming rpc message: processed"); } Err(_) => { tracing::trace!(%connection_id, "incoming rpc message: channel closed"); return Ok(()) } }, _ = create_timer(WRITE_TIMEOUT).fuse() => { tracing::trace!(%connection_id, "incoming rpc message: processing timed out"); Err(anyhow!("timed out processing incoming message"))? } } } break; }, _ = receive_timeout => { tracing::trace!(%connection_id, "receive timeout: delay between messages too long"); Err(anyhow!("delay between messages too long"))? } } } } }; let response_channels = connection_state.response_channels.clone(); let stream_response_channels = connection_state.stream_response_channels.clone(); self.connections .write() .insert(connection_id, connection_state); let incoming_rx = incoming_rx.filter_map(move |(incoming, received_at)| { let response_channels = response_channels.clone(); let stream_response_channels = stream_response_channels.clone(); async move { let message_id = incoming.id; tracing::trace!(?incoming, "incoming message future: start"); let _end = util::defer(move || { tracing::trace!(%connection_id, message_id, "incoming message future: end"); }); if let Some(responding_to) = incoming.responding_to { tracing::trace!( %connection_id, message_id, responding_to, "incoming response: received" ); let response_channel = response_channels.lock().as_mut()?.remove(&responding_to); let stream_response_channel = stream_response_channels .lock() .as_ref()? .get(&responding_to) .cloned(); if let Some(tx) = response_channel { let requester_resumed = oneshot::channel(); if let Err(error) = tx.send((incoming, received_at, requester_resumed.0)) { tracing::trace!( %connection_id, message_id, responding_to = responding_to, ?error, "incoming response: request future dropped", ); } tracing::trace!( %connection_id, message_id, responding_to, "incoming response: waiting to resume requester" ); let _ = requester_resumed.1.await; tracing::trace!( %connection_id, message_id, responding_to, "incoming response: requester resumed" ); } else if let Some(tx) = stream_response_channel { let requester_resumed = oneshot::channel(); if let Err(error) = tx.unbounded_send((Ok(incoming), requester_resumed.0)) { tracing::debug!( %connection_id, message_id, responding_to = responding_to, ?error, "incoming stream response: request future dropped", ); } tracing::debug!( %connection_id, message_id, responding_to, "incoming stream response: waiting to resume requester" ); let _ = requester_resumed.1.await; tracing::debug!( %connection_id, message_id, responding_to, "incoming stream response: requester resumed" ); } else { let message_type = proto::build_typed_envelope( connection_id.into(), received_at, incoming, ) .map(|p| p.payload_type_name()); tracing::warn!( %connection_id, message_id, responding_to, message_type, "incoming response: unknown request" ); } None } else { tracing::trace!(%connection_id, message_id, "incoming message: received"); proto::build_typed_envelope(connection_id.into(), received_at, incoming) .or_else(|| { tracing::error!( %connection_id, message_id, "unable to construct a typed envelope" ); None }) } } }); (connection_id, handle_io, incoming_rx.boxed()) } #[cfg(any(test, feature = "test-support"))] pub fn add_test_connection( self: &Arc<Self>, connection: Connection, executor: gpui::BackgroundExecutor, ) -> ( ConnectionId, impl Future<Output = anyhow::Result<()>> + Send, BoxStream<'static, Box<dyn AnyTypedEnvelope>>, ) { let executor = executor.clone(); self.add_connection(connection, move |duration| executor.timer(duration)) } pub fn disconnect(&self, connection_id: ConnectionId) { self.connections.write().remove(&connection_id); } #[cfg(any(test, feature = "test-support"))] pub fn reset(&self, epoch: u32) { self.next_connection_id.store(0, SeqCst); self.epoch.store(epoch, SeqCst); } pub fn teardown(&self) { self.connections.write().clear(); } /// Make a request and wait for a response. pub fn request<T: RequestMessage>( &self, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<T::Response>> { self.request_internal(None, receiver_id, request) .map_ok(|envelope| envelope.payload) } pub fn request_envelope<T: RequestMessage>( &self, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> { self.request_internal(None, receiver_id, request) } pub fn forward_request<T: RequestMessage>( &self, sender_id: ConnectionId, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<T::Response>> { self.request_internal(Some(sender_id), receiver_id, request) .map_ok(|envelope| envelope.payload) } fn request_internal<T: RequestMessage>( &self, original_sender_id: Option<ConnectionId>, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> { let envelope = request.into_envelope(0, None, original_sender_id.map(Into::into)); let response = self.request_dynamic(receiver_id, envelope, T::NAME); async move { let (response, received_at) = response.await?; Ok(TypedEnvelope { message_id: response.id, sender_id: receiver_id.into(), original_sender_id: response.original_sender_id, payload: T::Response::from_envelope(response) .ok_or_else(|| anyhow!("received response of the wrong type"))?, received_at, }) } } /// Make a request and wait for a response. /// /// The caller must make sure to deserialize the response into the request's /// response type. This interface is only useful in trait objects, where /// generics can't be used. If you have a concrete type, use `request`. pub fn request_dynamic( &self, receiver_id: ConnectionId, mut envelope: proto::Envelope, type_name: &'static str, ) -> impl Future<Output = Result<(proto::Envelope, Instant)>> { let (tx, rx) = oneshot::channel(); let send = self.connection_state(receiver_id).and_then(|connection| { envelope.id = connection.next_message_id.fetch_add(1, SeqCst); connection .response_channels .lock() .as_mut() .ok_or_else(|| anyhow!("connection was closed"))? .insert(envelope.id, tx); connection .outgoing_tx .unbounded_send(proto::Message::Envelope(envelope)) .map_err(|_| anyhow!("connection was closed"))?; Ok(()) }); async move { send?; let (response, received_at, _barrier) = rx.await.map_err(|_| anyhow!("connection was closed"))?; if let Some(proto::envelope::Payload::Error(error)) = &response.payload { return Err(RpcError::from_proto(error, type_name)); } Ok((response, received_at)) } } pub fn request_stream<T: RequestMessage>( &self, receiver_id: ConnectionId, request: T, ) -> impl Future<Output = Result<impl Unpin + Stream<Item = Result<T::Response>>>> { let (tx, rx) = mpsc::unbounded(); let send = self.connection_state(receiver_id).and_then(|connection| { let message_id = connection.next_message_id.fetch_add(1, SeqCst); let stream_response_channels = connection.stream_response_channels.clone(); stream_response_channels .lock() .as_mut() .ok_or_else(|| anyhow!("connection was closed"))? .insert(message_id, tx);