diff --git a/drivers/staging/android/ashmem.h b/drivers/staging/android/ashmem.h index 5836a654500e..fcacbfc9bd99 100644 --- a/drivers/staging/android/ashmem.h +++ b/drivers/staging/android/ashmem.h @@ -15,6 +15,10 @@ #include "uapi/ashmem.h" +#include +const gfp_t RUST_CONST_HELPER___GFP_FS = ___GFP_FS; +const gfp_t RUST_CONST_HELPER___GFP_IO = ___GFP_IO; + #define ASHMEM_NAME_PREFIX "dev/ashmem/" #define ASHMEM_NAME_PREFIX_LEN (sizeof(ASHMEM_NAME_PREFIX) - 1) #define ASHMEM_FULL_NAME_LEN (ASHMEM_NAME_LEN + ASHMEM_NAME_PREFIX_LEN) diff --git a/drivers/staging/android/ashmem_range.rs b/drivers/staging/android/ashmem_range.rs index 35527dbcc02e..669a7e02e64e 100644 --- a/drivers/staging/android/ashmem_range.rs +++ b/drivers/staging/android/ashmem_range.rs @@ -4,17 +4,27 @@ //! Keeps track of unpinned ranges in an ashmem file. -use crate::shmem::ShmemFile; -use core::{mem::MaybeUninit, pin::Pin}; +use crate::{ + ashmem_shrinker::{CountObjects, ScanObjects, ShrinkControl, Shrinker}, + shmem::ShmemFile, +}; +use core::{ + mem::MaybeUninit, + pin::Pin, + sync::atomic::{AtomicUsize, Ordering}, +}; use kernel::{ list::{List, ListArc, ListLinks}, + page::PAGE_SIZE, prelude::*, sync::{GlobalGuard, GlobalLockedBy, UniqueArc}, }; +// Only updated with ASHMEM_MUTEX held, but the shrinker will read it without the mutex. +pub(crate) static LRU_COUNT: AtomicUsize = AtomicUsize::new(0); + pub(crate) struct AshmemLru { lru_list: List, - lru_count: usize, } /// Represents ownership of the `ASHMEM_MUTEX` lock. @@ -47,15 +57,19 @@ impl AshmemGuard { // Only change the counter if the range is on the lru list. if !range.purged(self) { - self.lru_count -= old_size; - self.lru_count += new_size; + let mut lru_count = LRU_COUNT.load(Ordering::Relaxed); + lru_count -= old_size; + lru_count += new_size; + LRU_COUNT.store(lru_count, Ordering::Relaxed); } } fn insert_lru(&mut self, range: ListArc) { // Don't insert the range if it's already purged. if !range.purged(self) { - self.lru_count += range.size(self); + let mut lru_count = LRU_COUNT.load(Ordering::Relaxed); + lru_count += range.size(self); + LRU_COUNT.store(lru_count, Ordering::Relaxed); self.lru_list.push_front(range); } } @@ -67,7 +81,9 @@ impl AshmemGuard { // Only decrement lru_count if the range was actually in the list. if ret.is_some() { - self.lru_count -= range.size(self); + let mut lru_count = LRU_COUNT.load(Ordering::Relaxed); + lru_count -= range.size(self); + LRU_COUNT.store(lru_count, Ordering::Relaxed); } ret @@ -79,7 +95,6 @@ kernel::sync::global_lock! { // there are no calls to `lock` before `init` is called. pub(crate) unsafe(uninit) static ASHMEM_MUTEX: Mutex = AshmemLru { lru_list: List::new(), - lru_count: 0 }; } @@ -93,6 +108,7 @@ pub(crate) struct Range { lru: ListLinks<0>, #[pin] unpinned: ListLinks<1>, + file: ShmemFile, pub(crate) inner: GlobalLockedBy, } @@ -103,6 +119,10 @@ pub(crate) struct RangeInner { } impl Range { + pub(crate) fn set_purged(&self, guard: &mut AshmemGuard) { + self.inner.as_mut(guard).purged = true; + } + pub(crate) fn purged(&self, guard: &AshmemGuard) -> bool { self.inner.as_ref(guard).purged } @@ -344,7 +364,6 @@ impl Area { } pub(crate) struct NewRange<'a> { - #[allow(dead_code)] pub(crate) file: &'a ShmemFile, pub(crate) alloc: UniqueArc>, } @@ -354,6 +373,7 @@ impl<'a> NewRange<'a> { let new_range = self.alloc.pin_init_with(pin_init!(Range { lru <- ListLinks::new(), unpinned <- ListLinks::new(), + file: self.file.clone(), inner: GlobalLockedBy::new(inner), })); @@ -364,6 +384,62 @@ impl<'a> NewRange<'a> { } } +impl AshmemGuard { + pub(crate) fn free_lru(&mut self, stop_after: usize) -> usize { + let mut freed = 0; + while let Some(range) = self.lru_list.pop_back() { + let start = range.pgstart(self) * PAGE_SIZE; + let end = (range.pgend(self) + 1) * PAGE_SIZE; + range.set_purged(self); + self.remove_lru(&range); + freed += range.size(self); + + // C ashmem releases the mutex and uses a different mechanism to ensure mutual + // exclusion with `pin_unpin` operations, but we only hold `ASHMEM_MUTEX` here and in + // `pin_unpin`, so we don't need to release the mutex. A different mutex is used for + // all of the other ashmem operations. + range.file.punch_hole(start, end - start); + + if freed >= stop_after { + break; + } + + if super::shrinker_should_stop() { + break; + } + } + freed + } +} + +impl Shrinker for super::AshmemModule { + // Our shrinker data is in a global, so we don't need to set the private data. + type Ptr = (); + + fn count_objects(_: (), _sc: ShrinkControl<'_>) -> CountObjects { + let count = LRU_COUNT.load(super::Ordering::Relaxed); + if count == 0 { + CountObjects::EMPTY + } else { + CountObjects::new(count) + } + } + + fn scan_objects(_: (), sc: ShrinkControl<'_>) -> ScanObjects { + if !sc.reclaim_fs_allowed() { + return ScanObjects::STOP; + } + + let Some(guard) = super::ASHMEM_MUTEX.try_lock() else { + return ScanObjects::STOP; + }; + let mut guard = AshmemGuard(guard); + + let num_freed = guard.free_lru(sc.nr_to_scan()); + ScanObjects::from_count(num_freed) + } +} + #[cfg(test)] fn range_test() -> Result { fn get_random(max: usize) -> usize { diff --git a/drivers/staging/android/ashmem_rust.rs b/drivers/staging/android/ashmem_rust.rs index 1c0d756facee..fb50a5aceb9e 100644 --- a/drivers/staging/android/ashmem_rust.rs +++ b/drivers/staging/android/ashmem_rust.rs @@ -10,7 +10,11 @@ //! It is, in theory, a good memory allocator for low-memory devices, because it can discard shared //! memory units when under memory pressure. -use core::{ffi::c_int, pin::Pin}; +use core::{ + ffi::c_int, + pin::Pin, + sync::atomic::{AtomicUsize, Ordering}, +}; use kernel::{ bindings::{self, ASHMEM_GET_PIN_STATUS, ASHMEM_PIN, ASHMEM_UNPIN}, c_str, @@ -37,6 +41,9 @@ const PROT_EXEC: usize = bindings::PROT_EXEC as usize; const PROT_WRITE: usize = bindings::PROT_WRITE as usize; const PROT_MASK: usize = PROT_EXEC | PROT_READ | PROT_WRITE; +mod ashmem_shrinker; +use ashmem_shrinker::{ShrinkerBuilder, ShrinkerRegistration}; + mod ashmem_range; use ashmem_range::{Area, AshmemGuard, NewRange, ASHMEM_MUTEX}; @@ -50,6 +57,12 @@ fn read_implies_exec(task: &Task) -> bool { (personality & bindings::READ_IMPLIES_EXEC) != 0 } +static NUM_PIN_IOCTLS_WAITING: AtomicUsize = AtomicUsize::new(0); + +fn shrinker_should_stop() -> bool { + NUM_PIN_IOCTLS_WAITING.load(Ordering::Relaxed) > 0 +} + module! { type: AshmemModule, name: "ashmem_rust", @@ -60,6 +73,7 @@ module! { struct AshmemModule { _misc: Pin>>, + _shrinker: ShrinkerRegistration, } impl kernel::Module for AshmemModule { @@ -71,6 +85,9 @@ impl kernel::Module for AshmemModule { pr_info!("Using Rust implementation."); + let mut shrinker = ShrinkerBuilder::new(c_str!("android-ashmem"))?; + shrinker.set_seeks(4 * ashmem_shrinker::DEFAULT_SEEKS); + Ok(Self { _misc: Box::pin_init( MiscDeviceRegistration::register(MiscDeviceOptions { @@ -78,6 +95,7 @@ impl kernel::Module for AshmemModule { }), GFP_KERNEL, )?, + _shrinker: shrinker.register(()), }) } } @@ -356,9 +374,13 @@ impl Ashmem { Some(UniqueArc::new_uninit(GFP_KERNEL)?) }; + NUM_PIN_IOCTLS_WAITING.fetch_add(1, Ordering::Relaxed); let mut guard = AshmemGuard(ASHMEM_MUTEX.lock()); + NUM_PIN_IOCTLS_WAITING.fetch_sub(1, Ordering::Relaxed); + // C ashmem waits for in-flight shrinkers here using a separate mechanism, but we don't // release the lock when calling `punch_hole` in the shrinker, so we don't need to do that. + let asma = &mut *self.inner.lock(); let mut new_range = match asma.file.as_ref() { Some(file) => new_range.map(|alloc| NewRange { file, alloc }), diff --git a/drivers/staging/android/ashmem_shrinker.rs b/drivers/staging/android/ashmem_shrinker.rs new file mode 100644 index 000000000000..5c29d84c6a08 --- /dev/null +++ b/drivers/staging/android/ashmem_shrinker.rs @@ -0,0 +1,336 @@ +// SPDX-License-Identifier: GPL-2.0 + +// Copyright (C) 2024 Google LLC. + +#![allow(unreachable_pub, dead_code)] +//! Shrinker for handling memory pressure. +//! +//! C header: [`include/linux/shrinker.h`](srctree/include/linux/shrinker.h) + +use kernel::{alloc::AllocError, bindings, c_str, str::CStr, types::ForeignOwnable}; + +use core::{ + ffi::{c_int, c_long, c_ulong, c_void}, + marker::PhantomData, + ptr::NonNull, +}; + +const SHRINK_STOP: c_ulong = bindings::SHRINK_STOP as c_ulong; +const SHRINK_EMPTY: c_ulong = bindings::SHRINK_EMPTY as c_ulong; + +/// The default value for the number of seeks needed to recreate an object. +pub const DEFAULT_SEEKS: u32 = bindings::DEFAULT_SEEKS; + +/// An unregistered shrinker. +/// +/// This type can be used to modify the settings of the shrinker before it is registered. +/// +/// # Invariants +/// +/// The `shrinker` pointer references an unregistered shrinker. +pub struct ShrinkerBuilder { + shrinker: NonNull, +} + +// SAFETY: Moving an unregistered shrinker between threads is okay. +unsafe impl Send for ShrinkerBuilder {} +// SAFETY: An unregistered shrinker is thread safe. +unsafe impl Sync for ShrinkerBuilder {} + +impl ShrinkerBuilder { + /// Create a new shrinker. + pub fn new(name: &CStr) -> Result { + // TODO: Support numa/memcg aware shrinkers once list_lru is available. + let flags = 0; + + // SAFETY: Passing `0` as flags is okay. Using `%s` as the format string is okay when we + // pass a nul-terminated string as the string for `%s` to print. + let ptr = unsafe { + bindings::shrinker_alloc(flags, c_str!("%s").as_char_ptr(), name.as_char_ptr()) + }; + + let shrinker = NonNull::new(ptr).ok_or(AllocError)?; + + // INVARIANT: The allocated shrinker is unregistered. + Ok(Self { shrinker }) + } + + /// Create a new shrinker using format arguments for the name. + pub fn new_fmt(name: core::fmt::Arguments<'_>) -> Result { + // TODO: Support numa/memcg aware shrinkers once list_lru is available. + let flags = 0; + + // SAFETY: Passing `0` as flags is okay. Using `%pA` as the format string is okay when we + // pass a `fmt::Arguments` as the value to print. + let ptr = unsafe { + bindings::shrinker_alloc( + flags, + c_str!("%pA").as_char_ptr(), + &name as *const _ as *const c_void, + ) + }; + + let shrinker = NonNull::new(ptr).ok_or(AllocError)?; + + // INVARIANT: The allocated shrinker is unregistered. + Ok(Self { shrinker }) + } + + /// Set the number of seeks needed to recreate an object. + pub fn set_seeks(&mut self, seeks: u32) { + unsafe { (*self.shrinker.as_ptr()).seeks = seeks as c_int }; + } + + /// Set the batch size for reclaiming on this shrinker. + pub fn set_batch(&mut self, batch: usize) { + unsafe { (*self.shrinker.as_ptr()).batch = batch as c_long }; + } + + /// Register the shrinker. + /// + /// The provided pointer is used as the private data, and the type `T` determines the callbacks + /// that the shrinker will use. + pub fn register(self, private_data: T::Ptr) -> ShrinkerRegistration { + let shrinker = self.shrinker; + let ptr = shrinker.as_ptr(); + + // The destructor of `self` calls `shrinker_free`, so skip the destructor. + core::mem::forget(self); + + let private_data_ptr = ::into_foreign(private_data); + + // SAFETY: We own the private data, so we can assign to it. + unsafe { (*ptr).private_data = private_data_ptr.cast_mut() }; + // SAFETY: The shrinker is not yet registered, so we can update this field. + unsafe { (*ptr).count_objects = Some(rust_count_objects::) }; + // SAFETY: The shrinker is not yet registered, so we can update this field. + unsafe { (*ptr).scan_objects = Some(rust_scan_objects::) }; + + // SAFETY: The shrinker is unregistered, so it's safe to register it. + unsafe { bindings::shrinker_register(ptr) }; + + ShrinkerRegistration { + shrinker, + _phantom: PhantomData, + } + } +} + +impl Drop for ShrinkerBuilder { + fn drop(&mut self) { + // SAFETY: The shrinker is a valid but unregistered shrinker, and we will not use it + // anymore. + unsafe { bindings::shrinker_free(self.shrinker.as_ptr()) }; + } +} + +/// A shrinker that is registered with the kernel. +/// +/// # Invariants +/// +/// The `shrinker` pointer refers to a registered shrinker using `T` as the private data. +pub struct ShrinkerRegistration { + shrinker: NonNull, + _phantom: PhantomData, +} + +// SAFETY: This allows you to deregister the shrinker from a different thread, which means that +// private data could be dropped from any thread. +unsafe impl Send for ShrinkerRegistration where T::Ptr: Send {} +// SAFETY: The only thing you can do with an immutable reference is access the private data, which +// is okay to access in parallel as the `Shrinker` trait requires the private data to be `Sync`. +unsafe impl Sync for ShrinkerRegistration {} + +impl ShrinkerRegistration { + /// Access the private data in this shrinker. + pub fn private_data(&self) -> ::Borrowed<'_> { + // SAFETY: We own the private data, so we can access it. + let private = unsafe { (*self.shrinker.as_ptr()).private_data }; + // SAFETY: By the type invariants, the private data is `T`. This access could happen in + // parallel with a shrinker callback, but that's okay as the `Shrinker` trait ensures that + // `T::Ptr` is `Sync`. + unsafe { ::borrow(private) } + } +} + +impl Drop for ShrinkerRegistration { + fn drop(&mut self) { + // SAFETY: We own the private data, so we can access it. + let private = unsafe { (*self.shrinker.as_ptr()).private_data }; + // SAFETY: We will not access the shrinker after this call. + unsafe { bindings::shrinker_free(self.shrinker.as_ptr()) }; + // SAFETY: The above call blocked until the completion of any shrinker callbacks, so there + // are no longer any users of the private data. + drop(unsafe { ::from_foreign(private) }); + } +} + +/// Callbacks for a shrinker. +pub trait Shrinker { + /// The pointer type used to store the private data of the shrinker. + /// + /// Needs to be `Sync` because the shrinker callback could access this value immutably from + /// several thread in parallel. + type Ptr: ForeignOwnable + Sync; + + /// Count the number of freeable items in the cache. + fn count_objects( + me: ::Borrowed<'_>, + sc: ShrinkControl<'_>, + ) -> CountObjects; + + /// Free some objects in this cache. + fn scan_objects( + me: ::Borrowed<'_>, + sc: ShrinkControl<'_>, + ) -> ScanObjects; +} + +/// How many objects are there in the cache? +/// +/// This is used as the return value of [`Shrinker::count_objects`]. +pub struct CountObjects { + inner: c_ulong, +} + +impl CountObjects { + /// Indicates that the number of objects is zero. + pub const EMPTY: Self = Self { + inner: SHRINK_EMPTY, + }; + + /// The maximum possible number of freeable objects. + pub const MAX: Self = Self { + // The shrinker code assumes that it can multiply this value by two without overflow. + inner: c_ulong::MAX / 2, + }; + + /// Creates a new `CountObjects` with the given value. + /// + /// This should be the number of objects that were actually freed. Objects that were scanned + /// but not freed should be counted in `nr_scanned` but not here. + /// + /// If `count` is zero, then this indicates that the real count is unknown. Use + /// `CountObjects::EMPTY` to indicate that the shrinker is empty. + pub fn new(count: usize) -> Self { + if count > Self::MAX.inner as usize { + return Self::MAX; + } + + Self { + inner: count as c_ulong, + } + } +} + +/// How many objects were freed? +/// +/// This is used as the return value of [`Shrinker::scan_objects`]. +pub struct ScanObjects { + inner: c_ulong, +} + +impl ScanObjects { + /// Indicates that the shrinker should stop trying to free objects from this cache due to + /// potential deadlocks. + pub const STOP: Self = Self { inner: SHRINK_STOP }; + + /// The maximum possible number of freeable objects. + pub const MAX: Self = Self { + inner: SHRINK_STOP - 1, + }; + + /// Creates a new `CountObjects` with the given value. + pub fn from_count(count: usize) -> Self { + if count > Self::MAX.inner as usize { + return Self::MAX; + } + + Self { + inner: count as c_ulong, + } + } +} + +/// This struct is used to pass information from page reclaim to the shrinkers. +/// +/// # Invariants +/// +/// `ptr` has exclusive access to a valid `struct shrink_control`. +pub struct ShrinkControl<'a> { + ptr: NonNull, + _phantom: PhantomData<&'a bindings::shrink_control>, +} + +impl<'a> ShrinkControl<'a> { + /// Create a `ShrinkControl` from a raw pointer. + /// + /// # Safety + /// + /// The pointer should point at a valid `shrink_control` for the duration of 'a. + pub unsafe fn from_raw(ptr: *mut bindings::shrink_control) -> Self { + Self { + // SAFETY: Caller promises that this pointer is valid. + ptr: unsafe { NonNull::new_unchecked(ptr) }, + _phantom: PhantomData, + } + } + + /// Determines whether it is safe to call into filesystem code. + pub fn reclaim_fs_allowed(&self) -> bool { + // SAFETY: Okay by type invariants. + let mask = unsafe { (*self.ptr.as_ptr()).gfp_mask }; + + (mask & bindings::__GFP_FS) != 0 + } + + /// Determines whether it is safe to call into IO code. + pub fn reclaim_io_allowed(&self) -> bool { + // SAFETY: Okay by type invariants. + let mask = unsafe { (*self.ptr.as_ptr()).gfp_mask }; + + (mask & bindings::__GFP_IO) != 0 + } + + /// Returns the number of objects that `scan_objects` should try to reclaim. + pub fn nr_to_scan(&self) -> usize { + // SAFETY: Okay by type invariants. + unsafe { (*self.ptr.as_ptr()).nr_to_scan as usize } + } + + /// The callback should set this value to the number of objects inspected by the shrinker. + pub fn set_nr_scanned(&mut self, val: usize) { + // SAFETY: Okay by type invariants. + unsafe { (*self.ptr.as_ptr()).nr_scanned = val as c_ulong }; + } +} + +unsafe extern "C" fn rust_count_objects( + shrink: *mut bindings::shrinker, + sc: *mut bindings::shrink_control, +) -> c_ulong { + // SAFETY: We own the private data, so we can access it. + let private = unsafe { (*shrink).private_data }; + // SAFETY: This function is only used with shrinkers where `T` is the type of the private data. + let private = unsafe { ::borrow(private) }; + // SAFETY: The caller passes a valid `sc` pointer. + let sc = unsafe { ShrinkControl::from_raw(sc) }; + + let ret = T::count_objects(private, sc); + ret.inner +} + +unsafe extern "C" fn rust_scan_objects( + shrink: *mut bindings::shrinker, + sc: *mut bindings::shrink_control, +) -> c_ulong { + // SAFETY: We own the private data, so we can access it. + let private = unsafe { (*shrink).private_data }; + // SAFETY: This function is only used with shrinkers where `T` is the type of the private data. + let private = unsafe { ::borrow(private) }; + // SAFETY: The caller passes a valid `sc` pointer. + let sc = unsafe { ShrinkControl::from_raw(sc) }; + + let ret = T::scan_objects(private, sc); + ret.inner +} diff --git a/drivers/staging/android/shmem.rs b/drivers/staging/android/shmem.rs index 3bddaea0628d..f54efb6d6ca6 100644 --- a/drivers/staging/android/shmem.rs +++ b/drivers/staging/android/shmem.rs @@ -108,6 +108,25 @@ impl ShmemFile { } } + pub(crate) fn punch_hole(&self, start: usize, len: usize) { + use kernel::bindings::{FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE}; + + let f = self.inner.as_ptr(); + // SAFETY: f_op of a file is immutable, so okay to read. + let fallocate = unsafe { (*(*f).f_op).fallocate }; + + if let Some(fallocate) = fallocate { + unsafe { + fallocate( + f, + (FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE) as _, + start as _, + len as _, + ) + }; + } + } + pub(crate) fn inode_ino(&self) -> usize { // SAFETY: Accessing the ino is always okay. unsafe { (*(*self.inner.as_ptr()).f_inode).i_ino as usize }