Add some tests, and add a configurable soak test

backend/test_utils/src/server/channels.rs

@@ -1,4 +1,4 @@
-use std::time::Duration;
+use std::{ops::{Deref, DerefMut}, time::Duration};
 
 use crate::feed_message_de::FeedMessage;
 use crate::ws_client;
@@ -13,14 +13,7 @@ impl From<ws_client::Sender> for ShardSender {
     }
 }
 
-impl ShardSender {
-    /// Close this connection
-    pub async fn close(&mut self) -> Result<(), ws_client::SendError> {
-        self.0.close().await
-    }
-}
-
-impl Sink<ws_client::Message> for ShardSender {
+impl Sink<ws_client::SentMessage> for ShardSender {
     type Error = ws_client::SendError;
     fn poll_ready(
         mut self: std::pin::Pin<&mut Self>,
@@ -30,7 +23,7 @@ impl Sink<ws_client::Message> for ShardSender {
     }
     fn start_send(
         mut self: std::pin::Pin<&mut Self>,
-        item: ws_client::Message,
+        item: ws_client::SentMessage,
     ) -> Result<(), Self::Error> {
         self.0.start_send_unpin(item)
     }
@@ -49,19 +42,33 @@ impl Sink<ws_client::Message> for ShardSender {
 }
 
 impl ShardSender {
-    pub async fn send_json_binary(
+    /// Send JSON as a binary websocket message
+    pub fn send_json_binary(
         &mut self,
         json: serde_json::Value,
     ) -> Result<(), ws_client::SendError> {
         let bytes = serde_json::to_vec(&json).expect("valid bytes");
-        self.send(ws_client::Message::Binary(bytes)).await
+        self.unbounded_send(ws_client::SentMessage::Binary(bytes))
     }
-    pub async fn send_json_text(
+    /// Send JSON as a textual websocket message
+    pub fn send_json_text(
         &mut self,
         json: serde_json::Value,
     ) -> Result<(), ws_client::SendError> {
         let s = serde_json::to_string(&json).expect("valid string");
-        self.send(ws_client::Message::Text(s)).await
+        self.unbounded_send(ws_client::SentMessage::Text(s))
+    }
+}
+
+impl Deref for ShardSender {
+    type Target = ws_client::Sender;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+impl DerefMut for ShardSender {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
     }
 }
 
@@ -75,7 +82,7 @@ impl From<ws_client::Receiver> for ShardReceiver {
 }
 
 impl Stream for ShardReceiver {
-    type Item = Result<ws_client::Message, ws_client::RecvError>;
+    type Item = Result<ws_client::RecvMessage, ws_client::RecvError>;
     fn poll_next(
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
@@ -84,6 +91,18 @@ impl Stream for ShardReceiver {
     }
 }
 
+impl Deref for ShardReceiver {
+    type Target = ws_client::Receiver;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+impl DerefMut for ShardReceiver {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 /// Wrap a `ws_client::Sender` with convenient utility methods for feed connections
 pub struct FeedSender(ws_client::Sender);
 
@@ -93,7 +112,7 @@ impl From<ws_client::Sender> for FeedSender {
     }
 }
 
-impl Sink<ws_client::Message> for FeedSender {
+impl Sink<ws_client::SentMessage> for FeedSender {
     type Error = ws_client::SendError;
     fn poll_ready(
         mut self: std::pin::Pin<&mut Self>,
@@ -103,7 +122,7 @@ impl Sink<ws_client::Message> for FeedSender {
     }
     fn start_send(
         mut self: std::pin::Pin<&mut Self>,
-        item: ws_client::Message,
+        item: ws_client::SentMessage,
     ) -> Result<(), Self::Error> {
         self.0.start_send_unpin(item)
     }
@@ -121,21 +140,36 @@ impl Sink<ws_client::Message> for FeedSender {
     }
 }
 
+impl Deref for FeedSender {
+    type Target = ws_client::Sender;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for FeedSender {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 impl FeedSender {
-    pub async fn send_command<S: AsRef<str>>(
+    /// Send a command into the feed. A command consists of a string
+    /// "command" part, and another string "parameter" part.
+    pub fn send_command<S: AsRef<str>>(
         &mut self,
         command: S,
         param: S,
     ) -> Result<(), ws_client::SendError> {
-        self.send(ws_client::Message::Text(format!(
+        self.unbounded_send(ws_client::SentMessage::Text(format!(
             "{}:{}",
             command.as_ref(),
             param.as_ref()
         )))
-        .await
     }
 }
 
+
 /// Wrap a `ws_client::Receiver` with convenient utility methods for feed connections
 pub struct FeedReceiver(ws_client::Receiver);
 
@@ -146,7 +180,7 @@ impl From<ws_client::Receiver> for FeedReceiver {
 }
 
 impl Stream for FeedReceiver {
-    type Item = Result<ws_client::Message, ws_client::RecvError>;
+    type Item = Result<ws_client::RecvMessage, ws_client::RecvError>;
     fn poll_next(
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
@@ -155,6 +189,18 @@ impl Stream for FeedReceiver {
     }
 }
 
+impl Deref for FeedReceiver {
+    type Target = ws_client::Receiver;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+impl DerefMut for FeedReceiver {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 impl FeedReceiver {
     /// Wait for the next set of feed messages to arrive. Returns an error if the connection
     /// is closed, or the messages that come back cannot be properly decoded.
@@ -170,14 +216,14 @@ impl FeedReceiver {
             .ok_or_else(|| anyhow::anyhow!("Stream closed: no more messages"))??;
 
         match msg {
-            ws_client::Message::Binary(data) => {
+            ws_client::RecvMessage::Binary(data) => {
                 let messages = FeedMessage::from_bytes(&data)?;
                 Ok(messages)
-            }
-            ws_client::Message::Text(text) => {
+            },
+            ws_client::RecvMessage::Text(text) => {
                 let messages = FeedMessage::from_bytes(text.as_bytes())?;
                 Ok(messages)
-            }
+            },
         }
     }
 

backend/test_utils/src/server/server.rs

@@ -146,7 +146,7 @@ impl Server {
 
         // Since we're piping stdout from the child process, we need somewhere for it to go
         // else the process will get stuck when it tries to produce output:
-        utils::drain(child_stdout, tokio::io::sink());
+        utils::drain(child_stdout, tokio::io::stderr());
 
         let shard_uri = format!("http://127.0.0.1:{}/submit", shard_port)
             .parse()
@@ -184,7 +184,7 @@ impl Server {
 
         // Since we're piping stdout from the child process, we need somewhere for it to go
         // else the process will get stuck when it tries to produce output:
-        utils::drain(child_stdout, tokio::io::sink());
+        utils::drain(child_stdout, tokio::io::stderr());
 
         // URI for feeds to connect to the core:
         let feed_uri = format!("http://127.0.0.1:{}/feed", core_port)

backend/test_utils/src/ws_client.rs

@@ -7,28 +7,36 @@ use tokio_util::compat::TokioAsyncReadCompatExt;
 /// Send messages into the connection
 #[derive(Clone)]
 pub struct Sender {
-    inner: mpsc::UnboundedSender<SentMessage>,
+    inner: mpsc::UnboundedSender<SentMessageInternal>,
 }
 
 impl Sender {
     /// Ask the underlying Websocket connection to close.
     pub async fn close(&mut self) -> Result<(), SendError> {
-        self.inner.send(SentMessage::Close).await?;
+        self.inner.send(SentMessageInternal::Close).await?;
         Ok(())
     }
     /// Returns whether this channel is closed.
     pub fn is_closed(&mut self) -> bool {
         self.inner.is_closed()
     }
+    /// 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(SentMessageInternal::Message(msg))
+            .map_err(|e| e.into_send_error())?;
+        Ok(())
+    }
 }
 
 #[derive(thiserror::Error, Debug, Clone)]
 pub enum SendError {
     #[error("Failed to send message: {0}")]
-    ChannelError(#[from] mpsc::SendError),
+    ChannelError(#[from] mpsc::SendError)
 }
 
-impl Sink<Message> for Sender {
+impl Sink<SentMessage> for Sender {
     type Error = SendError;
     fn poll_ready(
         mut self: std::pin::Pin<&mut Self>,
@@ -36,9 +44,9 @@ impl Sink<Message> for Sender {
     ) -> 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: Message) -> Result<(), Self::Error> {
+    fn start_send(mut self: std::pin::Pin<&mut Self>, item: SentMessage) -> Result<(), Self::Error> {
         self.inner
-            .start_send_unpin(SentMessage::Message(item))
+            .start_send_unpin(SentMessageInternal::Message(item))
             .map_err(|e| e.into())
     }
     fn poll_flush(
@@ -57,7 +65,7 @@ impl Sink<Message> for Sender {
 
 /// Receive messages out of a connection
 pub struct Receiver {
-    inner: mpsc::UnboundedReceiver<Result<Message, RecvError>>,
+    inner: mpsc::UnboundedReceiver<Result<RecvMessage, RecvError>>,
 }
 
 #[derive(thiserror::Error, Debug)]
@@ -69,7 +77,7 @@ pub enum RecvError {
 }
 
 impl Stream for Receiver {
-    type Item = Result<Message, RecvError>;
+    type Item = Result<RecvMessage, RecvError>;
     fn poll_next(
         mut self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
@@ -78,15 +86,47 @@ impl Stream for Receiver {
     }
 }
 
-/// A message type that can be sent or received from the connection
-pub enum Message {
+/// A message that can be received from the connection
+#[derive(Debug, Clone)]
+pub enum RecvMessage {
+    /// Send an owned string into the socket.
     Text(String),
+    /// Send owned bytes into the socket.
+    Binary(Vec<u8>),
+}
+
+impl RecvMessage {
+    pub fn len(&self) -> usize {
+        match self {
+            RecvMessage::Binary(b) => b.len(),
+            RecvMessage::Text(s) => s.len(),
+        }
+    }
+}
+
+/// A message that can be sent into the connection
+#[derive(Debug, Clone)]
+pub enum SentMessage {
+    /// Being able to send static text is primarily useful for benchmarking,
+    /// so that we can avoid cloning an owned string and pass a static reference
+    /// (one such option here is using [`Box::leak`] to generate strings with
+    /// static lifetimes).
+    StaticText(&'static str),
+    /// Being able to send static bytes is primarily useful for benchmarking,
+    /// so that we can avoid cloning an owned string and pass a static reference
+    /// (one such option here is using [`Box::leak`] to generate bytes with
+    /// static lifetimes).
+    StaticBinary(&'static [u8]),
+    /// Send an owned string into the socket.
+    Text(String),
+    /// Send owned bytes into the socket.
     Binary(Vec<u8>),
 }
 
 /// Sent messages can be anything publically visible, or a close message.
-enum SentMessage {
-    Message(Message),
+#[derive(Debug, Clone)]
+enum SentMessageInternal {
+    Message(SentMessage),
     Close,
 }
 
@@ -151,9 +191,9 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
             };
 
             let msg = match message_data {
-                soketto::Data::Text(_) => Ok(Message::Binary(data)),
+                soketto::Data::Text(_) => Ok(RecvMessage::Binary(data)),
                 soketto::Data::Binary(_) => String::from_utf8(data)
-                    .map(|s| Message::Text(s))
+                    .map(|s| RecvMessage::Text(s))
                     .map_err(|e| e.into()),
             };
 
@@ -176,7 +216,7 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
     tokio::spawn(async move {
         while let Some(msg) = rx_from_external.next().await {
             match msg {
-                SentMessage::Message(Message::Text(s)) => {
+                SentMessageInternal::Message(SentMessage::Text(s)) => {
                     if let Err(e) = ws_to_connection.send_text_owned(s).await {
                         log::error!(
                             "Shutting down websocket connection: Failed to send text data: {}",
@@ -185,7 +225,7 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
                         break;
                     }
                 }
-                SentMessage::Message(Message::Binary(bytes)) => {
+                SentMessageInternal::Message(SentMessage::Binary(bytes)) => {
                     if let Err(e) = ws_to_connection.send_binary_mut(bytes).await {
                         log::error!(
                             "Shutting down websocket connection: Failed to send binary data: {}",
@@ -193,8 +233,26 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
                         );
                         break;
                     }
+                },
+                SentMessageInternal::Message(SentMessage::StaticText(s)) => {
+                    if let Err(e) = ws_to_connection.send_text(s).await {
+                        log::error!(
+                            "Shutting down websocket connection: Failed to send text data: {}",
+                            e
+                        );
+                        break;
+                    }
                 }
-                SentMessage::Close => {
+                SentMessageInternal::Message(SentMessage::StaticBinary(bytes)) => {
+                    if let Err(e) = ws_to_connection.send_binary(bytes).await {
+                        log::error!(
+                            "Shutting down websocket connection: Failed to send binary data: {}",
+                            e
+                        );
+                        break;
+                    }
+                },
+                SentMessageInternal::Close => {
                     if let Err(e) = ws_to_connection.close().await {
                         log::error!("Error attempting to close connection: {}", e);
                         break;