Flumify everything

backend/common/src/ws_client/connect.rs

@@ -14,8 +14,6 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 use super::on_close::OnClose;
-use futures::channel::mpsc;
-use futures::{SinkExt, StreamExt};
 use soketto::handshake::{Client, ServerResponse};
 use std::sync::Arc;
 use tokio::net::TcpStream;
@@ -73,7 +71,7 @@ impl Connection {
         let mut rx_closed2 = tx_closed1.subscribe();
 
         // Receive messages from the socket:
-        let (mut tx_to_external, rx_from_ws) = mpsc::unbounded();
+        let (tx_to_external, rx_from_ws) = flume::unbounded();
         tokio::spawn(async move {
             let mut send_to_external = true;
             loop {
@@ -112,7 +110,7 @@ impl Connection {
                         .map_err(|e| e.into()),
                 };
 
-                if let Err(e) = tx_to_external.send(msg).await {
+                if let Err(e) = tx_to_external.send_async(msg).await {
                     // Our external channel may have closed or errored, but the socket hasn't
                     // been closed, so keep receiving in order to allow the socket to continue to
                     // function properly (we may be happy just sending messages to it), but stop
@@ -124,12 +122,12 @@ impl Connection {
         });
 
         // Send messages to the socket:
-        let (tx_to_ws, mut rx_from_external) = mpsc::unbounded();
+        let (tx_to_ws, rx_from_external) = flume::unbounded::<SentMessage>();
         tokio::spawn(async move {
             loop {
                 // Wait for messages, or bail entirely if asked to close.
                 let msg = tokio::select! {
-                    msg = rx_from_external.next() => { msg },
+                    msg = rx_from_external.recv_async() => { msg },
                     _ = rx_closed2.recv() => {
                         // attempt to gracefully end the connection.
                         let _ = ws_to_connection.close().await;
@@ -141,8 +139,8 @@ impl Connection {
                 // needs to keep receiving data for the WS connection to stay open, there's no
                 // reason to keep this side of the loop open if our channel is closed.
                 let msg = match msg {
-                    None => break,
-                    Some(msg) => msg,
+                    Ok(msg) => msg,
+                    _ => break,
                 };
 
                 // We don't explicitly shut down the channel if we hit send errors. Why? Because the
@@ -207,7 +205,7 @@ impl Connection {
                 closer: Arc::clone(&on_close),
             },
             Receiver {
-                inner: rx_from_ws,
+                inner: rx_from_ws.into_stream(),
                 closer: on_close,
             },
         )

backend/common/src/ws_client/receiver.rs

@@ -15,13 +15,12 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 use super::on_close::OnClose;
-use futures::channel::mpsc;
 use futures::{Stream, StreamExt};
 use std::sync::Arc;
 
 /// Receive messages out of a connection
 pub struct Receiver {
-    pub(super) inner: mpsc::UnboundedReceiver<Result<RecvMessage, RecvError>>,
+    pub(super) inner: flume::r#async::RecvStream<'static, Result<RecvMessage, RecvError>>,
     pub(super) closer: Arc<OnClose>,
 }
 

backend/common/src/ws_client/sender.rs

@@ -15,8 +15,6 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 use super::on_close::OnClose;
-use futures::channel::mpsc;
-use futures::{Sink, SinkExt};
 use std::sync::Arc;
 
 /// A message that can be sent into the channel interface
@@ -41,62 +39,36 @@ pub enum SentMessage {
 /// Send messages into the connection
 #[derive(Clone)]
 pub struct Sender {
-    pub(super) inner: mpsc::UnboundedSender<SentMessage>,
+    pub(super) inner: flume::Sender<SentMessage>,
     pub(super) closer: Arc<OnClose>,
 }
 
 impl Sender {
     /// Ask the underlying Websocket connection to close.
-    pub async fn close(&mut self) -> Result<(), SendError> {
+    pub async fn close(&mut self) -> Result<(), SendError<SentMessage>> {
         self.closer.0.send(()).map_err(|_| SendError::CloseError)?;
         Ok(())
     }
     /// Returns whether this channel is closed.
     pub fn is_closed(&self) -> bool {
-        self.inner.is_closed()
+        self.inner.is_disconnected()
     }
     /// Unbounded send will always queue the message and doesn't
     /// need to be awaited.
-    pub fn unbounded_send(&self, msg: SentMessage) -> Result<(), SendError> {
-        self.inner
-            .unbounded_send(msg)
-            .map_err(|e| e.into_send_error())?;
+    pub fn unbounded_send(&self, msg: SentMessage) -> Result<(), flume::SendError<SentMessage>> {
+        self.inner.send(msg)?;
         Ok(())
     }
+    /// Convert this sender into a Sink
+    pub fn into_sink(self) -> impl futures::Sink<SentMessage> + std::marker::Unpin + Clone + 'static {
+        self.inner.into_sink()
+    }
 }
 
 #[derive(thiserror::Error, Debug, Clone)]
-pub enum SendError {
+pub enum SendError<T: std::fmt::Debug + 'static> {
     #[error("Failed to send message: {0}")]
-    ChannelError(#[from] mpsc::SendError),
+    ChannelError(#[from] flume::SendError<T>),
     #[error("Failed to send close message")]
     CloseError,
 }
-
-impl Sink<SentMessage> for Sender {
-    type Error = SendError;
-    fn poll_ready(
-        mut self: std::pin::Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Result<(), Self::Error>> {
-        self.inner.poll_ready_unpin(cx).map_err(|e| e.into())
-    }
-    fn start_send(
-        mut self: std::pin::Pin<&mut Self>,
-        item: SentMessage,
-    ) -> Result<(), Self::Error> {
-        self.inner.start_send_unpin(item).map_err(|e| e.into())
-    }
-    fn poll_flush(
-        mut self: std::pin::Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Result<(), Self::Error>> {
-        self.inner.poll_flush_unpin(cx).map_err(|e| e.into())
-    }
-    fn poll_close(
-        mut self: std::pin::Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Result<(), Self::Error>> {
-        self.inner.poll_close_unpin(cx).map_err(|e| e.into())
-    }
-}