Merge pull request #1 from ssrlive/mio

Mio

Author: ssrlive

.github/workflows/rust.yml

@@ -13,7 +13,7 @@ jobs:
         host_os:
           - ubuntu-latest
           - macos-latest
-          # - windows-latest
+          - windows-latest
 
     runs-on: ${{ matrix.host_os }}
 

Cargo.toml

@@ -11,17 +11,19 @@ categories = ["network-utilities"]
 readme = "README.md"
 
 [dependencies]
+cfg-if = "1.0"
 chrono = "0.4"
 clap = { version = "4.4", features = ["derive", "wrap_help"] }
 core_affinity = "0.8"
-ctrlc2 = "3.5"
+ctrlc2 = { version = "3.5", features = ["termination"] }
 env_logger = "0.10"
 log = { version = "0.4", features = ["std"] }
-mio = "0.6"
-nix = "0.20"
+mio = { version = "0.8", features = ["log", "os-poll", "net"] }
+nix = { version = "0.27", features = ["net"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 simple-error = "0.3"
+socket2 = { version = "0.5", features = ["all"] }
 uuid = { version = "1.6", features = ["v4"] }
 
 #configuration for cargo-deb

src/args.rs

@@ -33,49 +33,57 @@ pub struct Args {
 
     /// enable IPv6 on the server (on most hosts, this will allow both IPv4 and IPv6,
     /// but it might limit to just IPv6 on some)
-    #[arg(short = '6', long)]
+    #[arg(short = '6', long, conflicts_with = "client")]
     pub version6: bool,
 
     /// limit the number of concurrent clients that can be processed by a server;
     /// any over this count will be immediately disconnected
-    #[arg(long, value_name = "number", default_value = "0")]
+    #[arg(long, value_name = "number", default_value = "0", conflicts_with = "client")]
     pub client_limit: usize,
 
     /// run in client mode; value is the server's address
     #[arg(short, long, value_name = "host", conflicts_with = "server")]
     pub client: Option<std::net::IpAddr>,
 
     /// run in reverse-mode (server sends, client receives)
-    #[arg(short = 'R', long)]
+    #[arg(short = 'R', long, conflicts_with = "server")]
     pub reverse: bool,
 
     /// the format in which to deplay information (json, megabit/sec, megabyte/sec)
-    #[arg(short, long, value_enum, value_name = "format", default_value = "megabit")]
+    #[arg(
+        short,
+        long,
+        value_enum,
+        value_name = "format",
+        default_value = "megabit",
+        conflicts_with = "server"
+    )]
     pub format: Format,
 
     /// use UDP rather than TCP
-    #[arg(short, long)]
+    #[arg(short, long, conflicts_with = "server")]
     pub udp: bool,
 
     /// target bandwidth in bytes/sec; this value is applied to each stream,
     /// with a default target of 1 megabit/second for all protocols (note: megabit, not mebibit);
     /// the suffixes kKmMgG can also be used for xbit and xbyte, respectively
-    #[arg(short, long, default_value = "125000", value_name = "bytes/sec")]
+    #[arg(short, long, default_value = "125000", value_name = "bytes/sec", conflicts_with = "server")]
     pub bandwidth: String,
 
     /// the time in seconds for which to transmit
-    #[arg(short, long, default_value = "10.0", value_name = "seconds")]
+    #[arg(short, long, default_value = "10.0", value_name = "seconds", conflicts_with = "server")]
     pub time: f64,
 
     /// the interval at which to send batches of data, in seconds, between [0.0 and 1.0);
     /// this is used to evenly spread packets out over time
-    #[arg(long, default_value = "0.05", value_name = "seconds")]
+    #[arg(long, default_value = "0.05", value_name = "seconds", conflicts_with = "server")]
     pub send_interval: f64,
 
     /// length of the buffer to exchange; for TCP, this defaults to 32 kibibytes; for UDP, it's 1024 bytes
     #[arg(
         short,
         long,
+        conflicts_with = "server",
         default_value = "32768",
         default_value_if("udp", "true", Some("1024")),
         value_name = "bytes"
@@ -85,27 +93,27 @@ pub struct Args {
     /// send buffer, in bytes (only supported on some platforms;
     /// if set too small, a 'resource unavailable' error may occur;
     /// affects TCP window-size)
-    #[arg(long, default_value = "0", value_name = "bytes")]
+    #[arg(long, default_value = "0", value_name = "bytes", conflicts_with = "server")]
     pub send_buffer: usize,
 
     /// receive buffer, in bytes (only supported on some platforms;
     /// if set too small, a 'resource unavailable' error may occur; affects TCP window-size)
-    #[arg(long, default_value = "0", value_name = "bytes")]
+    #[arg(long, default_value = "0", value_name = "bytes", conflicts_with = "server")]
     pub receive_buffer: usize,
 
     /// the number of parallel data-streams to use
-    #[arg(short = 'P', long, value_name = "number", default_value = "1")]
+    #[arg(short = 'P', long, value_name = "number", default_value = "1", conflicts_with = "server")]
     pub parallel: usize,
 
     /// omit a number of seconds from the start of calculations,
     /// primarily to avoid including TCP ramp-up in averages;
     /// using this option may result in disagreement between bytes sent and received,
     /// since data can be in-flight across time-boundaries
-    #[arg(short, long, default_value = "0", value_name = "seconds")]
+    #[arg(short, long, default_value = "0", value_name = "seconds", conflicts_with = "server")]
     pub omit: usize,
 
     /// use no-delay mode for TCP tests, disabling Nagle's Algorithm
-    #[arg(short = 'N', long)]
+    #[arg(short = 'N', long, conflicts_with = "server")]
     pub no_delay: bool,
 
     /// an optional pool of IPv4 TCP ports over which data will be accepted;
@@ -127,6 +135,10 @@ pub struct Args {
     /// if omitted, any OS-assignable port is used; format: 1-10,19,21
     #[arg(long, value_name = "ports", default_value = "")]
     pub udp6_port_pool: String,
+
+    /// Verbosity level
+    #[arg(short, long, value_name = "level", value_enum, default_value = "info")]
+    pub verbosity: ArgVerbosity,
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, clap::ValueEnum, Default)]
@@ -146,3 +158,27 @@ impl std::fmt::Display for Format {
         }
     }
 }
