Fix async bug(s) in MGS serial console forwarding

Our `select!` loop was pulling enqueued data out of in-memory buffers,
but then if the future they went to was cancelled by the `select!` the
data was lost. This splits the forwarding up into three separate tokio
tasks, and leaves the `select!` loop only polling on cancel-safe futures
that don't lose any data if dropped.

Author: jgallagher

gateway/src/serial_console.rs

@@ -5,8 +5,8 @@
 // Copyright 2022 Oxide Computer Company
 
 use crate::error::Error;
-use futures::future::Fuse;
-use futures::FutureExt;
+use futures::stream::SplitSink;
+use futures::stream::SplitStream;
 use futures::SinkExt;
 use futures::StreamExt;
 use gateway_messages::SpComponent;
@@ -23,9 +23,8 @@ use slog::error;
 use slog::info;
 use slog::Logger;
 use std::borrow::Cow;
-use std::collections::VecDeque;
-use std::mem;
 use std::ops::Deref;
+use tokio::sync::mpsc;
 use tokio_tungstenite::tungstenite::handshake;
 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
 use tokio_tungstenite::tungstenite::protocol::CloseFrame;
@@ -143,41 +142,45 @@ struct SerialConsoleTask {
 
 impl SerialConsoleTask {
     async fn run(self, log: &Logger) -> Result<(), SerialTaskError> {
-        let (mut ws_sink, mut ws_stream) = self.ws_stream.split();
+        let (ws_sink, ws_stream) = self.ws_stream.split();
+
+        // Spawn a task to send any messages received from the SP to the client
+        // websocket.
+        //
+        // TODO-cleanup We have no way to apply backpressure to the SP, and are
+        // willing to buffer up an arbitray amount of data in memory. We should
+        // apply some form of backpressure (which the SP could only handle by
+        // discarding data).
+        let (ws_sink_tx, ws_sink_rx) = mpsc::unbounded_channel();
+        let mut ws_sink_handle =
+            tokio::spawn(Self::ws_sink_task(ws_sink, ws_sink_rx));
+
+        // Spawn a task to send any messages received from the client websocket
+        // to the SP.
         let (console_tx, mut console_rx) = self.console.split();
         let console_tx = DetachOnDrop::new(console_tx);
-
-        // TODO Currently we do not apply any backpressure on the SP and are
-        // willing to buffer up an arbitrary amount of data in memory. Is it
-        // reasonable to apply backpressure to the SP over UDP? Should we have
-        // caps on memory and start discarding data if we exceed them? We _do_
-        // apply backpressure to the websocket, delaying reading from it if we
-        // still have data waiting to be sent to the SP.
-        let mut data_from_sp: VecDeque<Vec<u8>> = VecDeque::new();
-        let mut data_to_sp: Vec<u8> = Vec::new();
+        let mut ws_recv_handle = tokio::spawn(Self::ws_recv_task(
+            ws_stream,
+            console_tx,
+            log.clone(),
+        ));
 
         loop {
-            let ws_send = if let Some(data) = data_from_sp.pop_front() {
-                ws_sink.send(Message::Binary(data)).fuse()
-            } else {
-                Fuse::terminated()
-            };
-
-            let (ws_recv, sp_send) = if data_to_sp.is_empty() {
-                (ws_stream.next().fuse(), Fuse::terminated())
-            } else {
-                // Steal `data_to_sp` and create a future to send it to the SP.
-                let mut to_send = Vec::new();
-                mem::swap(&mut to_send, &mut data_to_sp);
-                (Fuse::terminated(), console_tx.write(to_send).fuse())
-            };
-
             tokio::select! {
-                // Finished (or failed) sending data to the SP.
-                send_success = sp_send => {
-                    send_success
-                        .map_err(gateway_sp_comms::error::Error::from)
-                        .map_err(Error::from)?;
+                // Our ws_sink task completed; this is only possible if it
+                // fails, since it loops until we drop ws_sink_tx (which doesn't
+                // happen until we return!).
+                join_result = &mut ws_sink_handle => {
+                    let result = join_result.expect("ws sink task panicked");
+                    return result;
+                }
+
+                // Our ws_recv task completed; this is possible if the websocket
+                // connection fails or is closed by the client. In either case,
+                // we're also done.
+                join_result = &mut ws_recv_handle => {
+                    let result = join_result.expect("ws recv task panicked");
+                    return result;
                 }
 
                 // Receive a UDP packet from the SP.
@@ -188,7 +191,7 @@ impl SerialConsoleTask {
                                 log, "received serial console data from SP";
                                 "length" => data.len(),
                             );
-                            data_from_sp.push_back(data);
+                            let _ = ws_sink_tx.send(Message::Binary(data));
                         }
                         None => {
                             // Sender is closed; i.e., we've been detached.
@@ -198,47 +201,55 @@ impl SerialConsoleTask {
                                 code: CloseCode::Policy,
                                 reason: Cow::Borrowed("serial console was detached"),
                             };
-                            ws_sink.send(Message::Close(Some(close))).await?;
+                            let _ = ws_sink_tx.send(Message::Close(Some(close)));
                             return Ok(());
                         }
                     }
                 }
+            }
+        }
+    }
 
-                // Send a previously-received UDP packet of data to the websocket
-                // client
-                write_success = ws_send => {
-                    write_success?;
-                }
+    async fn ws_sink_task(
+        mut ws_sink: SplitSink<WebSocketStream<Upgraded>, Message>,
+        mut messages: mpsc::UnboundedReceiver<Message>,
+    ) -> Result<(), SerialTaskError> {
+        while let Some(message) = messages.recv().await {
+            ws_sink.send(message).await?;
+        }
+        Ok(())
+    }
 
-                // Receive from the websocket to send to the SP.
-                msg = ws_recv => {
-                    match msg {
-                        Some(Ok(Message::Binary(mut data))) => {
-                            // we only populate ws_recv when we have no data
-                            // currently queued up; sanity check that here
-                            assert!(data_to_sp.is_empty());
-                            data_to_sp.append(&mut data);
-                        }
-                        Some(Ok(Message::Close(_))) | None => {
-                            info!(
-                                log,
-                                "remote end closed websocket; terminating task",
-                            );
-                            return Ok(());
-                        }
-                        Some(other) => {
-                            let wrong_message = other?;
-                            error!(
-                                log,
-                                "bogus websocket message; terminating task";
-                                "message" => ?wrong_message,
-                            );
-                            return Ok(());
-                        }
-                    }
+    async fn ws_recv_task(
+        mut ws_stream: SplitStream<WebSocketStream<Upgraded>>,
+        console_tx: DetachOnDrop,
+        log: Logger,
+    ) -> Result<(), SerialTaskError> {
+        while let Some(message) = ws_stream.next().await {
+            match message {
+                Ok(Message::Binary(data)) => {
+                    console_tx
+                        .write(data)
+                        .await
+                        .map_err(gateway_sp_comms::error::Error::from)
+                        .map_err(Error::from)?;
+                }
+                Ok(Message::Close(_)) => {
+                    break;
+                }
+                Ok(other) => {
+                    error!(
+                        log,
+                        "bogus websocket message; terminating task";
+                        "message" => ?other,
+                    );
+                    return Ok(());
                 }
+                Err(err) => return Err(err.into()),
             }
         }
+        info!(log, "remote end closed websocket; terminating task",);
+        Ok(())
     }
 }