Rust speedtest logger

main.rs

/**
 * author: Florian Staudacher, 2016
 *
 * -----------------------------------------------------------------------------
 * "THE BEER-WARE LICENSE" (Revision 42):
 * <[email protected]> wrote this file. As long as you retain this
 * notice you can do whatever you want with this stuff. If we meet some day, and
 * you think this stuff is worth it, you can buy me a beer in return.
 * -----------------------------------------------------------------------------
 *
 * This program performs periodic speed tests against the closest speedtest.net
 * server and will give you the raw data. STDOUT is human-readable, but you can
 * produce a CSV file for performing further analysis with your results.
 *
 * Memory requirements should be minimal, since everything constructed during a
 * speed measurement will be destructed once it's done (thanks, Rust!).
 * You should be able to let this script run in the background without any leaks.
 * The default interval is set to 6 minutes, so you'll get 10 measurements in
 * an hour. The program tries to saturate your bandwith by using multiple
 * concurrent threads (safely ... thanks, Rust!) and down-/uploading with more
 * than one connection at the same time.
 *
 * invocation:
 *   ./speedtest
 *     This will output everything on STDOUT
 *
 *   ./speedtest log.csv
 *     This will still be quite noisy on STDOUT, but produce a
 *     CSV file for analysis
 *
 *  Disclaimer:
 *    The code is heavily based on the Python version
 *    https://github.com/sivel/speedtest-cli/
 *    and it is meant as a "toy project" for getting myself familiarized with
 *    Rust. It may or may not work for your purposes.
 *
 */

extern crate hyper;
extern crate xml;
extern crate time;
extern crate chrono;
extern crate threadpool;

use hyper::client;
use hyper::header::{UserAgent, ContentLength};
use hyper::status::StatusCode;
use xml::reader::{EventReader, XmlEvent};
use time::PreciseTime;
use std::net::ToSocketAddrs;
use std::env;
use std::fs::OpenOptions;
use std::io::{Read, Write, BufWriter, stdout};
use std::cmp::Ordering::Equal;
use std::fmt;
use std::path::Path;
use threadpool::ThreadPool;
use std::sync::mpsc;
use std::sync::Arc;
use chrono::{DateTime,UTC};

#[derive(Default,Debug)]
struct Client {
  ip: String,
  isp: String,
  position: Position,
}

#[derive(Default,Debug)]
struct Server {
  id: i32,
  name: String,
  country: String,
  cc: String,
  sponsor: String,
  host: String,
  url: String,
  url2: String,
  position: Position,
}

impl fmt::Display for Server {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "[{cc}:{id}] '{name}' by {sponsor}: {host}",
      name = self.name, host = self.host,
      sponsor = self.sponsor,
      cc = self.cc, id = self.id)
  }
}

impl Server {
  fn distance_to(&self, other: &Position) -> f64 {
    position_distance(self.position, *other)
  }
}

#[derive(Default,Debug,Copy,Clone)]
struct Position {
  lat: f32,
  lon: f32,
}

type TimeStamp = DateTime<UTC>;
type Bytes = u64;
type Milliseconds = u64;

#[derive(Debug)]
struct SpeedMeasurement {
  ts: TimeStamp,
  client_ip: String,
  client_isp: String,
  server_id: String,
  server_sponsor: String,
  server_host: String,
  server_ip: String,
  server_distance: f64,
  server_rtt: u64,
  download_bytes: Bytes,
  download_ms: Milliseconds,
  upload_bytes: Bytes,
  upload_ms: Milliseconds,
}

const USER_AGENT: &'static str = "Mozilla/5.0 (Linux; U; x64: en-us) Rust/1.10.0-beta.2 (KHTML, like Gecko) speedtest.rs/0.0.1-test";

