utils: move code to a dedicated mod; add AsyncRuntime (#2072)

easytier-gui/src-tauri/src/lib.rs

@@ -24,7 +24,7 @@ use easytier::{
     tunnel::TunnelListener,
     tunnel::ring::RingTunnelListener,
     tunnel::tcp::TcpTunnelListener,
-    utils::{self},
+    utils::panic::setup_panic_handler,
 };
 use std::ops::Deref;
 use std::sync::Arc;
@@ -1120,7 +1120,7 @@ pub fn run_gui() -> std::process::ExitCode {
         process::exit(0);
     }
 
-    utils::setup_panic_handler();
+    setup_panic_handler();
 
     let mut builder = tauri::Builder::default();
 

easytier-web/src/main.rs

@@ -17,7 +17,7 @@ use easytier::{
         network::{local_ipv4, local_ipv6},
     },
     tunnel::{TunnelListener, tcp::TcpTunnelListener, udp::UdpTunnelListener},
-    utils::setup_panic_handler,
+    utils::panic::setup_panic_handler,
 };
 
 use easytier::tunnel::IpScheme;

easytier/src/common/ifcfg/mod.rs

@@ -119,8 +119,8 @@ async fn run_shell_cmd(cmd: &str) -> Result<(), Error> {
             .creation_flags(CREATE_NO_WINDOW)
             .output()
             .await?;
-        stdout = crate::utils::utf8_or_gbk_to_string(cmd_out.stdout.as_slice());
-        stderr = crate::utils::utf8_or_gbk_to_string(cmd_out.stderr.as_slice());
+        stdout = crate::utils::string::utf8_or_gbk_to_string(cmd_out.stdout.as_slice());
+        stderr = crate::utils::string::utf8_or_gbk_to_string(cmd_out.stderr.as_slice());
     };
 
     #[cfg(not(target_os = "windows"))]

easytier/src/common/scoped_task.rs

@@ -4,40 +4,25 @@
 //! For example, if task A spawned task B but is doing something else, and task B is waiting for task C to join,
 //! aborting A will also abort both B and C.
 
+use derive_more::{Deref, DerefMut, From};
 use std::future::Future;
-use std::ops::Deref;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 use tokio::task::JoinHandle;
 
-#[derive(Debug)]
-pub struct ScopedTask<T> {
-    inner: JoinHandle<T>,
-}
+#[derive(Debug, From, Deref, DerefMut)]
+pub struct ScopedTask<T>(JoinHandle<T>);
 
 impl<T> Drop for ScopedTask<T> {
     fn drop(&mut self) {
-        self.inner.abort()
+        self.abort()
     }
 }
 
 impl<T> Future for ScopedTask<T> {
     type Output = <JoinHandle<T> as Future>::Output;
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        Pin::new(&mut self.inner).poll(cx)
-    }
-}
-
-impl<T> From<JoinHandle<T>> for ScopedTask<T> {
-    fn from(inner: JoinHandle<T>) -> Self {
-        Self { inner }
-    }
-}
-
-impl<T> Deref for ScopedTask<T> {
-    type Target = JoinHandle<T>;
-    fn deref(&self) -> &Self::Target {
-        &self.inner
+        Pin::new(&mut self.0).poll(cx)
     }
 }
 

easytier/src/core.rs

@@ -23,7 +23,7 @@ use crate::{
     launcher::add_proxy_network_to_config,
     proto::common::{CompressionAlgoPb, SecureModeConfig},
     rpc_service::ApiRpcServer,
-    utils::setup_panic_handler,
+    utils::panic::setup_panic_handler,
     web_client,
 };
 use anyhow::Context;

easytier/src/easytier-cli.rs

@@ -76,7 +76,7 @@ use easytier::{
         rpc_types::controller::BaseController,
     },
     tunnel::{TunnelScheme, tcp::TcpTunnelConnector},
-    utils::{PeerRoutePair, cost_to_str},
+    utils::{PeerRoutePair, string::cost_to_str},
 };
 
 rust_i18n::i18n!("locales", fallback = "en");

easytier/src/utils/mod.rs

