add stack pool

Author: loongs-zhang

core/src/coroutine/korosensei.rs

@@ -318,7 +318,7 @@ impl<'c, Param, Yield, Return> Coroutine<'c, Param, Yield, Return> {
             if remaining_stack >= red_zone {
                 return Ok(callback());
             }
-            return STACK_POOL.get_stack(stack_size).map(|stack| {
+            return STACK_POOL.allocate(stack_size).map(|stack| {
                 co.stack_infos.borrow_mut().push_back(StackInfo {
                     stack_top: stack.base().get(),
                     stack_bottom: stack.limit().get(),
@@ -339,7 +339,7 @@ impl<'c, Param, Yield, Return> Coroutine<'c, Param, Yield, Return> {
                 return Ok(callback());
             }
         }
-        STACK_POOL.get_stack(stack_size).map(|stack| {
+        STACK_POOL.allocate(stack_size).map(|stack| {
             STACK_INFOS.with(|s| {
                 s.borrow_mut().push_back(StackInfo {
                     stack_top: stack.base().get(),
@@ -382,7 +382,7 @@ where
         F: FnOnce(&Suspender<Param, Yield>, Param) -> Return + 'static,
     {
         let stack_size = stack_size.max(crate::common::page_size());
-        let stack = STACK_POOL.get_stack(stack_size)?;
+        let stack = STACK_POOL.allocate(stack_size)?;
         let stack_infos = RefCell::new(VecDeque::from([StackInfo {
             stack_top: stack.base().get(),
             stack_bottom: stack.limit().get(),

core/src/coroutine/stack_pool.rs

@@ -3,21 +3,21 @@ use std::cell::RefCell;
 use std::cmp::Ordering;
 use std::collections::BinaryHeap;
 use std::ops::Deref;
-use std::sync::Arc;
+use std::rc::Rc;
 
-pub(crate) struct PooledStack(Arc<DefaultStack>);
+pub(crate) struct PooledStack(Rc<DefaultStack>);
 
 impl Deref for PooledStack {
-    type Target = Arc<DefaultStack>;
+    type Target = Rc<DefaultStack>;
 
-    fn deref(&self) -> &Arc<DefaultStack> {
+    fn deref(&self) -> &Rc<DefaultStack> {
         &self.0
     }
 }
 
 impl PartialEq<Self> for PooledStack {
     fn eq(&self, other: &Self) -> bool {
-        Arc::strong_count(other).eq(&Arc::strong_count(self))
+        Rc::strong_count(other).eq(&Rc::strong_count(self))
     }
 }
 
@@ -32,7 +32,7 @@ impl PartialOrd<Self> for PooledStack {
 impl Ord for PooledStack {
     fn cmp(&self, other: &Self) -> Ordering {
         // BinaryHeap defaults to a large top heap, but we need a small top heap
-        Arc::strong_count(other).cmp(&Arc::strong_count(self))
+        Rc::strong_count(other).cmp(&Rc::strong_count(self))
     }
 }
 
@@ -56,28 +56,31 @@ unsafe impl Stack for PooledStack {
     #[cfg(windows)]
     #[inline]
     fn update_teb_fields(&mut self, stack_limit: usize, guaranteed_stack_bytes: usize) {
-        self.0
-            .update_teb_fields(stack_limit, guaranteed_stack_bytes)
+        while let Some(stack) = Rc::get_mut(&mut self) {
+            stack.update_teb_fields(stack_limit, guaranteed_stack_bytes);
+            return;
+        }
     }
 }
 
+/// todo `min_size` `max_size` `keep_alive_time`
 #[derive(Default)]
-pub(crate) struct StackPool(Arc<RefCell<BinaryHeap<PooledStack>>>);
+pub(crate) struct StackPool(RefCell<BinaryHeap<PooledStack>>);
 
 unsafe impl Send for StackPool {}
 
 unsafe impl Sync for StackPool {}
 
 impl StackPool {
-    pub(crate) fn get_stack(&self, stack_size: usize) -> std::io::Result<PooledStack> {
+    pub(crate) fn allocate(&self, stack_size: usize) -> std::io::Result<PooledStack> {
         loop {
             if let Ok(mut heap) = self.0.try_borrow_mut() {
                 if let Some(stack) = heap.peek() {
-                    if Arc::strong_count(stack) == 1 {
+                    if Rc::strong_count(stack) == 1 {
                         return Ok(PooledStack(stack.deref().clone()));
                     }
                 }
-                let stack = Arc::new(DefaultStack::new(stack_size)?);
+                let stack = Rc::new(DefaultStack::new(stack_size)?);
                 heap.push(PooledStack(stack.clone()));
                 return Ok(PooledStack(stack));
             }
@@ -98,13 +101,13 @@ mod tests {
     #[test]
     fn test_stack_pool() -> std::io::Result<()> {
         let pool = StackPool::default();
-        let stack = pool.get_stack(DEFAULT_STACK_SIZE)?;
-        assert_eq!(Arc::strong_count(&stack), 2);
+        let stack = pool.allocate(DEFAULT_STACK_SIZE)?;
+        assert_eq!(Rc::strong_count(&stack), 2);
         drop(stack);
-        let stack = pool.get_stack(DEFAULT_STACK_SIZE)?;
-        assert_eq!(Arc::strong_count(&stack), 2);
+        let stack = pool.allocate(DEFAULT_STACK_SIZE)?;
+        assert_eq!(Rc::strong_count(&stack), 2);
         assert_eq!(pool.len(), 1);
-        _ = pool.get_stack(DEFAULT_STACK_SIZE)?;
+        _ = pool.allocate(DEFAULT_STACK_SIZE)?;
         assert_eq!(pool.len(), 2);
         Ok(())
     }