Track kernel buffers per-process

src/interrupts.rs

@@ -332,7 +332,7 @@ extern "C" fn syscall_handler() {
                         retval = 1;
                     }
                 });
-                TASKING.record_buf_dealloc(buffer.len());
+                TASKING.record_curr_buf_dealloc(buffer.len());
             } else {
                 retval = 1;
             }
@@ -420,6 +420,12 @@ extern "C" fn syscall_handler() {
                     SECOND_PORT.write_u32s(&[
                         total_len, // Total block length
                     ]);
+                    TASKING.record_curr_buf_dealloc(buffer.len());
+                    #[expect(
+                        clippy::unwrap_used,
+                        reason = "The PID is known valid due to using it in message_queue_mut in the if-let condition"
+                    )]
+                    TASKING.record_buf_alloc(pid, buffer.len()).unwrap();
                     let buffer = Box::into_raw(buffer);
                     #[expect(
                         clippy::unwrap_used,
@@ -503,7 +509,7 @@ extern "C" fn syscall_handler() {
         }
         15 => {
             if let Some(buf) = get_buffer(regs.rcx) {
-                TASKING.record_buf_dealloc(buf.len());
+                TASKING.record_curr_buf_dealloc(buf.len());
             }
         }
         16 => {
@@ -511,7 +517,7 @@ extern "C" fn syscall_handler() {
             let rounded_size = size.next_multiple_of(4096);
             KERNEL_SPACE.lock().alloc_force_user = true;
             let mut buffer = Vec::with_capacity_in(rounded_size, &*KERNEL_SPACE);
-            TASKING.record_buf_alloc(rounded_size);
+            TASKING.record_curr_buf_alloc(rounded_size);
             buffer.resize(rounded_size, 0);
             let buffer = buffer.into_boxed_slice();
             let buffer = Box::into_raw(buffer);

src/tasking.rs

@@ -101,6 +101,7 @@ struct Process {
     data_buffers: Mutex<Slab<*mut [u8]>>,
     message_queue: Mutex<SegQueue<(usize, usize)>>,
     sleeping: RwLock<Option<SleepReason>>,
+    buf_allocated: AtomicUsize,
 }
 
 unsafe impl Send for Process {}
@@ -114,7 +115,6 @@ pub static TASKING: Lazy<Tasking> = Lazy::new(|| Tasking {
     ready_to_run: Mutex::new(VecDeque::new()),
     current_pid: RwLock::new(None),
     freeable_kstacks: Mutex::new(Vec::new()),
-    buf_allocated: AtomicUsize::new(0),
 });
 
 #[derive(Debug)]
@@ -123,7 +123,6 @@ pub struct Tasking {
     ready_to_run: Mutex<VecDeque<usize>>,
     current_pid: RwLock<Option<usize>>,
     freeable_kstacks: Mutex<Vec<Box<[usize], &'static ASpaceMutex>>>,
-    buf_allocated: AtomicUsize,
 }
 
 pub const KSTACK_SIZE: usize = 0x1_0000 / 8;
@@ -211,6 +210,7 @@ impl Tasking {
             message_queue: Mutex::new(SegQueue::new()),
             sleeping: RwLock::new(None),
             arguments: (user_arg_mem.cast(), arguments.len()),
+            buf_allocated: AtomicUsize::new(0),
         });
         self.ready_to_run.lock().push_back(pid);
         Ok(pid)
@@ -291,10 +291,13 @@ impl Tasking {
                     || ACTIVE_SPACE.lock().get_bytes_allocated(),
                     |space| space.get_bytes_allocated(),
                 );
+                let buf_bytes_used = current_process.buf_allocated.load(Ordering::Relaxed);
                 println!(
-                    "[TASKING] PID {current_pid} exiting, used {} ({}), this is being leaked.",
+                    "[TASKING] PID {current_pid} exiting, used {} ({}) + {} ({}) of kernel buffers, this is being leaked.",
                     SizeFormatter::new(bytes_used, BINARY),
-                    bytes_used / 4096
+                    bytes_used / 4096,
+                    SizeFormatter::new(buf_bytes_used, BINARY),
+                    buf_bytes_used / 4096,
                 );
                 self.freeable_kstacks.lock().push(processes.remove(current_pid).kernel_stack);
             }
@@ -392,32 +395,58 @@ impl Tasking {
         self.processes.read()[self.current_pid.read().unwrap()].arguments
     }
 
-    pub fn record_buf_alloc(&self, size: usize) {
+    pub fn record_curr_buf_alloc(&self, size: usize) {
-        self.buf_allocated.fetch_add(size, Ordering::Relaxed);
+        #[warn(clippy::unwrap_used, reason = "FIXME")]
+        #[warn(clippy::indexing_slicing, reason = "FIXME(?)")]
+        self.processes.read()[self.current_pid.read().unwrap()]
+            .buf_allocated
+            .fetch_add(size, Ordering::Relaxed);
     }
 
-    pub fn record_buf_dealloc(&self, size: usize) {
+    pub fn record_curr_buf_dealloc(&self, size: usize) {
-        self.buf_allocated.fetch_sub(size, Ordering::Relaxed);
+        #[warn(clippy::unwrap_used, reason = "FIXME")]
+        #[warn(clippy::indexing_slicing, reason = "FIXME(?)")]
+        self.processes.read()[self.current_pid.read().unwrap()]
+            .buf_allocated
+            .fetch_sub(size, Ordering::Relaxed);
+    }
+
+    pub fn record_buf_alloc(&self, pid: usize, size: usize) -> Result<(), InvalidPid> {
+        self.processes
+            .read()
+            .get(pid)
+            .ok_or(InvalidPid)?
+            .buf_allocated
+            .fetch_add(size, Ordering::Relaxed);
+        Ok(())
+    }
+
+    pub fn record_buf_dealloc(&self, pid: usize, size: usize) -> Result<(), InvalidPid> {
+        self.processes
+            .read()
+            .get(pid)
+            .ok_or(InvalidPid)?
+            .buf_allocated
+            .fetch_sub(size, Ordering::Relaxed);
+        Ok(())
     }
 
     pub fn print_stats(&self) {
         let mut total =
-            KERNEL_SPACE.lock().get_bytes_allocated() + self.buf_allocated.load(Ordering::Relaxed);
+            KERNEL_SPACE.lock().get_bytes_allocated();
         println!(
-            "[TASKING] Kernel misc used {}",
+            "[TASKING] Kernel used {}",
             SizeFormatter::new(KERNEL_SPACE.lock().get_bytes_allocated(), BINARY)
         );
-        println!(
-            "[TASKING] Kernel buffers used {}",
-            SizeFormatter::new(self.buf_allocated.load(Ordering::Relaxed), BINARY)
-        );
         for (i, process) in self.processes.read().iter() {
             let bytes_used = process.address_space.as_ref().map_or_else(
                 || ACTIVE_SPACE.lock().get_bytes_allocated(),
                 |space| space.get_bytes_allocated(),
             );
+            let buf_bytes_used = process.buf_allocated.load(Ordering::Relaxed);
             total += bytes_used;
-            println!("[TASKING] PID {} used {}", i, SizeFormatter::new(bytes_used, BINARY));
+            total += buf_bytes_used;
+            println!("[TASKING] PID {} used {} + {} of kernel buffers", i, SizeFormatter::new(bytes_used, BINARY), SizeFormatter::new(buf_bytes_used, BINARY));
         }
         println!("[TASKING] Total used {} ({})", SizeFormatter::new(total, BINARY), total / 4096);
     }