fix(kernel): fix kill wait-queue cleanup, RegionId tracking, ReleasePages ownership, and SwitchTo scheduler state

crates/domain/src/endpoint/mod.rs

@@ -61,6 +61,14 @@ impl Endpoint {
         !self.waiting_receivers.is_empty()
     }
 
+    pub fn remove_sender(&mut self, thread: ThreadId) -> bool {
+        self.waiting_senders.remove_thread(thread)
+    }
+
+    pub fn remove_receiver(&mut self, thread: ThreadId) -> bool {
+        self.waiting_receivers.remove_thread(thread)
+    }
+
     pub fn revoke(&mut self) {
         self.generation = self.generation.next();
     }
@@ -97,6 +105,15 @@ impl EndpointTable {
     pub fn destroy(&mut self, id: EndpointId) -> Option<Endpoint> {
         self.endpoints.get_mut(id.raw() as usize).and_then(|s| s.take())
     }
+
+    pub fn remove_thread_from_all(&mut self, thread: ThreadId) {
+        for slot in &mut self.endpoints {
+            if let Some(ep) = slot {
+                ep.remove_sender(thread);
+                ep.remove_receiver(thread);
+            }
+        }
+    }
 }
 
 pub trait MessagePassing {

crates/domain/src/memory/mod.rs

@@ -86,6 +86,21 @@ impl RegionTable {
         self.regions.get_mut(id.raw() as usize).and_then(|s| s.take())
     }
 
+    pub fn find_by_base(
+        &self,
+        owner: ProcessId,
+        base: PhysicalAddress,
+    ) -> Option<(RegionId, &MemoryRegion)> {
+        for (i, slot) in self.regions.iter().enumerate() {
+            if let Some(region) = slot {
+                if region.owner() == owner && region.base() == base {
+                    return Some((RegionId::from_raw(i as u16), region));
+                }
+            }
+        }
+        None
+    }
+
     pub fn regions_for_process(&self, process: ProcessId) -> impl Iterator<Item = &MemoryRegion> {
         self.regions
             .iter()

crates/domain/src/scheduler/multi_core.rs

@@ -34,6 +34,10 @@ impl MultiCoreScheduler {
         self.cores[core.raw() as usize].current()
     }
 
+    pub fn set_current(&mut self, core: CoreId, thread: ThreadId) {
+        self.cores[core.raw() as usize].set_current(thread);
+    }
+
     fn thread_count_on(&self, core: CoreId) -> usize {
         self.placements
             .iter()

crates/domain/src/scheduler/per_core.rs

@@ -43,6 +43,10 @@ impl PerCoreScheduler {
         self.current = None;
     }
 
+    pub fn set_current(&mut self, thread: ThreadId) {
+        self.current = Some(thread);
+    }
+
     pub fn highest_waiting_priority(&self) -> Option<Priority> {
         self.queue.highest_priority()
     }

crates/kernel/src/lib.rs

@@ -136,6 +136,31 @@ impl<P: Platform> Kernel<P> {
         self.processes.get(pid).is_ok()
     }
 
+    pub(crate) fn cleanup_threads(&mut self, thread_ids: &[Option<ThreadId>; 4]) {
+        use domain::scheduler::Scheduler;
+        for slot in thread_ids {
+            if let Some(tid) = slot {
+                self.endpoints.remove_thread_from_all(*tid);
+                self.cleanup_wait_any(*tid);
+                self.scheduler.remove(*tid);
+                self.pending_messages[tid.raw() as usize] = None;
+                self.pending_results[tid.raw() as usize] = None;
+                let _ = self.threads.destroy(*tid);
+            }
+        }
+    }
+
+    pub(crate) fn cleanup_wait_any(&mut self, tid: ThreadId) {
+        if let Some(wait_set) = self.wait_sets[tid.raw() as usize].take() {
+            for ep_slot in wait_set.endpoints() {
+                if let Some(_slot) = ep_slot {
+                    self.endpoints.remove_thread_from_all(tid);
+                    break;
+                }
+            }
+        }
+    }
+
     fn check_cap(
         &self,
         owner: ProcessId,

crates/kernel/src/syscall.rs

@@ -60,12 +60,7 @@ impl<P: Platform> Kernel<P> {
             }
         };
 
-        for slot in &thread_ids {
-            if let Some(tid) = slot {
-                self.scheduler.remove(*tid);
-                let _ = self.threads.destroy(*tid);
-            }
-        }
+        self.cleanup_threads(&thread_ids);
 
         let _ = self.processes.destroy(pid);
         self.process_memory[pid.raw() as usize] = None;
@@ -159,11 +154,22 @@ impl<P: Platform> Kernel<P> {
             count,
             types::generation::Generation::initial(),
         );
-        let _region_id = self.regions.insert(region);
+        let region_id = match self.regions.insert(region) {
+            Some(id) => id,
+            None => {
+                self.page_pool.free(addr, pages);
+                if let Some(pm) = &mut self.process_memory[pid.raw() as usize] {
+                    let _ = pm.release(pages);
+                }
+                return SyscallOutcome::Complete(SyscallResult::Error(
+                    domain::error::KernelError::OutOfMemory,
+                ));
+            }
+        };
 
         let slot = match self.grant_cap(
             pid,
-            domain::resource::Resource::Memory(types::id::RegionId::from_raw(0)),
+            domain::resource::Resource::Memory(region_id),
             types::permission::PermissionSet::empty()
                 .with(types::permission::Permission::Read)
                 .with(types::permission::Permission::Write),
@@ -184,12 +190,21 @@ impl<P: Platform> Kernel<P> {
         let pages = count.value() as usize;
         let phys = types::address::PhysicalAddress::from_raw(addr.raw());
 
+        let region_id = match self.regions.find_by_base(pid, phys) {
+            Some((rid, _)) => rid,
+            None => return SyscallOutcome::Complete(SyscallResult::Error(
+                domain::error::KernelError::OutOfMemory,
+            )),
+        };
+
         if !self.page_pool.free(phys, pages) {
             return SyscallOutcome::Complete(SyscallResult::Error(
                 domain::error::KernelError::OutOfMemory,
             ));
         }
 
+        let _ = self.regions.remove(region_id);
+
         if let Some(pm) = &mut self.process_memory[pid.raw() as usize] {
             let _ = pm.release(pages);
         }
@@ -315,12 +330,7 @@ impl<P: Platform> Kernel<P> {
             }
         };
 
-        for slot in &thread_ids {
-            if let Some(tid) = slot {
-                self.scheduler.remove(*tid);
-                let _ = self.threads.destroy(*tid);
-            }
-        }
+        self.cleanup_threads(&thread_ids);
 
         let _ = self.processes.destroy(target_pid);
         self.process_memory[target_pid.raw() as usize] = None;
@@ -409,14 +419,14 @@ impl<P: Platform> Kernel<P> {
             self.pending_results[receiver_tid.raw() as usize] =
                 Some(SyscallResult::Message(message));
 
+            let core = self.platform.core_id();
+
             let receiver = match self.threads.get_mut(receiver_tid) {
                 Ok(t) => t,
                 Err(e) => return SyscallOutcome::Complete(SyscallResult::Error(e)),
             };
-            let priority = receiver.priority();
             let _ = receiver.unblock();
-            let core = self.platform.core_id();
-            self.scheduler.add(receiver_tid, priority, core);
+            let _ = receiver.set_running(core);
 
             let caller = match self.threads.get_mut(caller_tid) {
                 Ok(t) => t,
@@ -426,8 +436,7 @@ impl<P: Platform> Kernel<P> {
             let _ = caller.preempt();
             self.scheduler.add(caller_tid, caller_priority, core);
 
-            self.scheduler.next(core);
-            let _ = self.threads.get_mut(receiver_tid).map(|t| t.set_running(core));
+            self.scheduler.set_current(core, receiver_tid);
 
             SyscallOutcome::SwitchTo(receiver_tid)
         } else {