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

If a test connection has been killed, never return a message

Browse source
This commit is contained in:
Nathan Sobo 2021-09-09 08:25:58 -06:00
parent 3c61a3e826
commit e2b56e8764
1 changed files with 30 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
zrpc/src/conn.rs
View file

@ -1,5 +1,6 @@
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage}; use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
use futures::{SinkExt as _, StreamExt as _}; use futures::{channel::mpsc, SinkExt as _, Stream, StreamExt as _};
use std::{io, task::Poll};
pub struct Conn { pub struct Conn {
pub(crate) tx: pub(crate) tx:
@ -53,10 +54,10 @@ impl Conn {
Box<dyn Send + Unpin + futures::Sink<WebSocketMessage, Error = WebSocketError>>, Box<dyn Send + Unpin + futures::Sink<WebSocketMessage, Error = WebSocketError>>,
Box<dyn Send + Unpin + futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>>, Box<dyn Send + Unpin + futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>>,
) { ) {
use futures::{future, stream, SinkExt as _, StreamExt as _}; use futures::{future, SinkExt as _};
use std::io::{Error, ErrorKind}; use io::{Error, ErrorKind};
let (tx, rx) = futures::channel::mpsc::unbounded::<WebSocketMessage>(); let (tx, rx) = mpsc::unbounded::<WebSocketMessage>();
let tx = tx let tx = tx
.sink_map_err(|e| WebSocketError::from(Error::new(ErrorKind::Other, e))) .sink_map_err(|e| WebSocketError::from(Error::new(ErrorKind::Other, e)))
.with({ .with({
@ -69,19 +70,32 @@ impl Conn {
} }
} }
}); });
let rx = stream::select( let rx = KillableReceiver { kill_rx, rx };
rx.map(Ok),
kill_rx.filter_map(|kill| {
if kill.is_none() {
future::ready(None)
} else {
future::ready(Some(Err(
Error::new(ErrorKind::Other, "connection killed").into()
)))
}
}),
);
(Box::new(tx), Box::new(rx)) (Box::new(tx), Box::new(rx))
} }
} }
struct KillableReceiver {
rx: mpsc::UnboundedReceiver<WebSocketMessage>,
kill_rx: postage::watch::Receiver<Option<()>>,
}
impl Stream for KillableReceiver {
type Item = Result<WebSocketMessage, WebSocketError>;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
if let Poll::Ready(Some(Some(()))) = self.kill_rx.poll_next_unpin(cx) {
Poll::Ready(Some(Err(io::Error::new(
io::ErrorKind::Other,
"connection killed",
)
.into())))
} else {
self.rx.poll_next_unpin(cx).map(|value| value.map(Ok))
}
}
}