diff --git a/src/relay/loadbalancing/server/ping.rs b/src/relay/loadbalancing/server/ping.rs index 74747879..27a42d48 100644 --- a/src/relay/loadbalancing/server/ping.rs +++ b/src/relay/loadbalancing/server/ping.rs @@ -1,8 +1,9 @@ use std::{ - io, + collections::VecDeque, + fmt, io, sync::{ - atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, - Arc, + atomic::{AtomicU64, AtomicUsize, Ordering}, + Arc, Mutex, }, time::{Duration, Instant}, }; @@ -10,6 +11,7 @@ use std::{ use crate::{ config::ServerConfig, context::SharedContext, + relay::tcprelay::EncryptedWrite, relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient}, }; @@ -17,13 +19,13 @@ use futures::{Future, Stream}; use log::{debug, error, info}; use tokio::{ self, + io::read_exact, timer::{Interval, Timeout}, }; struct Server { config: Arc, elapsed: AtomicU64, - available: AtomicBool, } impl Server { @@ -31,20 +33,80 @@ impl Server { Server { config: Arc::new(config), elapsed: AtomicU64::new(0), - available: AtomicBool::new(true), } } - fn is_available(&self) -> bool { - self.available.load(Ordering::Acquire) - } - fn delay(&self) -> u64 { self.elapsed.load(Ordering::Acquire) } } +const MAX_LATENCY_QUEUE_SIZE: usize = 17; + +struct ServerLatencyInner { + latency_queue: VecDeque, +} + +impl ServerLatencyInner { + fn new() -> ServerLatencyInner { + ServerLatencyInner { + latency_queue: VecDeque::with_capacity(MAX_LATENCY_QUEUE_SIZE), + } + } + + fn push(&mut self, lat: u64) -> u64 { + self.latency_queue.push_back(lat); + if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE { + self.latency_queue.pop_front(); + } + + self.median() + } + + fn median(&self) -> u64 { + if self.latency_queue.is_empty() { + return 0; + } + + let mut v = self.latency_queue.iter().cloned().collect::>(); + v.sort(); + + let mid = v.len() / 2; + if (v.len() & 1) == 0 { + (v[mid - 1] + v[mid]) / 2 + } else { + v[mid] + } + } +} + +#[derive(Clone)] +struct ServerLatency { + inner: Arc>, +} + +impl ServerLatency { + fn new() -> ServerLatency { + ServerLatency { + inner: Arc::new(Mutex::new(ServerLatencyInner::new())), + } + } + + fn push(&self, lat: u64) -> u64 { + let mut inner = self.inner.lock().unwrap(); + inner.push(lat) + } +} + +impl fmt::Debug for ServerLatency { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let inner = self.inner.lock().unwrap(); + write!(f, "ServerLatency {{ latency: {:?} }}", inner.latency_queue) + } +} + const DEFAULT_CHECK_INTERVAL_SEC: u64 = 10; +const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5; struct Inner { servers: Vec>, @@ -70,8 +132,8 @@ impl Inner { continue; } - let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80); let context = context.clone(); + let latency = ServerLatency::new(); tokio::spawn( // Check every 10 seconds @@ -81,21 +143,20 @@ impl Inner { ) .for_each(move |_| { let sc = sc.clone(); + let lat = latency.clone(); - let fut1 = Inner::check_delay(sc.clone(), context.clone(), addr.clone()); - let fut2 = Inner::check_delay(sc.clone(), context.clone(), addr.clone()); - let fut3 = Inner::check_delay(sc.clone(), context.clone(), addr.clone()); - - fut1.join3(fut2, fut3).then(move |res| { - match res { - Ok((d1, d2, d3)) => { - sc.available.store(true, Ordering::Release); - sc.elapsed.store((d1 + d2 + d3) / 3, Ordering::Release); - } - Err(..) => { - sc.available.store(false, Ordering::Release); - } - } + Inner::check_delay(sc.clone(), context.clone()).then(move |res| { + let md = match res { + Ok(d) => lat.push(d), + Err(..) => lat.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty + }; + debug!( + "updated remote server {} (median: {} ms) {:?}", + sc.config.addr(), + md, + lat + ); + sc.elapsed.store(md, Ordering::Release); Ok(()) }) @@ -110,24 +171,27 @@ impl Inner { } } - fn check_delay( - sc: Arc, - context: SharedContext, - addr: Address, - ) -> impl Future { + fn check_delay(sc: Arc, context: SharedContext) -> impl Future { + static GET_BODY: &[u8] = b"GET / HTTP/1.1\r\nHost: www.example.com\r\nConnection: close\r\nAccept: */*\r\n\r\n"; + + let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80); + let start = Instant::now(); - let fut = ServerClient::connect(context, addr, sc.config.clone()); - Timeout::new(fut, Duration::from_secs(5)).then(move |res| { + + // Send HTTP GET and read the first byte + let fut = ServerClient::connect(context, addr, sc.config.clone()).and_then(|ServerClient { r, w }| { + w.write_all(GET_BODY) + .and_then(|(mut w, _)| w.flush()) + .and_then(move |_| read_exact(r, [0u8, 1])) + }); + + Timeout::new(fut, Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC)).then(move |res| { let elapsed = Instant::now() - start; let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms match res { Ok(..) => { - // Connected ... record its time - debug!( - "checked remote server {} connected with {} ms", - sc.config.addr(), - elapsed - ); + // Got the result ... record its time + debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed); Ok(elapsed) } Err(err) => { @@ -140,8 +204,8 @@ impl Inner { } None => { // Timeout - error!( - "checked remote server {} connect timeout, elapsed {} ms", + debug!( + "checked remote server {} latency timeout, elapsed {} ms", sc.config.addr(), elapsed ); @@ -195,30 +259,21 @@ impl PingBalancer { for (idx, svr) in inner.servers.iter().enumerate() { let choosen_svr = &inner.servers[svr_idx]; - if svr.is_available() && (!choosen_svr.is_available() || svr.delay() < choosen_svr.delay()) { + if svr.delay() < choosen_svr.delay() { svr_idx = idx; } } let choosen_svr = &inner.servers[svr_idx]; - if svr_idx == 0 && !choosen_svr.is_available() { - // Cannot find any usable servers, use the first one (svr_idx = 0) - error!( - "cannot find any usable servers, picked {} (delay: {} ms)", - choosen_svr.config.addr(), - choosen_svr.delay() - ); - } else { - debug!( - "chosen the best server {} (delay: {} ms)", - choosen_svr.config.addr(), - choosen_svr.delay() - ); - } + debug!( + "chosen the best server {} (delay: {} ms)", + choosen_svr.config.addr(), + choosen_svr.delay() + ); if inner.best_idx() != svr_idx { info!( - "switched server from {} (delay: {} ms) to {} (delay: {} ms)", + "switched server from {} (latency: {} ms) to {} (latency: {} ms)", inner.best_server().config.addr(), inner.best_server().delay(), choosen_svr.config.addr(), diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index d769befa..c29ccc93 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -26,9 +26,7 @@ use tokio::{ }; use tokio_io::{ io::{read_exact, write_all, ReadHalf, WriteHalf}, - try_nb, - AsyncRead, - AsyncWrite, + try_nb, AsyncRead, AsyncWrite, }; pub use self::crypto_io::{DecryptedRead, EncryptedWrite};