fn parse_client_cfg(client: &mut Client, http_client: &client::Client) {
  let cfg_resp = http_client.get("http://www.speedtest.net/speedtest-config.php")
                   .header(UserAgent(USER_AGENT.to_owned()))
                   .send().unwrap();
  let parser = EventReader::new(cfg_resp);
  for event in parser {
    match event {
      Ok(XmlEvent::StartElement { name, attributes, .. }) => {
        if name.local_name == "client" {
          for attr in attributes {
            match &*attr.name.local_name {
              "ip" => client.ip = attr.value,
              "isp" => client.isp = attr.value,
              "lat" => client.position.lat = attr.value.to_string().parse::<f32>().unwrap(),
              "lon" => client.position.lon = attr.value.to_string().parse::<f32>().unwrap(),
              _ => {}
            }
          }
        }
      }
      Err(e) => {
        println!("Error: {}", e);
        break;
      }
      _ => {}
    }
  }
}

fn parse_server_list(servers: &mut Vec<Server>, http_client: &client::Client) {
  let urls = vec!(
    "http://www.speedtest.net/speedtest-servers-static.php",
    "http://c.speedtest.net/speedtest-servers-static.php",
    "http://www.speedtest.net/speedtest-servers.php",
    "http://c.speedtest.net/speedtest-servers-static.php"
  );

  let conn_test = urls.iter().find(|&&url| {
    print!("\n  trying {}...", url);
    match http_client.head(url)
      .header(UserAgent(USER_AGENT.to_owned()))
      .send() {
      Ok(resp) => {
        if **resp.headers.get::<ContentLength>().unwrap() == 0 {
          false
        } else {
          true
        }
      },
      Err(_) => false
    }
  });

  let url = match conn_test {
    Some(server) => *server,
    None => {
      //writeln!(std::io::stderr(), "server list could not be retrieved!");
      return
    }
  };

  let srv_resp = http_client.get(url)
                   .header(UserAgent(USER_AGENT.to_owned()))
                   .send().unwrap();

  let parser = EventReader::new(srv_resp);
  for event in parser {
    match event {
      Ok(XmlEvent::StartElement { name, attributes, .. }) => {
        if name.local_name == "server" {
          let mut server = Server::default();
          for attr in attributes {
            match &*attr.name.local_name {
              "id"      => server.id = attr.value.to_string().parse::<i32>().unwrap(),
              "name"    => server.name = attr.value,
              "country" => server.country = attr.value,
              "cc"      => server.cc = attr.value,
              "sponsor" => server.sponsor = attr.value,
              "host"    => server.host = attr.value,
              "url"     => server.url = attr.value,
              "url2"    => server.url2 = attr.value,
              "lat"     => server.position.lat = attr.value.to_string().parse::<f32>().unwrap(),
              "lon"     => server.position.lon = attr.value.to_string().parse::<f32>().unwrap(),
              _         => {}
            }
          }
          servers.push(server);
        }
      }
      Err(e) => {
        println!("Error: {}", e);
        break;
      }
      _ => {}
    }
  }
}

fn position_distance(p1: Position, p2: Position) -> f64 {
  let r = 6371.0; // km
  let dlat = (p2.lat - p1.lat).to_radians();
  let dlon = (p2.lon - p1.lon).to_radians();

  let a = (dlat / 2.0).sin() * (dlat / 2.0).sin() +
          (p1.lat).to_radians().cos() *
          (p2.lat).to_radians().cos() * (dlon / 2.0).sin() *
          (dlon / 2.0).sin();
  let c = 2.0 * a.sqrt().atan2((1.0-a).sqrt());

  r * c as f64
}

fn lowest_latency(servers: &mut [Server], http_client: &client::Client) -> (i32, Milliseconds) {
  let mut recommend_id = 0;
  let mut best_latency: u64 = 3600;

  for server in servers {
    let passes = 3;
    let mut latency: u64 = 0;
    let chk_url = Path::new(&server.url).parent().unwrap().join("latency.txt");

    for _ in 0..passes {
      let start = PreciseTime::now();
      let mut resp = http_client.get(chk_url.to_str().unwrap())
                      .header(UserAgent(USER_AGENT.to_owned()))
                      .send().unwrap();
      let duration = start.to(PreciseTime::now()).num_milliseconds() as u64;

      if resp.status != StatusCode::Ok {
        latency += 3600;
        continue;
      }

      let mut body = [0u8; 9];
      resp.read_exact(&mut body).unwrap();

      let body: Vec<u8> = body.to_vec();
      let body = String::from_utf8(body).unwrap();

      if body == "test=test" {
        latency += duration;
      } else {
        latency += 3600;
      }
    }

    let latency = latency / 3;
    if latency < best_latency {
      recommend_id = server.id;
      best_latency = latency;
    }
  }

  (recommend_id, best_latency)
}

