diff --git a/vnt/src/handle/tun_tap/mod.rs b/vnt/src/handle/tun_tap/mod.rs index 1d1bd33..9dd4e9e 100644 --- a/vnt/src/handle/tun_tap/mod.rs +++ b/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 bool + Send>>>>, + f: Arc>>>, stopped: Arc>, } impl DeviceStop { pub fn set_stop_fn(&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) { diff --git a/vnt/src/handle/tun_tap/unix.rs b/vnt/src/handle/tun_tap/unix.rs index 4931857..5c3b1f3 100644 --- a/vnt/src/handle/tun_tap/unix.rs +++ b/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); diff --git a/vnt/src/handle/tun_tap/windows.rs b/vnt/src/handle/tun_tap/windows.rs index e2469eb..9b2d5e7 100644 --- a/vnt/src/handle/tun_tap/windows.rs +++ b/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(()) diff --git a/vnt/src/util/notify.rs b/vnt/src/util/notify.rs index defa700..ce4f20b 100644 --- a/vnt/src/util/notify.rs +++ b/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) } }