feat: unified ServerHandle

crates/shadowsocks-service/src/lib.rs

@@ -70,6 +70,7 @@ pub mod net;
 #[cfg(feature = "server")]
 pub mod server;
 mod sys;
+mod utils;
 
 /// Default UDP association's expire duration
 #[allow(dead_code)]

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

@@ -1,27 +1,24 @@
 //! Shadowsocks Local Server
 
 use std::{
-    future::Future,
     io::{self, ErrorKind},
-    pin::Pin,
     sync::Arc,
-    task::{Context, Poll},
     time::Duration,
 };
 
-use futures::{future, ready};
+use futures::future;
 use log::trace;
 use shadowsocks::{
     config::Mode,
     net::{AcceptOpts, ConnectOpts},
 };
-use tokio::task::JoinHandle;
 
 #[cfg(feature = "local-flow-stat")]
 use crate::{config::LocalFlowStatAddress, net::FlowStat};
 use crate::{
     config::{Config, ConfigType, ProtocolType},
     dns::build_dns_resolver,
+    utils::ServerHandle,
 };
 
 use self::{
@@ -70,27 +67,6 @@ pub mod utils;
 /// This is borrowed from Go's `net` library's default setting
 pub(crate) const LOCAL_DEFAULT_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(15);
 
-struct ServerHandle(JoinHandle<io::Result<()>>);
-
-impl Drop for ServerHandle {
-    #[inline]
-    fn drop(&mut self) {
-        self.0.abort();
-    }
-}
-
-impl Future for ServerHandle {
-    type Output = io::Result<()>;
-
-    #[inline]
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match ready!(Pin::new(&mut self.0).poll(cx)) {
-            Ok(res) => res.into(),
-            Err(err) => Err(io::Error::new(ErrorKind::Other, err)).into(),
-        }
-    }
-}
-
 /// Local Server instance
 pub struct Server {
     balancer: PingBalancer,

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

@@ -1,22 +1,15 @@
 //! Shadowsocks server
 
-use std::{
-    future::Future,
-    io::{self, ErrorKind},
-    pin::Pin,
-    sync::Arc,
-    task::{Context, Poll},
-    time::Duration,
-};
+use std::{io, sync::Arc, time::Duration};
 
-use futures::{future, ready};
+use futures::future;
 use log::trace;
 use shadowsocks::net::{AcceptOpts, ConnectOpts};
-use tokio::task::JoinHandle;
 
 use crate::{
     config::{Config, ConfigType},
     dns::build_dns_resolver,
+    utils::ServerHandle,
 };
 
 pub use self::{
@@ -170,24 +163,3 @@ pub async fn run(config: Config) -> io::Result<()> {
     let (res, ..) = future::select_all(vfut).await;
     res
 }
-
-struct ServerHandle(JoinHandle<io::Result<()>>);
-
-impl Drop for ServerHandle {
-    #[inline]
-    fn drop(&mut self) {
-        self.0.abort();
-    }
-}
-
-impl Future for ServerHandle {
-    type Output = io::Result<()>;
-
-    #[inline]
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match ready!(Pin::new(&mut self.0).poll(cx)) {
-            Ok(res) => res.into(),
-            Err(err) => Err(io::Error::new(ErrorKind::Other, err)).into(),
-        }
-    }
-}

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

@@ -2,15 +2,12 @@
 
 use std::{
     collections::HashMap,
-    future::Future,
     io::{self, ErrorKind},
-    pin::Pin,
     sync::Arc,
-    task::{Context, Poll},
     time::Duration,
 };
 
-use futures::{future, ready};
+use futures::future;
 use log::{error, trace};
 use shadowsocks::{
     config::{ManagerAddr, ServerConfig},
@@ -19,9 +16,9 @@ use shadowsocks::{
     plugin::{Plugin, PluginMode},
     ManagerClient,
 };
-use tokio::{task::JoinHandle, time};
+use tokio::time;
 
-use crate::{acl::AccessControl, config::SecurityConfig, net::FlowStat};
+use crate::{acl::AccessControl, config::SecurityConfig, net::FlowStat, utils::ServerHandle};
 
 use super::{context::ServiceContext, tcprelay::TcpServer, udprelay::UdpServer};
 
@@ -162,27 +159,6 @@ impl ServerBuilder {
     }
 }
 
-struct ServerHandle(JoinHandle<io::Result<()>>);
-
-impl Drop for ServerHandle {
-    #[inline]
-    fn drop(&mut self) {
-        self.0.abort();
-    }
-}
-
-impl Future for ServerHandle {
-    type Output = io::Result<()>;
-
-    #[inline]
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match ready!(Pin::new(&mut self.0).poll(cx)) {
-            Ok(res) => res.into(),
-            Err(err) => Err(io::Error::new(ErrorKind::Other, err)).into(),
-        }
-    }
-}
-
 /// Shadowsocks Server instance
 pub struct Server {
     context: Arc<ServiceContext>,

crates/shadowsocks-service/src/utils.rs

@@ -0,0 +1,36 @@
+//! Service Utilities
+
+use std::{
+    future::Future,
+    io::{self, ErrorKind},
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use futures::ready;
+use tokio::task::JoinHandle;
+
+/// Wrapper of `tokio::task::JoinHandle`, which links to a server instance.
+///
+/// `ServerHandle` implements `Future` which will join the `JoinHandle` and get the result.
+/// When `ServerHandle` drops, it will abort the task.
+pub struct ServerHandle(pub JoinHandle<io::Result<()>>);
+
+impl Drop for ServerHandle {
+    #[inline]
+    fn drop(&mut self) {
+        self.0.abort();
+    }
+}
+
+impl Future for ServerHandle {
+    type Output = io::Result<()>;
+
+    #[inline]
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match ready!(Pin::new(&mut self.0).poll(cx)) {
+            Ok(res) => res.into(),
+            Err(err) => Err(io::Error::new(ErrorKind::Other, err)).into(),
+        }
+    }
+}