Revert "aborting tasks with JoinHandle instead of futures abortable"

This reverts commit 9ab1f8b.

Author: zonyitoo

Cargo.lock

@@ -2,6 +2,7 @@
 
 use std::{
     fmt::{self, Debug, Display},
+    future::Future,
     io,
     net::{Ipv4Addr, SocketAddr},
     sync::{
@@ -12,7 +13,7 @@ use std::{
 };
 
 use byte_string::ByteStr;
-use futures::future;
+use futures::future::{self, AbortHandle};
 use log::{debug, info, trace};
 use shadowsocks::{
     config::Mode,
@@ -25,7 +26,6 @@ use shadowsocks::{
 };
 use tokio::{
     io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
-    task::JoinHandle,
     time,
 };
 
@@ -72,7 +72,7 @@ impl PingBalancerBuilder {
         self.servers.push(Arc::new(server));
     }
 
-    pub async fn build(self) -> PingBalancer {
+    pub async fn build(self) -> (PingBalancer, impl Future<Output = ()>) {
         assert!(!self.servers.is_empty(), "build PingBalancer without any servers");
 
         let balancer_context = PingBalancerContext {
@@ -87,17 +87,21 @@ impl PingBalancerBuilder {
 
         let shared_context = Arc::new(balancer_context);
 
-        let abortable = {
+        let (checker, abortable) = {
             let shared_context = shared_context.clone();
-            tokio::spawn(async move { shared_context.checker_task().await })
+            future::abortable(async move { shared_context.checker_task().await })
+        };
+        let checker = async move {
+            let _ = checker.await;
         };
 
-        PingBalancer {
+        let balancer = PingBalancer {
             inner: Arc::new(PingBalancerInner {
                 context: shared_context,
                 abortable,
             }),
-        }
+        };
+        (balancer, checker)
     }
 }
 
@@ -256,7 +260,7 @@ impl PingBalancerContext {
 
 struct PingBalancerInner {
     context: Arc<PingBalancerContext>,
-    abortable: JoinHandle<()>,
+    abortable: AbortHandle,
 }
 
 impl Drop for PingBalancerInner {

crates/shadowsocks-service/src/local/loadbalancing/ping_balancer.rs

@@ -189,7 +189,10 @@ pub async fn run(mut config: Config) -> io::Result<()> {
         for server in config.server {
             balancer_builder.add_server(ServerIdent::new(server));
         }
-        balancer_builder.build().await
+        let (balancer, checker) = balancer_builder.build().await;
+        tokio::spawn(checker);
+
+        balancer
     };
 
     #[cfg(feature = "local-flow-stat")]

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

@@ -9,6 +9,7 @@ use std::{
 
 use async_trait::async_trait;
 use bytes::Bytes;
+use futures::future::{self, AbortHandle};
 use log::{debug, error, trace, warn};
 use lru_time_cache::LruCache;
 use shadowsocks::{
@@ -23,7 +24,6 @@ use spin::Mutex as SpinMutex;
 use tokio::{
     net::UdpSocket,
     sync::{mpsc, Mutex},
-    task::JoinHandle,
     time,
 };
 
@@ -53,8 +53,8 @@ where
     respond_writer: W,
     context: Arc<ServiceContext>,
     assoc_map: SharedAssociationMap<W>,
-    cleanup_abortable: JoinHandle<()>,
-    keepalive_abortable: JoinHandle<()>,
+    cleanup_abortable: AbortHandle,
+    keepalive_abortable: AbortHandle,
     keepalive_tx: mpsc::Sender<SocketAddr>,
     balancer: PingBalancer,
 }
@@ -89,25 +89,29 @@ where
 
         let cleanup_abortable = {
             let assoc_map = assoc_map.clone();
-            tokio::spawn(async move {
+            let (cleanup_task, cleanup_abortable) = future::abortable(async move {
                 loop {
                     time::sleep(time_to_live).await;
 
                     // cleanup expired associations. iter() will remove expired elements
                     let _ = assoc_map.lock().await.iter();
                 }
-            })
+            });
+            tokio::spawn(cleanup_task);
+            cleanup_abortable
         };
 
-        let (keepalive_tx, mut keepalive_rx) = mpsc::channel(64);
+        let (keepalive_tx, mut keepalive_rx) = mpsc::channel(256);
 
         let keepalive_abortable = {
             let assoc_map = assoc_map.clone();
-            tokio::spawn(async move {
+            let (keepalive_task, keepalive_abortable) = future::abortable(async move {
                 while let Some(peer_addr) = keepalive_rx.recv().await {
                     assoc_map.lock().await.get(&peer_addr);
                 }
-            })
+            });
+            tokio::spawn(keepalive_task);
+            keepalive_abortable
         };
 
         UdpAssociationManager {
@@ -195,7 +199,7 @@ enum UdpAssociationBypassState {
     Empty,
     Connected {
         socket: Arc<UdpSocket>,
-        abortable: JoinHandle<io::Result<()>>,
+        abortable: AbortHandle,
     },
     Aborted,
 }
@@ -213,7 +217,7 @@ impl UdpAssociationBypassState {
         UdpAssociationBypassState::Empty
     }
 
-    fn set_connected(&mut self, socket: Arc<UdpSocket>, abortable: JoinHandle<io::Result<()>>) {
+    fn set_connected(&mut self, socket: Arc<UdpSocket>, abortable: AbortHandle) {
         *self = UdpAssociationBypassState::Connected { socket, abortable };
     }
 
@@ -226,7 +230,7 @@ enum UdpAssociationProxyState {
     Empty,
     Connected {
         socket: Arc<MonProxySocket>,
-        abortable: JoinHandle<io::Result<()>>,
+        abortable: AbortHandle,
     },
     Aborted,
 }
@@ -247,7 +251,7 @@ impl UdpAssociationProxyState {
         *self = UdpAssociationProxyState::Empty;
     }
 
-    fn set_connected(&mut self, socket: Arc<MonProxySocket>, abortable: JoinHandle<io::Result<()>>) {
+    fn set_connected(&mut self, socket: Arc<MonProxySocket>, abortable: AbortHandle) {
         self.abort_inner();
         *self = UdpAssociationProxyState::Connected { socket, abortable };
     }
@@ -390,12 +394,13 @@ where
                         ShadowUdpSocket::connect_any_with_opts(&target_addr, self.context.connect_opts_ref()).await?;
                     let socket: Arc<UdpSocket> = Arc::new(socket.into());
 
-                    // CLIENT <- REMOTE
-                    let r2l_abortable = {
+                    let (r2l_fut, r2l_abortable) = {
                         let assoc = self.clone();
-                        tokio::spawn(assoc.copy_bypassed_r2l(socket.clone()))
+                        future::abortable(assoc.copy_bypassed_r2l(socket.clone()))
                     };
 
+                    // CLIENT <- REMOTE
+                    tokio::spawn(r2l_fut);
                     debug!(
                         "created udp association for {} (bypassed) with {:?}",
                         self.peer_addr,
@@ -444,12 +449,13 @@ where
                         ShadowUdpSocket::connect_any_with_opts(&target_addr, self.context.connect_opts_ref()).await?;
                     let socket: Arc<UdpSocket> = Arc::new(socket.into());
 
-                    // CLIENT <- REMOTE
-                    let r2l_abortable = {
+                    let (r2l_fut, r2l_abortable) = {
                         let assoc = self.clone();
-                        tokio::spawn(assoc.copy_bypassed_r2l(socket.clone()))
+                        future::abortable(assoc.copy_bypassed_r2l(socket.clone()))
                     };
 
+                    // CLIENT <- REMOTE
+                    tokio::spawn(r2l_fut);
                     debug!(
                         "created udp association for {} (bypassed) with {:?}",
                         self.peer_addr,
@@ -509,12 +515,14 @@ where
                         let socket = MonProxySocket::from_socket(socket, self.context.flow_stat());
                         let socket = Arc::new(socket);
 
-                        // CLIENT <- REMOTE
-                        let r2l_abortable = {
+                        let (r2l_fut, r2l_abortable) = {
                             let assoc = self.clone();
-                            tokio::spawn(assoc.copy_proxied_r2l(socket.clone()))
+                            future::abortable(assoc.copy_proxied_r2l(socket.clone()))
                         };
 
+                        // CLIENT <- REMOTE
+                        tokio::spawn(r2l_fut);
+
                         debug!(
                             "created udp association for {} <-> {} (proxied) with {:?}",
                             self.peer_addr,

crates/shadowsocks-service/src/local/net/udp/association.rs

@@ -8,6 +8,7 @@ use std::{
 };
 
 use async_trait::async_trait;
+use futures::future::{self, AbortHandle};
 use log::{error, info, trace, warn};
 use lru_time_cache::LruCache;
 use shadowsocks::{
@@ -16,7 +17,7 @@ use shadowsocks::{
     relay::{socks5::Address, udprelay::MAXIMUM_UDP_PAYLOAD_SIZE},
     ServerAddr,
 };
-use tokio::{sync::Mutex, task::JoinHandle};
+use tokio::sync::Mutex;
 
 use crate::{
     config::RedirType,
@@ -40,7 +41,7 @@ const INBOUND_SOCKET_CACHE_CAPACITY: usize = 256;
 
 struct UdpRedirInboundCache {
     cache: Arc<Mutex<LruCache<SocketAddr, Arc<UdpRedirSocket>>>>,
-    watcher: JoinHandle<()>,
+    watcher: AbortHandle,
 }
 
 impl Drop for UdpRedirInboundCache {
@@ -56,15 +57,16 @@ impl UdpRedirInboundCache {
             INBOUND_SOCKET_CACHE_CAPACITY,
         )));
 
-        let watcher = {
+        let (cleanup_fut, watcher) = {
             let cache = cache.clone();
-            tokio::spawn(async move {
+            future::abortable(async move {
                 loop {
                     tokio::time::sleep(INBOUND_SOCKET_CACHE_EXPIRATION).await;
                     let _ = cache.lock().await.iter();
                 }
             })
         };
+        tokio::spawn(cleanup_fut);
 
         UdpRedirInboundCache { cache, watcher }
     }