fn measure_download(server: &Server, http_client: &Arc<client::Client>) -> (Bytes, Milliseconds) {
  let cconnections = 5;
  let thread_pool = ThreadPool::new(cconnections);
//   let sizes = vec!(350, 500, 750, 1000, 1500, 2000, 2500, 3000, 3500, 4000);
  let img_sizes = vec!(350, 500, 750, 1000, 1500, 2000);
  let sizes: Vec<_> = img_sizes.iter().cycle().take(12).collect();
  let (tchan_tx, tchan_rx) = mpsc::channel();
  let start = PreciseTime::now();

  for s in sizes {
    let down_url = Path::new(&server.url).parent().unwrap().join(format!("random{0}x{0}.jpg", s));
    let thread_client = http_client.clone();
    let thread_chan = tchan_tx.clone();

    thread_pool.execute(move || {
      print!(" .");
      std::io::stdout().flush().unwrap();

      let mut buf = Vec::new();
      let mut resp = thread_client.get(down_url.to_str().unwrap())
                                  .header(UserAgent(USER_AGENT.to_owned()))
                                  .send().unwrap();
      match resp.read_to_end(&mut buf) {
        Ok(bytes) => {
          thread_chan.send(bytes).unwrap();
        },
        Err(e) => panic!("{:?}", e)
      };
    });
  }

  drop(tchan_tx); // main thread doesn't count

  let sum_bytes = tchan_rx.iter().fold(0, |sum, val| sum + val) as u64;
  let duration = start.to(PreciseTime::now()).num_milliseconds() as u64;

  (sum_bytes, duration)
}

fn measure_upload(server: &Server, http_client: &Arc<client::Client>) -> (Bytes, Milliseconds) {
  let cconnections = 3;
  let thread_pool = ThreadPool::new(cconnections);
  //let chunk_sizes = vec!(250*1000, 500*1000, 1000*1000, 2*1000*1000, 4*1000*1000, 8*1000*1000,16*1000*1000);
  let chunk_sizes: Vec<i32> = vec!(250*1000, 500*1000, 1000*1000, 2*1000*1000);
  let sizes: Vec<_> = chunk_sizes.iter().cycle().take(7).collect();
  let (tchan_tx, tchan_rx) = mpsc::channel();
  let start = PreciseTime::now();
  let char_data = "01213456789ABCDEFGHIJKLMNOPQRSTUVWXYZ".to_string();

  for s in sizes {
    let up_url = server.url.clone();
    let thread_client = http_client.clone();
    let thread_chan = tchan_tx.clone();
    let data: String = char_data.chars().cycle().take((s-9) as usize).collect();
    let data = "content1=".to_string() + &data;

    thread_pool.execute(move || {
      print!(" .");
      std::io::stdout().flush().unwrap();

      let mut buf = Vec::new();
      let mut resp = thread_client.post(&up_url)
                  .header(UserAgent(USER_AGENT.to_owned()))
                  .body(&data)
                  .send().unwrap();

      match resp.read_to_end(&mut buf) {
        Ok(_) => {
          thread_chan.send(data.len()).unwrap();
        },
        Err(e) => panic!("{:?}", e)
      };
    });
  }

  drop(tchan_tx); // main thread doesn't count

  let sum_bytes = tchan_rx.iter().fold(0, |sum, val| sum + val) as u64;
  let duration = start.to(PreciseTime::now()).num_milliseconds() as u64;

  (sum_bytes, duration)
}

