allow customizing servers' weight for balancer

- fixes #510

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,
         }
     ],
 

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<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     mode: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    tcp_weight: Option<f32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    udp_weight: Option<f32>,
 }
 
 /// Server config type
@@ -1195,7 +1199,7 @@ impl Config {
                     Some(mode) => match mode.parse::<Mode>() {
                         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
+                        },
                     });
                 }
 

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,
         }
     }

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

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::<ServerAddr>() {