1use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
  2use futures::{SinkExt as _, Stream, StreamExt as _};
  3use std::{io, task::Poll};
  4
  5pub struct Connection {
  6    pub(crate) tx:
  7        Box<dyn 'static + Send + Unpin + futures::Sink<WebSocketMessage, Error = WebSocketError>>,
  8    pub(crate) rx: Box<
  9        dyn 'static
 10            + Send
 11            + Unpin
 12            + futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>,
 13    >,
 14}
 15
 16impl Connection {
 17    pub fn new<S>(stream: S) -> Self
 18    where
 19        S: 'static
 20            + Send
 21            + Unpin
 22            + futures::Sink<WebSocketMessage, Error = WebSocketError>
 23            + futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>,
 24    {
 25        let (tx, rx) = stream.split();
 26        Self {
 27            tx: Box::new(tx),
 28            rx: Box::new(rx),
 29        }
 30    }
 31
 32    pub async fn send(&mut self, message: WebSocketMessage) -> Result<(), WebSocketError> {
 33        self.tx.send(message).await
 34    }
 35
 36    #[cfg(any(test, feature = "test-support"))]
 37    pub fn in_memory(
 38        executor: std::sync::Arc<gpui::executor::Background>,
 39    ) -> (Self, Self, postage::watch::Sender<Option<()>>) {
 40        let (kill_tx, mut kill_rx) = postage::watch::channel_with(None);
 41        postage::stream::Stream::try_recv(&mut kill_rx).unwrap();
 42
 43        let (a_tx, a_rx) = Self::channel(kill_rx.clone(), executor.clone());
 44        let (b_tx, b_rx) = Self::channel(kill_rx, executor);
 45        (
 46            Self { tx: a_tx, rx: b_rx },
 47            Self { tx: b_tx, rx: a_rx },
 48            kill_tx,
 49        )
 50    }
 51
 52    #[cfg(any(test, feature = "test-support"))]
 53    fn channel(
 54        kill_rx: postage::watch::Receiver<Option<()>>,
 55        executor: std::sync::Arc<gpui::executor::Background>,
 56    ) -> (
 57        Box<dyn Send + Unpin + futures::Sink<WebSocketMessage, Error = WebSocketError>>,
 58        Box<dyn Send + Unpin + futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>>,
 59    ) {
 60        use futures::channel::mpsc;
 61        use io::{Error, ErrorKind};
 62        use std::sync::Arc;
 63
 64        let (tx, rx) = mpsc::unbounded::<WebSocketMessage>();
 65        let tx = tx
 66            .sink_map_err(|e| WebSocketError::from(Error::new(ErrorKind::Other, e)))
 67            .with({
 68                let executor = Arc::downgrade(&executor);
 69                let kill_rx = kill_rx.clone();
 70                move |msg| {
 71                    let kill_rx = kill_rx.clone();
 72                    let executor = executor.clone();
 73                    Box::pin(async move {
 74                        if let Some(executor) = executor.upgrade() {
 75                            executor.simulate_random_delay().await;
 76                        }
 77                        if kill_rx.borrow().is_none() {
 78                            Ok(msg)
 79                        } else {
 80                            Err(Error::new(ErrorKind::Other, "connection killed").into())
 81                        }
 82                    })
 83                }
 84            });
 85        let rx = rx.then(move |msg| {
 86            let executor = Arc::downgrade(&executor);
 87            Box::pin(async move {
 88                if let Some(executor) = executor.upgrade() {
 89                    executor.simulate_random_delay().await;
 90                }
 91                msg
 92            })
 93        });
 94        let rx = KillableReceiver { kill_rx, rx };
 95
 96        (Box::new(tx), Box::new(rx))
 97    }
 98}
 99
100struct KillableReceiver<S> {
101    rx: S,
102    kill_rx: postage::watch::Receiver<Option<()>>,
103}
104
105impl<S: Unpin + Stream<Item = WebSocketMessage>> Stream for KillableReceiver<S> {
106    type Item = Result<WebSocketMessage, WebSocketError>;
107
108    fn poll_next(
109        mut self: std::pin::Pin<&mut Self>,
110        cx: &mut std::task::Context<'_>,
111    ) -> Poll<Option<Self::Item>> {
112        if let Poll::Ready(Some(Some(()))) = self.kill_rx.poll_next_unpin(cx) {
113            Poll::Ready(Some(Err(io::Error::new(
114                io::ErrorKind::Other,
115                "connection killed",
116            )
117            .into())))
118        } else {
119            self.rx.poll_next_unpin(cx).map(|value| value.map(Ok))
120        }
121    }
122}