stack overflow handler (#51)

* use global_asm

* initfn abi fix

* add armv7, riscv64 support

* remove target

* add more platforms

* update ci

* filename update

* cargo fmt && filename update

* fix align

* overflow check

* update deps

* update

* make buffer bigger

* overflow check

* windows fix

* windows fix

* windows fix

* windows

* cargo fmt

* remove explicit align

* process sig mask

* refactor

* windows impl

* improve test

Cargo.toml

@@ -20,9 +20,10 @@ exclude = [
 
 
 [target.'cfg(windows)'.dependencies.windows]
-version = "0.48"
+version = "0.51.1"
 features = [
     "Win32_System_Memory",
+    "Win32_System_Kernel",
     "Win32_Foundation",
     "Win32_System_SystemInformation",
     "Win32_System_Diagnostics_Debug"

src/detail/aarch64_unix.rs

@@ -23,13 +23,13 @@ extern "C" {
     pub fn swap_registers(out_regs: *mut Registers, in_regs: *const Registers);
 }
 
-#[repr(C, align(16))]
+#[repr(C)]
 #[derive(Debug)]
 pub struct Registers {
     // We save the 13 callee-saved registers:
     //  x19--x28, fp (x29), lr (x30), sp
     // and the 8 callee-saved floating point registers:
-    //  v8--v15
+    //  d8--d15
     gpr: [usize; 32],
 }
 

src/detail/asm/asm_aarch64_aapcs_elf.S

@@ -1,7 +1,7 @@
 .text
 .globl prefetch
 .type prefetch,@function
-.align 16
+.align 2
 prefetch:
     prfm pldl1keep, [x0]
     ret
@@ -10,7 +10,7 @@ prefetch:
 .text
 .globl bootstrap_green_task
 .type bootstrap_green_task,@function
-.align 16
+.align 2
 bootstrap_green_task:
     mov x0, x19   // arg0
     mov x1, x20   // arg1
@@ -21,7 +21,7 @@ bootstrap_green_task:
 .text
 .globl swap_registers
 .type swap_registers,@function
-.align 16
+.align 2
 swap_registers:
     stp x19, x20, [x0, #0]
     stp x21, x22, [x0, #16]

src/detail/asm/asm_aarch64_aapcs_macho.S

@@ -1,13 +1,13 @@
 .text
 .globl _prefetch
-.align 8
+.align 2
 _prefetch:
     prfm pldl1keep, [x0]
     ret
 
 .text
 .globl _bootstrap_green_task
-.align 8
+.align 2
 _bootstrap_green_task:
     mov x0, x19  // arg0
     mov x1, x20  // arg1
@@ -16,7 +16,7 @@ _bootstrap_green_task:
 
 .text
 .globl _swap_registers
-.align 8
+.align 2
 _swap_registers:
     stp x19, x20, [x0, #0]
     stp x21, x22, [x0, #16]

src/detail/asm/asm_arm_aapcs_elf.S

@@ -1,7 +1,7 @@
 .text
 .globl prefetch
 .type prefetch, %function
-.align 16
+.align 2
 prefetch:
     pld [r0]
     bx lr
@@ -10,7 +10,7 @@ prefetch:
 .text
 .globl bootstrap_green_task
 .type bootstrap_green_task, %function
-.align 16
+.align 2
 bootstrap_green_task:
      mov r0, r4 // arg0
      mov r1, r5 // arg1
@@ -21,7 +21,7 @@ bootstrap_green_task:
 .text
 .globl swap_registers
 .type swap_registers, %function
-.align 16
+.align 2
 swap_registers:
      // Android doesn't like to use sp directly
      stmia r0!, {{v1-v7, fp}}

src/detail/asm/asm_loongarch64_sysv_elf.S

@@ -1,7 +1,7 @@
 .text
 .globl prefetch
 .type prefetch,@function
-.align 16
+.align 2
 prefetch:
     preld 0, $a0, 0
     ret
@@ -10,7 +10,7 @@ prefetch:
 .text
 .globl bootstrap_green_task
 .type bootstrap_green_task,@function
-.align 16
+.align 2
 bootstrap_green_task:
     move $a0, $s0   // arg0
     move $a1, $s1   // arg1
@@ -21,7 +21,7 @@ bootstrap_green_task:
 .text
 .globl swap_registers
 .type swap_registers,@function
-.align 16
+.align 2
 swap_registers:
     st.d $ra, $a0, 0
     st.d $sp, $a0, 8

src/detail/asm/asm_riscv64_c_elf.S

@@ -3,15 +3,15 @@
 .text
 .globl prefetch
 .type prefetch,@function
-.align 16
+.align 1
 prefetch:
     ret
 .size prefetch,.-prefetch
 
 .text
 .globl bootstrap_green_task
 .type bootstrap_green_task,@function
-.align 16
+.align 1
 bootstrap_green_task:
      mv a0, s2   // arg0
      mv a1, s3   // arg1
@@ -22,7 +22,7 @@ bootstrap_green_task:
 .text
 .globl swap_registers
 .type swap_registers,@function
-.align 16
+.align 1
 swap_registers:
      sd s2, 0*8(a0)
      sd s3, 1*8(a0)

src/detail/gen.rs

@@ -1,5 +1,5 @@
 use crate::rt::ContextStack;
-use crate::stack::Func;
+use crate::stack::{overflow, Func};
 use crate::yield_::yield_now;
 use crate::Error;
 use std::any::Any;
@@ -29,6 +29,8 @@ fn catch_unwind_filter<F: FnOnce() -> R + panic::UnwindSafe, R>(f: F) -> std::th
 /// the init function passed to reg_context
 #[inline]
 pub fn gen_init_impl(_: usize, f: *mut usize) -> ! {
+    overflow::init_once();
+
     let clo = move || {
         // consume self.f
         let f: &mut Option<Func> = unsafe { &mut *(f as *mut _) };

src/detail/mod.rs

@@ -1,3 +1,5 @@
+#![allow(clippy::eq_op)]
+
 // Register contexts used in various architectures
 //
 // These structures all represent a context of one task throughout its

src/detail/x86_64_windows.rs

@@ -174,10 +174,11 @@ pub use self::asm_impl::*;
 */
 
 // windows need to restore xmm6~xmm15, for most cases only use two xmm registers
+// so we use sysv64
 #[repr(C)]
 #[derive(Debug)]
 pub struct Registers {
-    gpr: [usize; 16],
+    pub(crate) gpr: [usize; 16],
 }
 
 impl Registers {

src/gen_impl.rs

@@ -342,6 +342,8 @@ impl<'a, A, T> GeneratorImpl<'a, A, T> {
 
         self.f = Some(func);
 
+        let guard = (self.stack.begin() as usize, self.stack.end() as usize);
+        self.context.stack_guard = guard;
         self.context.regs.init_with(
             gen_init,
             0,
@@ -460,10 +462,10 @@ impl<'a, A, T> GeneratorImpl<'a, A, T> {
         // so that we can stop the inner func
         self.context._ref = 2;
         // save the old panic hook, we don't want to print anything for the Cancel
-        let old = ::std::panic::take_hook();
-        ::std::panic::set_hook(Box::new(|_| {}));
+        let old = panic::take_hook();
+        panic::set_hook(Box::new(|_| {}));
         self.resume_gen();
-        ::std::panic::set_hook(old);
+        panic::set_hook(old);
     }
 
     /// cancel the generator
@@ -520,7 +522,7 @@ impl<'a, A, T> Drop for GeneratorImpl<'a, A, T> {
             // set_stack_size::<F>(used_stack);
         } else {
             error!("stack overflow detected!");
-            std::panic::panic_any(Error::StackErr);
+            panic::panic_any(Error::StackErr);
         }
     }
 }

src/reg_context.rs

@@ -4,7 +4,7 @@ use crate::stack::Stack;
 #[derive(Debug)]
 pub struct RegContext {
     /// Hold the registers while the task or scheduler is suspended
-    regs: Registers,
+    pub(crate) regs: Registers,
 }
 
 impl RegContext {

src/rt.rs

@@ -62,6 +62,8 @@ pub struct Context {
     pub local_data: *mut u8,
     /// propagate panic
     pub err: Option<Box<dyn Any + Send>>,
+    /// cached stack guard for fast path
+    pub stack_guard: (usize, usize),
 }
 
 impl Context {
@@ -76,6 +78,7 @@ impl Context {
             child: ptr::null_mut(),
             parent: ptr::null_mut(),
             local_data: ptr::null_mut(),
+            stack_guard: (0, 0),
         }
     }
 
@@ -170,7 +173,7 @@ impl Context {
 
 /// Coroutine managing environment
 pub struct ContextStack {
-    root: *mut Context,
+    pub(crate) root: *mut Context,
 }
 
 #[cfg(nightly)]
@@ -292,6 +295,22 @@ pub fn get_local_data() -> *mut u8 {
     ptr::null_mut()
 }
 
+pub mod guard {
+    use crate::is_generator;
+    use crate::rt::ContextStack;
+    use crate::stack::sys::page_size;
+    use std::ops::Range;
+
+    pub type Guard = Range<usize>;
+
+    pub fn current() -> Guard {
+        assert!(is_generator());
+        let guard = unsafe { (*(*ContextStack::current().root).child).stack_guard };
+
+        guard.0 - page_size()..guard.1
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::is_generator;
@@ -301,4 +320,32 @@ mod test {
         // this is the root context
         assert!(!is_generator());
     }
+
+    #[test]
+    fn test_overflow() {
+        use crate::*;
+        use std::panic::catch_unwind;
+
+        // test signal mask
+        for _ in 0..2 {
+            let result = catch_unwind(|| {
+                let mut g = Gn::new_scoped(move |_s: Scope<(), ()>| {
+                    let guard = super::guard::current();
+
+                    // make sure the compiler does not apply any optimization on it
+                    std::hint::black_box(unsafe { *(guard.start as *const usize) });
+
+                    eprintln!("entered unreachable code");
+                    std::process::abort();
+                });
+
+                g.next();
+            });
+
+            assert!(matches!(
+                result.map_err(|err| *err.downcast::<Error>().unwrap()),
+                Err(Error::StackErr)
+            ));
+        }
+    }
 }

src/stack/mod.rs

@@ -13,6 +13,8 @@ use std::ptr;
 #[cfg_attr(windows, path = "windows.rs")]
 pub mod sys;
 
+pub use sys::overflow;
+
 // must align with StackBoxHeader
 const ALIGN: usize = std::mem::size_of::<StackBoxHeader>();
 const HEADER_SIZE: usize = std::mem::size_of::<StackBoxHeader>() / std::mem::size_of::<usize>();
@@ -58,7 +60,7 @@ impl<T> StackBox<T> {
             header.data_size = data_size;
             header.need_drop = need_drop;
             header.stack = stack.shadow_clone();
-            std::mem::MaybeUninit::new(StackBox { ptr })
+            MaybeUninit::new(StackBox { ptr })
         }
     }
 
@@ -312,13 +314,8 @@ impl Stack {
     /// Allocate a new stack of `size`. If size = 0, this is a `dummy_stack`
     pub fn new(size: usize) -> Stack {
         let track = (size & 1) != 0;
-        let mut bytes = size * std::mem::size_of::<usize>();
-        // the minimal size
-        let min_size = SysStack::min_size();
 
-        if bytes < min_size {
-            bytes = min_size;
-        }
+        let bytes = usize::max(size * std::mem::size_of::<usize>(), SysStack::min_size());
 
         let buf = SysStack::allocate(bytes, true).expect("failed to alloc sys stack");