Merge pull request hatoo#695 from thomasgl-orange/randomize-connect-to

Randomize --connect-to if multiple matching options

README.md

@@ -192,6 +192,7 @@ Options:
           Accept invalid certs.
       --connect-to <CONNECT_TO>
           Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443'
+          Note: if used several times for the same host:port:target_host:target_port, a random choice is made
       --disable-color
           Disable the color scheme.
       --unix-socket <UNIX_SOCKET>

src/client.rs

@@ -79,11 +79,14 @@ impl Dns {
             .port_or_known_default()
             .ok_or(ClientError::PortNotFound)?;
 
-        // Try to find an override (passed via `--connect-to`) that applies to this (host, port)
+        // Try to find an override (passed via `--connect-to`) that applies to this (host, port),
+        // choosing one randomly if several match.
         let (host, port) = if let Some(entry) = self
             .connect_to
             .iter()
-            .find(|entry| entry.requested_port == port && entry.requested_host == host)
+            .filter(|entry| entry.requested_port == port && entry.requested_host == host)
+            .collect::<Vec<_>>()
+            .choose(rng)
         {
             (entry.target_host.as_str(), entry.target_port)
         } else {

src/main.rs

@@ -236,7 +236,8 @@ Note: If qps is specified, burst will be ignored",
     #[arg(help = "Accept invalid certs.", long = "insecure")]
     insecure: bool,
     #[arg(
-        help = "Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443'",
+        help = "Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443'
+Note: if used several times for the same host:port:target_host:target_port, a random choice is made",
         long = "connect-to"
     )]
     connect_to: Vec<ConnectToEntry>,

tests/tests.rs

@@ -302,6 +302,62 @@ async fn burst_10_req_delay_2s_rate_4(iteration: u8, args: &[&str]) -> usize {
     count
 }
 
+// Randomly spread 100 requests on two matching --connect-to targets, and return a count for each
+async fn distribution_on_two_matching_connect_to(host: &'static str) -> (i32, i32) {
+    let (tx1, rx1) = flume::unbounded();
+    let (tx2, rx2) = flume::unbounded();
+
+    let app1 = Router::new().route(
+        "/",
+        get(move || async move {
+            tx1.send(()).unwrap();
+            "Success1"
+        }),
+    );
+
+    let app2 = Router::new().route(
+        "/",
+        get(move || async move {
+            tx2.send(()).unwrap();
+            "Success2"
+        }),
+    );
+
+    let (listener1, port1) = bind_port().await;
+    tokio::spawn(async { axum::serve(listener1, app1).await });
+
+    let (listener2, port2) = bind_port().await;
+    tokio::spawn(async { axum::serve(listener2, app2).await });
+
+    tokio::task::spawn_blocking(move || {
+        Command::cargo_bin("oha")
+            .unwrap()
+            .args(["-n", "100", "--no-tui"])
+            .arg(format!("http://{host}/"))
+            .arg("--connect-to")
+            .arg(format!("{host}:80:localhost:{port1}"))
+            .arg("--connect-to")
+            .arg(format!("{host}:80:localhost:{port2}"))
+            .assert()
+            .success();
+    })
+    .await
+    .unwrap();
+
+    let mut count1 = 0;
+    let mut count2 = 0;
+    loop {
+        if rx1.try_recv().is_ok() {
+            count1 += 1;
+        } else if rx2.try_recv().is_ok() {
+            count2 += 1;
+        } else {
+            break;
+        }
+    }
+    (count1, count2)
+}
+
 #[tokio::test]
 async fn test_enable_compression_default() {
     let req = get_req("/", &[]).await;
@@ -596,6 +652,13 @@ async fn test_connect_to() {
     )
 }
 
+#[tokio::test]
+async fn test_connect_to_randomness() {
+    let (count1, count2) = distribution_on_two_matching_connect_to("invalid.example.org").await;
+    assert!(count1 >= 10 && count2 >= 10); // should not be too flaky with 100 coin tosses
+    assert!(count1 + count2 == 100);
+}
+
 #[tokio::test]
 async fn test_connect_to_ipv6_target() {
     assert_eq!(