harness: improve robustness of tests (#1194)

* improve robustness of `should_update_load_epoch_when_balancer_restart_2_monoliths`

* slightly improve `should_prioritize_same_region_ws` robustness

* refactor harness discoverer to make it easier to debug

* fix harness discoverer sometimes listening on a port that is already being used

* improve some error messages

* use ipv6 localhost to check port availability

* minor refactor

* automatically try to restart the balancer if starting it the first time fails

* make `route_ws_to_correct_monolith_race` much less flaky

* improve reliability of `should_not_unload_rooms_when_balancer_restart`

* format

Author: dyc3

crates/harness-tests/src/routing.rs

@@ -138,6 +138,7 @@ async fn route_ws_to_correct_monolith_race(ctx: &mut TestRunner) {
 
     for i in 0..100 {
         let room_name = format!("foo{}", i);
+        println!("iteration: {}", room_name);
         m.load_room(room_name.clone()).await;
 
         let mut client = Client::new(ctx).unwrap();
@@ -153,19 +154,30 @@ async fn route_ws_to_correct_monolith_race(ctx: &mut TestRunner) {
                 },
                 result = tokio::time::timeout(Duration::from_secs(1), dummy.wait_recv()) => {
                     result.expect("msg recv timeout");
+                    println!("dummy received message");
+                    tokio::time::timeout(Duration::from_millis(100), dummy.wait_recv()).await.expect("dummy never received unload message"); // wait for unload message
                     continue; // because we are waiting for the client to reconnect
                 },
                 _ = client.wait_for_disconnect() => {
-                    println!("client disconnected, retrying");
+                    println!("client disconnected, retrying =====================================");
                     client.join(room_name.clone()).await;
                     continue;
                 }
             };
         }
 
         let recvd = m.collect_recv();
-        assert_eq!(recvd.len(), 1);
-        assert!(matches!(recvd[0], MsgB2M::Join(_)));
+        assert_eq!(
+            recvd.len(),
+            1,
+            "expected exactly one message, got {:?}",
+            recvd
+        );
+        if let MsgB2M::Join(m) = &recvd[0] {
+            assert_eq!(m.room, room_name.into());
+        } else {
+            panic!("expected join message, got {:?}", recvd[0])
+        }
         m.clear_recv();
         dummy.clear_recv();
     }
@@ -237,7 +249,12 @@ async fn unicast_messaging(ctx: &mut TestRunner) {
     let res = vec![c1.recv().await, c2.recv().await];
     let oks: Vec<_> = res.iter().filter(|r| r.is_ok()).collect();
 
-    assert_eq!(oks.len(), 1);
+    assert_eq!(
+        oks.len(),
+        1,
+        "expected exactly one ok result, got these messages: {:?}",
+        oks
+    );
     assert_eq!(oks[0].as_ref().unwrap().to_string(), "{}");
 }
 
