diff --git a/README.md b/README.md index 3ef5a641..b3f43c45 100644 --- a/README.md +++ b/README.md @@ -394,6 +394,13 @@ Example configuration: "plugin": "...", "plugin_opts": "...", "timeout": 7200, + + // Customized weight for local server's balancer + // + // Weight must be in [0, 1], default is 1.0. + // The higher weight, the server may rank higher. + "tcp_weight": 1.0, + "udp_weight": 1.0, } ], diff --git a/crates/shadowsocks-service/src/config.rs b/crates/shadowsocks-service/src/config.rs index b2a17fbf..21539dda 100644 --- a/crates/shadowsocks-service/src/config.rs +++ b/crates/shadowsocks-service/src/config.rs @@ -62,7 +62,7 @@ use serde::{Deserialize, Serialize}; #[cfg(any(feature = "local-tunnel", feature = "local-dns"))] use shadowsocks::relay::socks5::Address; use shadowsocks::{ - config::{ManagerAddr, Mode, ServerAddr, ServerConfig}, + config::{ManagerAddr, Mode, ServerAddr, ServerConfig, ServerWeight}, crypto::v1::CipherKind, plugin::PluginConfig, }; @@ -212,6 +212,10 @@ struct SSServerExtConfig { id: Option, #[serde(skip_serializing_if = "Option::is_none")] mode: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tcp_weight: Option, + #[serde(skip_serializing_if = "Option::is_none")] + udp_weight: Option, } /// Server config type @@ -1195,7 +1199,7 @@ impl Config { Some(mode) => match mode.parse::() { Ok(mode) => nsvr.set_mode(mode), Err(..) => { - let err = Error::new(ErrorKind::Malformed, "invalid `mode`", None); + let err = Error::new(ErrorKind::Invalid, "invalid `mode`", None); return Err(err); } }, @@ -1227,6 +1231,23 @@ impl Config { nsvr.set_id(id); } + if svr.tcp_weight.is_some() || svr.udp_weight.is_some() { + let tcp_weight = svr.tcp_weight.unwrap_or(1.0); + if tcp_weight < 0.0 || tcp_weight > 1.0 { + let err = Error::new(ErrorKind::Invalid, "invalid `tcp_weight`, must be in [0, 1]", None); + return Err(err); + } + let udp_weight = svr.udp_weight.unwrap_or(1.0); + if udp_weight < 0.0 || udp_weight > 1.0 { + let err = Error::new(ErrorKind::Invalid, "invalid `udp_weight`, must be in [0, 1]", None); + return Err(err); + } + let mut weight = ServerWeight::new(); + weight.set_tcp_weight(tcp_weight); + weight.set_udp_weight(udp_weight); + nsvr.set_weight(weight); + } + nconfig.server.push(nsvr); } } @@ -1732,6 +1753,16 @@ impl fmt::Display for Config { remarks: svr.remarks().map(ToOwned::to_owned), id: svr.id().map(ToOwned::to_owned), mode: Some(svr.mode().to_string()), + tcp_weight: if svr.weight().tcp_weight() != 1.0 { + Some(svr.weight().tcp_weight()) + } else { + None + }, + udp_weight: if svr.weight().udp_weight() != 1.0 { + Some(svr.weight().udp_weight()) + } else { + None + }, }); } diff --git a/crates/shadowsocks-service/src/local/loadbalancing/server_data.rs b/crates/shadowsocks-service/src/local/loadbalancing/server_data.rs index b28d54fa..0e84ac51 100644 --- a/crates/shadowsocks-service/src/local/loadbalancing/server_data.rs +++ b/crates/shadowsocks-service/src/local/loadbalancing/server_data.rs @@ -16,19 +16,13 @@ pub struct ServerScore { score: AtomicU32, } -impl Default for ServerScore { - fn default() -> Self { - ServerScore { - stat_data: Mutex::new(ServerStat::new()), - score: AtomicU32::new(0), - } - } -} - impl ServerScore { /// Create a `ServerScore` - pub fn new() -> ServerScore { - ServerScore::default() + pub fn new(user_weight: f32) -> ServerScore { + ServerScore { + stat_data: Mutex::new(ServerStat::new(user_weight)), + score: AtomicU32::new(0), + } } /// Get server's current statistic scores @@ -70,8 +64,8 @@ impl ServerIdent { /// Create a ServerIdent` pub fn new(svr_cfg: ServerConfig) -> ServerIdent { ServerIdent { - tcp_score: ServerScore::new(), - udp_score: ServerScore::new(), + tcp_score: ServerScore::new(svr_cfg.weight().tcp_weight()), + udp_score: ServerScore::new(svr_cfg.weight().udp_weight()), 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 fec6ced1..a71ed56b 100644 --- a/crates/shadowsocks-service/src/local/loadbalancing/server_stat.rs +++ b/crates/shadowsocks-service/src/local/loadbalancing/server_stat.rs @@ -35,6 +35,8 @@ pub struct ServerStat { latency_stdev: f64, /// Score's average latency_mean: f64, + /// User's customized weight + user_weight: f32, } fn max_latency_stdev() -> f64 { @@ -46,22 +48,19 @@ fn max_latency_stdev() -> f64 { (diff1 + diff2).sqrt() } -impl Default for ServerStat { - fn default() -> Self { +impl ServerStat { + pub fn new(user_weight: f32) -> ServerStat { + assert!(user_weight >= 0.0 && user_weight <= 1.0); + ServerStat { rtt: MAX_SERVER_RTT, fail_rate: 1.0, latency_queue: VecDeque::new(), latency_stdev: 0.0, latency_mean: 0.0, + user_weight, } } -} - -impl ServerStat { - pub fn new() -> ServerStat { - ServerStat::default() - } fn score(&self) -> u32 { // Normalize rtt @@ -73,14 +72,19 @@ impl ServerStat { const SCORE_RTT_WEIGHT: f64 = 1.0; const SCORE_FAIL_WEIGHT: f64 = 3.0; const SCORE_STDEV_WEIGHT: f64 = 1.0; + const SCORE_USER_WEIGHT: f64 = 5.0; - // Score = (norm_lat * 1.0 + prop_err * 3.0 + stdev * 1.0) / 5.0 + // Score = (norm_lat * 1.0 + prop_err * 3.0 + stdev * 1.0 + user_weight * 5.0) / 10.0 // // 1. The lower latency, the better // 2. The lower errored count, the better // 3. The lower latency's stdev, the better - let score = (nrtt * SCORE_RTT_WEIGHT + self.fail_rate * SCORE_FAIL_WEIGHT + nstdev * SCORE_STDEV_WEIGHT) - / (SCORE_RTT_WEIGHT + SCORE_FAIL_WEIGHT + SCORE_STDEV_WEIGHT); + // 4. The higher user's weight, the better + let score = (nrtt * SCORE_RTT_WEIGHT + + self.fail_rate * SCORE_FAIL_WEIGHT + + nstdev * SCORE_STDEV_WEIGHT + + self.user_weight as f64 * SCORE_USER_WEIGHT) + / (SCORE_RTT_WEIGHT + SCORE_FAIL_WEIGHT + SCORE_STDEV_WEIGHT + SCORE_USER_WEIGHT); // Times 10000 converts to u32, for 0.0001 precision (score * 10000.0) as u32 diff --git a/crates/shadowsocks/src/config.rs b/crates/shadowsocks/src/config.rs index ff2c9e54..44b3448a 100644 --- a/crates/shadowsocks/src/config.rs +++ b/crates/shadowsocks/src/config.rs @@ -97,6 +97,53 @@ impl FromStr for Mode { } } +/// Server's weight +/// +/// Commonly for using in balancer +#[derive(Debug, Clone)] +pub struct ServerWeight { + tcp_weight: f32, + udp_weight: f32, +} + +impl Default for ServerWeight { + fn default() -> Self { + ServerWeight::new() + } +} + +impl ServerWeight { + /// Creates a default weight for server, which will have 1.0 for both TCP and UDP + pub fn new() -> ServerWeight { + ServerWeight { + tcp_weight: 1.0, + udp_weight: 1.0, + } + } + + /// Weight for TCP balancer + pub fn tcp_weight(&self) -> f32 { + self.tcp_weight + } + + /// Set weight for TCP balancer in `[0, 1]` + pub fn set_tcp_weight(&mut self, weight: f32) { + assert!(weight >= 0.0 && weight <= 1.0); + self.tcp_weight = weight; + } + + /// Weight for UDP balancer + pub fn udp_weight(&self) -> f32 { + self.udp_weight + } + + /// Set weight for UDP balancer in `[0, 1]` + pub fn set_udp_weight(&mut self, weight: f32) { + assert!(weight >= 0.0 && weight <= 1.0); + self.udp_weight = weight; + } +} + /// Configuration for a server #[derive(Clone, Debug)] pub struct ServerConfig { @@ -123,6 +170,9 @@ pub struct ServerConfig { /// Mode mode: Mode, + + /// Weight + weight: ServerWeight, } impl ServerConfig { @@ -148,6 +198,7 @@ impl ServerConfig { remarks: None, id: None, mode: Mode::TcpOnly, + weight: ServerWeight::new(), } } @@ -264,6 +315,16 @@ impl ServerConfig { self.mode = mode; } + /// Get server's balancer weight + pub fn weight(&self) -> &ServerWeight { + &self.weight + } + + /// Set server's balancer weight + pub fn set_weight(&mut self, weight: ServerWeight) { + self.weight = weight; + } + /// Get URL for QRCode /// ```plain /// ss:// + base64(method:password@host:port) @@ -327,7 +388,7 @@ impl ServerConfig { let mut sp2 = account.splitn(2, ':'); let (method, pwd) = match (sp2.next(), sp2.next()) { (Some(m), Some(p)) => (m, p), - _ => return Err(UrlParseError::InvalidUserInfo) + _ => return Err(UrlParseError::InvalidUserInfo), }; let addr = match addr.parse::() {