mirrors/zed
1
0
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2025-01-15 14:47:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 9

Prevent requests from hanging when shutting down a connection

Browse source 
When closing a connection (either due to an error or simply because the
user wanted to), we will now  *take* `response_channels` as opposed to
clearing them. This ensures that `Peer::request` can't succeed in both
adding the oneshot channel in `response_channels` map _and_ submit the
message onto the `outgoing_tx` channel.

This also streamlines how we close a connection by unifying all the exit
code paths of the IO handling future.
This commit is contained in:
Antonio Scandurra 2021-09-23 16:18:56 +02:00
parent c70b4a99c9
commit 6b0ed820bf
1 changed files with 14 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
zrpc/src/peer.rs
View file

@ -87,7 +87,7 @@ pub struct Peer {
struct ConnectionState { struct ConnectionState {
outgoing_tx: mpsc::Sender<proto::Envelope>, outgoing_tx: mpsc::Sender<proto::Envelope>,
next_message_id: Arc<AtomicU32>, next_message_id: Arc<AtomicU32>,
response_channels: Arc<Mutex<HashMap<u32, mpsc::Sender<proto::Envelope>>>>, response_channels: Arc<Mutex<Option<HashMap<u32, mpsc::Sender<proto::Envelope>>>>>,
} }
impl Peer { impl Peer {
@ -123,7 +123,7 @@ impl Peer {
let this = self.clone(); let this = self.clone();
let response_channels = connection_state.response_channels.clone(); let response_channels = connection_state.response_channels.clone();
let handle_io = async move { let handle_io = async move {
loop { let result = 'outer: loop {
let read_message = reader.read_message().fuse(); let read_message = reader.read_message().fuse();
futures::pin_mut!(read_message); futures::pin_mut!(read_message);
loop { loop {
@ -131,7 +131,7 @@ impl Peer {
incoming = read_message => match incoming { incoming = read_message => match incoming {
Ok(incoming) => { Ok(incoming) => {
if let Some(responding_to) = incoming.responding_to { if let Some(responding_to) = incoming.responding_to {
let channel = response_channels.lock().await.remove(&responding_to); let channel = response_channels.lock().await.as_mut().unwrap().remove(&responding_to);
if let Some(mut tx) = channel { if let Some(mut tx) = channel {
tx.send(incoming).await.ok(); tx.send(incoming).await.ok();
} else { } else {
@ -140,9 +140,7 @@ impl Peer {
} else { } else {
if let Some(envelope) = proto::build_typed_envelope(connection_id, incoming) { if let Some(envelope) = proto::build_typed_envelope(connection_id, incoming) {
if incoming_tx.send(envelope).await.is_err() { if incoming_tx.send(envelope).await.is_err() {
response_channels.lock().await.clear(); break 'outer Ok(())
this.connections.write().await.remove(&connection_id);
return Ok(())
} }
} else { } else {
log::error!("unable to construct a typed envelope"); log::error!("unable to construct a typed envelope");
@ -152,28 +150,24 @@ impl Peer {
break; break;
} }
Err(error) => { Err(error) => {
response_channels.lock().await.clear(); break 'outer Err(error).context("received invalid RPC message")
this.connections.write().await.remove(&connection_id);
Err(error).context("received invalid RPC message")?;
} }
}, },
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 { if let Err(result) = writer.write_message(&outgoing).await {
response_channels.lock().await.clear(); break 'outer Err(result).context("failed to write RPC message")
this.connections.write().await.remove(&connection_id);
Err(result).context("failed to write RPC message")?;
} }
} }
None => { None => break 'outer Ok(()),
response_channels.lock().await.clear();
this.connections.write().await.remove(&connection_id);
return Ok(())
}
} }
} }
} }
} };
response_channels.lock().await.take();
this.connections.write().await.remove(&connection_id);
result
}; };
self.connections self.connections
@ -226,6 +220,8 @@ impl Peer {
.response_channels .response_channels
.lock() .lock()
.await .await
.as_mut()
.ok_or_else(|| anyhow!("connection was closed"))?
.insert(message_id, tx); .insert(message_id, tx);
connection connection
.outgoing_tx .outgoing_tx