Impose a timeout on writing RPC messages

Co-Authored-By: Nathan Sobo <nathan@zed.dev>
This commit is contained in:
Max Brunsfeld 2021-11-22 13:23:55 -08:00
parent 447f710570
commit 7474813a17
4 changed files with 82 additions and 105 deletions

11
Cargo.lock generated
View file

@ -4308,6 +4308,7 @@ dependencies = [
"rsa", "rsa",
"serde 1.0.125", "serde 1.0.125",
"smol", "smol",
"smol-timeout",
"tempdir", "tempdir",
"zstd", "zstd",
] ]
@ -4867,6 +4868,16 @@ dependencies = [
"once_cell", "once_cell",
] ]
[[package]]
name = "smol-timeout"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "847d777e2c6c166bad26264479e80a9820f3d364fcb4a0e23cd57bbfa8e94961"
dependencies = [
"async-io",
"pin-project-lite 0.1.12",
]
[[package]] [[package]]
name = "socket2" name = "socket2"
version = "0.3.19" version = "0.3.19"

View file

@ -20,6 +20,7 @@ prost = "0.8"
rand = "0.8" rand = "0.8"
rsa = "0.4" rsa = "0.4"
serde = { version = "1", features = ["derive"] } serde = { version = "1", features = ["derive"] }
smol-timeout = "0.6"
zstd = "0.9" zstd = "0.9"
[build-dependencies] [build-dependencies]

View file

