Message-ID: <CAH5fLgjRuGZWMRWTbmS5w+eLG1kPQNo95c6E6db9y8nUt2RH2A@mail.gmail.com>
Date: Mon, 16 Sep 2024 12:22:40 +0200
From: Alice Ryhl <aliceryhl@...gle.com>
To: Dave Chinner <david@...morbit.com>
Cc: Miguel Ojeda <ojeda@...nel.org>, Andrew Morton <akpm@...ux-foundation.org>, 
	Qi Zheng <zhengqi.arch@...edance.com>, Roman Gushchin <roman.gushchin@...ux.dev>, 
	Muchun Song <muchun.song@...ux.dev>, Boqun Feng <boqun.feng@...il.com>, 
	Gary Guo <gary@...yguo.net>, Björn Roy Baron <bjorn3_gh@...tonmail.com>, 
	Benno Lossin <benno.lossin@...ton.me>, Andreas Hindborg <a.hindborg@...sung.com>, 
	Trevor Gross <tmgross@...ch.edu>, linux-kernel@...r.kernel.org, linux-mm@...ck.org, 
	rust-for-linux@...r.kernel.org
Subject: Re: [PATCH] rust: shrinker: add shrinker abstraction

On Mon, Sep 16, 2024 at 4:59 AM Dave Chinner <david@...morbit.com> wrote:
>
> On Thu, Sep 12, 2024 at 09:54:01AM +0000, Alice Ryhl wrote:
> > Rust Binder holds incoming transactions in a read-only mmap'd region
> > where it manually manages the pages. These pages are only in use until
> > the incoming transaction is consumed by userspace, but the kernel will
> > keep the pages around for future transactions. Rust Binder registers a
> > shrinker with the kernel so that it can give back these pages if the
> > system comes under memory pressure.
> >
> > Separate types are provided for registered and unregistered shrinkers.
> > The unregistered shrinker type can be used to configure the shrinker
> > before registering it. Separating it into two types also enables the
> > user to construct the private data between the calls to `shrinker_alloc`
> > and `shrinker_register` and avoid constructing the private data if
> > allocating the shrinker fails.
>
> This seems a bit nasty. It appears to me that the code does an
> unsafe copy of the internal shrinker state between the unregistered
> and registered types. Shrinkers have additional internal state when
> SHRINKER_MEMCG_AWARE and/or SHRINKER_NUMA_AWARE flags are set,
> and this abstraction doesn't seem to handle either those flags or
> the internal state they use at all.

Doing an unsafe copy of the internal shrinker state is certainly not
my intent. The UnregisteredShrinker and RegisteredShrinker types just
wrap a non-null pointer, so the moves inside `register` are not copies
of the whole `struct shrinker` but just copies of a pointer to the
shrinker.

> > +impl UnregisteredShrinker {
> > +    /// Create a new shrinker.
> > +    pub fn alloc(name: &CStr) -> Result<Self, AllocError> {
> > +        // 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(0, c_str!("%s").as_char_ptr(), name.as_char_ptr()) };
>
> This passes flags as 0, so doesn't support SHRINKER_MEMCG_AWARE or
> SHRINKER_NUMA_AWARE shrinker variants. These flags should be
> passed through here.

I don't have a user for memcg/numa aware shrinkers in Rust, which is
why I did not extend these abstractions to support them. However, if
you would like me to, I'm happy to do so. It is easy to add a flags
argument to this method, but I imagine they also need a few other
additions elsewhere in the API to really be supported?

Now that I think about it, perhaps Binder's shrinker *could* be memcg
aware? It uses the list_lru abstraction, and currently calls
`list_lru_count` in `count_objects`, but maybe we could just use
`list_lru_shrink_count` instead and magically become memcg aware ...?

> > +    /// 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 };
> > +    }
>
> Seems kinda weird to have a separate function for setting seeks,
> when...
>
> > +
> > +    /// 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<T: Shrinker>(self, private_data: T::Ptr) -> RegisteredShrinker<T> {
>
> .... all the other data needed to set up a shrinker is provided to
> this function....
>
> > +        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 = <T::Ptr as ForeignOwnable>::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::<T>) };
> > +        // SAFETY: The shrinker is not yet registered, so we can update this field.
> > +        unsafe { (*ptr).scan_objects = Some(rust_scan_objects::<T>) };
>
> .... and implemented exactly the same way.

Well .. this design makes setting `seeks` optional, but `private_data`
mandatory. We *have* to set the two function pointers, but you don't
have to set `seeks` explicitly.

It's not obvious, so it's probably worth mentioning that this design
doesn't actually require you to make an allocation and store a real
pointer in the private data. When implementing the Shrinker trait for
your custom shrinker, you have to choose which pointer type to use. If
you choose the unit type `()` as the pointer type, then no real
pointer is actually stored.

