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
 63        let (tx, rx) = mpsc::unbounded::<WebSocketMessage>();
 64        let tx = tx
 65            .sink_map_err(|e| WebSocketError::from(Error::new(ErrorKind::Other, e)))
 66            .with({
 67                let kill_rx = kill_rx.clone();
 68                let executor = executor.clone();
 69                move |msg| {
 70                    let kill_rx = kill_rx.clone();
 71                    let executor = executor.clone();
 72                    Box::pin(async move {
 73                        executor.simulate_random_delay().await;
 74                        if kill_rx.borrow().is_none() {
 75                            Ok(msg)
 76                        } else {
 77                            Err(Error::new(ErrorKind::Other, "connection killed").into())
 78                        }
 79                    })
 80                }
 81            });
 82        let rx = rx.then(move |msg| {
 83            let executor = executor.clone();
 84            Box::pin(async move {
 85                executor.simulate_random_delay().await;
 86                msg
 87            })
 88        });
 89        let rx = KillableReceiver { kill_rx, rx };
 90
 91        (Box::new(tx), Box::new(rx))
 92    }
 93}
 94
 95struct KillableReceiver<S> {
 96    rx: S,
 97    kill_rx: postage::watch::Receiver<Option<()>>,
 98}
 99
100impl<S: Unpin + Stream<Item = WebSocketMessage>> Stream for KillableReceiver<S> {
101    type Item = Result<WebSocketMessage, WebSocketError>;
102
103    fn poll_next(
104        mut self: std::pin::Pin<&mut Self>,
105        cx: &mut std::task::Context<'_>,
106    ) -> Poll<Option<Self::Item>> {
107        if let Poll::Ready(Some(Some(()))) = self.kill_rx.poll_next_unpin(cx) {
108            Poll::Ready(Some(Err(io::Error::new(
109                io::ErrorKind::Other,
110                "connection killed",
111            )
112            .into())))
113        } else {
114            self.rx.poll_next_unpin(cx).map(|value| value.map(Ok))
115        }
116    }
117}