Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Log the time incoming RPC messages were queued (#8909)

Browse source 
Release Notes:

- N/A

Co-authored-by: Conrad <conrad@zed.dev>
This commit is contained in:
Max Brunsfeld 2024-03-05 14:40:09 -08:00 • committed by GitHub
parent bca98caa07
commit 35c516fda9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 47 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

3
crates/rpc/src/macros.rs
View file

@ -1,7 +1,7 @@
#[macro_export]
macro_rules! messages {
($(($name:ident, $priority:ident)),* $(,)?) => {
pub fn build_typed_envelope(sender_id: ConnectionId, envelope: Envelope) -> Option<Box<dyn AnyTypedEnvelope>> {
pub fn build_typed_envelope(sender_id: ConnectionId, received_at: Instant, envelope: Envelope) -> Option<Box<dyn AnyTypedEnvelope>> {
match envelope.payload {
$(Some(envelope::Payload::$name(payload)) => {
Some(Box::new(TypedEnvelope {
@ -12,6 +12,7 @@ macro_rules! messages {
}),
message_id: envelope.id,
payload,
received_at,
}))
}, )*
_ => None

52
crates/rpc/src/peer.rs
View file

@ -13,7 +13,7 @@ use futures::{
};
use parking_lot::{Mutex, RwLock};
use serde::{ser::SerializeStruct, Serialize};
use std::{fmt, sync::atomic::Ordering::SeqCst};
use std::{fmt, sync::atomic::Ordering::SeqCst, time::Instant};
use std::{
future::Future,
marker::PhantomData,
@ -79,6 +79,7 @@ pub struct TypedEnvelope<T> {
pub original_sender_id: Option<PeerId>,
pub message_id: u32,
pub payload: T,
pub received_at: Instant,
}
impl<T> TypedEnvelope<T> {
@ -111,8 +112,16 @@ pub struct ConnectionState {
next_message_id: Arc<AtomicU32>,
#[allow(clippy::type_complexity)]
#[serde(skip)]
response_channels:
Arc<Mutex<Option<HashMap<u32, oneshot::Sender<(proto::Envelope, oneshot::Sender<()>)>>>>>,
response_channels: Arc<
Mutex<
Option<
HashMap<
u32,
oneshot::Sender<(proto::Envelope, std::time::Instant, oneshot::Sender<()>)>,
>,
>,
>,
>,
}
const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(1);
@ -154,7 +163,7 @@ impl Peer {
#[cfg(any(test, feature = "test-support"))]
const INCOMING_BUFFER_SIZE: usize = 1;
#[cfg(not(any(test, feature = "test-support")))]
const INCOMING_BUFFER_SIZE: usize = 64;
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();
@ -238,10 +247,10 @@ impl Peer {
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) = incoming {
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).fuse() => match result {
result = incoming_tx.send((incoming, received_at)).fuse() => match result {
Ok(_) => {
tracing::trace!(%connection_id, "incoming rpc message: processed");
}
@ -272,7 +281,7 @@ impl Peer {
.write()
.insert(connection_id, connection_state);
let incoming_rx = incoming_rx.filter_map(move |incoming| {
let incoming_rx = incoming_rx.filter_map(move |(incoming, received_at)| {
let response_channels = response_channels.clone();
async move {
let message_id = incoming.id;
@ -291,7 +300,7 @@ impl Peer {
let channel = response_channels.lock().as_mut()?.remove(&responding_to);
if let Some(tx) = channel {
let requester_resumed = oneshot::channel();
if let Err(error) = tx.send((incoming, requester_resumed.0)) {
if let Err(error) = tx.send((incoming, received_at, requester_resumed.0)) {
tracing::trace!(
%connection_id,
message_id,
@ -315,8 +324,9 @@ impl Peer {
"incoming response: requester resumed"
);
} else {
let message_type = proto::build_typed_envelope(connection_id, incoming)
.map(|p| p.payload_type_name());
let message_type =
proto::build_typed_envelope(connection_id, received_at, incoming)
.map(|p| p.payload_type_name());
tracing::warn!(
%connection_id,
message_id,
@ -329,14 +339,16 @@ impl Peer {
None
} else {
tracing::trace!(%connection_id, message_id, "incoming message: received");
proto::build_typed_envelope(connection_id, incoming).or_else(|| {
tracing::error!(
%connection_id,
message_id,
"unable to construct a typed envelope"
);
None
})
proto::build_typed_envelope(connection_id, received_at, incoming).or_else(
|| {
tracing::error!(
%connection_id,
message_id,
"unable to construct a typed envelope"
);
None
},
)
}
}
});
@ -425,7 +437,8 @@ impl Peer {
});
async move {
send?;
let (response, _barrier) = rx.await.map_err(|_| anyhow!("connection was closed"))?;
let (response, received_at, _barrier) =
rx.await.map_err(|_| anyhow!("connection was closed"))?;
if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
Err(RpcError::from_proto(&error, T::NAME))
@ -436,6 +449,7 @@ impl Peer {
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,
})
}
}

10
crates/rpc/src/proto.rs
View file

@ -8,6 +8,7 @@ use futures::{SinkExt as _, StreamExt as _};
use prost::Message as _;
use serde::Serialize;
use std::any::{Any, TypeId};
use std::time::Instant;
use std::{
cmp,
fmt::Debug,
@ -515,8 +516,9 @@ impl<S> MessageStream<S>
where
S: futures::Stream<Item = Result<WebSocketMessage, anyhow::Error>> + Unpin,
{
pub async fn read(&mut self) -> Result<Message, anyhow::Error> {
pub async fn read(&mut self) -> Result<(Message, Instant), anyhow::Error> {
while let Some(bytes) = self.stream.next().await {
let received_at = Instant::now();
match bytes? {
WebSocketMessage::Binary(bytes) => {
zstd::stream::copy_decode(bytes.as_slice(), &mut self.encoding_buffer).unwrap();
@ -525,10 +527,10 @@ where
self.encoding_buffer.clear();
self.encoding_buffer.shrink_to(MAX_BUFFER_LEN);
return Ok(Message::Envelope(envelope));
return Ok((Message::Envelope(envelope), received_at));
}
WebSocketMessage::Ping(_) => return Ok(Message::Ping),
WebSocketMessage::Pong(_) => return Ok(Message::Pong),
WebSocketMessage::Ping(_) => return Ok((Message::Ping, received_at)),
WebSocketMessage::Pong(_) => return Ok((Message::Pong, received_at)),
WebSocketMessage::Close(_) => break,
_ => {}
}