diff --git a/Cargo.lock b/Cargo.lock index ef62984a..c262574a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -525,9 +525,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.3.4" +version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7f3675cfef6a30c8031cf9e6493ebdc3bb3272a3fea3923c4210d1830e6a472" +checksum = "6c06815895acec637cd6ed6e9662c935b866d20a106f8361892893a7d9234964" dependencies = [ "bytes", "fnv", @@ -969,9 +969,9 @@ dependencies = [ [[package]] name = "nix" -version = "0.22.1" +version = "0.22.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7555d6c7164cc913be1ce7f95cbecdabda61eb2ccd89008524af306fb7f5031" +checksum = "d3bb9a13fa32bc5aeb64150cd3f32d6cf4c748f8f8a417cce5d2eb976a8370ba" dependencies = [ "bitflags", "cc", @@ -1750,9 +1750,9 @@ checksum = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a" [[package]] name = "syn" -version = "1.0.77" +version = "1.0.78" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5239bc68e0fef57495900cfea4e8dc75596d9a319d7e16b1e0a440d24e6fe0a0" +checksum = "a4eac2e6c19f5c3abc0c229bea31ff0b9b091c7b14990e8924b92902a303a0c0" dependencies = [ "proc-macro2", "quote", @@ -1905,9 +1905,9 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "1.3.0" +version = "1.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54473be61f4ebe4efd09cec9bd5d16fa51d70ea0192213d754d2d500457db110" +checksum = "154794c8f499c2619acd19e839294703e9e32e7630ef5f46ea80d4ef0fbee5eb" dependencies = [ "proc-macro2", "quote", @@ -1989,9 +1989,9 @@ dependencies = [ [[package]] name = "tracing-attributes" -version = "0.1.16" +version = "0.1.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98863d0dd09fa59a1b79c6750ad80dbda6b75f4e71c437a6a1a8cb91a8bcbd77" +checksum = "c4f915eb6abf914599c200260efced9203504c4c37380af10cdf3b7d36970650" dependencies = [ "proc-macro2", "quote", @@ -2000,9 +2000,9 @@ dependencies = [ [[package]] name = "tracing-core" -version = "0.1.20" +version = "0.1.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46125608c26121c81b0c6d693eab5a420e416da7e43c426d2e8f7df8da8a3acf" +checksum = "1f4ed65637b8390770814083d20756f87bfa2c21bf2f110babdc5438351746e4" dependencies = [ "lazy_static", ] diff --git a/README.md b/README.md index 81a9fa88..fde8ec47 100644 --- a/README.md +++ b/README.md @@ -558,7 +558,16 @@ Example configuration: "nofile": 10240, // Try to resolve domain name to IPv6 (AAAA) addresses first - "ipv6_first": false + "ipv6_first": false, + + // Balancer customization + "balancer": { + // MAX Round-Trip-Time (RTT) of servers + // The timeout seconds of each individual checks + "max_server_rtt": 5, + // Interval seconds between each check + "check_interval": 10, + } } ``` diff --git a/crates/shadowsocks-service/src/config.rs b/crates/shadowsocks-service/src/config.rs index 055c5052..492a7aa3 100644 --- a/crates/shadowsocks-service/src/config.rs +++ b/crates/shadowsocks-service/src/config.rs @@ -95,6 +95,14 @@ struct SSSecurityReplayAttackConfig { policy: Option, } +#[derive(Serialize, Deserialize, Debug, Default)] +struct SSBalancerConfig { + #[serde(skip_serializing_if = "Option::is_none")] + max_server_rtt: Option, + #[serde(skip_serializing_if = "Option::is_none")] + check_interval: Option, +} + #[derive(Serialize, Deserialize, Debug, Default)] struct SSConfig { #[serde(skip_serializing_if = "Option::is_none")] @@ -148,6 +156,8 @@ struct SSConfig { fast_open: Option, #[serde(skip_serializing_if = "Option::is_none")] security: Option, + #[serde(skip_serializing_if = "Option::is_none")] + balancer: Option, } #[derive(Serialize, Deserialize, Debug, Default)] @@ -837,6 +847,15 @@ pub struct SecurityReplayAttackConfig { pub policy: ReplayAttackPolicy, } +/// Balancer Config +#[derive(Clone, Debug, Default)] +pub struct BalancerConfig { + /// MAX rtt of servers, which is the timeout duration of each check requests + pub max_server_rtt: Option, + /// Interval between each checking + pub check_interval: Option, +} + /// Configuration #[derive(Clone, Debug)] pub struct Config { @@ -916,6 +935,9 @@ pub struct Config { /// Replay attack policy pub security: SecurityConfig, + + /// Balancer config of local server + pub balancer: BalancerConfig, } /// Configuration parsing error kind @@ -1020,6 +1042,8 @@ impl Config { stat_path: None, security: SecurityConfig::default(), + + balancer: BalancerConfig::default(), } } @@ -1550,6 +1574,13 @@ impl Config { } } + if let Some(balancer) = config.balancer { + nconfig.balancer = BalancerConfig { + max_server_rtt: balancer.max_server_rtt.map(Duration::from_secs), + check_interval: balancer.check_interval.map(Duration::from_secs), + }; + } + Ok(nconfig) } @@ -1729,6 +1760,21 @@ impl Config { ); return Err(err); } + + // Balancer related checks + if let Some(rtt) = self.balancer.max_server_rtt { + if rtt.as_secs() == 0 { + let err = Error::new(ErrorKind::Invalid, "balancer.max_server_rtt must be > 0", None); + return Err(err); + } + } + + if let Some(intv) = self.balancer.check_interval { + if intv.as_secs() == 0 { + let err = Error::new(ErrorKind::Invalid, "balancer.check_interval must be > 0", None); + return Err(err); + } + } } if self.config_type.is_server() && self.server.is_empty() { @@ -2081,6 +2127,14 @@ impl fmt::Display for Config { }); } + // Balancer + if self.balancer.max_server_rtt.is_some() || self.balancer.check_interval.is_some() { + jconf.balancer = Some(SSBalancerConfig { + max_server_rtt: self.balancer.max_server_rtt.as_ref().map(Duration::as_secs), + check_interval: self.balancer.check_interval.as_ref().map(Duration::as_secs), + }); + } + write!(f, "{}", json5::to_string(&jconf).unwrap()) } } diff --git a/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs b/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs index 8eac4cac..9cb3683b 100644 --- a/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs +++ b/crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs @@ -58,6 +58,8 @@ pub struct PingBalancerBuilder { servers: Vec>, context: Arc, mode: Mode, + max_server_rtt: Duration, + check_interval: Duration, } impl PingBalancerBuilder { @@ -66,11 +68,22 @@ impl PingBalancerBuilder { servers: Vec::new(), context, mode, + max_server_rtt: Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC), + check_interval: Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC), } } - pub fn add_server(&mut self, server: ServerIdent) { - self.servers.push(Arc::new(server)); + pub fn add_server(&mut self, server: ServerConfig) { + let ident = ServerIdent::new(server, self.max_server_rtt); + self.servers.push(Arc::new(ident)); + } + + pub fn max_server_rtt(&mut self, rtt: Duration) { + self.max_server_rtt = rtt; + } + + pub fn check_interval(&mut self, intv: Duration) { + self.check_interval = intv; } pub async fn build(self) -> PingBalancer { @@ -131,6 +144,8 @@ impl PingBalancerBuilder { best_udp_idx: AtomicUsize::new(best_udp_idx), context: self.context, mode: self.mode, + max_server_rtt: self.max_server_rtt, + check_interval: self.check_interval, }; balancer_context.init_score().await; @@ -157,6 +172,8 @@ struct PingBalancerContext { best_udp_idx: AtomicUsize, context: Arc, mode: Mode, + max_server_rtt: Duration, + check_interval: Duration, } impl PingBalancerContext { @@ -231,6 +248,7 @@ impl PingBalancerContext { server: server.clone(), server_type: ServerType::Tcp, context: self.context.clone(), + max_server_rtt: self.max_server_rtt, }; vfut_tcp.push(checker.check_update_score()); } @@ -240,6 +258,7 @@ impl PingBalancerContext { server: server.clone(), server_type: ServerType::Udp, context: self.context.clone(), + max_server_rtt: self.max_server_rtt, }; vfut_udp.push(checker.check_update_score()); } @@ -336,7 +355,7 @@ impl PingBalancerContext { async fn checker_task_real(&self) { loop { - time::sleep(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await; + time::sleep(self.check_interval).await; // Sleep before check. // PingBalancer already checked once when constructing @@ -404,6 +423,7 @@ struct PingChecker { server: Arc, server_type: ServerType, context: Arc, + max_server_rtt: Duration, } impl PingChecker { @@ -563,8 +583,7 @@ impl PingChecker { let start = Instant::now(); // Send HTTP GET and read the first byte - let timeout = Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC); - let res = time::timeout(timeout, self.check_request()).await; + let res = time::timeout(self.max_server_rtt, self.check_request()).await; let elapsed = Instant::now() - start; let elapsed = elapsed.as_secs() as u32 * 1000 + elapsed.subsec_millis(); // Converted to ms diff --git a/crates/shadowsocks-service/src/local/loadbalancing/server_data.rs b/crates/shadowsocks-service/src/local/loadbalancing/server_data.rs index 96d2f696..a0f6f4e9 100644 --- a/crates/shadowsocks-service/src/local/loadbalancing/server_data.rs +++ b/crates/shadowsocks-service/src/local/loadbalancing/server_data.rs @@ -3,6 +3,7 @@ use std::{ fmt::{self, Debug}, sync::atomic::{AtomicU32, Ordering}, + time::Duration, }; use shadowsocks::ServerConfig; @@ -18,9 +19,12 @@ pub struct ServerScore { impl ServerScore { /// Create a `ServerScore` - pub fn new(user_weight: f32) -> ServerScore { + pub fn new(user_weight: f32, max_server_rtt: Duration) -> ServerScore { + let max_server_rtt = max_server_rtt.as_millis() as u32; + assert!(max_server_rtt > 0); + ServerScore { - stat_data: Mutex::new(ServerStat::new(user_weight)), + stat_data: Mutex::new(ServerStat::new(user_weight, max_server_rtt)), score: AtomicU32::new(u32::MAX), } } @@ -61,11 +65,11 @@ pub struct ServerIdent { } impl ServerIdent { - /// Create a ServerIdent` - pub fn new(svr_cfg: ServerConfig) -> ServerIdent { + /// Create a `ServerIdent` + pub fn new(svr_cfg: ServerConfig, max_server_rtt: Duration) -> ServerIdent { ServerIdent { - tcp_score: ServerScore::new(svr_cfg.weight().tcp_weight()), - udp_score: ServerScore::new(svr_cfg.weight().udp_weight()), + tcp_score: ServerScore::new(svr_cfg.weight().tcp_weight(), max_server_rtt), + udp_score: ServerScore::new(svr_cfg.weight().udp_weight(), max_server_rtt), svr_cfg, } } diff --git a/crates/shadowsocks-service/src/local/loadbalancing/server_stat.rs b/crates/shadowsocks-service/src/local/loadbalancing/server_stat.rs index 4e3ef3a6..ac93a991 100644 --- a/crates/shadowsocks-service/src/local/loadbalancing/server_stat.rs +++ b/crates/shadowsocks-service/src/local/loadbalancing/server_stat.rs @@ -7,8 +7,7 @@ pub const DEFAULT_CHECK_INTERVAL_SEC: u64 = 10; /// Timeout of each check pub const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5; // A common connection timeout of 5 seconds. -const MAX_SERVER_RTT: u32 = DEFAULT_CHECK_TIMEOUT_SEC as u32 * 1000; -const MAX_LATENCY_QUEUE_SIZE: usize = 59; // Account for the last 10 minutes. +const MAX_LATENCY_QUEUE_SIZE: usize = 67; /// Statistic score #[derive(Debug, Copy, Clone)] @@ -27,20 +26,24 @@ pub struct ServerStat { /// Use median instead of average time, /// because probing result may have some really bad cases rtt: u32, + /// MAX server's RTT, normally the check timeout milliseconds + max_server_rtt: u32, /// Total_Fail / Total_Probe fail_rate: f64, /// Recently probe data latency_queue: VecDeque, /// Score's standard deviation latency_stdev: f64, + /// Score's standard deviation MAX + max_latency_stdev: f64, /// Score's average latency_mean: f64, /// User's customized weight user_weight: f32, } -fn max_latency_stdev() -> f64 { - let mrtt = MAX_SERVER_RTT as f64; +fn max_latency_stdev(max_server_rtt: u32) -> f64 { + let mrtt = max_server_rtt as f64; let avg = (0.0 + mrtt) / 2.0; let diff1 = (0.0 - avg) * (0.0 - avg); let diff2 = (mrtt - avg) * (mrtt - avg); @@ -49,14 +52,16 @@ fn max_latency_stdev() -> f64 { } impl ServerStat { - pub fn new(user_weight: f32) -> ServerStat { + pub fn new(user_weight: f32, max_server_rtt: u32) -> ServerStat { assert!((0.0..=1.0).contains(&user_weight)); ServerStat { - rtt: MAX_SERVER_RTT, + rtt: max_server_rtt, + max_server_rtt, fail_rate: 1.0, latency_queue: VecDeque::new(), latency_stdev: 0.0, + max_latency_stdev: max_latency_stdev(max_server_rtt), latency_mean: 0.0, user_weight, } @@ -64,10 +69,10 @@ impl ServerStat { fn score(&self) -> u32 { // Normalize rtt - let nrtt = self.rtt as f64 / MAX_SERVER_RTT as f64; + let nrtt = self.rtt as f64 / self.max_server_rtt as f64; // Normalize stdev - let nstdev = self.latency_stdev / max_latency_stdev(); + let nstdev = self.latency_stdev / self.max_latency_stdev; const SCORE_RTT_WEIGHT: f64 = 1.0; const SCORE_FAIL_WEIGHT: f64 = 3.0; diff --git a/crates/shadowsocks-service/src/local/mod.rs b/crates/shadowsocks-service/src/local/mod.rs index c594243c..2af86239 100644 --- a/crates/shadowsocks-service/src/local/mod.rs +++ b/crates/shadowsocks-service/src/local/mod.rs @@ -27,10 +27,7 @@ use crate::{ dns::build_dns_resolver, }; -use self::{ - context::ServiceContext, - loadbalancing::{PingBalancerBuilder, ServerIdent}, -}; +use self::{context::ServiceContext, loadbalancing::PingBalancerBuilder}; pub mod context; #[cfg(feature = "local-dns")] @@ -202,9 +199,20 @@ pub async fn run(mut config: Config) -> io::Result<()> { } let mut balancer_builder = PingBalancerBuilder::new(context.clone(), mode); - for server in config.server { - balancer_builder.add_server(ServerIdent::new(server)); + + // max_server_rtt have to be set before add_server + if let Some(rtt) = config.balancer.max_server_rtt { + balancer_builder.max_server_rtt(rtt); } + + if let Some(intv) = config.balancer.check_interval { + balancer_builder.check_interval(intv); + } + + for server in config.server { + balancer_builder.add_server(server); + } + balancer_builder.build().await };