refactor memory pool and io_uring macro

Author: loongs-zhang

core/src/common/memory_pool.rs

@@ -0,0 +1,192 @@
+use crate::common::now;
+use crossbeam_utils::atomic::AtomicCell;
+use std::collections::BinaryHeap;
+use std::num::NonZeroUsize;
+use std::sync::atomic::{AtomicU64, AtomicUsize};
+
+/// A trait for reusable memory.
+///
+/// # Safety
+/// The memory requires exclusive use when writing.
+pub unsafe trait ReusableMemory: Ord + Clone {
+    /// Create a new instance of `ReusableMemory`.
+    fn new(size: usize) -> std::io::Result<Self>
+    where
+        Self: Sized;
+
+    /// Get the reference count of the memory.
+    fn reference_count(&self) -> usize;
+
+    /// Get the top of the memory.
+    fn top(&self) -> NonZeroUsize;
+
+    /// Get the bottom of the memory.
+    fn bottom(&self) -> NonZeroUsize;
+
+    /// Get the size of the memory.
+    fn size(&self) -> usize {
+        self.top()
+            .get()
+            .checked_sub(self.bottom().get())
+            .expect("the `bottom` is bigger than `top`")
+    }
+
+    /// Get the creation time of the memory.
+    fn create_time(&self) -> u64;
+
+    /// Callback when the memory is reused.
+    fn on_reuse(&mut self) -> std::io::Result<()>;
+}
+
+/// A memory pool for reusing.
+#[repr(C)]
+#[derive(educe::Educe)]
+#[educe(Debug)]
+pub struct MemoryPool<M: ReusableMemory> {
+    #[educe(Debug(ignore))]
+    pool: AtomicCell<BinaryHeap<M>>,
+    len: AtomicUsize,
+    //最小内存数，即核心内存数
+    min_count: AtomicUsize,
+    //非核心内存的最大存活时间，单位ns
+    keep_alive_time: AtomicU64,
+}
+
+unsafe impl<M: ReusableMemory> Send for MemoryPool<M> {}
+
+unsafe impl<M: ReusableMemory> Sync for MemoryPool<M> {}
+
+impl<M: ReusableMemory> Default for MemoryPool<M> {
+    fn default() -> Self {
+        Self::new(0, 0)
+    }
+}
+
+#[allow(missing_docs)]
+impl<M: ReusableMemory> MemoryPool<M> {
+    /// Create a new instance of `MemoryPool`.
+    #[must_use]
+    pub fn new(min_count: usize, keep_alive_time: u64) -> Self {
+        Self {
+            pool: AtomicCell::new(BinaryHeap::default()),
+            len: AtomicUsize::default(),
+            min_count: AtomicUsize::new(min_count),
+            keep_alive_time: AtomicU64::new(keep_alive_time),
+        }
+    }
+
+    pub fn allocate(&self, memory_size: usize) -> std::io::Result<M> {
+        let heap = unsafe {
+            self.pool
+                .as_ptr()
+                .as_mut()
+                .expect("MemoryPool is not unique")
+        };
+        // find min memory
+        let mut not_use = Vec::new();
+        while let Some(memory) = heap.peek() {
+            if memory.reference_count() > 1 {
+                // can't use the memory
+                break;
+            }
+            if let Some(mut memory) = heap.pop() {
+                self.sub_len();
+                if memory_size <= memory.size() {
+                    for s in not_use {
+                        heap.push(s);
+                        self.add_len();
+                    }
+                    heap.push(memory.clone());
+                    self.add_len();
+                    return memory.on_reuse().map(|()| memory);
+                }
+                if self.min_count() < self.len()
+                    && now() <= memory.create_time().saturating_add(self.keep_alive_time())
+                {
+                    // clean the expired memory
+                    continue;
+                }
+                not_use.push(memory);
+            }
+        }
+        let memory = M::new(memory_size)?;
+        heap.push(memory.clone());
+        self.add_len();
+        Ok(memory)
+    }
+
+    pub fn set_keep_alive_time(&self, keep_alive_time: u64) -> &Self {
+        self.keep_alive_time
+            .store(keep_alive_time, std::sync::atomic::Ordering::Release);
+        self
+    }
+
+    pub fn keep_alive_time(&self) -> u64 {
+        self.keep_alive_time
+            .load(std::sync::atomic::Ordering::Acquire)
+    }
+
+    pub fn set_min_count(&self, min_count: usize) -> &Self {
+        self.min_count
+            .store(min_count, std::sync::atomic::Ordering::Release);
+        self
+    }
+
+    pub fn min_count(&self) -> usize {
+        self.min_count.load(std::sync::atomic::Ordering::Acquire)
+    }
+
+    pub fn is_empty(&self) -> bool {
+        0 == self.len()
+    }
+
+    pub fn len(&self) -> usize {
+        self.len.load(std::sync::atomic::Ordering::Acquire)
+    }
+
+    fn add_len(&self) {
+        self.len.store(
+            self.len().saturating_add(1),
+            std::sync::atomic::Ordering::Release,
+        );
+    }
+
+    fn sub_len(&self) {
+        self.len.store(
+            self.len().saturating_sub(1),
+            std::sync::atomic::Ordering::Release,
+        );
+    }
+
+    /// Clean the expired memory.
+    #[allow(dead_code)]
+    pub fn clean(&self) {
+        let heap = unsafe {
+            self.pool
+                .as_ptr()
+                .as_mut()
+                .expect("MemoryPool is not unique")
+        };
+        let mut maybe_free = Vec::new();
+        while let Some(memory) = heap.peek() {
+            if memory.reference_count() > 1 {
+                // can't free the memory
+                break;
+            }
+            if let Some(memory) = heap.pop() {
+                self.sub_len();
+                maybe_free.push(memory);
+            }
+        }
+        for memory in maybe_free {
+            if self.min_count() < self.len()
+                && now() <= memory.create_time().saturating_add(self.keep_alive_time())
+            {
+                // free the memory
+                continue;
+            }
+            heap.push(memory);
+            self.add_len();
+        }
+    }
+}

