From 0e61744443bb25f156ec1e316f7a7145f2f955cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Thu, 11 Nov 2021 17:50:52 +0100
Subject: [PATCH] http load test: connection open interval setting, other
 improvements

---
 Cargo.lock                            |  1 +
 aquatic_http_load_test/Cargo.toml     |  1 +
 aquatic_http_load_test/src/config.rs  |  9 ++++-
 aquatic_http_load_test/src/network.rs | 47 +++++++++++++++++++--------
 4 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c1aa85a..a27992b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -139,6 +139,7 @@ dependencies = [
  "futures-lite",
  "glommio",
  "hashbrown 0.11.2",
+ "log",
  "mimalloc",
  "quickcheck",
  "quickcheck_macros",
diff --git a/aquatic_http_load_test/Cargo.toml b/aquatic_http_load_test/Cargo.toml
index a33fdeb..889de16 100644
--- a/aquatic_http_load_test/Cargo.toml
+++ b/aquatic_http_load_test/Cargo.toml
@@ -20,6 +20,7 @@ aquatic_http_protocol = "0.1.0"
 futures-lite = "1"
 hashbrown = "0.11.2"
 glommio = { git = "https://github.com/DataDog/glommio.git", rev = "4e6b14772da2f4325271fbcf12d24cf91ed466e5" }
+log = "0.4"
 mimalloc = { version = "0.1", default-features = false }
 rand = { version = "0.8", features = ["small_rng"] }
 rand_distr = "0.4"
diff --git a/aquatic_http_load_test/src/config.rs b/aquatic_http_load_test/src/config.rs
index 78b3615..3352957 100644
--- a/aquatic_http_load_test/src/config.rs
+++ b/aquatic_http_load_test/src/config.rs
@@ -9,7 +9,13 @@ pub struct Config {
     pub server_address: SocketAddr,
     pub log_level: LogLevel,
     pub num_workers: usize,
+    /// Maximum number of connections to keep open
     pub num_connections: usize,
+    /// How often to check if num_connections connections are open, and
+    /// open a new one otherwise. A value of 0 means that connections are
+    /// opened as quickly as possible, which is useful when the tracker
+    /// doesn't keep connections alive.
+    pub connection_creation_interval_ms: u64,
     pub duration: usize,
     pub torrents: TorrentConfig,
     #[cfg(feature = "cpu-pinning")]
@@ -46,7 +52,8 @@ impl Default for Config {
             server_address: "127.0.0.1:3000".parse().unwrap(),
             log_level: LogLevel::Error,
             num_workers: 1,
-            num_connections: 8,
+            num_connections: 128,
+            connection_creation_interval_ms: 10,
             duration: 0,
             torrents: TorrentConfig::default(),
             #[cfg(feature = "cpu-pinning")]
diff --git a/aquatic_http_load_test/src/network.rs b/aquatic_http_load_test/src/network.rs
index aaa15c0..fdae866 100644
--- a/aquatic_http_load_test/src/network.rs
+++ b/aquatic_http_load_test/src/network.rs
@@ -24,14 +24,36 @@ pub async fn run_socket_thread(
     let config = Rc::new(config);
     let num_active_connections = Rc::new(RefCell::new(0usize));
 
-    TimerActionRepeat::repeat(move || {
-        periodically_open_connections(
-            config.clone(),
-            tls_config.clone(),
-            load_test_state.clone(),
-            num_active_connections.clone(),
-        )
-    });
+    let interval = config.connection_creation_interval_ms;
+
+    if interval == 0 {
+        loop {
+            if *num_active_connections.borrow() < config.num_connections {
+                if let Err(err) = Connection::run(
+                    config.clone(),
+                    tls_config.clone(),
+                    load_test_state.clone(),
+                    num_active_connections.clone(),
+                )
+                .await
+                {
+                    ::log::error!("connection creation error: {:?}", err);
+                }
+            }
+        }
+    } else {
+        let interval = Duration::from_millis(interval);
+
+        TimerActionRepeat::repeat(move || {
+            periodically_open_connections(
+                config.clone(),
+                interval,
+                tls_config.clone(),
+                load_test_state.clone(),
+                num_active_connections.clone(),
+            )
+        });
+    }
 
     futures_lite::future::pending::<bool>().await;
 
@@ -40,6 +62,7 @@ pub async fn run_socket_thread(
 
 async fn periodically_open_connections(
     config: Rc<Config>,
+    interval: Duration,
     tls_config: Arc<rustls::ClientConfig>,
     load_test_state: LoadTestState,
     num_active_connections: Rc<RefCell<usize>>,
@@ -49,13 +72,13 @@ async fn periodically_open_connections(
             if let Err(err) =
                 Connection::run(config, tls_config, load_test_state, num_active_connections).await
             {
-                eprintln!("connection creation error: {:?}", err);
+                ::log::error!("connection creation error: {:?}", err);
             }
         })
         .detach();
     }
 
-    Some(Duration::from_secs(1))
+    Some(interval)
 }
 
 struct Connection {
@@ -97,10 +120,8 @@ impl Connection {
 
         *num_active_connections.borrow_mut() += 1;
 
-        println!("run connection");
-
         if let Err(err) = connection.run_connection_loop().await {
-            eprintln!("connection error: {:?}", err);
+            ::log::info!("connection error: {:?}", err);
         }
 
         *num_active_connections.borrow_mut() -= 1;