active cleanup of expired udp associations, fixed #351

crates/shadowsocks-service/src/local/redir/udprelay/mod.rs

@@ -40,6 +40,13 @@ pub struct UdpRedir {
     context: Arc<ServiceContext>,
     redir_ty: RedirType,
     assoc_map: Arc<Mutex<LruCache<SocketAddr, Arc<UdpAssociation>>>>,
+    cleanup_abortable: AbortHandle,
+}
+
+impl Drop for UdpRedir {
+    fn drop(&mut self) {
+        self.cleanup_abortable.abort();
+    }
 }
 
 impl UdpRedir {
@@ -50,14 +57,31 @@ impl UdpRedir {
         capacity: Option<usize>,
     ) -> UdpRedir {
         let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION);
+        let assoc_map = Arc::new(Mutex::new(match capacity {
+            Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity),
+            None => LruCache::with_expiry_duration(time_to_live),
+        }));
+
+        let cleanup_abortable = {
+            let assoc_map = assoc_map.clone();
+            let (cleanup_task, cleanup_abortable) = future::abortable(async move {
+                let mut interval = time::interval(time_to_live);
+                loop {
+                    interval.tick().await;
+
+                    // iter() will trigger a cleanup of expired associations
+                    let _ = assoc_map.lock().await.iter();
+                }
+            });
+            tokio::spawn(cleanup_task);
+            cleanup_abortable
+        };
 
         UdpRedir {
             context,
             redir_ty,
-            assoc_map: Arc::new(Mutex::new(match capacity {
-                Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity),
-                None => LruCache::with_expiry_duration(time_to_live),
-            })),
+            assoc_map,
+            cleanup_abortable,
         }
     }
 

crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/bsd.rs

@@ -52,19 +52,7 @@ impl UdpRedirSocket {
 
     /// Send data to the socket to the given target address
     pub async fn send_to(&mut self, buf: &[u8], target: SocketAddr) -> io::Result<usize> {
-        // poll_fn(|cx| self.poll_send_to(cx, buf, target)).await
-
-        loop {
-            let mut write_guard = self.io.writable().await?;
-
-            match self.io.get_ref().send_to(buf, target) {
-                Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
-                    write_guard.clear_ready();
-                    // false positive
-                }
-                x => return x,
-            }
-        }
+        poll_fn(|cx| self.poll_send_to(cx, buf, target)).await
     }
 
     fn poll_send_to(&self, cx: &mut Context<'_>, buf: &[u8], target: SocketAddr) -> Poll<io::Result<usize>> {

crates/shadowsocks-service/src/local/redir/udprelay/sys/unix/linux.rs

@@ -8,7 +8,7 @@ use std::{
     task::{Context, Poll},
 };
 
-use futures::ready;
+use futures::{future::poll_fn, ready};
 use socket2::{Domain, Protocol, SockAddr, Socket, Type};
 use tokio::io::unix::AsyncFd;
 

crates/shadowsocks-service/src/local/socks/server/socks5/udprelay.rs

@@ -37,6 +37,13 @@ use crate::{
 pub struct Socks5UdpServer {
     context: Arc<ServiceContext>,
     assoc_map: Arc<Mutex<LruCache<SocketAddr, Arc<UdpAssociation>>>>,
+    cleanup_abortable: AbortHandle,
+}
+
+impl Drop for Socks5UdpServer {
+    fn drop(&mut self) {
+        self.cleanup_abortable.abort();
+    }
 }
 
 impl Socks5UdpServer {
@@ -46,13 +53,30 @@ impl Socks5UdpServer {
         capacity: Option<usize>,
     ) -> Socks5UdpServer {
         let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION);
+        let assoc_map = Arc::new(Mutex::new(match capacity {
+            Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity),
+            None => LruCache::with_expiry_duration(time_to_live),
+        }));
+
+        let cleanup_abortable = {
+            let assoc_map = assoc_map.clone();
+            let (cleanup_task, cleanup_abortable) = future::abortable(async move {
+                let mut interval = time::interval(time_to_live);
+                loop {
+                    interval.tick().await;
+
+                    // iter() will trigger a cleanup of expired associations
+                    let _ = assoc_map.lock().await.iter();
+                }
+            });
+            tokio::spawn(cleanup_task);
+            cleanup_abortable
+        };
 
         Socks5UdpServer {
             context,
-            assoc_map: Arc::new(Mutex::new(match capacity {
-                Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity),
-                None => LruCache::with_expiry_duration(time_to_live),
-            })),
+            assoc_map,
+            cleanup_abortable,
         }
     }
 

