调整任务停止逻辑

vnt/src/handle/tun_tap/mod.rs

@@ -18,22 +18,20 @@ pub(crate) use windows::*;
 /// 仅仅是停止tun，不停止vnt
 #[derive(Clone, Default)]
 pub struct DeviceStop {
-    f: Arc<Mutex<Option<Box<dyn FnOnce() -> bool + Send>>>>,
+    f: Arc<Mutex<Option<Box<dyn FnOnce() + Send>>>>,
     stopped: Arc<AtomicCell<bool>>,
 }
 
 impl DeviceStop {
     pub fn set_stop_fn<F>(&self, f: F)
     where
-        F: FnOnce() -> bool + Send + 'static,
+        F: FnOnce() + Send + 'static,
     {
         self.f.lock().replace(Box::new(f));
     }
-    pub fn stop(&self) -> bool {
+    pub fn stop(&self) {
         if let Some(f) = self.f.lock().take() {
             f()
-        } else {
-            false
         }
     }
     pub fn stopped(&self) {

vnt/src/handle/tun_tap/unix.rs

@@ -35,36 +35,23 @@ pub(crate) fn start_simple(
     compressor: Compressor,
     device_stop: DeviceStop,
 ) -> anyhow::Result<()> {
-    let stop_all = Arc::new(AtomicCell::new(true));
     let poll = Poll::new()?;
     let waker = Arc::new(Waker::new(poll.registry(), STOP)?);
     let _waker = waker.clone();
-    let device_cell = Arc::new(AtomicCell::new(Some(waker)));
     let worker = {
-        let device_cell = device_cell.clone();
         stop_manager.add_listener("tun_device".into(), move || {
-            if let Some(waker) = device_cell.take() {
-                if let Err(e) = waker.wake() {
-                    log::warn!("{:?}", e);
-                }
+            if let Err(e) = waker.wake() {
+                log::warn!("{:?}", e);
             }
         })?
     };
-    {
-        let stop_all = stop_all.clone();
-        device_stop.set_stop_fn(move || {
-            if let Some(waker) = device_cell.take() {
-                stop_all.store(false);
-                if let Err(e) = waker.wake() {
-                    log::warn!("{:?}", e);
-                    return false;
-                }
-                true
-            } else {
-                false
-            }
-        });
-    }
+    let worker_cell = Arc::new(AtomicCell::new(Some(worker)));
+    let _worker_cell = worker_cell.clone();
+    device_stop.set_stop_fn(move || {
+        if let Some(worker) = _worker_cell.take() {
+            worker.stop_self()
+        }
+    });
     if let Err(e) = start_simple0(
         poll,
         context,
@@ -82,7 +69,7 @@ pub(crate) fn start_simple(
         log::error!("{:?}", e);
     };
     device_stop.stopped();
-    if stop_all.load() {
+    if let Some(worker) = worker_cell.take() {
         worker.stop_all();
     }
     drop(_waker);

vnt/src/handle/tun_tap/windows.rs

@@ -28,30 +28,21 @@ pub(crate) fn start_simple(
     compressor: Compressor,
     device_stop: DeviceStop,
 ) -> anyhow::Result<()> {
-    let device_cell = Arc::new(AtomicCell::new(Some(device.clone())));
-    let stop_all = Arc::new(AtomicCell::new(true));
     let worker = {
-        let device_cell = device_cell.clone();
+        let device = device.clone();
         stop_manager.add_listener("tun_device".into(), move || {
-            if let Some(device) = device_cell.take() {
-                if let Err(e) = device.shutdown() {
-                    log::warn!("{:?}", e);
-                }
+            if let Err(e) = device.shutdown() {
+                log::warn!("{:?}", e);
             }
         })?
     };
+    let worker_cell = Arc::new(AtomicCell::new(Some(worker)));
+
     {
-        let stop_all = stop_all.clone();
+        let worker_cell = worker_cell.clone();
         device_stop.set_stop_fn(move || {
-            if let Some(device) = device_cell.take() {
-                stop_all.store(false);
-                if let Err(e) = device.shutdown() {
-                    log::warn!("{:?}", e);
-                    return false;
-                }
-                true
-            } else {
-                false
+            if let Some(worker) = worker_cell.take() {
+                worker.stop_self()
             }
         });
     }
@@ -71,7 +62,7 @@ pub(crate) fn start_simple(
         log::error!("{:?}", e);
     }
     device_stop.stopped();
-    if stop_all.load() {
+    if let Some(worker) = worker_cell.take() {
         worker.stop_all();
     }
     Ok(())

vnt/src/util/notify.rs

@@ -28,7 +28,7 @@ impl StopManager {
         self.inner.add_listener(name, f)
     }
     pub fn stop(&self) {
-        self.inner.stop("");
+        self.inner.stop();
     }
     pub fn wait(&self) {
         self.inner.wait();
@@ -81,14 +81,11 @@ impl StopManagerInner {
         guard.1.push((name.clone(), Box::new(f)));
         Ok(Worker::new(name, self.clone()))
     }
-    fn stop(&self, skip_name: &str) {
+    fn stop(&self) {
         self.state.store(true, Ordering::Release);
         let mut guard = self.listeners.lock();
         guard.0 = true;
-        for (name, listener) in guard.1.drain(..) {
-            if &name == skip_name {
-                continue;
-            }
+        for (_name, listener) in guard.1.drain(..) {
             listener();
         }
     }
@@ -136,6 +133,19 @@ impl Worker {
     }
     fn release0(&self) {
         let inner = &self.inner;
+        let worker_name = &self.name;
+        {
+            let mut mutex_guard = inner.listeners.lock();
+            if let Some(pos) = mutex_guard
+                .1
+                .iter()
+                .position(|(name, _)| name == worker_name)
+            {
+                let (_, listener) = mutex_guard.1.remove(pos);
+                listener();
+            }
+        }
+
         let count = inner.worker_num.fetch_sub(1, Ordering::AcqRel);
         if count == 1 {
             for x in inner.park_threads.lock().drain(..) {
@@ -145,7 +155,10 @@ impl Worker {
         }
     }
     pub fn stop_all(self) {
-        self.inner.stop(&self.name)
+        self.inner.stop()
+    }
+    pub fn stop_self(self) {
+        drop(self)
     }
 }