diff --git a/vnt/src/core/conn.rs b/vnt/src/core/conn.rs index 8dc6b5b..b63c55f 100644 --- a/vnt/src/core/conn.rs +++ b/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, +} +impl Vnt { + #[cfg(feature = "integrated_tun")] + pub fn new(config: Config, callback: Call) -> anyhow::Result { + let inner = Arc::new(VntInner::new(config, callback)?); + Ok(Self { inner }) + } + #[cfg(not(feature = "integrated_tun"))] + pub fn new_device( + config: Config, + callback: Call, + device: Device, + ) -> anyhow::Result { + 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>, @@ -48,10 +75,10 @@ pub struct Vnt { external_route: ExternalRoute, } -impl Vnt { +impl VntInner { #[cfg(feature = "integrated_tun")] pub fn new(config: Config, callback: Call) -> anyhow::Result { - Vnt::new_device0(config, callback, DeviceAdapter::default()) + VntInner::new_device0(config, callback, DeviceAdapter::default()) } #[cfg(not(feature = "integrated_tun"))] pub fn new_device( @@ -59,7 +86,7 @@ impl Vnt { callback: Call, device: Device, ) -> anyhow::Result { - Vnt::new_device0(config, callback, device) + VntInner::new_device0(config, callback, device) } fn new_device0( config: Config, @@ -384,7 +411,7 @@ pub fn start( ) } -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(&self, name: String, f: F) -> anyhow::Result where F: FnOnce() + Send + 'static, @@ -478,3 +508,8 @@ impl Vnt { } } } +impl Drop for VntInner { + fn drop(&mut self) { + self.stop(); + } +} diff --git a/vnt/src/util/notify.rs b/vnt/src/util/notify.rs index ce4f20b..3495f8c 100644 --- a/vnt/src/util/notify.rs +++ b/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 { - self.inner.state.load(Ordering::Acquire) + pub fn is_stopped(&self) -> bool { + 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(); } diff --git a/vnt/src/util/scheduler.rs b/vnt/src/util/scheduler.rs index 5d5485e..808cafd 100644 --- a/vnt/src/util/scheduler.rs +++ b/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, 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 { 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, + state: Arc>, } + +#[derive(Copy, Clone, Eq, PartialEq)] +enum SchedulerState { + Running, + ShutdownNow, // 立即停止任务执行,队列中剩余的任务不再执行 + _Shutdown, //执行完队列中剩余的任务再停止 +} + impl Scheduler { pub fn new(stop_manager: StopManager) -> anyhow::Result { let (sender, receiver) = sync_channel::(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, s_inner: Scheduler) { + +fn run(receiver: Receiver, s_inner: &Scheduler) { let mut binary_heap = BinaryHeap::::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, 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, s_inner: Scheduler) { } } } + fn add_task(op: Op, binary_heap: &mut BinaryHeap) -> bool { return match op { Op::Task(task) => {