Refine ping ranking logic

zonyitoo · zonyitoo · commit 96e3e99d082c · 2019-12-15T23:45:27.000+08:00
diff --git a/src/relay/loadbalancing/server/ping.rs b/src/relay/loadbalancing/server/ping.rs
@@ -1,9 +1,11 @@
 use std::{
     collections::VecDeque,
-    fmt, io,
+    fmt,
+    io,
     sync::{
         atomic::{AtomicU64, AtomicUsize, Ordering},
-        Arc, Mutex,
+        Arc,
+        Mutex,
     },
 };
 
@@ -40,8 +42,14 @@ impl Server {
 
 const MAX_LATENCY_QUEUE_SIZE: usize = 37;
 
+#[derive(Debug, Copy, Clone)]
+enum Score {
+    Latency(u64),
+    Errored,
+}
+
 struct ServerLatencyInner {
-    latency_queue: VecDeque<u64>,
+    latency_queue: VecDeque<Score>,
 }
 
 impl ServerLatencyInner {
@@ -51,7 +59,7 @@ impl ServerLatencyInner {
         }
     }
 
-    fn push(&mut self, lat: u64) -> u64 {
+    fn push(&mut self, lat: Score) -> u64 {
         self.latency_queue.push_back(lat);
         if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
             self.latency_queue.pop_front();
@@ -62,18 +70,45 @@ impl ServerLatencyInner {
 
     fn score(&self) -> u64 {
         if self.latency_queue.is_empty() {
-            return u64::max_value();
+            // Never checked, assume it is the worst of all
+            return 2 * 1000;
+        }
+
+        // 1. Mid Latency
+        // 2. Proportion of Errors
+        let mut vec_lat = Vec::with_capacity(self.latency_queue.len());
+        let mut acc_err = 0;
+        for lat in &self.latency_queue {
+            match lat {
+                Score::Latency(l) => vec_lat.push(l),
+                Score::Errored => acc_err += 1,
+            }
         }
 
-        let mut v = self.latency_queue.iter().cloned().collect::<Vec<u64>>();
-        v.sort();
+        let max_lat = DEFAULT_CHECK_TIMEOUT_SEC * 1000;
 
-        let mid = v.len() / 2;
-        if (v.len() & 1) == 0 {
-            (v[mid - 1] + v[mid]) / 2
+        // Find the mid of latencies
+        let mid_lat = if vec_lat.is_empty() {
+            // The whole array are errors
+            max_lat
         } else {
-            v[mid]
-        }
+            vec_lat.sort();
+            let mid = vec_lat.len() / 2;
+            if vec_lat.len() % 2 == 0 {
+                (vec_lat[mid] + vec_lat[mid - 1]) / 2
+            } else {
+                *vec_lat[mid]
+            }
+        };
+
+        // Score = norm_lat + prop_err
+        //
+        // 1. The lower latency, the better
+        // 2. The lower errored count, the better
+        let norm_lat = mid_lat as f64 / max_lat as f64;
+        let prop_err = acc_err as f64 / self.latency_queue.len() as f64;
+
+        ((norm_lat + prop_err) * 1000.0) as u64
     }
 }
 
@@ -89,7 +124,7 @@ impl ServerLatency {
         }
     }
 
-    fn push(&self, lat: u64) -> u64 {
+    fn push(&self, lat: Score) -> u64 {
         let mut inner = self.inner.lock().unwrap();
         inner.push(lat)
     }
@@ -135,20 +170,22 @@ impl Inner {
             tokio::spawn(
                 // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
                 async move {
-                    // Wait until the server is ready (plugins are started)
-                    time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
-
                     let mut interval = time::interval(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC));
 
                     while context.server_running() {
-                        interval.tick().await;
+                        // First round may be failed, plugins are started asynchronously
+
                         let score = match Inner::check_delay(&*sc, &*context).await {
-                            Ok(d) => latency.push(d),
-                            Err(..) => latency.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
+                            Ok(d) => latency.push(Score::Latency(d)),
+                            Err(..) => latency.push(Score::Errored), // Penalty
                         };
                         debug!("updated remote server {} (score: {})", sc.config.addr(), score);
                         sc.score.store(score, Ordering::Release);
+
+                        interval.tick().await;
                     }
+
+                    debug!("server {} latency ping task stopped", sc.config.addr());
                 },
             );
         }
@@ -236,12 +273,10 @@ impl PingBalancer {
             tokio::spawn(async move {
                 let inner = cloned_inner;
 
-                // Wait until the server is ready (plugins are started)
-                time::delay_for(Duration::from_secs(2)).await;
-
                 let mut interval = time::interval(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC));
 
                 while context.server_running() {
+                    // Must be here, wait for the first checking round
                     interval.tick().await;
 
                     assert!(!inner.servers.is_empty());
@@ -257,23 +292,17 @@ impl PingBalancer {
                     }
 
                     let choosen_svr = &inner.servers[svr_idx];
-                    debug!(
-                        "chosen the best server {} (score: {})",
-                        choosen_svr.config.addr(),
-                        choosen_svr.score()
-                    );
-
-                    if inner.best_idx() != svr_idx {
+                    if choosen_svr.score() != inner.servers[inner.best_idx()].score() {
                         info!(
                             "switched server from {} (score: {}) to {} (score: {})",
                             inner.best_server().config.addr(),
                             inner.best_server().score(),
                             choosen_svr.config.addr(),
                             choosen_svr.score(),
                         );
-                    }
 
-                    inner.set_best_idx(svr_idx);
+                        inner.set_best_idx(svr_idx);
+                    }
                 }
             });
         }
diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs
@@ -3,8 +3,10 @@
 use std::{
     io::{self, Cursor},
     net::{IpAddr, Ipv4Addr, SocketAddr},
-    sync::atomic::{AtomicBool, Ordering},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
     time::Duration,
 };
 
@@ -57,7 +59,7 @@ impl UdpAssociation {
         src_addr: SocketAddr,
         mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
     ) -> io::Result<UdpAssociation> {
-        debug!("Created UDP Association for {}", src_addr);
+        debug!("created UDP Association for {}", src_addr);
 
         // Create a socket for receiving packets
         let local_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
@@ -82,7 +84,7 @@ impl UdpAssociation {
                 if let Err(err) =
                     UdpAssociation::relay_l2r(&*context, src_addr, &mut sender, &pkt[..], timeout, &*c_svr_cfg).await
                 {
-                    error!("Failed to relay packet, {} -> ..., error: {}", src_addr, err);
+                    error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);
 
                     // FIXME: Ignore? Or how to deal with it?
                 }
@@ -101,10 +103,10 @@ impl UdpAssociation {
                 match UdpAssociation::relay_r2l(src_addr, &mut receiver, timeout, &mut response_tx, &*svr_cfg).await {
                     Ok(..) => {}
                     Err(ref err) if err.kind() == ErrorKind::TimedOut => {
-                        trace!("Receive packet timeout, {} <- ...", src_addr);
+                        trace!("receive packet timeout, {} <- ...", src_addr);
                     }
                     Err(err) => {
-                        error!("Failed to receive packet, {} <- .., error: {}", src_addr, err);
+                        error!("railed to receive packet, {} <- .., error: {}", src_addr, err);
 
                         // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                         // break;
@@ -132,12 +134,12 @@ impl UdpAssociation {
         let decrypted_pkt = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), pkt) {
             Ok(Some(pkt)) => pkt,
             Ok(None) => {
-                error!("Failed to decrypt pkt in UDP relay, packet too short");
+                error!("failed to decrypt pkt in UDP relay, packet too short");
                 let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                 return Err(err);
             }
             Err(err) => {
-                error!("Failed to decrypt pkt in UDP relay: {}", err);
+                error!("failed to decrypt pkt in UDP relay: {}", err);
                 let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                 return Err(err);
             }
@@ -232,7 +234,7 @@ impl UdpAssociation {
 
         // Send back to src_addr
         if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
-            error!("Failed to send packet into response channel, error: {}", err);
+            error!("failed to send packet into response channel, error: {}", err);
 
             // FIXME: What to do? Ignore?
         }
@@ -244,7 +246,7 @@ impl UdpAssociation {
         match self.tx.send(pkt.to_vec()).await {
             Ok(..) => true,
             Err(err) => {
-                error!("Failed to send packet, error: {}", err);
+                error!("failed to send packet, error: {}", err);
                 false
             }
         }
@@ -290,7 +292,7 @@ async fn listen(context: SharedContext, svr_cfg: Arc<ServerConfig>) -> io::Resul
         // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
         let pkt = &pkt_buf[..recv_len];
 
-        trace!("Received UDP packet from {}, length {} bytes", src, recv_len);
+        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
 
         // Check or (re)create an association
         loop {
@@ -329,7 +331,7 @@ pub async fn run(context: SharedContext) -> io::Result<()> {
 
     match vec_fut.into_future().await.0 {
         Some(res) => {
-            error!("One of TCP servers exited unexpectly, result: {:?}", res);
+            error!("one of TCP servers exited unexpectly, result: {:?}", res);
             let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
             Err(err)
         }
diff --git a/src/relay/udprelay/socks5_local.rs b/src/relay/udprelay/socks5_local.rs
@@ -3,8 +3,10 @@
 use std::{
     io::{self, Cursor, ErrorKind, Read},
     net::{IpAddr, Ipv4Addr, SocketAddr},
-    sync::atomic::{AtomicBool, Ordering},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
     time::Duration,
 };
 
@@ -45,7 +47,7 @@ async fn parse_packet(pkt: &[u8]) -> io::Result<(Address, Vec<u8>)> {
     let header = UdpAssociateHeader::read_from(&mut cur).await?;
 
     if header.frag != 0 {
-        error!("Received UDP associate with frag != 0, which is not supported by ShadowSocks");
+        error!("received UDP associate with frag != 0, which is not supported by ShadowSocks");
         let err = io::Error::new(ErrorKind::Other, "unsupported UDP fragmentation");
         return Err(err);
     }
@@ -81,7 +83,7 @@ impl UdpAssociation {
         src_addr: SocketAddr,
         mut response_tx: mpsc::Sender<(SocketAddr, Vec<u8>)>,
     ) -> io::Result<UdpAssociation> {
-        debug!("Created UDP Association for {}", src_addr);
+        debug!("created UDP Association for {}", src_addr);
 
         // Create a socket for receiving packets
         let local_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
@@ -106,7 +108,7 @@ impl UdpAssociation {
                 if let Err(err) =
                     UdpAssociation::relay_l2r(&*context, src_addr, &mut sender, &pkt[..], timeout, &*c_svr_cfg).await
                 {
-                    error!("Failed to send packet {} -> ..., error: {}", src_addr, err);
+                    error!("failed to send packet {} -> ..., error: {}", src_addr, err);
 
                     // FIXME: Ignore? Or how to deal with it?
                 }
@@ -123,10 +125,10 @@ impl UdpAssociation {
                 match UdpAssociation::relay_r2l(src_addr, &mut receiver, timeout, &mut response_tx, &*svr_cfg).await {
                     Ok(..) => {}
                     Err(ref err) if err.kind() == ErrorKind::TimedOut => {
-                        trace!("Receive packet timeout, {} <- ...", src_addr);
+                        trace!("receive packet timeout, {} <- ...", src_addr);
                     }
                     Err(err) => {
-                        error!("Failed to receive packet, {} <- .., error: {}", src_addr, err);
+                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);
 
                         // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                         // break;
@@ -238,7 +240,7 @@ impl UdpAssociation {
 
         // Send back to src_addr
         if let Err(err) = response_tx.send((src_addr, payload)).await {
-            error!("Failed to send packet into response channel, error: {}", err);
+            error!("failed to send packet into response channel, error: {}", err);
 
             // FIXME: What to do? Ignore?
         }
@@ -250,7 +252,7 @@ impl UdpAssociation {
         match self.tx.send(pkt.to_vec()).await {
             Ok(..) => true,
             Err(err) => {
-                error!("Failed to send packet, error: {}", err);
+                error!("failed to send packet, error: {}", err);
                 false
             }
         }
@@ -295,7 +297,7 @@ async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> {
         // Copy bytes, because udp_associate runs in another tokio Task
         let pkt = &pkt_buf[..recv_len];
 
-        trace!("Received UDP packet from {}, length {} bytes", src, recv_len);
+        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
 
         // Pick a server
         let svr_cfg = balancer.pick_server();
diff --git a/src/relay/udprelay/tunnel_local.rs b/src/relay/udprelay/tunnel_local.rs