@@ -297,7 +314,7 @@ async fn should_prioritize_same_region_ws(ctx: &mut TestRunner) {
     let mut c1 = Client::new(ctx).unwrap();
     c1.join("foo").await;
 
-    tokio::time::timeout(Duration::from_millis(100), m1.wait_recv())
+    tokio::time::timeout(Duration::from_millis(200), m1.wait_recv())
         .await
         .expect("timeout waiting for join message");
 

crates/harness-tests/src/state.rs

@@ -66,15 +66,19 @@ async fn should_not_unload_rooms_when_balancer_restart(ctx: &mut TestRunner) {
 
     m.show().await;
     c1.join("foo").await;
-    m.wait_recv().await;
+    tokio::time::timeout(Duration::from_millis(200), m.wait_recv())
+        .await
+        .expect("m did not receive join message (pre restart)");
     m.clear_recv();
 
     ctx.restart_balancer().await;
 
     m.wait_for_balancer_connect().await;
     c1.disconnect().await;
     c1.join("foo").await;
-    m.wait_recv().await;
+    tokio::time::timeout(Duration::from_millis(200), m.wait_recv())
+        .await
+        .expect("m did not receive join message (post restart)");
     m.gossip().await;
     let _ = tokio::time::timeout(Duration::from_millis(200), m.wait_recv()).await;
 
@@ -103,15 +107,19 @@ async fn should_update_load_epoch_when_balancer_restart_2_monoliths(ctx: &mut Te
 
     m.show().await;
     c1.join("foo").await;
-    m.wait_recv().await;
+    tokio::time::timeout(Duration::from_millis(200), m.wait_recv())
+        .await
+        .expect("m did not receive join message (pre restart)");
     m.clear_recv();
 
     ctx.restart_balancer().await;
+    c1.disconnect().await;
 
     m.wait_for_balancer_connect().await;
-    c1.disconnect().await;
     c1.join("foo").await;
-    m.wait_recv().await;
+    tokio::time::timeout(Duration::from_millis(200), m.wait_recv())
+        .await
+        .expect("m did not receive join message (post restart)");
     m.gossip().await;
     let _ = tokio::time::timeout(Duration::from_millis(200), m.wait_recv()).await;
 
@@ -121,7 +129,9 @@ async fn should_update_load_epoch_when_balancer_restart_2_monoliths(ctx: &mut Te
         .await;
     m2.show().await;
     m2.load_room("foo").await;
-    m2.wait_recv().await;
+    tokio::time::timeout(Duration::from_millis(200), m2.wait_recv())
+        .await
+        .expect("m2 did not receive any message");
 
     let recv = m2.collect_recv();
     for msg in recv {

crates/harness/src/test_runner.rs

@@ -49,11 +49,13 @@ impl TestRunner {
             warn!("restart_balancer: Failed to kill balancer: {:?}", result);
         }
 
-        self.child = Self::spawn_balancer(&self.spawn_options).await;
+        self.child = Self::spawn_balancer(&self.spawn_options)
+            .await
+            .expect("failed to respawn balancer");
     }
 
     /// Spawn a new balancer and wait for it to be ready.
-    async fn spawn_balancer(opts: &BalancerSpawnOptions) -> Child {
+    async fn spawn_balancer(opts: &BalancerSpawnOptions) -> anyhow::Result<Child> {
         let mut envs: HashMap<_, _> = HashMap::from_iter([
             ("BALANCER_PORT", format!("{}", opts.port)),
             (
@@ -65,7 +67,12 @@ impl TestRunner {
             envs.insert("BALANCER_REGION", region.clone());
         }
 
-        let child = Command::new("cargo")
+        let client = reqwest::Client::builder()
+            .timeout(std::time::Duration::from_millis(100))
+            .build()
+            .expect("failed to build request client");
+
+        let mut child = Command::new("cargo")
             .args([
                 "run",
                 "-p",
@@ -77,12 +84,8 @@ impl TestRunner {
             .envs(envs)
             .spawn()
             .expect("Failed to start balancer");
-
         loop {
-            let client = reqwest::Client::builder()
-                .timeout(std::time::Duration::from_millis(100))
-                .build()
-                .expect("failed to build request client");
+            println!("waiting for balancer to start");
             match client
                 .get(&format!("http://localhost:{}/api/status", opts.port))
                 .send()
@@ -94,12 +97,19 @@ impl TestRunner {
                     }
                 }
                 Err(_) => {
+                    match child.try_wait() {
+                        Ok(Some(status)) => {
+                            anyhow::bail!("Exited with status {}", status);
+                        }
+                        Ok(None) => {} // process is still running
+                        Err(e) => anyhow::bail!("Error waiting for balancer to start: {}", e),
+                    }
                     tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
                 }
             }
         }
 
-        child
+        Ok(child)
     }
 
     /// Create a URL that points to the balancer. This creates URLs that clients should connect to when making HTTP requests.
@@ -117,22 +127,38 @@ impl AsyncTestContext for TestRunner {
     async fn setup() -> Self {
         let mut port;
         let mut harness_port;
+
+        let mut opts;
+        let child;
+        let mut attempts = 0;
         loop {
+            if attempts > 5 {
+                panic!("Failed to find an unused port after 5 attempts");
+            }
             port = random_unused_port();
             harness_port = random_unused_port();
             // Ensure that the harness port is different from the balancer port.
-            if port != harness_port {
-                break;
+            if port == harness_port {
+                continue;
             }
-        }
-
-        let opts = BalancerSpawnOptions {
-            port,
-            harness_port,
-            region: None,
-        };
+            attempts += 1;
+
+            opts = BalancerSpawnOptions {
+                port,
+                harness_port,
+                region: None,
+            };
+
+            child = match Self::spawn_balancer(&opts).await {
+                Ok(child) => child,
+                Err(e) => {
+                    println!("Failed to spawn balancer: {}", e);
+                    continue;
+                }
+            };
 
-        let child = Self::spawn_balancer(&opts).await;
+            break;
+        }
 
         let (_provider_task, monolith_add_tx, monolith_remove_tx) =
             crate::provider::DiscoveryProvider::connect(harness_port).await;

crates/harness/src/util.rs

@@ -11,5 +11,5 @@ pub(crate) fn random_unused_port() -> u16 {
 }
 
 pub(crate) fn port_is_available(port: u16) -> bool {
-    std::net::TcpListener::bind(("localhost", port)).is_ok()
+    std::net::TcpListener::bind(("::1", port)).is_ok()
 }