task: reduce the amount of OOMs

Signed-off-by: Andy-Python-Programmer <[email protected]>

Author: Andy-Python-Programmer

aero.py

@@ -470,9 +470,11 @@ def run_in_emulator(args, iso_path):
         print("Running with KVM acceleration enabled")
 
         if platform.system() == 'Darwin':
-            qemu_args += ['-accel', 'hvf', '-cpu', 'qemu64,+la57' if args.la57 else 'qemu64']
+            qemu_args += ['-accel', 'hvf', '-cpu',
+                          'qemu64,+la57' if args.la57 else 'qemu64']
         else:
-            qemu_args += ['-enable-kvm', '-cpu', 'host,+la57' if args.la57 else 'host']
+            qemu_args += ['-enable-kvm', '-cpu',
+                          'host,+la57' if args.la57 else 'host']
     else:
         qemu_args += ["-cpu", "qemu64,+la57" if args.la57 else "qemu64"]
 

src/aero_kernel/src/arch/x86_64/task.rs

@@ -114,8 +114,10 @@ const USERLAND_STACK_BOTTOM: VirtAddr = USERLAND_STACK_TOP.const_sub_u64(USERLAN
 
 pub struct ArchTask {
     context: Unique<Context>,
+
     address_space: AddressSpace,
     context_switch_rsp: VirtAddr,
+    user: bool,
 
     fs_base: VirtAddr,
     gs_base: VirtAddr,
@@ -130,29 +132,32 @@ impl ArchTask {
             // Since the IDLE task is a special kernel task, we use the kernel's
             // address space here and we also use the kernel privilage level here.
             address_space: AddressSpace::this(),
+            user: false,
 
             fs_base: VirtAddr::zero(),
             gs_base: VirtAddr::zero(),
         }
     }
 
     pub fn new_kernel(entry_point: VirtAddr, enable_interrupts: bool) -> Self {
+        let switch_stack = Self::alloc_switch_stack().unwrap().as_mut_ptr::<u8>();
+
         let task_stack = unsafe {
             let layout = Layout::from_size_align_unchecked(4096 * 16, 0x1000);
             alloc_zeroed(layout).add(layout.size())
         };
 
         let address_space = AddressSpace::this();
 
-        let mut stack_ptr = task_stack as u64;
+        let mut stack_ptr = switch_stack as u64;
         let mut stack = StackHelper::new(&mut stack_ptr);
 
         let kframe = unsafe { stack.offset::<InterruptErrorStack>() };
 
         kframe.stack.iret.ss = 0x10; // kernel stack segment
         kframe.stack.iret.cs = 0x08; // kernel code segment
         kframe.stack.iret.rip = entry_point.as_u64();
-        kframe.stack.iret.rsp = unsafe { task_stack.sub(8) as u64 };
+        kframe.stack.iret.rsp = task_stack as u64;
         kframe.stack.iret.rflags = if enable_interrupts { 0x200 } else { 0x00 };
 
         extern "C" {
@@ -168,7 +173,8 @@ impl ArchTask {
         Self {
             context: unsafe { Unique::new_unchecked(context) },
             address_space,
-            context_switch_rsp: VirtAddr::new(task_stack as u64),
+            context_switch_rsp: VirtAddr::new(switch_stack as u64),
+            user: false,
 
             fs_base: VirtAddr::zero(),
             gs_base: VirtAddr::zero(),
@@ -184,6 +190,8 @@ impl ArchTask {
     }
 
     pub fn fork(&self) -> Result<Self, MapToError<Size4KiB>> {
+        assert!(self.user, "cannot fork a kernel task");
+
         let new_address_space = AddressSpace::this().offset_page_table().fork()?;
 
         // Since the fork function marks all of the userspace entries in both the forked
@@ -193,15 +201,7 @@ impl ArchTask {
             asm!("mov cr3, {}", in(reg) controlregs::read_cr3_raw(), options(nostack));
         }
 
-        let switch_stack = unsafe {
-            let frame: PhysFrame = FRAME_ALLOCATOR.allocate_frame().unwrap();
-
-            frame
-                .start_address()
-                .as_hhdm_virt()
-                .as_mut_ptr::<u8>()
-                .add(Size4KiB::SIZE as usize)
-        };
+        let switch_stack = Self::alloc_switch_stack()?.as_mut_ptr::<u8>();
 
         let mut old_stack_ptr = self.context_switch_rsp.as_u64();
         let mut old_stack = StackHelper::new(&mut old_stack_ptr);
@@ -232,6 +232,7 @@ impl ArchTask {
             context: unsafe { Unique::new_unchecked(context) },
             context_switch_rsp: VirtAddr::new(switch_stack as u64),
             address_space: new_address_space,
+            user: false,
 
             // The FS and GS bases are inherited from the parent process.
             fs_base: self.fs_base.clone(),
@@ -247,9 +248,18 @@ impl ArchTask {
         argv: Option<ExecArgs>,
         envv: Option<ExecArgs>,
     ) -> Result<(), MapToError<Size4KiB>> {
-        let address_space = AddressSpace::new()?;
+        let address_space = if self.user {
+            self.unref_pt();
+            AddressSpace::new()?
+        } else {
+            AddressSpace::new()?
+        };
+
         let loaded_binary = vm.load_bin(executable).expect("exec: failed to load ELF");
 
+        // a kernel task can only execute a user executable
+        self.user = true;
+
         // mmap the userland stack...
         vm.mmap(
             USERLAND_STACK_BOTTOM,
@@ -335,6 +345,49 @@ impl ArchTask {
         Ok(())
     }
 
+    /// Allocates a new context switch stack for the process and returns the stack
+    /// top address. See the module level documentation for more information.
+    fn alloc_switch_stack() -> Result<VirtAddr, MapToError<Size4KiB>> {
+        let frame: PhysFrame<Size4KiB> = FRAME_ALLOCATOR
+            .allocate_frame()
+            .ok_or(MapToError::FrameAllocationFailed)?;
+
+        Ok(frame.start_address().as_hhdm_virt() + Size4KiB::SIZE)
+    }
+
+    fn unref_pt(&mut self) {
+        self.address_space
+            .offset_page_table()
+            .page_table()
+            .for_each_entry(
+                PageTableFlags::PRESENT | PageTableFlags::USER_ACCESSIBLE,
+                |entry| {
+                    entry.unref_vm_frame();
+                    entry.set_unused();
+
+                    Ok(())
+                },
+            )
+            .expect("dealloc: failed to unref the page table");
+    }
+
+    /// Deallocates the architecture-specific task resources. This function is called
+    /// when the process is turned into a zombie.
+    pub fn dealloc(&mut self) {
+        if self.user {
+            self.unref_pt();
+        }
+
+        // deallocate the switch stack
+        {
+            let frame: PhysFrame<Size4KiB> = PhysFrame::containing_address(
+                (self.context_switch_rsp - Size4KiB::SIZE).as_hhdm_phys(),
+            );
+
+            FRAME_ALLOCATOR.deallocate_frame(frame);
+        }
+    }
+
     /// Returns the saved GS base for this task.
     pub fn get_gs_base(&self) -> VirtAddr {
         self.gs_base

src/aero_kernel/src/mem/paging/addr.rs

@@ -111,6 +111,11 @@ impl VirtAddr {
         }
     }
 
+    /// Converts this HHDM (Higher Half Direct Map) virtual address to its physical address.
+    pub fn as_hhdm_phys(&self) -> PhysAddr {
+        unsafe { PhysAddr::new(self.clone() - crate::PHYSICAL_MEMORY_OFFSET) }
+    }
+
     /// Returns if the address is valid to read `sizeof(T)` bytes at the address.
     fn validate_read<T: Sized>(&self) -> bool {
         (*self + core::mem::size_of::<T>()) <= crate::arch::task::userland_last_address()

src/aero_kernel/src/mem/paging/frame.rs

@@ -61,6 +61,10 @@ unsafe impl FrameAllocator<Size4KiB> for LockedFrameAllocator {
                     .map(|f| PhysFrame::containing_address(f))
             })
             .unwrap_or(None)
+            .map(|frame| {
+                frame.as_slice_mut().fill(0);
+                frame
+            })
     }
 
     fn deallocate_frame(&self, frame: PhysFrame<Size4KiB>) {
@@ -81,6 +85,10 @@ unsafe impl FrameAllocator<Size2MiB> for LockedFrameAllocator {
                     .map(|f| PhysFrame::containing_address(f))
             })
             .unwrap_or(None)
+            .map(|frame| {
+                frame.as_slice_mut().fill(0);
+                frame
+            })
     }
 
     fn deallocate_frame(&self, frame: PhysFrame<Size2MiB>) {
@@ -548,8 +556,9 @@ impl VmFrame {
     pub fn dec_ref_count(&self) {
         let mut this = self.lock.lock();
 
-        assert!(this.use_count > 0);
-        this.use_count -= 1;
+        if this.use_count > 0 {
+            this.use_count -= 1;
+        }
     }
 
     pub fn inc_ref_count(&self) {
@@ -574,6 +583,14 @@ mod tests {
         let mut offset_table = address_space.offset_page_table();
 
         let frame: PhysFrame = unsafe { FRAME_ALLOCATOR.allocate_frame().unwrap() };
+
+        assert!(!FRAME_ALLOCATOR
+            .0
+            .get()
+            .unwrap()
+            .lock()
+            .is_free(frame.start_address(), 0));
+
         let page = Page::<Size4KiB>::containing_address(VirtAddr::new(0xcafebabedea));
 
         let vm_frame = frame.start_address().as_vm_frame().unwrap();

src/aero_kernel/src/mem/paging/mapper.rs

@@ -1045,6 +1045,10 @@ impl<'a> OffsetPageTable<'a> {
             inner: MappedPageTable::new(page_table, phys_offset),
         }
     }
+
+    pub fn page_table(&mut self) -> &mut PageTable {
+        self.inner.page_table
+    }
 }
 
 #[derive(Debug)]

src/aero_kernel/src/mem/paging/page_table.rs

@@ -283,6 +283,23 @@ impl PageTable {
         }
     }
 
+    pub fn for_each_entry(
+        &mut self,
+        flags: PageTableFlags,
+        mut fun: impl FnMut(&mut PageTableEntry) -> Result<(), MapToError<Size4KiB>>,
+    ) -> Result<(), MapToError<Size4KiB>> {
+        self.for_entries_mut(flags, |_, _entry, table| {
+            table.for_entries_mut(flags, |_, _entry, table| {
+                table.for_entries_mut(flags, |_, _entry, table| {
+                    table.for_entries_mut(flags, |_, entry, _| fun(entry))?;
+                    Ok(())
+                })?;
+                Ok(())
+            })?;
+            Ok(())
+        })
+    }
+
     pub fn for_entries_mut(
         &mut self,
         flags: PageTableFlags,

src/aero_kernel/src/userland/task.rs

@@ -543,7 +543,7 @@ impl Task {
     }
 
     pub(super) fn into_zombie(&self) {
-        // TODO: Deallocate the arch task's resources.
+        self.arch_task_mut().dealloc();
 
         if let Some(parent) = self.get_parent() {
             parent.remove_child(self);