Add/test max_nodes_per_connection feature

backend/test_utils/src/server/channels.rs

@@ -208,12 +208,13 @@ impl FeedReceiver {
     /// Prefer [`FeedReceiver::recv_feed_messages`]; tests should generally be
     /// robust in assuming that messages may not all be delivered at once (unless we are
     /// specifically testing which messages are buffered together).
-    pub async fn recv_feed_messages_once(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> {
-        let msg = self
-            .0
-            .next()
-            .await
-            .ok_or_else(|| anyhow::anyhow!("Stream closed: no more messages"))??;
+    pub async fn recv_feed_messages_once_timeout(&mut self, timeout: Duration) -> Result<Vec<FeedMessage>, anyhow::Error> {
+        let msg = match tokio::time::timeout(timeout, self.0.next()).await {
+            // Timeout elapsed; no messages back:
+            Err(_) => return Ok(Vec::new()),
+            // Something back; Complain if error no stream closed:
+            Ok(res) => res.ok_or_else(|| anyhow::anyhow!("Stream closed: no more messages"))??
+        };
 
         match msg {
             ws_client::RecvMessage::Binary(data) => {
@@ -227,11 +228,26 @@ impl FeedReceiver {
         }
     }
 
+    /// Wait for the next set of feed messages to arrive.
+    /// See `recv_feed_messages_once_timeout`.
+    pub async fn recv_feed_messages_once(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> {
+        // Default to a timeout of 30 seconds, meaning that the test will eventually end,
+        self.recv_feed_messages_once_timeout(Duration::from_secs(30)).await
+    }
+
     /// Wait for feed messages to be sent back, building up a list of output messages until
     /// the channel goes quiet for a short while.
-    pub async fn recv_feed_messages(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> {
+    ///
+    /// If no new messages are received within the timeout given, bail with whatever we have so far.
+    /// This differs from `recv_feed_messages` and `recv_feed_messages_once`, which will block indefinitely
+    /// waiting for something to arrive
+    pub async fn recv_feed_messages_timeout(&mut self, timeout: Duration) -> Result<Vec<FeedMessage>, anyhow::Error> {
         // Block as long as needed for messages to start coming in:
-        let mut feed_messages = self.recv_feed_messages_once().await?;
+        let mut feed_messages = match tokio::time::timeout(timeout, self.recv_feed_messages_once()).await {
+            Ok(msgs) => msgs?,
+            Err(_) => return Ok(Vec::new()),
+        };
+
         // Then, loop a little to make sure we catch any additional messages that are sent soon after:
         loop {
             match tokio::time::timeout(Duration::from_millis(250), self.recv_feed_messages_once())
@@ -250,4 +266,11 @@ impl FeedReceiver {
             }
         }
     }
+
+    /// Wait for feed messages until nothing else arrives in a timely fashion.
+    /// See `recv_feed_messages_timeout`.
+    pub async fn recv_feed_messages(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> {
+        // Default to a timeout of 30 seconds, meaning that the test will eventually end,
+        self.recv_feed_messages_timeout(Duration::from_secs(30)).await
+    }
 }

backend/test_utils/src/workspace/start_server.rs

@@ -1,7 +1,7 @@
 use super::commands;
 use crate::server::{self, Server, Command};
 
-/// Additional options to pass to the feed command.
+/// Additional options to pass to the core command.
 pub struct CoreOpts {
     pub feed_timeout: Option<u64>
 }
@@ -14,6 +14,19 @@ impl Default for CoreOpts {
     }
 }
 
+/// Additional options to pass to the shard command.
+pub struct ShardOpts {
+    pub max_nodes_per_connection: Option<usize>
+}
+
+impl Default for ShardOpts {
+    fn default() -> Self {
+        Self {
+            max_nodes_per_connection: None
+        }
+    }
+}
+
 /// Start a telemetry server. We'll use `cargo run` by default, but you can also provide
 /// env vars to configure the binary that runs for the shard and core process. Either:
 ///
@@ -31,7 +44,7 @@ impl Default for CoreOpts {
 /// - `TELEMETRY_SUBMIT_HOSTS` - hosts (comma separated) to connect to for telemetry `/submit`s.
 /// - `TELEMETRY_FEED_HOST` - host to connect to for feeds (eg 127.0.0.1:3000)
 ///
-pub async fn start_server(release_mode: bool, core_opts: CoreOpts) -> Server {
+pub async fn start_server(release_mode: bool, core_opts: CoreOpts, shard_opts: ShardOpts) -> Server {
     // Start to a single process:
     if let Ok(bin) = std::env::var("TELEMETRY_BIN") {
         return Server::start(server::StartOpts::SingleProcess {
@@ -52,10 +65,17 @@ pub async fn start_server(release_mode: bool, core_opts: CoreOpts) -> Server {
     }
 
     // Build the shard command
-    let shard_command = std::env::var("TELEMETRY_SHARD_BIN")
+    let mut shard_command = std::env::var("TELEMETRY_SHARD_BIN")
         .map(|val| Command::new(val))
         .unwrap_or_else(|_| commands::cargo_run_telemetry_shard(release_mode).expect("must be in rust workspace to run shard command"));
 
+    // Append additional opts to the shard command
+    if let Some(max_nodes_per_connection) = shard_opts.max_nodes_per_connection {
+        shard_command = shard_command
+            .arg("--max-nodes-per-connection")
+            .arg(max_nodes_per_connection.to_string());
+    }
+
     // Build the core command
     let mut core_command = std::env::var("TELEMETRY_CORE_BIN")
         .map(|val| Command::new(val))
@@ -77,10 +97,10 @@ pub async fn start_server(release_mode: bool, core_opts: CoreOpts) -> Server {
 
 /// Start a telemetry core server in debug mode. see [`start_server`] for details.
 pub async fn start_server_debug() -> Server {
-    start_server(false, CoreOpts::default()).await
+    start_server(false, CoreOpts::default(), ShardOpts::default()).await
 }
 
 /// Start a telemetry core server in release mode. see [`start_server`] for details.
 pub async fn start_server_release() -> Server {
-    start_server(true, CoreOpts::default()).await
+    start_server(true, CoreOpts::default(), ShardOpts::default()).await
 }