@ -7,6 +7,7 @@ use postage::{
mpsc, mpsc,
prelude::{Sink as _, Stream as _}, prelude::{Sink as _, Stream as _},
}; };
use smol_timeout::TimeoutExt as _;
use std::{ use std::{
collections::HashMap, collections::HashMap,
fmt, fmt,
@ -16,6 +17,7 @@ use std::{
atomic::{self, AtomicU32}, atomic::{self, AtomicU32},
Arc, Arc,
}, },
time::Duration,
}; };
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
@ -90,6 +92,8 @@ struct ConnectionState {
response_channels: Arc<Mutex<Option<HashMap<u32, mpsc::Sender<proto::Envelope>>>>>, response_channels: Arc<Mutex<Option<HashMap<u32, mpsc::Sender<proto::Envelope>>>>>,
} }
const WRITE_TIMEOUT: Duration = Duration::from_secs(10);
impl Peer { impl Peer {
pub fn new() -> Arc<Self> { pub fn new() -> Arc<Self> {
Arc::new(Self { Arc::new(Self {
@ -155,8 +159,10 @@ impl Peer {
}, },
outgoing = outgoing_rx.recv().fuse() => match outgoing { outgoing = outgoing_rx.recv().fuse() => match outgoing {
Some(outgoing) => { Some(outgoing) => {
if let Err(result) = writer.write_message(&outgoing).await { match writer.write_message(&outgoing).timeout(WRITE_TIMEOUT).await {
break 'outer Err(result).context("failed to write RPC message") None => break 'outer Err(anyhow!("timed out writing RPC message")),
Some(Err(result)) => break 'outer Err(result).context("failed to write RPC message"),
_ => {}
} }
} }
None => break 'outer Ok(()), None => break 'outer Ok(()),

View file

@ -6,9 +6,10 @@ use super::{
AppState, AppState,
}; };
use anyhow::anyhow; use anyhow::anyhow;
use async_std::{sync::RwLock, task}; use async_std::task;
use async_tungstenite::{tungstenite::protocol::Role, WebSocketStream}; use async_tungstenite::{tungstenite::protocol::Role, WebSocketStream};
use futures::{future::BoxFuture, FutureExt}; use futures::{future::BoxFuture, FutureExt};
use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use postage::{mpsc, prelude::Sink as _, prelude::Stream as _}; use postage::{mpsc, prelude::Sink as _, prelude::Stream as _};
use rpc::{ use rpc::{
proto::{self, AnyTypedEnvelope, EnvelopedMessage}, proto::{self, AnyTypedEnvelope, EnvelopedMessage},
@ -23,7 +24,7 @@ use std::{
sync::Arc, sync::Arc,
time::Instant, time::Instant,
}; };
use store::{JoinedWorktree, Store, Worktree}; use store::{Store, Worktree};
use surf::StatusCode; use surf::StatusCode;
use tide::log; use tide::log;
use tide::{ use tide::{
@ -116,9 +117,7 @@ impl Server {
async move { async move {
let (connection_id, handle_io, mut incoming_rx) = let (connection_id, handle_io, mut incoming_rx) =
this.peer.add_connection(connection).await; this.peer.add_connection(connection).await;
this.state_mut() this.state_mut().add_connection(connection_id, user_id);
.await
.add_connection(connection_id, user_id);
if let Err(err) = this.update_collaborators_for_users(&[user_id]).await { if let Err(err) = this.update_collaborators_for_users(&[user_id]).await {
log::error!("error updating collaborators for {:?}: {}", user_id, err); log::error!("error updating collaborators for {:?}: {}", user_id, err);
} }
@ -168,7 +167,7 @@ impl Server {
async fn sign_out(self: &mut Arc<Self>, connection_id: ConnectionId) -> tide::Result<()> { async fn sign_out(self: &mut Arc<Self>, connection_id: ConnectionId) -> tide::Result<()> {
self.peer.disconnect(connection_id).await; self.peer.disconnect(connection_id).await;
let removed_connection = self.state_mut().await.remove_connection(connection_id)?; let removed_connection = self.state_mut().remove_connection(connection_id)?;
for (worktree_id, worktree) in removed_connection.hosted_worktrees { for (worktree_id, worktree) in removed_connection.hosted_worktrees {
if let Some(share) = worktree.share { if let Some(share) = worktree.share {
@ -213,10 +212,7 @@ impl Server {
request: TypedEnvelope<proto::OpenWorktree>, request: TypedEnvelope<proto::OpenWorktree>,
) -> tide::Result<()> { ) -> tide::Result<()> {
let receipt = request.receipt(); let receipt = request.receipt();
let host_user_id = self let host_user_id = self.state().user_id_for_connection(request.sender_id)?;
.state()
.await
.user_id_for_connection(request.sender_id)?;
let mut collaborator_user_ids = HashSet::new(); let mut collaborator_user_ids = HashSet::new();
collaborator_user_ids.insert(host_user_id); collaborator_user_ids.insert(host_user_id);
@ -236,7 +232,7 @@ impl Server {
} }
let collaborator_user_ids = collaborator_user_ids.into_iter().collect::<Vec<_>>(); let collaborator_user_ids = collaborator_user_ids.into_iter().collect::<Vec<_>>();
let worktree_id = self.state_mut().await.add_worktree(Worktree { let worktree_id = self.state_mut().add_worktree(Worktree {
host_connection_id: request.sender_id, host_connection_id: request.sender_id,
collaborator_user_ids: collaborator_user_ids.clone(), collaborator_user_ids: collaborator_user_ids.clone(),
root_name: request.payload.root_name, root_name: request.payload.root_name,
@ -259,7 +255,6 @@ impl Server {
let worktree_id = request.payload.worktree_id; let worktree_id = request.payload.worktree_id;
let worktree = self let worktree = self
.state_mut() .state_mut()
.await
.remove_worktree(worktree_id, request.sender_id)?; .remove_worktree(worktree_id, request.sender_id)?;
if let Some(share) = worktree.share { if let Some(share) = worktree.share {
@ -294,7 +289,6 @@ impl Server {
let collaborator_user_ids = let collaborator_user_ids =
self.state_mut() self.state_mut()
.await
.share_worktree(worktree.id, request.sender_id, entries); .share_worktree(worktree.id, request.sender_id, entries);
if let Some(collaborator_user_ids) = collaborator_user_ids { if let Some(collaborator_user_ids) = collaborator_user_ids {
self.peer self.peer
@ -322,7 +316,6 @@ impl Server {
let worktree_id = request.payload.worktree_id; let worktree_id = request.payload.worktree_id;
let worktree = self let worktree = self
.state_mut() .state_mut()
.await
.unshare_worktree(worktree_id, request.sender_id)?; .unshare_worktree(worktree_id, request.sender_id)?;
broadcast(request.sender_id, worktree.connection_ids, |conn_id| { broadcast(request.sender_id, worktree.connection_ids, |conn_id| {
@ -341,22 +334,17 @@ impl Server {
request: TypedEnvelope<proto::JoinWorktree>, request: TypedEnvelope<proto::JoinWorktree>,
) -> tide::Result<()> { ) -> tide::Result<()> {
let worktree_id = request.payload.worktree_id; let worktree_id = request.payload.worktree_id;
let user_id = self
.state()
.await
.user_id_for_connection(request.sender_id)?;
let mut state = self.state_mut().await; let user_id = self.state().user_id_for_connection(request.sender_id)?;
match state.join_worktree(request.sender_id, user_id, worktree_id) { let response_data = self
Ok(JoinedWorktree { .state_mut()
replica_id, .join_worktree(request.sender_id, user_id, worktree_id)
worktree, .and_then(|joined| {
}) => { let share = joined.worktree.share()?;
let share = worktree.share()?;
let peer_count = share.guest_connection_ids.len(); let peer_count = share.guest_connection_ids.len();
let mut peers = Vec::with_capacity(peer_count); let mut peers = Vec::with_capacity(peer_count);
peers.push(proto::Peer { peers.push(proto::Peer {
peer_id: worktree.host_connection_id.0, peer_id: joined.worktree.host_connection_id.0,
replica_id: 0, replica_id: 0,
}); });
for (peer_conn_id, peer_replica_id) in &share.guest_connection_ids { for (peer_conn_id, peer_replica_id) in &share.guest_connection_ids {
@ -370,16 +358,19 @@ impl Server {
let response = proto::JoinWorktreeResponse { let response = proto::JoinWorktreeResponse {
worktree: Some(proto::Worktree { worktree: Some(proto::Worktree {
id: worktree_id, id: worktree_id,
root_name: worktree.root_name.clone(), root_name: joined.worktree.root_name.clone(),
entries: share.entries.values().cloned().collect(), entries: share.entries.values().cloned().collect(),
}), }),
replica_id: replica_id as u32, replica_id: joined.replica_id as u32,
peers, peers,
}; };
let connection_ids = worktree.connection_ids(); let connection_ids = joined.worktree.connection_ids();
let collaborator_user_ids = worktree.collaborator_user_ids.clone(); let collaborator_user_ids = joined.worktree.collaborator_user_ids.clone();
drop(state); Ok((response, connection_ids, collaborator_user_ids))
});
match response_data {
Ok((response, connection_ids, collaborator_user_ids)) => {
broadcast(request.sender_id, connection_ids, |conn_id| { broadcast(request.sender_id, connection_ids, |conn_id| {
self.peer.send( self.peer.send(
conn_id, conn_id,
@ -398,7 +389,6 @@ impl Server {
.await?; .await?;
} }
Err(error) => { Err(error) => {
drop(state);
self.peer self.peer
.respond_with_error( .respond_with_error(
request.receipt(), request.receipt(),
@ -419,10 +409,7 @@ impl Server {
) -> tide::Result<()> { ) -> tide::Result<()> {
let sender_id = request.sender_id; let sender_id = request.sender_id;
let worktree_id = request.payload.worktree_id; let worktree_id = request.payload.worktree_id;
let worktree = self let worktree = self.state_mut().leave_worktree(sender_id, worktree_id);
.state_mut()
.await
.leave_worktree(sender_id, worktree_id);
if let Some(worktree) = worktree { if let Some(worktree) = worktree {
broadcast(sender_id, worktree.connection_ids, |conn_id| { broadcast(sender_id, worktree.connection_ids, |conn_id| {
self.peer.send( self.peer.send(
@ -444,7 +431,7 @@ impl Server {
mut self: Arc<Server>, mut self: Arc<Server>,
request: TypedEnvelope<proto::UpdateWorktree>, request: TypedEnvelope<proto::UpdateWorktree>,
) -> tide::Result<()> { ) -> tide::Result<()> {
let connection_ids = self.state_mut().await.update_worktree( let connection_ids = self.state_mut().update_worktree(
request.sender_id, request.sender_id,
request.payload.worktree_id, request.payload.worktree_id,
&request.payload.removed_entries, &request.payload.removed_entries,
@ -467,7 +454,6 @@ impl Server {
let receipt = request.receipt(); let receipt = request.receipt();
let host_connection_id = self let host_connection_id = self
.state() .state()
.await
.worktree_host_connection_id(request.sender_id, request.payload.worktree_id)?; .worktree_host_connection_id(request.sender_id, request.payload.worktree_id)?;
let response = self let response = self
.peer .peer
@ -483,7 +469,6 @@ impl Server {
) -> tide::Result<()> { ) -> tide::Result<()> {
let host_connection_id = self let host_connection_id = self
.state() .state()
.await
.worktree_host_connection_id(request.sender_id, request.payload.worktree_id)?; .worktree_host_connection_id(request.sender_id, request.payload.worktree_id)?;
self.peer self.peer
.forward_send(request.sender_id, host_connection_id, request.payload) .forward_send(request.sender_id, host_connection_id, request.payload)
@ -498,7 +483,7 @@ impl Server {
let host; let host;
let guests; let guests;
{ {
let state = self.state().await; let state = self.state();
host = state host = state
.worktree_host_connection_id(request.sender_id, request.payload.worktree_id)?; .worktree_host_connection_id(request.sender_id, request.payload.worktree_id)?;
guests = state guests = state
@ -532,16 +517,13 @@ impl Server {
self: Arc<Server>, self: Arc<Server>,
request: TypedEnvelope<proto::UpdateBuffer>, request: TypedEnvelope<proto::UpdateBuffer>,
) -> tide::Result<()> { ) -> tide::Result<()> {
broadcast( let receiver_ids = self
request.sender_id, .state()
self.state() .worktree_connection_ids(request.sender_id, request.payload.worktree_id)?;
.await broadcast(request.sender_id, receiver_ids, |connection_id| {
.worktree_connection_ids(request.sender_id, request.payload.worktree_id)?, self.peer
|connection_id| { .forward_send(request.sender_id, connection_id, request.payload.clone())
self.peer })
.forward_send(request.sender_id, connection_id, request.payload.clone())
},
)
.await?; .await?;
self.peer.respond(request.receipt(), proto::Ack {}).await?; self.peer.respond(request.receipt(), proto::Ack {}).await?;
Ok(()) Ok(())
@ -551,17 +533,13 @@ impl Server {
self: Arc<Server>, self: Arc<Server>,
request: TypedEnvelope<proto::BufferSaved>, request: TypedEnvelope<proto::BufferSaved>,
) -> tide::Result<()> { ) -> tide::Result<()> {
broadcast( let receiver_ids = self
request.sender_id, .state()
self.store .worktree_connection_ids(request.sender_id, request.payload.worktree_id)?;
.read() broadcast(request.sender_id, receiver_ids, |connection_id| {
.await self.peer
.worktree_connection_ids(request.sender_id, request.payload.worktree_id)?, .forward_send(request.sender_id, connection_id, request.payload.clone())
|connection_id| { })
self.peer
.forward_send(request.sender_id, connection_id, request.payload.clone())
},
)
.await?; .await?;
Ok(()) Ok(())
} }
@ -570,10 +548,7 @@ impl Server {
self: Arc<Server>, self: Arc<Server>,
request: TypedEnvelope<proto::GetChannels>, request: TypedEnvelope<proto::GetChannels>,
) -> tide::Result<()> { ) -> tide::Result<()> {
let user_id = self let user_id = self.state().user_id_for_connection(request.sender_id)?;
.state()
.await
.user_id_for_connection(request.sender_id)?;
let channels = self.app_state.db.get_accessible_channels(user_id).await?; let channels = self.app_state.db.get_accessible_channels(user_id).await?;
self.peer self.peer
.respond( .respond(
@ -622,20 +597,20 @@ impl Server {
) -> tide::Result<()> { ) -> tide::Result<()> {
let mut send_futures = Vec::new(); let mut send_futures = Vec::new();
let state = self.state().await; {
for user_id in user_ids { let state = self.state();
let collaborators = state.collaborators_for_user(*user_id); for user_id in user_ids {
for connection_id in state.connection_ids_for_user(*user_id) { let collaborators = state.collaborators_for_user(*user_id);
send_futures.push(self.peer.send( for connection_id in state.connection_ids_for_user(*user_id) {
connection_id, send_futures.push(self.peer.send(
proto::UpdateCollaborators { connection_id,
collaborators: collaborators.clone(), proto::UpdateCollaborators {
}, collaborators: collaborators.clone(),
)); },
));
}
} }
} }
drop(state);
futures::future::try_join_all(send_futures).await?; futures::future::try_join_all(send_futures).await?;
Ok(()) Ok(())
@ -645,10 +620,7 @@ impl Server {
mut self: Arc<Self>, mut self: Arc<Self>,
request: TypedEnvelope<proto::JoinChannel>, request: TypedEnvelope<proto::JoinChannel>,
) -> tide::Result<()> { ) -> tide::Result<()> {
let user_id = self let user_id = self.state().user_id_for_connection(request.sender_id)?;
.state()
.await
.user_id_for_connection(request.sender_id)?;
let channel_id = ChannelId::from_proto(request.payload.channel_id); let channel_id = ChannelId::from_proto(request.payload.channel_id);
if !self if !self
.app_state .app_state
@ -659,9 +631,7 @@ impl Server {
Err(anyhow!("access denied"))?; Err(anyhow!("access denied"))?;
} }
self.state_mut() self.state_mut().join_channel(request.sender_id, channel_id);
.await
.join_channel(request.sender_id, channel_id);
let messages = self let messages = self
.app_state .app_state
.db .db
@ -692,10 +662,7 @@ impl Server {
mut self: Arc<Self>, mut self: Arc<Self>,
request: TypedEnvelope<proto::LeaveChannel>, request: TypedEnvelope<proto::LeaveChannel>,
) -> tide::Result<()> { ) -> tide::Result<()> {
let user_id = self let user_id = self.state().user_id_for_connection(request.sender_id)?;
.state()
.await
.user_id_for_connection(request.sender_id)?;
let channel_id = ChannelId::from_proto(request.payload.channel_id); let channel_id = ChannelId::from_proto(request.payload.channel_id);
if !self if !self
.app_state .app_state
@ -707,7 +674,6 @@ impl Server {
} }
self.state_mut() self.state_mut()
.await
.leave_channel(request.sender_id, channel_id); .leave_channel(request.sender_id, channel_id);
Ok(()) Ok(())
@ -722,7 +688,7 @@ impl Server {
let user_id; let user_id;
let connection_ids; let connection_ids;
{ {
let state = self.state().await; let state = self.state();
user_id = state.user_id_for_connection(request.sender_id)?; user_id = state.user_id_for_connection(request.sender_id)?;
if let Some(ids) = state.channel_connection_ids(channel_id) { if let Some(ids) = state.channel_connection_ids(channel_id) {
connection_ids = ids; connection_ids = ids;
@ -809,10 +775,7 @@ impl Server {
self: Arc<Self>, self: Arc<Self>,
request: TypedEnvelope<proto::GetChannelMessages>, request: TypedEnvelope<proto::GetChannelMessages>,
) -> tide::Result<()> { ) -> tide::Result<()> {
let user_id = self let user_id = self.state().user_id_for_connection(request.sender_id)?;
.state()
.await
.user_id_for_connection(request.sender_id)?;
let channel_id = ChannelId::from_proto(request.payload.channel_id); let channel_id = ChannelId::from_proto(request.payload.channel_id);
if !self if !self
.app_state .app_state
@ -853,15 +816,11 @@ impl Server {
Ok(()) Ok(())
} }
fn state<'a>( fn state<'a>(self: &'a Arc<Self>) -> RwLockReadGuard<'a, Store> {
self: &'a Arc<Self>,
) -> impl Future<Output = async_std::sync::RwLockReadGuard<'a, Store>> {
self.store.read() self.store.read()
} }
fn state_mut<'a>( fn state_mut<'a>(self: &'a mut Arc<Self>) -> RwLockWriteGuard<'a, Store> {
self: &'a mut Arc<Self>,
) -> impl Future<Output = async_std::sync::RwLockWriteGuard<'a, Store>> {
self.store.write() self.store.write()
} }
} }
@ -961,7 +920,7 @@ mod tests {
github, AppState, Config, github, AppState, Config,
}; };
use ::rpc::Peer; use ::rpc::Peer;
use async_std::{sync::RwLockReadGuard, task}; use async_std::task;
use gpui::{ModelHandle, TestAppContext}; use gpui::{ModelHandle, TestAppContext};
use parking_lot::Mutex; use parking_lot::Mutex;
use postage::{mpsc, watch}; use postage::{mpsc, watch};
@ -2372,7 +2331,7 @@ mod tests {
} }
async fn state<'a>(&'a self) -> RwLockReadGuard<'a, Store> { async fn state<'a>(&'a self) -> RwLockReadGuard<'a, Store> {
self.server.store.read().await self.server.store.read()
} }
async fn condition<F>(&mut self, mut predicate: F) async fn condition<F>(&mut self, mut predicate: F)
@ -2380,7 +2339,7 @@ mod tests {
F: FnMut(&Store) -> bool, F: FnMut(&Store) -> bool,
{ {
async_std::future::timeout(Duration::from_millis(500), async { async_std::future::timeout(Duration::from_millis(500), async {
while !(predicate)(&*self.server.store.read().await) { while !(predicate)(&*self.server.store.read()) {
self.notifications.recv().await; self.notifications.recv().await;
} }
}) })