> > +/// 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 unknown.
> > +    pub const UNKNOWN: Self = Self { inner: 0 };
> > +
> > +    /// 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.
> > +    pub fn from_count(count: usize) -> Self {
> > +        if count == 0 {
> > +            return Self::EMPTY;
> > +        }
>
> No. A return count of 0 is not the same as a return value of
> SHRINK_EMPTY.
>
> A return value of 0 means "no reclaim work can be done right now".
>
> This implies that there are objects that can be reclaimed in the near
> future, but right now they are unavailable for reclaim. This can be
> due to a trylock protecting the count operation failing, the all the
> objects in the cache being dirty and needing work to be done before
> they can be reclaimed, etc.
>
> A return value of SHRINK_EMPTY means "there are no reclaimable
> objects at all".
>
> This implies that the cache is empty - it has absolutely nothing in
> it that can be reclaimed either now or in the near future.
>
>
> These two return values are treated very differently by the high
> level code. SHRINK_EMPTY is used by shrink_slab_memcg() to maintain
> a "shrinker needs to run" bit in the memcg shrinker search bitmap.
> The bit is cleared when SHRINK_EMPTY is returned, meaning the
> shrinker will not get called again until a new object is queued
> to it's LRU. This sets the "shrinker needs to run" bit again, and
> so the shrinker will run next time shrink_slab_memcg() is called.
>
> In constrast, a return value of zero (i.e. no work to be done right
> now) does not change the "shrinker needs to run" state bit, and
> hence it will always be called the next time the shrink_slab_memcg()
> is run to try to reclaim objects from that memcg shrinker...

This is a part of the API that I was pretty unsure about, so very
happy to get your input. For both of the shrinkers I am familiar with,
they know the exact count of objects and there's no error case for any
lock. The shrinkers just return the exact count directly from
count_objects, but that seemed incorrect to me, as this means that the
shrinker returns 0 when it really should return SHRINK_EMPTY. That's
why I implemented `from_count` this way - the intent is that you use
`CountObjects::from_count` when you know the exact count, so if you
pass 0 to that method, then that means the shrinker really is empty.
If you don't know what the count is, then you return
`CountObjects::UNKNOWN` instead.

That said, I guess there is some reason that the C shrinker API is
designed to use the value 0 for unknown rather than empty, so perhaps
I should not try to do it differently.

> > +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 {
> > +        // The shrinker code assumes that it can multiply this value by two without overflow.
> > +        inner: c_ulong::MAX / 2,
>
> No it doesn't. This is purely a count of objects freed by the
> shrinker.

In mm/shrinker.c it multiplies by two here:
total_scan = min(total_scan, (2 * freeable));

And I guess there is even a multiplication by four a bit earlier in
the same function if priority is zero.

> > +    /// Determines whether it is safe to recurse into filesystem code.
> > +    pub fn gfp_fs(&self) -> bool {
> > +        // SAFETY: Okay by type invariants.
> > +        let mask = unsafe { (*self.ptr.as_ptr()).gfp_mask };
> > +
> > +        (mask & bindings::__GFP_FS) != 0
> > +    }
>
> This needs a check for __GFP_IO as well, as there are block layer
> shrinkers that need to be GFP_NOIO aware...

Would you prefer that I add additional methods, or that I modify this
method to also check __GFP_IO, or that I just expose a way to get the
`gfp_mask` value directly?

> > +    /// 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 processed.
> > +    pub fn set_nr_scanned(&mut self, val: usize) {
> > +        let mut val = val as c_ulong;
> > +        // SAFETY: Okay by type invariants.
> > +        let max = unsafe { (*self.ptr.as_ptr()).nr_to_scan };
> > +        if val > max {
> > +            val = max;
> > +        }
>
> No. Shrinkers are allowed to scan more objects in a batch than
> the high level code asked them to scan. If they do this, they
> *must* report back the count of all the objects they scanned so the
> batched scan accounting can adjust the remaining amount of work that
> needs to be done appropriately.

Acknowledged.

This is related to another question that I had, actually. Other than
returning SHRINK_STOP, in what situations should `nr_scanned` be
different from the return value of `scan_objects`?

> > +
> > +        // SAFETY: Okay by type invariants.
> > +        unsafe { (*self.ptr.as_ptr()).nr_scanned = val };
> > +    }
> > +}
> > +
> > +unsafe extern "C" fn rust_count_objects<T: Shrinker>(
> > +    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 { <T::Ptr as ForeignOwnable>::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
> > +}
>
> Why are there two assignments to "private" here? And for the one
> that is borrowing a reference, why is it needed as the shrinker is
> supposed to hold a reference, right? Also, where is that reference
> dropped - it's not at all obvious to me (as a relative n00b to rust)
> what is going on here.

It's just to avoid having one long line. The first line obtains a void
pointer to the private data. The second line converts it into a
borrowed form of whichever pointer type the private data has chosen to
use.

For example:
* If the private data is `Box<T>`, that is, a pointer with exclusive
ownership over an allocation created with kmalloc, then the borrowed
form is a normal shared Rust reference `&T`.
* If the private data is `Arc<T>`, that is, a pointer that owns a
refcount on an allocation created with kmalloc, then the borrowed form
is `ArcBorrow<T>`, which is like `&T` but also gives you the ability
to increment the refcount and create a new `Arc<T>` if you would like
one. An `ArcBorrow<T>` does not hold a refcount itself.

Notably, the call to `borrow` does *not* increment any refcounts.
Think of it as a pointer cast from void pointer to whichever pointer
type is appropriate for the kind of ownership used with the private
data. If you look at the definition of ForeignOwnable in
rust/kernel/types.rs, which is a trait for anything that can be stored
in a void pointer, you'll find that it has three methods:

* `into_foreign` - converts an owned version of the pointer into a void pointer.
* `from_foreign` - converts the void pointer back into an owned
version, effectively taking ownership of the void pointer
* `borrow` - let's you look inside a void pointer without taking ownership of it

Here, `from_foreign` requires that there are no active calls to
`borrow` when you take back ownership.

Alice

Powered by blists - more mailing lists