netspeed 0.1.1

measure tcp throughput
Documentation
use crate::{
    command::{Command, DeclineReason, Operator},
    Result,
};
use anyhow::{anyhow, Context};
use log::{debug, info};
use std::{
    fmt,
    io::{self, Write},
    net::{TcpStream, ToSocketAddrs},
    str::FromStr,
    time::Duration,
};

#[derive(Default, Debug)]
struct Throughput {
    bytes: u64,
    duration: Duration,
}

#[derive(Default, Debug)]
struct NetworkSpec {
    downstream: Throughput,
    upstream: Throughput,
}

pub struct Client {
    operator: Operator,
    spec: NetworkSpec,
}

impl Client {
    pub fn new(addr: impl ToSocketAddrs + fmt::Debug) -> Result<Self> {
        info!("Connecting to {:?}", addr);

        Ok(Self {
            operator: Operator::new(
                TcpStream::connect_timeout(
                    &addr.to_socket_addrs().unwrap().next().unwrap(),
                    Duration::from_secs(3),
                )
                .context(format!("Addr:{:?}", addr))?,
            ),
            spec: NetworkSpec::default(),
        })
    }

    pub fn duration(mut self, duration: Option<&str>) -> Self {
        let duration =
            Duration::from_secs(u64::from_str(duration.unwrap_or("3").as_ref()).unwrap());
        self.spec.downstream.duration = duration;
        self.spec.upstream.duration = duration;
        self
    }

    pub fn run(mut self) -> Result<()> {
        self.check_server_status()
            .and_then(|_| self.ping_pon())
            .and_then(|_| self.downstream())
            .and_then(|_| self.upstream())
            .and_then(|_| self.print_result(io::stdout()))
    }

    fn check_server_status(&mut self) -> Result<()> {
        let cmd = self.operator.read()?;
        match cmd {
            Command::Ready => {
                debug!("Receive server ready");
                Ok(())
            }
            Command::Decline => match self.operator.read_decline_reason()? {
                DeclineReason::Unknown => Err(anyhow!("Server decline speed test :(")),
                DeclineReason::MaxThreadsExceed(max_threads) => Err(anyhow!(
                    "Server decline speed test. Cause: max threads exceeded({})",
                    max_threads
                )),
            },
            _ => Err(anyhow!("Unexpected command {:?}", cmd)),
        }
    }

    fn ping_pon(&mut self) -> Result<()> {
        self.operator.ping_write_then_read().map(|r| {
            debug!("Successfully ping to remote server");
            r
        })
    }

    fn downstream(&mut self) -> Result<()> {
        info!(
            "Start downstream duration: {} seconds",
            self.spec.downstream.duration.as_secs()
        );
        self.operator
            .request_downstream(self.spec.downstream.duration)?;
        self.spec.downstream.bytes = self.operator.read_loop()?;
        Ok(())
    }

    fn upstream(&mut self) -> Result<()> {
        info!(
            "Start upstream duration: {} seconds",
            self.spec.upstream.duration.as_secs()
        );
        self.operator
            .request_upstream(self.spec.upstream.duration)?;
        self.spec.upstream.bytes = self.operator.write_loop(self.spec.upstream.duration)?;
        Ok(())
    }

    fn print_result<W: Write>(&mut self, mut writer: W) -> Result<()> {
        writeln!(
            writer,
            "Downstream: {}",
            self.format_throughput(&self.spec.downstream)
        )
        .and(writeln!(
            writer,
            "  Upstream: {}",
            self.format_throughput(&self.spec.upstream)
        ))
        .map_err(anyhow::Error::from)
    }

    fn format_throughput(&self, throughput: &Throughput) -> String {
        use crate::util::*;
        format_bps(to_bps(throughput.bytes, throughput.duration))
    }
}