core/src/common/mod.rs

@@ -13,6 +13,9 @@ pub mod constants;
 /// Check <https://www.rustwiki.org.cn/en/reference/introduction.html> for help information.
 pub(crate) mod macros;
 
+/// Reuse memory.
+pub mod memory_pool;
+
 /// `BeanFactory` impls.
 pub mod beans;
 

core/src/coroutine/korosensei.rs

@@ -2,7 +2,7 @@ use crate::catch;
 use crate::common::constants::CoroutineState;
 use crate::coroutine::listener::Listener;
 use crate::coroutine::local::CoroutineLocal;
-use crate::coroutine::stack_pool::{MemoryPool, PooledStack};
+use crate::coroutine::stack_pool::{PooledStack, StackPool};
 use crate::coroutine::suspender::Suspender;
 use crate::coroutine::StackInfo;
 use corosensei::stack::Stack;
@@ -326,17 +326,14 @@ impl<'c, Param, Yield, Return> Coroutine<'c, Param, Yield, Return> {
         stack_size: usize,
         callback: F,
     ) -> std::io::Result<R> {
-        let stack_pool = MemoryPool::get_instance();
+        let stack_pool = StackPool::get_instance();
         if let Some(co) = Self::current() {
             let remaining_stack = unsafe { co.remaining_stack() };
             if remaining_stack >= red_zone {
                 return Ok(callback());
             }
             return stack_pool.allocate(stack_size).map(|stack| {
-                co.stack_infos_mut().push_back(StackInfo {
-                    stack_top: stack.base().get(),
-                    stack_bottom: stack.limit().get(),
-                });
+                co.stack_infos_mut().push_back(StackInfo::from(&stack));
                 let r = corosensei::on_stack(stack, callback);
                 _ = co.stack_infos_mut().pop_back();
                 r
@@ -357,10 +354,7 @@ impl<'c, Param, Yield, Return> Coroutine<'c, Param, Yield, Return> {
         }
         stack_pool.allocate(stack_size).map(|stack| {
             STACK_INFOS.with(|s| {
-                s.borrow_mut().push_back(StackInfo {
-                    stack_top: stack.base().get(),
-                    stack_bottom: stack.limit().get(),
-                });
+                s.borrow_mut().push_back(StackInfo::from(&stack));
             });
             let r = corosensei::on_stack(stack, callback);
             _ = STACK_INFOS.with(|s| s.borrow_mut().pop_back());
@@ -398,7 +392,7 @@ where
         F: FnOnce(&Suspender<Param, Yield>, Param) -> Return + 'static,
     {
         let stack_size = stack_size.max(crate::common::page_size());
-        let stack = MemoryPool::get_instance().allocate(stack_size)?;
+        let stack = StackPool::get_instance().allocate(stack_size)?;
         let stack_infos = UnsafeCell::new(VecDeque::from([StackInfo {
             stack_top: stack.base().get(),
             stack_bottom: stack.limit().get(),
@@ -469,3 +463,12 @@ where
         }
     }
 }
+
+impl<S: Stack> From<&S> for StackInfo {
+    fn from(stack: &S) -> Self {
+        Self {
+            stack_top: stack.base().get(),
+            stack_bottom: stack.limit().get(),
+        }
+    }
+}