feat(tun): poll thread park only if tx channel is empty

Author: zonyitoo

crates/shadowsocks-service/src/local/tun/tcp.rs

@@ -239,6 +239,7 @@ pub struct TcpTun {
     balancer: PingBalancer,
     iface_rx: mpsc::UnboundedReceiver<Vec<u8>>,
     iface_tx: mpsc::UnboundedSender<Vec<u8>>,
+    iface_tx_avail: Arc<AtomicBool>,
 }
 
 impl Drop for TcpTun {
@@ -255,7 +256,7 @@ impl TcpTun {
         capabilities.medium = Medium::Ip;
         capabilities.max_transmission_unit = mtu as usize;
 
-        let (mut device, iface_rx, iface_tx) = VirtTunDevice::new(capabilities);
+        let (mut device, iface_rx, iface_tx, iface_tx_avail) = VirtTunDevice::new(capabilities);
 
         let mut iface_config = InterfaceConfig::default();
         iface_config.random_seed = rand::random();
@@ -444,11 +445,13 @@ impl TcpTun {
                         socket_set.remove(socket_handle);
                     }
 
-                    let next_duration = iface
-                        .poll_delay(before_poll, &socket_set)
-                        .unwrap_or(SmolDuration::from_millis(5));
-                    if next_duration != SmolDuration::ZERO {
-                        thread::park_timeout(Duration::from(next_duration));
+                    if !device.recv_available() {
+                        let next_duration = iface
+                            .poll_delay(before_poll, &socket_set)
+                            .unwrap_or(SmolDuration::from_millis(5));
+                        if next_duration != SmolDuration::ZERO {
+                            thread::park_timeout(Duration::from(next_duration));
+                        }
                     }
                 }
 
@@ -467,6 +470,7 @@ impl TcpTun {
             balancer,
             iface_rx,
             iface_tx,
+            iface_tx_avail,
         }
     }
 
@@ -525,6 +529,7 @@ impl TcpTun {
         }
 
         // Wake up and poll the interface.
+        self.iface_tx_avail.store(true, Ordering::Release);
         self.manager_notify.notify();
     }
 

crates/shadowsocks-service/src/local/tun/virt_device.rs

@@ -1,6 +1,12 @@
 //! Virtual Device for receiving packets from tun
 
-use std::marker::PhantomData;
+use std::{
+    marker::PhantomData,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+};
 
 use smoltcp::{
     phy::{self, Device, DeviceCapabilities},
@@ -12,25 +18,39 @@ pub struct VirtTunDevice {
     capabilities: DeviceCapabilities,
     in_buf: mpsc::UnboundedReceiver<Vec<u8>>,
     out_buf: mpsc::UnboundedSender<Vec<u8>>,
+    in_buf_avail: Arc<AtomicBool>,
 }
 
 impl VirtTunDevice {
     pub fn new(
         capabilities: DeviceCapabilities,
-    ) -> (Self, mpsc::UnboundedReceiver<Vec<u8>>, mpsc::UnboundedSender<Vec<u8>>) {
+    ) -> (
+        Self,
+        mpsc::UnboundedReceiver<Vec<u8>>,
+        mpsc::UnboundedSender<Vec<u8>>,
+        Arc<AtomicBool>,
+    ) {
         let (iface_tx, iface_output) = mpsc::unbounded_channel();
         let (iface_input, iface_rx) = mpsc::unbounded_channel();
+        let in_buf_avail = Arc::new(AtomicBool::new(false));
 
         (
             Self {
                 capabilities,
                 in_buf: iface_rx,
                 out_buf: iface_tx,
+                in_buf_avail: in_buf_avail.clone(),
             },
             iface_output,
             iface_input,
+            in_buf_avail,
         )
     }
+
+    #[inline]
+    pub fn recv_available(&self) -> bool {
+        self.in_buf_avail.load(Ordering::Acquire)
+    }
 }
 
 impl Device for VirtTunDevice {
@@ -46,6 +66,7 @@ impl Device for VirtTunDevice {
             let tx = VirtTxToken(self);
             return Some((rx, tx));
         }
+        self.in_buf_avail.store(false, Ordering::Release);
         None
     }
 

crates/shadowsocks-service/src/manager/server.rs

@@ -11,12 +11,22 @@ use shadowsocks::{
     crypto::CipherKind,
     dns_resolver::DnsResolver,
     manager::protocol::{
-        self, AddRequest, AddResponse, ErrorResponse, ListResponse, ManagerRequest, PingResponse, RemoveRequest,
-        RemoveResponse, ServerUserConfig, StatRequest,
+        self,
+        AddRequest,
+        AddResponse,
+        ErrorResponse,
+        ListResponse,
+        ManagerRequest,
+        PingResponse,
+        RemoveRequest,
+        RemoveResponse,
+        ServerUserConfig,
+        StatRequest,
     },
     net::{AcceptOpts, ConnectOpts},
     plugin::PluginConfig,
-    ManagerListener, ServerAddr,
+    ManagerListener,
+    ServerAddr,
 };
 use tokio::{sync::Mutex, task::JoinHandle};