[mio] 增加统计、监听器、定时器

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<U64AdderInner>,
+    index: Option<usize>,
+}
+
+struct U64AdderInner {
+    global: AtomicU64,
+    base: Vec<AtomicU64>,
+}
+
+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<usize> = 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<U64AdderInner>,
+}
+
+impl WatchU64Adder {
+    pub fn get(&self) -> u64 {
+        self.inner.get()
+    }
+}

vnt/src/util/counter/mod.rs

@@ -0,0 +1,2 @@
+mod adder;
+pub use adder::*;

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::*;

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<StopManagerInner>,
+}
+
+impl StopManager {
+    pub fn new<F>(f: F) -> Self
+    where
+        F: FnOnce() + Send + 'static,
+    {
+        Self {
+            inner: Arc::new(StopManagerInner::new(f)),
+        }
+    }
+    pub fn add_listener<F>(&self, name: String, f: F) -> io::Result<Worker>
+    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<dyn FnOnce() + Send>)>)>,
+    park_threads: Mutex<Vec<Thread>>,
+    worker_num: AtomicUsize,
+    state: AtomicBool,
+    stop_call: Mutex<Option<Box<dyn FnOnce() + Send>>>,
+}
+
+impl StopManagerInner {
+    fn new<F>(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<F>(self: &Arc<Self>, name: String, f: F) -> io::Result<Worker>
+    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<StopManagerInner>,
+}
+
+impl Worker {
+    fn new(name: String, inner: Arc<StopManagerInner>) -> 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);
+    }
+}

vnt/src/util/result_convert.rs

@@ -0,0 +1,10 @@
+use std::fmt::Display;
+use std::io;
+
+#[inline]
+pub fn io_convert<T, R: Display, F: FnOnce(&io::Error) -> R>(
+    rs: io::Result<T>,
+    f: F,
+) -> io::Result<T> {
+    rs.map_err(|e| io::Error::new(e.kind(), format!("{},internal error:{:?}", f(&e), e)))
+}

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<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>,
+}
+impl Scheduler {
+    pub fn new(stop_manager: StopManager) -> io::Result<Self> {
+        let (sender, receiver) = sync_channel::<Op>(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<F>(&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<Op>, s_inner: Scheduler) {
+    let mut binary_heap = BinaryHeap::<DelayedTask>::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<DelayedTask>) -> bool {
+    return match op {
+        Op::Task(task) => {
+            binary_heap.push(task);
+            true
+        }
+        Op::Stop => false,
+    };
+}