From 71afddf9b1dab678fb195a509fa1e1e4de7e03b5 Mon Sep 17 00:00:00 2001
From: Erin Power <erin.power@embark-studios.com>
Date: Sun, 15 Oct 2023 11:15:31 +0200
Subject: [PATCH] a

---
 benches/throughput.rs     | 17 ++++++++++++++++-
 src/utils/net.rs          |  1 -
 tests/local_rate_limit.rs | 12 ++++++++----
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/benches/throughput.rs b/benches/throughput.rs
index 2cd558d259..066dc7ca2a 100644
--- a/benches/throughput.rs
+++ b/benches/throughput.rs
@@ -56,6 +56,9 @@ static THROUGHPUT_SERVER_INIT: Lazy<()> = Lazy::new(|| {
 static FEEDBACK_LOOP: Lazy<()> = Lazy::new(|| {
     std::thread::spawn(|| {
         let socket = UdpSocket::bind(FEEDBACK_LOOP_ADDR).unwrap();
+        socket
+            .set_read_timeout(Some(std::time::Duration::from_millis(500)))
+            .unwrap();
 
         loop {
             let mut packet = [0; MESSAGE_SIZE];
@@ -74,6 +77,9 @@ fn throughput_benchmark(c: &mut Criterion) {
     // Sleep to give the servers some time to warm-up.
     std::thread::sleep(std::time::Duration::from_millis(500));
     let socket = UdpSocket::bind(BENCH_LOOP_ADDR).unwrap();
+    socket
+        .set_read_timeout(Some(std::time::Duration::from_millis(500)))
+        .unwrap();
     let mut packet = [0; MESSAGE_SIZE];
 
     let mut group = c.benchmark_group("throughput");
@@ -125,6 +131,9 @@ fn write_feedback(addr: SocketAddr) -> mpsc::Sender<Vec<u8>> {
     let (write_tx, write_rx) = mpsc::channel::<Vec<u8>>();
     std::thread::spawn(move || {
         let socket = UdpSocket::bind(addr).unwrap();
+        socket
+            .set_read_timeout(Some(std::time::Duration::from_millis(500)))
+            .unwrap();
         let mut packet = [0; MESSAGE_SIZE];
         let (_, source) = socket.recv_from(&mut packet).unwrap();
         while let Ok(packet) = write_rx.recv() {
@@ -142,6 +151,9 @@ fn readwrite_benchmark(c: &mut Criterion) {
     let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>();
     std::thread::spawn(move || {
         let socket = UdpSocket::bind(READ_LOOP_ADDR).unwrap();
+        socket
+            .set_read_timeout(Some(std::time::Duration::from_millis(500)))
+            .unwrap();
         let mut packet = [0; MESSAGE_SIZE];
         loop {
             let (length, _) = socket.recv_from(&mut packet).unwrap();
@@ -164,9 +176,12 @@ fn readwrite_benchmark(c: &mut Criterion) {
     Lazy::force(&WRITE_SERVER_INIT);
 
     // Sleep to give the servers some time to warm-up.
-    std::thread::sleep(std::time::Duration::from_millis(500));
+    std::thread::sleep(std::time::Duration::from_millis(150));
 
     let socket = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
+    socket
+        .set_read_timeout(Some(std::time::Duration::from_millis(500)))
+        .unwrap();
 
     // prime the direct write connection
     socket.send_to(PACKETS[0], direct_write_addr).unwrap();
diff --git a/src/utils/net.rs b/src/utils/net.rs
index 860b993a8e..8d41afe8da 100644
--- a/src/utils/net.rs
+++ b/src/utils/net.rs
@@ -223,4 +223,3 @@ pub fn to_canonical(addr: &mut SocketAddr) {
 
     addr.set_ip(ip);
 }
-
diff --git a/tests/local_rate_limit.rs b/tests/local_rate_limit.rs
index 7b93007643..55068c5ba3 100644
--- a/tests/local_rate_limit.rs
+++ b/tests/local_rate_limit.rs
@@ -35,7 +35,8 @@ period: 1
 ";
     let echo = t.run_echo_server(&AddressType::Random).await;
 
-    let server_addr = available_addr(&AddressType::Random).await;
+    let mut server_addr = available_addr(&AddressType::Random).await;
+    quilkin::test_utils::map_addr_to_localhost(&mut server_addr);
     let server_proxy = quilkin::cli::Proxy {
         port: server_addr.port(),
         ..<_>::default()
@@ -53,20 +54,21 @@ period: 1
         .map(std::sync::Arc::new)
         .unwrap(),
     );
+    tracing::trace!("spawning server");
     t.run_server(server_config, server_proxy, None);
-    tokio::time::sleep(Duration::from_millis(50)).await;
 
     let msg = "hello";
     let (mut rx, socket) = t.open_socket_and_recv_multiple_packets().await;
 
     for _ in 0..3 {
+        tracing::trace!(%server_addr, %msg, "sending");
         socket.send_to(msg.as_bytes(), &server_addr).await.unwrap();
     }
 
     for _ in 0..2 {
         assert_eq!(
             msg,
-            timeout(Duration::from_secs(5), rx.recv())
+            timeout(Duration::from_millis(500), rx.recv())
                 .await
                 .unwrap()
                 .unwrap()
@@ -76,5 +78,7 @@ period: 1
     // Allow enough time to have received any response.
     tokio::time::sleep(Duration::from_millis(100)).await;
     // Check that we do not get any response.
-    assert!(timeout(Duration::from_secs(1), rx.recv()).await.is_err());
+    assert!(timeout(Duration::from_millis(500), rx.recv())
+        .await
+        .is_err());
 }