+
+#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, clap::ValueEnum)]
+pub enum ArgVerbosity {
+    Off,
+    Error,
+    Warn,
+    #[default]
+    Info,
+    Debug,
+    Trace,
+}
+
+impl std::fmt::Display for ArgVerbosity {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ArgVerbosity::Off => write!(f, "off"),
+            ArgVerbosity::Error => write!(f, "error"),
+            ArgVerbosity::Warn => write!(f, "warn"),
+            ArgVerbosity::Info => write!(f, "info"),
+            ArgVerbosity::Debug => write!(f, "debug"),
+            ArgVerbosity::Trace => write!(f, "trace"),
+        }
+    }
+}

src/client.rs

@@ -27,10 +27,10 @@ use crate::{
         results::{IntervalResultBox, IntervalResultKind, TcpTestResults, TestResults, UdpTestResults},
     },
     stream::{tcp, udp, TestStream},
+    BoxResult,
 };
 use mio::net::TcpStream;
 use std::{
-    error::Error,
     net::{IpAddr, Shutdown, ToSocketAddrs},
     sync::{
         atomic::{AtomicBool, Ordering},
@@ -41,8 +41,6 @@ use std::{
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
-type BoxResult<T> = Result<T, Box<dyn Error>>;
-
 /// when false, the system is shutting down
 static ALIVE: AtomicBool = AtomicBool::new(true);
 
@@ -60,23 +58,30 @@ fn connect_to_server(address: &str, port: &u16) -> BoxResult<TcpStream> {
     if server_addr.is_none() {
         return Err(Box::new(simple_error::simple_error!("unable to resolve {}", address)));
     }
-    let raw_stream = match std::net::TcpStream::connect_timeout(&server_addr.unwrap(), CONNECT_TIMEOUT) {
+    let stream = match std::net::TcpStream::connect_timeout(&server_addr.unwrap(), CONNECT_TIMEOUT) {
         Ok(s) => s,
         Err(e) => return Err(Box::new(simple_error::simple_error!("unable to connect: {}", e))),
     };
-    let stream = match TcpStream::from_stream(raw_stream) {
-        Ok(s) => s,
-        Err(e) => {
-            return Err(Box::new(simple_error::simple_error!(
-                "unable to prepare TCP control-channel: {}",
-                e
-            )))
+
+    let stream = {
+        let socket: socket2::Socket = socket2::Socket::from(stream);
+        let keepalive = socket2::TcpKeepalive::new()
+            .with_time(KEEPALIVE_DURATION)
+            .with_interval(KEEPALIVE_DURATION);
+        cfg_if::cfg_if! {
+            if #[cfg(unix)] {
+                let keepalive = keepalive.with_retries(4);
+            }
         }
+        socket.set_tcp_keepalive(&keepalive)?;
+        socket.set_nodelay(true)?;
+
+        let stream: std::net::TcpStream = socket.into();
+
+        TcpStream::from_std(stream)
     };
-    log::info!("connected to server");
 
-    stream.set_nodelay(true).expect("cannot disable Nagle's algorithm");
-    stream.set_keepalive(Some(KEEPALIVE_DURATION)).expect("unable to set TCP keepalive");
+    log::info!("connected to server");
 
     Ok(stream)
 }

src/lib.rs

@@ -0,0 +1,14 @@
+pub mod args;
+pub mod client;
+pub(crate) mod protocol;
+pub mod server;
+pub(crate) mod stream;
+pub(crate) mod utils;
+
+pub type BoxResult<T> = Result<T, Box<dyn std::error::Error + Send + Sync + 'static>>;
+
+/// a global token generator
+pub(crate) fn get_global_token() -> mio::Token {
+    mio::Token(TOKEN_SEED.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1)
+}
+static TOKEN_SEED: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);

src/main.rs

@@ -18,42 +18,34 @@
  * along with rperf.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-mod args;
-mod client;
-mod protocol;
-mod server;
-mod stream;
-mod utils;
+use rperf::{args, client, server, BoxResult};
 
-fn main() {
+fn main() -> BoxResult<()> {
     use clap::Parser;
     let args = args::Args::parse();
 
-    let mut env = env_logger::Env::default().filter_or("RUST_LOG", "info");
+    let default = args.verbosity.to_string();
+    let mut env = env_logger::Env::default().filter_or("RUST_LOG", &default);
     if args.debug {
         env = env.filter_or("RUST_LOG", "debug");
     }
     env_logger::init_from_env(env);
 
     if args.server {
         log::debug!("registering SIGINT handler...");
-        ctrlc2::set_handler(move || {
+        let exiting = ctrlc2::set_handler(move || {
             if server::kill() {
                 log::warn!("shutdown requested; please allow a moment for any in-progress tests to stop");
             } else {
                 log::warn!("forcing shutdown immediately");
                 std::process::exit(3);
             }
             true
-        })
-        .expect("unable to set SIGINT handler");
+        })?;
 
         log::debug!("beginning normal operation...");
-        let service = server::serve(&args);
-        if service.is_err() {
-            log::error!("unable to run server: {}", service.unwrap_err());
-            std::process::exit(4);
-        }
+        server::serve(&args)?;
+        exiting.join().expect("unable to join SIGINT handler thread");
     } else if args.client.is_some() {
         log::debug!("registering SIGINT handler...");
         ctrlc2::set_handler(move || {
@@ -64,19 +56,14 @@ fn main() {
                 std::process::exit(3);
             }
             true
-        })
-        .expect("unable to set SIGINT handler");
+        })?;
 
         log::debug!("connecting to server...");
-        let execution = client::execute(&args);
-        if execution.is_err() {
-            log::error!("unable to run client: {}", execution.unwrap_err());
-            std::process::exit(4);
-        }
+        client::execute(&args)?;
     } else {
         use clap::CommandFactory;
         let mut cmd = args::Args::command();
         cmd.print_help().unwrap();
-        std::process::exit(2);
     }
+    Ok(())
 }