fn speedtest() -> Option<SpeedMeasurement> {
  let mut client = Client::default();
  let http_client = client::Client::new();
  parse_client_cfg(&mut client, &http_client);
  println!("hello, {ip}!\n you're located at about {lat}°N {lon}°E and your ISP is '{isp}'.",
    ip = client.ip, lat = client.position.lat, lon = client.position.lon, isp = client.isp);

  print!(" retrieving server list...");
  let mut servers = Vec::<Server>::new();
  parse_server_list(&mut servers, &http_client);
  println!(" ({} entries)", servers.len());

  if servers.len() == 0 {
    return None;
  }

  let num_servers = 5;
  servers.sort_by(|a, b| {
    a.distance_to(&client.position).partial_cmp(&b.distance_to(&client.position)).unwrap_or(Equal)
  });

  println!(" chosing a close server based on latency...");
  let (best_id, best_latency) = lowest_latency(&mut servers[0..num_servers], &http_client);
  let best_idx = servers.iter().position(|s| s.id == best_id).unwrap();

  let server = servers.remove(best_idx);
  let server_ip = &server.host.to_socket_addrs().unwrap().next().unwrap().ip();
  let server_distance = server.distance_to(&client.position);
  println!(" using server:\n* {} ({})\n  (dist: {:.2}km, rtt: {}ms)",
    server, server_ip, server_distance, best_latency);

  println!(" other candidates:");
  for s in &servers[0..(num_servers-1)] {
    println!("- {}", s);
  }

  let http_client = Arc::new(http_client);

  print!(" testing download");
  let (dsum_bytes, dduration) = measure_download(&server, &http_client);
  let dmbit = (dsum_bytes as f64) * 8.0 / (dduration as f64) / 1000.0;
  println!("\n> result: {}Byte in {}ms (avg. {:.3}Mbit/s, {:.3}MByte/s)",
    dsum_bytes, dduration, dmbit, dmbit/8.0);

  print!(" testing upload");
  let (usum_bytes, uduration) = measure_upload(&server, &http_client);
  let umbit = (usum_bytes as f64) * 8.0 / (uduration as f64) / 1000.0;
  println!("\n> result: {}Byte in {}ms (avg. {:.3}Mbit/s, {:.3}MByte/s)",
    usum_bytes, uduration, umbit, umbit/8.0);

  Some(SpeedMeasurement {
    ts: UTC::now(),
    client_ip: client.ip,
    client_isp: client.isp,
    server_id: format!("{cc}:{id}", cc = server.cc, id = server.id),
    server_sponsor: server.sponsor,
    server_host: server.host,
    server_ip: format!("{}", server_ip), // hacky...
    server_distance: server_distance,
    server_rtt: best_latency,
    download_bytes: dsum_bytes,
    download_ms: dduration,
    upload_bytes: usum_bytes,
    upload_ms: uduration,
  })
}

fn output_result(result: &SpeedMeasurement) {
  let mut out: BufWriter<Box<Write>> = BufWriter::new(
    if let Some(filename) = env::args().nth(1) {
      Box::new(OpenOptions::new().create(true).append(true).open(filename).unwrap())
    } else {
      Box::new(stdout())
    }
  );

  out.write_fmt(format_args!(
    "\"{ts}\",\"{cip}\",\"{cisp}\",\"{sid}\",\"{ssponsor}\",\"{shost}\",\"{sip}\",\"{sdist}\",\"{srtt}\",\"{db}\",\"{dms}\",\"{ub}\",\"{ums}\"\n",
    ts       = result.ts,
    cip      = result.client_ip,
    cisp     = result.client_isp,
    sid      = result.server_id,
    ssponsor = result.server_sponsor,
    shost    = result.server_host,
    sip      = result.server_ip,
    sdist    = result.server_distance,
    srtt     = result.server_rtt,
    db       = result.download_bytes,
    dms      = result.download_ms,
    ub       = result.upload_bytes,
    ums      = result.upload_ms,
  )).unwrap();
}

pub fn main() {
  loop {
    {
      match speedtest() {
        Some(result) => {
          //println!("\n\n### RESULT ###\n\n{:?}", result);
          output_result(&result);
        },
        None => {
          writeln!(std::io::stderr(), "! aborting this iteration, server list empty!").unwrap();
        }
      }
    }

    std::thread::sleep(time::Duration::minutes(6).to_std().unwrap());
    println!(". repeating...");
  }
}