crates/shadowsocks-service/src/local/tunnel/udprelay.rs

@@ -30,18 +30,42 @@ use crate::{
 pub struct UdpTunnel {
     context: Arc<ServiceContext>,
     assoc_map: Arc<Mutex<LruCache<SocketAddr, UdpAssociation>>>,
+    cleanup_abortable: AbortHandle,
+}
+
+impl Drop for UdpTunnel {
+    fn drop(&mut self) {
+        self.cleanup_abortable.abort();
+    }
 }
 
 impl UdpTunnel {
     pub fn new(context: Arc<ServiceContext>, time_to_live: Option<Duration>, capacity: Option<usize>) -> UdpTunnel {
         let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION);
+        let assoc_map = Arc::new(Mutex::new(match capacity {
+            Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity),
+            None => LruCache::with_expiry_duration(time_to_live),
+        }));
+
+        let cleanup_abortable = {
+            let assoc_map = assoc_map.clone();
+            let (cleanup_task, cleanup_abortable) = future::abortable(async move {
+                let mut interval = time::interval(time_to_live);
+                loop {
+                    interval.tick().await;
+
+                    // iter() will trigger a cleanup of expired associations
+                    let _ = assoc_map.lock().await.iter();
+                }
+            });
+            tokio::spawn(cleanup_task);
+            cleanup_abortable
+        };
 
         UdpTunnel {
             context,
-            assoc_map: Arc::new(Mutex::new(match capacity {
-                Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity),
-                None => LruCache::with_expiry_duration(time_to_live),
-            })),
+            assoc_map,
+            cleanup_abortable,
         }
     }
 

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

@@ -29,18 +29,42 @@ use super::context::ServiceContext;
 pub struct UdpServer {
     context: Arc<ServiceContext>,
     assoc_map: Arc<Mutex<LruCache<SocketAddr, Arc<UdpAssociation>>>>,
+    cleanup_abortable: AbortHandle,
+}
+
+impl Drop for UdpServer {
+    fn drop(&mut self) {
+        self.cleanup_abortable.abort();
+    }
 }
 
 impl UdpServer {
     pub fn new(context: Arc<ServiceContext>, time_to_live: Option<Duration>, capacity: Option<usize>) -> UdpServer {
         let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION);
+        let assoc_map = Arc::new(Mutex::new(match capacity {
+            Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity),
+            None => LruCache::with_expiry_duration(time_to_live),
+        }));
+
+        let cleanup_abortable = {
+            let assoc_map = assoc_map.clone();
+            let (cleanup_task, cleanup_abortable) = future::abortable(async move {
+                let mut interval = time::interval(time_to_live);
+                loop {
+                    interval.tick().await;
+
+                    // iter() will trigger a cleanup of expired associations
+                    let _ = assoc_map.lock().await.iter();
+                }
+            });
+            tokio::spawn(cleanup_task);
+            cleanup_abortable
+        };
 
         UdpServer {
             context,
-            assoc_map: Arc::new(Mutex::new(match capacity {
-                Some(capacity) => LruCache::with_expiry_duration_and_capacity(time_to_live, capacity),
-                None => LruCache::with_expiry_duration(time_to_live),
-            })),
+            assoc_map,
+            cleanup_abortable,
         }
     }