From ec3aa01b9ca60b56bc6dd6f0537f3138a74e4c32 Mon Sep 17 00:00:00 2001 From: lubeilin <1791778603@qq.com> Date: Thu, 29 Feb 2024 22:17:27 +0800 Subject: [PATCH] =?UTF-8?q?[mio]=20=E5=A2=9E=E5=8A=A0=E7=BB=9F=E8=AE=A1?= =?UTF-8?q?=E3=80=81=E7=9B=91=E5=90=AC=E5=99=A8=E3=80=81=E5=AE=9A=E6=97=B6?= =?UTF-8?q?=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- vnt/src/util/counter/adder.rs | 96 ++++++++++++++++++++++ vnt/src/util/counter/mod.rs | 2 + vnt/src/util/mod.rs | 10 ++- vnt/src/util/notify.rs | 143 +++++++++++++++++++++++++++++++++ vnt/src/util/result_convert.rs | 10 +++ vnt/src/util/scheduler.rs | 132 ++++++++++++++++++++++++++++++ 6 files changed, 392 insertions(+), 1 deletion(-) create mode 100644 vnt/src/util/counter/adder.rs create mode 100644 vnt/src/util/counter/mod.rs create mode 100644 vnt/src/util/notify.rs create mode 100644 vnt/src/util/result_convert.rs create mode 100644 vnt/src/util/scheduler.rs diff --git a/vnt/src/util/counter/adder.rs b/vnt/src/util/counter/adder.rs new file mode 100644 index 0000000..e539969 --- /dev/null +++ b/vnt/src/util/counter/adder.rs @@ -0,0 +1,96 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; + +/// 并发计数器，销毁计数器并不会释放计数槽，这不适用计数器会多次创建销毁的场景 +pub struct U64Adder { + inner: Arc, + index: Option, +} + +struct U64AdderInner { + global: AtomicU64, + base: Vec, +} + +impl U64AdderInner { + pub fn get(&self) -> u64 { + let mut count = self.global.load(Ordering::Relaxed); + for counter in self.base.iter() { + let num = counter.load(Ordering::Relaxed); + if num > 1 { + count = count + num - 1; + } + } + count + } +} + +impl U64Adder { + /// 计数槽容量 + pub fn with_capacity(capacity: usize) -> Self { + let mut base = Vec::with_capacity(capacity); + base.push(AtomicU64::new(1)); + for _ in 1..capacity { + base.push(AtomicU64::new(0)) + } + let inner = Arc::new(U64AdderInner { + global: AtomicU64::new(0), + base, + }); + U64Adder { + inner, + index: Some(0), + } + } + pub fn add(&mut self, num: u64) { + if let Some(index) = self.index { + let counter = &self.inner.base[index]; + let i = counter.load(Ordering::Relaxed); + counter.store(i + num, Ordering::Relaxed); + } else { + self.inner.global.fetch_add(num, Ordering::Relaxed); + } + } + pub fn get(&self) -> u64 { + self.inner.get() + } + pub fn watch(&self) -> WatchU64Adder { + WatchU64Adder { + inner: self.inner.clone(), + } + } +} + +impl Clone for U64Adder { + fn clone(&self) -> Self { + let mut index: Option = None; + for (i, counter) in self.inner.base.iter().enumerate() { + //占用一个空闲的计数槽 + if counter.load(Ordering::Acquire) == 0 { + if counter + .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Relaxed) + .is_ok() + { + index = Some(i); + break; + } + } + } + + Self { + inner: self.inner.clone(), + index, + } + } +} + +#[derive(Clone)] +pub struct WatchU64Adder { + inner: Arc, +} + +impl WatchU64Adder { + pub fn get(&self) -> u64 { + self.inner.get() + } +} diff --git a/vnt/src/util/counter/mod.rs b/vnt/src/util/counter/mod.rs new file mode 100644 index 0000000..0ebd885 --- /dev/null +++ b/vnt/src/util/counter/mod.rs @@ -0,0 +1,2 @@ +mod adder; +pub use adder::*; diff --git a/vnt/src/util/mod.rs b/vnt/src/util/mod.rs index 773f6d7..ce8a1f9 100644 --- a/vnt/src/util/mod.rs +++ b/vnt/src/util/mod.rs @@ -1 +1,9 @@ -pub mod wait; +mod notify; +mod result_convert; +pub use result_convert::io_convert; +mod scheduler; +pub use notify::StopManager; +pub use scheduler::Scheduler; + +mod counter; +pub use counter::*; diff --git a/vnt/src/util/notify.rs b/vnt/src/util/notify.rs new file mode 100644 index 0000000..bc9d242 --- /dev/null +++ b/vnt/src/util/notify.rs @@ -0,0 +1,143 @@ +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::Arc; +use std::thread::Thread; +use std::{io, thread}; + +use parking_lot::Mutex; + +#[derive(Clone)] +pub struct StopManager { + inner: Arc, +} + +impl StopManager { + pub fn new(f: F) -> Self + where + F: FnOnce() + Send + 'static, + { + Self { + inner: Arc::new(StopManagerInner::new(f)), + } + } + pub fn add_listener(&self, name: String, f: F) -> io::Result + where + F: FnOnce() + Send + 'static, + { + self.inner.add_listener(name, f) + } + pub fn stop(&self) { + self.inner.stop(""); + } + pub fn wait(&self) { + self.inner.wait(); + } + pub fn is_stop(&self) -> bool { + self.inner.state.load(Ordering::Acquire) + } +} + +struct StopManagerInner { + listeners: Mutex<(bool, Vec<(String, Box)>)>, + park_threads: Mutex>, + worker_num: AtomicUsize, + state: AtomicBool, + stop_call: Mutex>>, +} + +impl StopManagerInner { + fn new(f: F) -> Self + where + F: FnOnce() + Send + 'static, + { + Self { + listeners: Mutex::new((false, Vec::with_capacity(32))), + park_threads: Mutex::new(Vec::with_capacity(4)), + worker_num: AtomicUsize::new(0), + state: AtomicBool::new(false), + stop_call: Mutex::new(Some(Box::new(f))), + } + } + fn add_listener(self: &Arc, name: String, f: F) -> io::Result + where + F: FnOnce() + Send + 'static, + { + if name.is_empty() { + return Err(io::Error::new(io::ErrorKind::Other, "name cannot be empty")); + } + let mut guard = self.listeners.lock(); + if guard.0 { + return Err(io::Error::new(io::ErrorKind::Other, "stopped")); + } + for (n, _) in &guard.1 { + if &name == n { + return Err(io::Error::new( + io::ErrorKind::Other, + format!("stop add_listener {:?} name already exists", name), + )); + } + } + guard.1.push((name.clone(), Box::new(f))); + Ok(Worker::new(name, self.clone())) + } + fn stop(&self, skip_name: &str) { + 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; + } + listener(); + } + } + fn wait(&self) { + { + let mut guard = self.park_threads.lock(); + guard.push(thread::current()); + drop(guard); + } + loop { + if self.worker_num.load(Ordering::Acquire) == 0 { + return; + } + thread::park() + } + } + fn stop_call(&self) { + if let Some(call) = self.stop_call.lock().take() { + call(); + } + } +} + +pub struct Worker { + name: String, + inner: Arc, +} + +impl Worker { + fn new(name: String, inner: Arc) -> Self { + let _ = inner.worker_num.fetch_add(1, Ordering::AcqRel); + Self { name, inner } + } + fn release0(&self) { + let inner = &self.inner; + let count = inner.worker_num.fetch_sub(1, Ordering::AcqRel); + if count == 1 { + for x in inner.park_threads.lock().drain(..) { + x.unpark(); + } + self.inner.stop_call(); + } + } + pub fn stop_all(self) { + self.inner.stop(&self.name) + } +} + +impl Drop for Worker { + fn drop(&mut self) { + self.release0(); + log::info!("stop {}", self.name); + } +} diff --git a/vnt/src/util/result_convert.rs b/vnt/src/util/result_convert.rs new file mode 100644 index 0000000..981d8c2 --- /dev/null +++ b/vnt/src/util/result_convert.rs @@ -0,0 +1,10 @@ +use std::fmt::Display; +use std::io; + +#[inline] +pub fn io_convert R>( + rs: io::Result, + f: F, +) -> io::Result { + rs.map_err(|e| io::Error::new(e.kind(), format!("{},internal error:{:?}", f(&e), e))) +} diff --git a/vnt/src/util/scheduler.rs b/vnt/src/util/scheduler.rs new file mode 100644 index 0000000..1b4bfdd --- /dev/null +++ b/vnt/src/util/scheduler.rs @@ -0,0 +1,132 @@ +use crate::util::StopManager; +use std::collections::BinaryHeap; +use std::{ + cmp::Ordering, + io, + sync::mpsc::{sync_channel, Receiver, SyncSender}, + time::{Duration, Instant}, +}; + +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, +} +impl Scheduler { + pub fn new(stop_manager: StopManager) -> io::Result { + let (sender, receiver) = sync_channel::(32); + let s = Self { sender }; + let s_inner = s.clone(); + let worker = { + let scheduler = s.clone(); + stop_manager.add_listener("Scheduler".into(), move || { + scheduler.shutdown(); + })? + }; + std::thread::Builder::new() + .name("Scheduler".into()) + .spawn(move || { + run(receiver, s_inner); + worker.stop_all(); + }) + .unwrap(); + Ok(s) + } + pub fn timeout(&self, time: Duration, f: F) -> bool + where + F: FnOnce(&Scheduler) + Send + 'static, + { + let task = DelayedTask { + f: Box::new(f), + next: Instant::now().checked_add(time).unwrap(), + }; + self.sender.send(Op::Task(task)).is_ok() + } + pub fn shutdown(self) { + let _ = self.sender.send(Op::Stop); + } +} +fn run(receiver: Receiver, s_inner: Scheduler) { + let mut binary_heap = BinaryHeap::::with_capacity(32); + loop { + while let Some(task) = binary_heap.peek() { + let now = Instant::now(); + if now < task.next { + //需要等待对应时间 + match receiver.recv_timeout(task.next - now) { + Ok(op) => { + if add_task(op, &mut binary_heap) { + continue; + } + return; + } + Err(e) => match e { + std::sync::mpsc::RecvTimeoutError::Timeout => continue, + std::sync::mpsc::RecvTimeoutError::Disconnected => return, + }, + } + } else { + if let Some(task) = binary_heap.pop() { + (task.f)(&s_inner); + } + } + } + //取出所有任务 + loop { + match receiver.try_recv() { + Ok(op) => { + if add_task(op, &mut binary_heap) { + continue; + } + return; + } + Err(e) => match e { + std::sync::mpsc::TryRecvError::Empty => break, + std::sync::mpsc::TryRecvError::Disconnected => return, + }, + } + } + + if binary_heap.is_empty() { + //任务队列为空时陷入等待 + if let Ok(op) = receiver.recv() { + if add_task(op, &mut binary_heap) { + continue; + } + } + return; + } + } +} +fn add_task(op: Op, binary_heap: &mut BinaryHeap) -> bool { + return match op { + Op::Task(task) => { + binary_heap.push(task); + true + } + Op::Stop => false, + }; +}