优化停止逻辑

vnt/src/core/conn.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::net::Ipv4Addr;
+use std::ops::Deref;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -33,6 +34,32 @@ use crate::{nat, VntCallback};
 
 #[derive(Clone)]
 pub struct Vnt {
+    inner: Arc<VntInner>,
+}
+impl Vnt {
+    #[cfg(feature = "integrated_tun")]
+    pub fn new<Call: VntCallback>(config: Config, callback: Call) -> anyhow::Result<Self> {
+        let inner = Arc::new(VntInner::new(config, callback)?);
+        Ok(Self { inner })
+    }
+    #[cfg(not(feature = "integrated_tun"))]
+    pub fn new_device<Call: VntCallback, Device: DeviceWrite>(
+        config: Config,
+        callback: Call,
+        device: Device,
+    ) -> anyhow::Result<Self> {
+        let inner = Arc::new(VntInner::new_device(config, callback, device)?);
+        Ok(Self { inner })
+    }
+}
+impl Deref for Vnt {
+    type Target = VntInner;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+pub struct VntInner {
     stop_manager: StopManager,
     config: Config,
     current_device: Arc<AtomicCell<CurrentDeviceInfo>>,
@@ -48,10 +75,10 @@ pub struct Vnt {
     external_route: ExternalRoute,
 }
 
-impl Vnt {
+impl VntInner {
     #[cfg(feature = "integrated_tun")]
     pub fn new<Call: VntCallback>(config: Config, callback: Call) -> anyhow::Result<Self> {
-        Vnt::new_device0(config, callback, DeviceAdapter::default())
+        VntInner::new_device0(config, callback, DeviceAdapter::default())
     }
     #[cfg(not(feature = "integrated_tun"))]
     pub fn new_device<Call: VntCallback, Device: DeviceWrite>(
@@ -59,7 +86,7 @@ impl Vnt {
         callback: Call,
         device: Device,
     ) -> anyhow::Result<Self> {
-        Vnt::new_device0(config, callback, device)
+        VntInner::new_device0(config, callback, device)
     }
     fn new_device0<Call: VntCallback, Device: DeviceWrite>(
         config: Config,
@@ -384,7 +411,7 @@ pub fn start<Call: VntCallback>(
     )
 }
 
-impl Vnt {
+impl VntInner {
     pub fn name(&self) -> &str {
         &self.config.name
     }
@@ -449,6 +476,9 @@ impl Vnt {
         let _ = self.context.lock().take();
         self.stop_manager.stop()
     }
+    pub fn is_stopped(&self) -> bool {
+        self.stop_manager.is_stopped()
+    }
     pub fn add_stop_listener<F>(&self, name: String, f: F) -> anyhow::Result<crate::util::Worker>
     where
         F: FnOnce() + Send + 'static,
@@ -478,3 +508,8 @@ impl Vnt {
         }
     }
 }
+impl Drop for VntInner {
+    fn drop(&mut self) {
+        self.stop();
+    }
+}

vnt/src/util/notify.rs

@@ -36,8 +36,8 @@ impl StopManager {
     pub fn wait_timeout(&self, dur: Duration) -> bool {
         self.inner.wait_timeout(dur)
     }
-    pub fn is_stop(&self) -> bool {
+    pub fn is_stopped(&self) -> bool {
-        self.inner.state.load(Ordering::Acquire)
+        self.inner.is_stopped()
     }
 }
 
@@ -89,6 +89,9 @@ impl StopManagerInner {
             listener();
         }
     }
+    pub fn is_stopped(&self) -> bool {
+        self.worker_num.load(Ordering::Acquire) == 0
+    }
     fn wait(&self) {
         {
             let mut guard = self.park_threads.lock();
@@ -115,6 +118,7 @@ impl StopManagerInner {
         self.worker_num.load(Ordering::Acquire) == 0
     }
     fn stop_call(&self) {
+        self.stop();
         if let Some(call) = self.stop_call.lock().take() {
             call();
         }

vnt/src/util/scheduler.rs

@@ -1,5 +1,8 @@
 use crate::util::StopManager;
+use crossbeam_utils::atomic::AtomicCell;
 use std::collections::BinaryHeap;
+use std::sync::mpsc::TrySendError;
+use std::sync::Arc;
 use std::{
     cmp::Ordering,
     sync::mpsc::{sync_channel, Receiver, SyncSender},
@@ -10,45 +13,62 @@ struct DelayedTask {
     f: Box<dyn FnOnce(&Scheduler) + Send>,
     next: Instant,
 }
+
 impl Eq for DelayedTask {}
+
 impl PartialEq for DelayedTask {
     fn eq(&self, other: &Self) -> bool {
         self.next.eq(&other.next)
     }
 }
+
 impl PartialOrd for DelayedTask {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         self.next.partial_cmp(&other.next).map(|ord| ord.reverse())
     }
 }
+
 impl Ord for DelayedTask {
     fn cmp(&self, other: &Self) -> Ordering {
         self.next.cmp(&other.next).reverse()
     }
 }
+
 enum Op {
     Task(DelayedTask),
     Stop,
 }
+
 #[derive(Clone)]
 pub struct Scheduler {
     sender: SyncSender<Op>,
+    state: Arc<AtomicCell<SchedulerState>>,
 }
+
+#[derive(Copy, Clone, Eq, PartialEq)]
+enum SchedulerState {
+    Running,
+    ShutdownNow, // 立即停止任务执行，队列中剩余的任务不再执行
+    _Shutdown,   //执行完队列中剩余的任务再停止
+}
+
 impl Scheduler {
     pub fn new(stop_manager: StopManager) -> anyhow::Result<Self> {
         let (sender, receiver) = sync_channel::<Op>(32);
-        let s = Self { sender };
+        let state = Arc::new(AtomicCell::new(SchedulerState::Running));
+        let s = Self { sender, state };
         let s_inner = s.clone();
         let worker = {
             let scheduler = s.clone();
             stop_manager.add_listener("Scheduler".into(), move || {
-                scheduler.shutdown();
+                scheduler.shutdown_now();
             })?
         };
         std::thread::Builder::new()
             .name("Scheduler".into())
             .spawn(move || {
-                run(receiver, s_inner);
+                run(receiver, &s_inner);
+                s_inner.shutdown_now();
                 worker.stop_all();
             })
             .expect("Scheduler");
@@ -58,20 +78,44 @@ impl Scheduler {
     where
         F: FnOnce(&Scheduler) + Send + 'static,
     {
+        if self.state.load() != SchedulerState::Running {
+            log::error!("定时任务执行停止");
+            return false;
+        }
         let task = DelayedTask {
             f: Box::new(f),
             next: Instant::now().checked_add(time).unwrap(),
         };
-        self.sender.send(Op::Task(task)).is_ok()
+        // 如果是任务中调用此方法，那这里用send可能会导致整个定时任务阻塞
+        // 任务总数不能大于或等于通道长度，所以改成try_send快速失败
+        match self.sender.try_send(Op::Task(task)) {
+            Ok(_) => true,
+            Err(e) => {
+                match e {
+                    TrySendError::Full(_) => {
+                        log::error!("定时任务队列达到上限");
+                    }
+                    TrySendError::Disconnected(_) => {
+                        log::error!("定时任务执行停止 通道关闭");
+                    }
+                }
+                false
+            }
+        }
     }
-    pub fn shutdown(self) {
+    pub fn shutdown_now(&self) {
+        self.state.store(SchedulerState::ShutdownNow);
         let _ = self.sender.send(Op::Stop);
     }
 }
-fn run(receiver: Receiver<Op>, s_inner: Scheduler) {
+
+fn run(receiver: Receiver<Op>, s_inner: &Scheduler) {
     let mut binary_heap = BinaryHeap::<DelayedTask>::with_capacity(32);
     loop {
         while let Some(task) = binary_heap.peek() {
+            if s_inner.state.load() == SchedulerState::ShutdownNow {
+                return;
+            }
             let now = Instant::now();
             if now < task.next {
                 //需要等待对应时间
@@ -89,7 +133,7 @@ fn run(receiver: Receiver<Op>, s_inner: Scheduler) {
                 }
             } else {
                 if let Some(task) = binary_heap.pop() {
-                    (task.f)(&s_inner);
+                    (task.f)(s_inner);
                 }
             }
         }
@@ -120,6 +164,7 @@ fn run(receiver: Receiver<Op>, s_inner: Scheduler) {
         }
     }
 }
+
 fn add_task(op: Op, binary_heap: &mut BinaryHeap<DelayedTask>) -> bool {
     return match op {
         Op::Task(task) => {