feat(mm): add page reclaim for file-backed memory pressure (rebased)

os/StarryOS/kernel/src/entry.rs

@@ -30,6 +30,8 @@ pub fn init(args: &[String], envs: &[String]) {
     pseudofs::mount_all().expect("Failed to mount pseudofs");
     spawn_alarm_task();
 
+    ax_alloc::register_page_reclaim_fn(ax_fs::page_cache_reclaim);
+
     let loc = FS_CONTEXT
         .lock()
         .resolve(&args[0])

os/arceos/modules/axalloc/src/buddy_slab.rs

@@ -165,10 +165,11 @@ impl GlobalAllocator {
     /// Allocate arbitrary number of bytes. Returns the left bound of the
     /// allocated region.
     pub fn alloc(&self, layout: Layout) -> AllocResult<NonNull<u8>> {
-        let result = {
-            let inner = self.inner.lock();
-            inner.alloc(layout).map_err(crate::AllocError::from)
-        };
+        let result = self
+            .inner
+            .lock()
+            .alloc(layout)
+            .map_err(crate::AllocError::from);
         if result.is_ok() {
             self.usages.lock().alloc(UsageKind::RustHeap, layout.size());
         }
@@ -177,13 +178,10 @@ impl GlobalAllocator {
 
     /// Gives back the allocated region to the byte allocator.
     pub fn dealloc(&self, pos: NonNull<u8>, layout: Layout) {
-        {
-            let inner = self.inner.lock();
-            unsafe { inner.dealloc(pos, layout) };
-        }
         self.usages
             .lock()
             .dealloc(UsageKind::RustHeap, layout.size());
+        unsafe { self.inner.lock().dealloc(pos, layout) };
     }
 
     /// Allocates contiguous pages.
@@ -193,16 +191,24 @@ impl GlobalAllocator {
         alignment: usize,
         kind: UsageKind,
     ) -> AllocResult<usize> {
-        let result = {
-            let inner = self.inner.lock();
-            inner
-                .alloc_pages(num_pages, alignment)
-                .map_err(crate::AllocError::from)
-        };
-        if result.is_ok() {
-            self.usages.lock().alloc(kind, num_pages * PAGE_SIZE);
+        let mut result = self.inner.lock().alloc_pages(num_pages, alignment);
+        if result.is_err() {
+            for _ in 0..4 {
+                let reclaimed = crate::try_page_reclaim(num_pages.max(16));
+                // Retry allocation regardless of whether reclaim ran;
+                // concurrent reclaim may have freed pages.
+                result = self.inner.lock().alloc_pages(num_pages, alignment);
+                if result.is_ok() {
+                    break;
+                }
+                if reclaimed == 0 {
+                    break;
+                }
+            }
         }
-        result
+        let addr = result.map_err(crate::AllocError::from)?;
+        self.usages.lock().alloc(kind, num_pages * PAGE_SIZE);
+        Ok(addr)
     }
 
     /// Allocates contiguous low-memory pages (physical address < 4 GiB).
@@ -212,16 +218,22 @@ impl GlobalAllocator {
         alignment: usize,
         kind: UsageKind,
     ) -> AllocResult<usize> {
-        let result = {
-            let inner = self.inner.lock();
-            inner
-                .alloc_pages_lowmem(num_pages, alignment)
-                .map_err(crate::AllocError::from)
-        };
-        if result.is_ok() {
-            self.usages.lock().alloc(kind, num_pages * PAGE_SIZE);
+        let mut result = self.inner.lock().alloc_pages_lowmem(num_pages, alignment);
+        if result.is_err() {
+            for _ in 0..4 {
+                let reclaimed = crate::try_page_reclaim(num_pages.max(16));
+                result = self.inner.lock().alloc_pages_lowmem(num_pages, alignment);
+                if result.is_ok() {
+                    break;
+                }
+                if reclaimed == 0 {
+                    break;
+                }
+            }
         }
-        result
+        let addr = result.map_err(crate::AllocError::from)?;
+        self.usages.lock().alloc(kind, num_pages * PAGE_SIZE);
+        Ok(addr)
     }
 
     /// Allocates contiguous pages starting from the given address.
@@ -237,11 +249,8 @@ impl GlobalAllocator {
 
     /// Gives back the allocated pages starts from `pos` to the page allocator.
     pub fn dealloc_pages(&self, pos: usize, num_pages: usize, kind: UsageKind) {
-        {
-            let inner = self.inner.lock();
-            inner.dealloc_pages(pos, num_pages);
-        }
         self.usages.lock().dealloc(kind, num_pages * PAGE_SIZE);
+        self.inner.lock().dealloc_pages(pos, num_pages);
     }
 
     /// Returns the number of allocated bytes in the allocator backend.

os/arceos/modules/axalloc/src/lib.rs

@@ -19,6 +19,29 @@ use strum::{IntoStaticStr, VariantArray};
 
 const PAGE_SIZE: usize = 0x1000;
 
+/// A function that tries to reclaim physical pages (e.g. by evicting
+/// clean file-backed page cache pages). Returns the number of pages freed.
+pub type PageReclaimFn = fn(num_pages: usize) -> usize;
+
+static PAGE_RECLAIM_FN: ax_kspin::SpinNoIrq<Option<PageReclaimFn>> = ax_kspin::SpinNoIrq::new(None);
+
+/// Register a callback that the allocator will invoke when a page allocation
+/// cannot be satisfied.
+pub fn register_page_reclaim_fn(f: PageReclaimFn) {
+    *PAGE_RECLAIM_FN.lock() = Some(f);
+}
+
+/// Try to reclaim physical pages by invoking the registered callback.
+/// Returns the number of pages actually freed.
+///
+/// The `SpinNoIrq` guard is released before calling into the reclaim
+/// function so that the reclaim path (and any evict listeners it
+/// triggers) runs with interrupts enabled.
+pub fn try_page_reclaim(num_pages: usize) -> usize {
+    let reclaim_fn = { *PAGE_RECLAIM_FN.lock() };
+    reclaim_fn.map_or(0, |f| f(num_pages))
+}
+
 mod page;
 pub use page::GlobalPage;
 

os/arceos/modules/axfs-ng/src/highlevel/file.rs

@@ -5,7 +5,7 @@ use alloc::{
 };
 use core::{
     num::NonZeroUsize,
-    sync::atomic::{AtomicU8, Ordering},
+    sync::atomic::{AtomicBool, AtomicU8, Ordering},
     task::Context,
 };
 
@@ -493,6 +493,105 @@ impl CachedFileShared {
             evict_listeners: Mutex::new(LinkedList::default()),
         }
     }
+
+    /// Scan the LRU and evict up to `max` clean pages.
+    ///
+    /// Collects non-dirty pages from the LRU cache (least-recently-used
+    /// first), pops them, notifies all registered evict listeners, and
+    /// drops the pages (freeing their backing physical memory).
+    /// Returns the number of pages evicted.
+    ///
+    /// # Lock ordering: `page_cache` -> `evict_listeners`
+    ///
+    /// This function holds the `page_cache` lock while acquiring
+    /// `evict_listeners`.  Listeners must never acquire `page_cache`
+    /// (or any lock taken before `evict_listeners`), or a deadlock will
+    /// occur.  The current callers (address-space backend page-table
+    /// invalidation) do not touch the file's page cache, so this
+    /// ordering is safe in practice.
+    ///
+    /// # Context
+    ///
+    /// When called from the page-reclaim path, the reclaim callback is
+    /// invoked with interrupts enabled: `try_page_reclaim` releases its
+    /// `SpinNoIrq` guard before invoking the callback, so `might_sleep()`
+    /// in the blocking `evict_listeners.lock()` is safe.
+    fn try_evict_clean_pages(&self, max: usize) -> usize {
+        let Some(mut cache) = self.page_cache.try_lock() else {
+            return 0;
+        };
+        let mut to_evict = [0u32; 256];
+        let limit = max.min(256);
+        let mut cnt = 0;
+        for (&pn, page) in cache.iter().rev() {
+            if !page.dirty && cnt < limit {
+                to_evict[cnt] = pn;
+                cnt += 1;
+            }
+        }
+        for &pn in to_evict[..cnt].iter() {
+            if let Some(page) = cache.pop(&pn) {
+                for listener in self.evict_listeners.lock().iter() {
+                    (listener.listener)(pn, &page);
+                }
+            }
+        }
+        cnt
+    }
+}
+
+struct ReclaimGuard;
+
+impl Drop for ReclaimGuard {
+    fn drop(&mut self) {
+        RECLAIM_IN_PROGRESS.store(false, Ordering::Release);
+    }
+}
+
+static GLOBAL_CACHED_FILES: spin::RwLock<alloc::vec::Vec<Weak<CachedFileShared>>> =
+    spin::RwLock::new(alloc::vec::Vec::new());
+
+static RECLAIM_IN_PROGRESS: AtomicBool = AtomicBool::new(false);
+
+pub fn page_cache_reclaim(num_pages: usize) -> usize {
+    if RECLAIM_IN_PROGRESS.swap(true, Ordering::AcqRel) {
+        return 0;
+    }
+    let _guard = ReclaimGuard;
+
+    let mut reclaimed = 0;
+    let target = num_pages.max(16) * 2;
+    let mut file_count = 0;
+
+    if let Some(guard) = GLOBAL_CACHED_FILES.try_read() {
+        for weak in guard.iter() {
+            if let Some(file) = weak.upgrade() {
+                let freed = file.try_evict_clean_pages(target - reclaimed);
+                reclaimed += freed;
+                file_count += 1;
+                if reclaimed >= target {
+                    break;
+                }
+            }
+        }
+    } else {
+        return 0;
+    }
+
+    if reclaimed > 0 {
+        debug!(
+            "page_cache_reclaim: evicted {} clean pages across {} files",
+            reclaimed, file_count
+        );
+    }
+
+    reclaimed
+}
+
+fn register_cached_file(file: &Arc<CachedFileShared>) {
+    let mut guard = GLOBAL_CACHED_FILES.write();
+    guard.retain(|w| w.upgrade().is_some());
+    guard.push(Arc::downgrade(file));
 }
 
 /// A file handle with an LRU page cache for buffered I/O.
@@ -536,22 +635,29 @@ impl CachedFile {
         let in_memory = location.filesystem().name() == "tmpfs";
 
         let mut guard = location.user_data();
-        let shared = if let Some(shared) = guard.get::<FileUserData>().and_then(|it| it.get()) {
-            shared
-        } else {
-            let (shared, user_data) = if in_memory {
-                let shared = Arc::new(CachedFileShared::new_unbounded());
-                (shared.clone(), FileUserData::Strong(shared))
+        let (shared, is_new) =
+            if let Some(shared) = guard.get::<FileUserData>().and_then(|it| it.get()) {
+                (shared, false)
             } else {
-                let shared = Arc::new(CachedFileShared::new());
-                let user_data = FileUserData::Weak(Arc::downgrade(&shared));
-                (shared, user_data)
+                let (shared, user_data) = if in_memory {
+                    let shared = Arc::new(CachedFileShared::new_unbounded());
+                    (shared.clone(), FileUserData::Strong(shared))
+                } else {
+                    let shared = Arc::new(CachedFileShared::new());
+                    let user_data = FileUserData::Weak(Arc::downgrade(&shared));
+                    (shared, user_data)
+                };
+                guard.insert(user_data);
+                (shared, true)
             };
-            guard.insert(user_data);
-            shared
-        };
         drop(guard);
 
+        // In-memory files (tmpfs) have no backing store, so evicting clean
+        // pages would lose data. Only register disk-backed files for reclaim.
+        if is_new && !in_memory {
+            register_cached_file(&shared);
+        }
+
         Self {
             inner: location,
             shared,

os/arceos/modules/axsync/src/mutex.rs

@@ -7,9 +7,9 @@ use ax_task::{WaitQueue, current, might_sleep};
 /// A [`lock_api::RawMutex`] implementation.
 ///
 /// When the mutex is locked, the current task will block and be put into the
-/// wait queue. When the mutex is unlocked, ownership is released before waking
-/// at most one waiting task; the woken task then acquires the mutex with the
-/// normal compare-exchange path.
+/// wait queue. When the mutex is unlocked, ownership is handed off to at most
+/// one task waiting on the queue; if no tasks are waiting, the mutex simply
+/// becomes unlocked.
 pub struct RawMutex {
     wq: WaitQueue,
     owner_id: AtomicU64,
@@ -129,8 +129,10 @@ unsafe impl lock_api::RawMutex for RawMutex {
         );
         #[cfg(feature = "lockdep")]
         crate::lockdep::release(self);
-        self.owner_id.store(0, Ordering::Release);
-        self.wq.notify_one(true);
+        // wake up one waiting thread.
+        self.wq.notify_one_with(true, |id: u64| {
+            self.owner_id.swap(id, Ordering::Release);
+        });
     }
 
     #[inline(always)]
@@ -166,10 +168,7 @@ impl RawMutex {
                         owner_id, current_id,
                         "Thread({current_id}) tried to acquire mutex it already owns.",
                     );
-                    // Wait until the lock is released. The woken waiter
-                    // competes through the normal CAS path, avoiding a state
-                    // where the owner id names a task that has not returned a
-                    // guard yet.
+                    // Wait until someone hands off lock to me or lock is released
                     self.wq
                         .wait_until(|| self.is_owner(current_id) || !self.is_locked_inner());
                     // This check is necessary: some newcomers may race with a wakened one.
@@ -184,7 +183,7 @@ impl RawMutex {
     #[inline(always)]
     #[track_caller]
     #[cfg(feature = "lockdep")]
-    fn lock_nested(&self, subclass: crate::lockdep::LockSubclass) {
+    fn lock_nested(&self, subclass: LockSubclass) {
         might_sleep();
         let current_id = current().id().as_u64();
 
@@ -197,7 +196,8 @@ impl RawMutex {
     #[track_caller]
     #[cfg(not(feature = "lockdep"))]
     fn try_lock_plain(&self) -> bool {
-        might_sleep();
+        // try_lock is a single atomic CAS — it never blocks or sleeps,
+        // so it is safe to call from atomic context (cf. Linux mutex_trylock).
         let current_id = current().id().as_u64();
         self.try_lock_after_prepare(current_id)
     }
@@ -210,8 +210,9 @@ impl RawMutex {
     #[inline(always)]
     #[track_caller]
     #[cfg(feature = "lockdep")]
-    fn try_lock_nested(&self, subclass: crate::lockdep::LockSubclass) -> bool {
-        might_sleep();
+    fn try_lock_nested(&self, subclass: LockSubclass) -> bool {
+        // try_lock is a single atomic CAS — it never blocks or sleeps,
+        // so it is safe to call from atomic context.
         let current_id = current().id().as_u64();
 
         let lockdep = crate::lockdep::LockdepAcquire::prepare_nested(self, true, subclass);