@@ -0,0 +1,30 @@
+pub mod panic;
+pub mod string;
+pub mod task;
+
+use std::net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener};
+use std::sync::{Arc, Weak};
+
+pub type PeerRoutePair = crate::proto::api::instance::PeerRoutePair;
+
+pub fn check_tcp_available(port: u16) -> bool {
+    let s = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port);
+    TcpListener::bind(s).is_ok()
+}
+
+pub fn find_free_tcp_port(mut range: std::ops::Range<u16>) -> Option<u16> {
+    range.find(|&port| check_tcp_available(port))
+}
+
+pub fn weak_upgrade<T>(weak: &Weak<T>) -> anyhow::Result<Arc<T>> {
+    weak.upgrade()
+        .ok_or_else(|| anyhow::anyhow!("{} not available", std::any::type_name::<T>()))
+}
+
+pub trait BoxExt: Sized {
+    fn boxed(self) -> Box<Self> {
+        Box::new(self)
+    }
+}
+
+impl<T> BoxExt for T {}

easytier/src/utils/panic.rs

@@ -1,43 +1,14 @@
 use crate::common::log;
 use indoc::formatdoc;
-use std::sync::Arc;
-use std::{fs::OpenOptions, str::FromStr};
-
-pub type PeerRoutePair = crate::proto::api::instance::PeerRoutePair;
-
-pub fn cost_to_str(cost: i32) -> String {
-    if cost == 1 {
-        "p2p".to_string()
-    } else {
-        format!("relay({})", cost)
-    }
-}
-
-pub fn float_to_str(f: f64, precision: usize) -> String {
-    format!("{:.1$}", f, precision)
-}
-
-#[cfg(target_os = "windows")]
-pub fn utf8_or_gbk_to_string(s: &[u8]) -> String {
-    use encoding::{DecoderTrap, Encoding, all::GBK};
-    if let Ok(utf8_str) = String::from_utf8(s.to_vec()) {
-        utf8_str
-    } else {
-        // 如果解码失败，则尝试使用GBK解码
-        if let Ok(gbk_str) = GBK.decode(s, DecoderTrap::Strict) {
-            gbk_str
-        } else {
-            String::from_utf8_lossy(s).to_string()
-        }
-    }
-}
+use std::fs::OpenOptions;
+use std::str::FromStr;
+use std::{backtrace, io::Write};
 
 thread_local! {
     static PANIC_COUNT : std::cell::RefCell<u32> = const { std::cell::RefCell::new(0) };
 }
 
 pub fn setup_panic_handler() {
-    use std::{backtrace, io::Write};
     std::panic::set_hook(Box::new(|info| {
         let mut stderr = std::io::stderr();
         let sep = format!("{}\n", "=======".repeat(10));
@@ -126,26 +97,3 @@ pub fn setup_panic_handler() {
         std::process::exit(1);
     }));
 }
-
-pub fn check_tcp_available(port: u16) -> bool {
-    use std::net::TcpListener;
-    let s = std::net::SocketAddr::new(std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED), port);
-    TcpListener::bind(s).is_ok()
-}
-
-pub fn find_free_tcp_port(mut range: std::ops::Range<u16>) -> Option<u16> {
-    range.find(|&port| check_tcp_available(port))
-}
-
-pub fn weak_upgrade<T>(weak: &std::sync::Weak<T>) -> anyhow::Result<Arc<T>> {
-    weak.upgrade()
-        .ok_or_else(|| anyhow::anyhow!("{} not available", std::any::type_name::<T>()))
-}
-
-pub trait BoxExt: Sized {
-    fn boxed(self) -> Box<Self> {
-        Box::new(self)
-    }
-}
-
-impl<T> BoxExt for T {}

easytier/src/utils/string.rs

@@ -0,0 +1,26 @@
+pub fn cost_to_str(cost: i32) -> String {
+    if cost == 1 {
+        "p2p".to_string()
+    } else {
+        format!("relay({})", cost)
+    }
+}
+
+pub fn float_to_str(f: f64, precision: usize) -> String {
+    format!("{:.1$}", f, precision)
+}
+
+#[cfg(target_os = "windows")]
+pub fn utf8_or_gbk_to_string(s: &[u8]) -> String {
+    use encoding::{DecoderTrap, Encoding, all::GBK};
+    if let Ok(utf8_str) = String::from_utf8(s.to_vec()) {
+        utf8_str
+    } else {
+        // 如果解码失败，则尝试使用GBK解码
+        if let Ok(gbk_str) = GBK.decode(s, DecoderTrap::Strict) {
+            gbk_str
+        } else {
+            String::from_utf8_lossy(s).to_string()
+        }
+    }
+}

easytier/src/utils/task.rs

