Adjust process & thread exit & kill behavior

h2o/kernel/src/cpu/x86_64/intr.rs

@@ -137,7 +137,7 @@ unsafe fn exception(frame_ptr: *mut Frame, vec: def::ExVec) {
                     });
                 }
                 // Kill the fucking task.
-                SCHED.exit_current(sv_call::EFAULT.into_retval())
+                SCHED.exit_current(sv_call::EFAULT.into_retval(), true)
             }
             // unreachable!()
         }

h2o/kernel/src/sched/imp.rs

@@ -166,7 +166,7 @@ impl Scheduler {
     /// # Panics
     ///
     /// Panics if the scheduler unexpectedly returns.
-    pub fn exit_current(&self, retval: usize) -> ! {
+    pub fn exit_current(&self, retval: usize, kill_all: bool) -> ! {
         self.canary.assert();
         let pree = PREEMPT.lock();
 
@@ -181,6 +181,10 @@ impl Scheduler {
             SCHED_INFO[self.cpu]
                 .expected_runtime
                 .fetch_sub(current.time_slice.as_micros() as u64, Release);
+
+            if kill_all {
+                current.space().try_stop(&current.tid);
+            }
         }
 
         let _ = self.schedule_impl(Instant::now(), pree, None, |task| {
@@ -221,6 +225,25 @@ impl Scheduler {
             None => return Some(pree),
         };
         log::trace!("Checking task {:?}'s pending signal", cur.tid.raw());
+
+        if cur.space().has_to_stop() {
+            log::trace!(
+                "Killing task {:?}, P{} due to main task stopped",
+                cur.tid.raw(),
+                PREEMPT.raw()
+            );
+
+            SCHED_INFO[self.cpu]
+                .expected_runtime
+                .fetch_sub(cur.time_slice.as_micros() as u64, Release);
+
+            let _ = self.schedule_impl(cur_time, pree, None, |task| {
+                task::Ready::exit(task, sv_call::EKILLED.into_retval());
+                Ok(())
+            });
+            unreachable!("Dead task");
+        }
+
         let ti = &*cur.tid;
 
         if ti.ty() == task::Type::Kernel {
@@ -234,6 +257,7 @@ impl Scheduler {
                 SCHED_INFO[self.cpu]
                     .expected_runtime
                     .fetch_sub(cur.time_slice.as_micros() as u64, Release);
+                cur.space().try_stop(&cur.tid);
 
                 let _ = self.schedule_impl(cur_time, pree, None, |task| {
                     task::Ready::exit(task, sv_call::EKILLED.into_retval());

h2o/kernel/src/sched/task.rs

@@ -81,6 +81,7 @@ fn exec_inner(
         .unwrap();
 
     let tid = tid::allocate(ti).map_err(|_| sv_call::EBUSY)?;
+    space.set_main(&tid);
 
     let entry = ctx::Entry {
         entry: s.entry,
@@ -128,6 +129,7 @@ fn create(name: Option<String>, space: Arc<Space>, init_chan: sv_call::Handle) -
         .unwrap();
 
     let tid = tid::allocate(ti).map_err(|_| sv_call::EBUSY)?;
+    space.set_main(&tid);
 
     let mut kstack = ctx::Kstack::new(None, ty);
     kstack.task_frame_mut().set_args(init_chan.raw() as _, 0);

h2o/kernel/src/sched/task/idle.rs

@@ -40,6 +40,7 @@ pub(super) static IDLE: Lazy<Tid> = Lazy::new(|| {
     let kstack = ctx::Kstack::new(Some(entry), Type::Kernel);
 
     let tid = tid::allocate(ti).expect("Tid exhausted");
+    space.set_main(&tid);
 
     let init = Init::new(tid.clone(), space, kstack, ctx::ExtFrame::zeroed());
     crate::sched::SCHED.unblock(init, true);

h2o/kernel/src/sched/task/space.rs

@@ -1,4 +1,5 @@
 use alloc::sync::Arc;
+use core::sync::atomic::{AtomicU32, Ordering::*};
 
 use sv_call::Feature;
 
@@ -16,6 +17,7 @@ pub struct Space {
     mem: Arc<mem::space::Space>,
     handles: HandleMap,
     futexes: Futexes,
+    main: AtomicU32,
 }
 
 unsafe impl Send for Space {}
@@ -28,6 +30,7 @@ impl Space {
             mem,
             handles: HandleMap::new(),
             futexes: Default::default(),
+            main: AtomicU32::new(0),
         })?)
     }
 
@@ -36,6 +39,7 @@ impl Space {
             mem: mem::space::with_current(Arc::clone),
             handles: HandleMap::new(),
             futexes: Default::default(),
+            main: AtomicU32::new(0),
         })
     }
 
@@ -44,6 +48,21 @@ impl Space {
         &self.mem
     }
 
+    #[inline]
+    pub fn set_main(&self, tid: &Tid) {
+        let _ = self.main.compare_exchange(0, tid.raw(), AcqRel, Acquire);
+    }
+
+    #[inline]
+    pub fn try_stop(&self, tid: &Tid) {
+        let _ = self.main.compare_exchange(tid.raw(), 0, AcqRel, Acquire);
+    }
+
+    #[inline]
+    pub fn has_to_stop(&self) -> bool {
+        self.main.load(Acquire) == 0
+    }
+
     #[inline]
     pub fn handles(&self) -> &HandleMap {
         &self.handles

h2o/kernel/src/sched/task/syscall.rs

@@ -48,8 +48,8 @@ unsafe impl DefaultFeature for SuspendToken {
 }
 
 #[syscall]
-fn task_exit(retval: usize) -> Result {
-    SCHED.exit_current(retval);
+fn task_exit(retval: usize, kill_all: bool) -> Result {
+    SCHED.exit_current(retval, kill_all);
     #[allow(unreachable_code)]
     Err(EKILLED)
 }
@@ -156,7 +156,7 @@ fn task_new(
 
     let mut sus_slot = Arsc::try_new_uninit()?;
 
-    let (task, hdl) = super::create(name, Arc::clone(&space), init_chan)?;
+    let (task, hdl) = super::create(name, space, init_chan)?;
 
     let task = super::Ready::block(
         super::IntoReady::into_ready(task, unsafe { crate::cpu::id() }, MIN_TIME_GRAN),

h2o/kernel/syscall/task.json

@@ -12,6 +12,10 @@
                 {
                     "name": "retval",
                     "ty": "usize"
+                },
+                {
+                    "name": "kill_all",
+                    "ty": "bool"
                 }
             ]
         },

h2o/tinit/src/test/ipc.rs

@@ -120,7 +120,7 @@ pub unsafe fn test(virt: &Virt, stack: (*mut u8, *mut u8, Handle)) {
                 .expect("Failed to send the response");
 
             ::log::trace!("Finished");
-            sv_task_exit(0).into_res().expect("Failed to exit the task");
+            sv_task_exit(0, false).into_res().expect("Failed to exit the task");
         }
 
         let other = {

h2o/tinit/src/test/task.rs

@@ -36,7 +36,7 @@ unsafe extern "C" fn func(_: Handle, arg: u32) {
         },
         _ => {}
     }
-    sv_task_exit(12345)
+    sv_task_exit(12345, false)
         .into_res()
         .expect("Failed to exit the task");
 }

src/lib/h2o_rs/src/task.rs

@@ -188,8 +188,8 @@ impl SuspendToken {
 ///
 /// This function doesn't clean up the current self-maintained context, and the
 /// caller must ensure it is destroyed before calling this function.
-pub unsafe fn exit(retval: usize) -> ! {
-    let _ = sv_call::sv_task_exit(retval);
+pub unsafe fn exit(retval: usize, kill_all: bool) -> ! {
+    let _ = sv_call::sv_task_exit(retval, kill_all);
     unreachable!("The task failed to exit");
 }
 

src/lib/h2o_std/core/src/thread/imp.rs

@@ -60,7 +60,7 @@ impl Thread {
                 __libc_allocate_tcb();
                 Box::from_raw(arg as *mut Box<dyn FnOnce()>)();
                 __libc_deallocate_tcb();
-                exit(0);
+                exit(0, false);
             }
         }
     }

src/lib/h2o_std/src/rt.rs

@@ -110,7 +110,7 @@ macro_rules! entry {
                 unsafe { solvent::prelude::Object::from_raw(init_chan) },
                 $main,
             );
-            unsafe { solvent::task::exit(0) };
+            unsafe { solvent::task::exit(0, false) };
         }
     };
 }

src/lib/libc/src/ffi/stdlib.rs

@@ -22,7 +22,7 @@ pub extern "C" fn exit(s: i32) -> ! {
 /// caller must ensure it is destroyed before calling this function.
 #[no_mangle]
 pub unsafe extern "C" fn _Exit(s: i32) -> ! {
-    solvent::task::exit(s as usize);
+    solvent::task::exit(s as usize, true);
 }
 
 /// # Safety