@@ -0,0 +1,140 @@
+use crate::common::scoped_task::ScopedTask;
+use derivative::Derivative;
+use derive_more::{Deref, DerefMut};
+use parking_lot::Mutex;
+use std::future::Future;
+use std::mem::take;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::Notify;
+use tokio::task::{AbortHandle, JoinError};
+use tokio_util::sync::CancellationToken;
+
+#[derive(Derivative, Debug)]
+#[derivative(Default(bound = ""))]
+enum AsyncRuntimeState<R: Send + 'static> {
+    #[derivative(Default)]
+    Idle,
+    Running {
+        id: tokio::task::Id,
+        task: ScopedTask<R>,
+        token: CancellationToken,
+    },
+    Stopping(AbortHandle),
+}
+
+#[derive(Derivative, Debug)]
+#[derivative(Default(bound = ""))]
+pub struct AsyncRuntimeInner<R: Send + 'static = ()> {
+    state: Mutex<AsyncRuntimeState<R>>,
+    idle: Notify,
+}
+
+#[derive(Derivative, Deref, DerefMut)]
+#[derivative(Debug = "transparent", Default(bound = ""), Clone(bound = ""))]
+pub struct AsyncRuntime<R: Send + 'static = ()>(Arc<AsyncRuntimeInner<R>>);
+
+impl<R: Send + 'static> AsyncRuntime<R> {
+    pub fn token(&self) -> Option<CancellationToken> {
+        if let AsyncRuntimeState::Running { token, .. } = &*self.state.lock() {
+            Some(token.clone())
+        } else {
+            None
+        }
+    }
+
+    pub fn start<F, Fut>(&self, token: Option<CancellationToken>, factory: F) -> anyhow::Result<()>
+    where
+        F: FnOnce(CancellationToken) -> Fut,
+        Fut: Future<Output = R> + Send + 'static,
+    {
+        let mut state = self.state.lock();
+        if !matches!(*state, AsyncRuntimeState::Idle) {
+            return Err(anyhow::anyhow!("task is already running/stopping"));
+        }
+
+        let token = token.unwrap_or_default();
+
+        let task = {
+            let f = factory(token.clone());
+            let this = (*self).clone();
+            tokio::spawn(async move {
+                let result = f.await;
+                let mut state = this.state.lock();
+                if let AsyncRuntimeState::Running { id, .. } = &*state
+                    && *id == tokio::task::id()
+                {
+                    take(&mut *state);
+                }
+                result
+            })
+        };
+
+        *state = AsyncRuntimeState::Running {
+            id: task.id(),
+            task: task.into(),
+            token,
+        };
+
+        Ok(())
+    }
+
+    pub async fn stop(&self, timeout: Duration) -> Option<Result<R, JoinError>> {
+        let state = {
+            let mut state = self.state.lock();
+            match &*state {
+                AsyncRuntimeState::Running { .. } => {
+                    let AsyncRuntimeState::Running { task, token, .. } = take(&mut *state) else {
+                        unreachable!()
+                    };
+                    *state = AsyncRuntimeState::Stopping(task.abort_handle());
+                    Ok((task, token))
+                }
+                AsyncRuntimeState::Stopping(_) => Err(self.idle.notified()),
+                AsyncRuntimeState::Idle => return None,
+            }
+        };
+
+        let (mut task, token) = match state {
+            Ok(running) => running,
+            Err(stopping) => {
+                stopping.await;
+                return None;
+            }
+        };
+
+        token.cancel();
+        let result = if let Ok(result) = tokio::time::timeout(timeout, &mut task).await {
+            result
+        } else {
+            task.abort();
+            tracing::warn!("task stop timeout after {:?}, aborted", timeout);
+            task.await
+        };
+
+        {
+            let mut state = self.state.lock();
+            if matches!(*state, AsyncRuntimeState::Stopping(_)) {
+                *state = AsyncRuntimeState::Idle;
+                drop(state);
+                self.idle.notify_waiters();
+            }
+        }
+
+        Some(result)
+    }
+
+    pub fn abort(&self) {
+        let mut state = self.state.lock();
+        match &*state {
+            AsyncRuntimeState::Running { task, .. } => {
+                task.abort();
+                *state = AsyncRuntimeState::Idle;
+                drop(state);
+                self.idle.notify_waiters();
+            }
+            AsyncRuntimeState::Stopping(handle) => handle.abort(),
+            _ => {}
+        }
+    }
+}