Additional allocation functions for special contexts

[geofft]

We've bound the Rust standard allocator to kmalloc(GFP_KERNEL), the common-case GFP ("get free pages") flag. This is mostly fine, but there are some cases where you don't want to do that:

Fallible allocations. If the kernel is out of memory, you get a Rust panic. That turns into a BUG(), which is usually acceptable (though not great), because if it happens in a syscall it kills the process with SIGKILL but leaves the kernel running (i.e., it's not a kernel panic). But if you're in some other context, it turns into a kernel panic, which is no good.

There has been some ongoing work in upstream Rust for fallible allocations - the most current thing is Vec::try_reserve() and friends, which is available in nightly but stabilization seems like it's blocked on complicated things (rust-lang/rust#48043). Also, that helps for resizable containers like Vec, but not so much for Box (unless you feel like using a one-element Vec instead).

Non-blocking allocations. GFP_KERNEL can block if it needs to free memory by switching to some other thread of execution. There are other flags like GFP_ATOMIC, GFP_NOWAIT, etc. for doing allocations from an interrupt or some other context where you can't block. See https://www.kernel.org/doc/html/latest/core-api/memory-allocation.html and https://www.kernel.org/doc/html/latest/core-api/mm-api.html#useful-gfp-flag-combinations for the various options.

While there is some work in Rust on custom allocators (https://github.com/rust-lang/wg-allocators is probably the best starting point), it turns out it's not quite what we want. That work is for type-level customization of the allocator (analogous to C++'s custom allocators), e.g., making a Vec<u8, GFPAtomic> or something. First, strictly speaking, one allocated, you don't need to keep track of how it was allocated - the same function kfree can be used regardless of what flags were specified. Second and more importantly, you might want to change the allocation type for a specific allocation: perhaps you have a Vec or BTreeMap or something, and you want to add something to it from atomic context, and it needs to allocate to do so - you'd want that allocation to be done with GFP_ATOMIC, but there's no need for your allocations in general to use anything other than GFP_KERNEL. Once you're out of atomic context, the same collection can use GFP_KERNEL for its next allocation.

So, I think we can address both of these by providing our own APIs for allocations, which run the right kmalloc variant but then cast their result back to a normal Box/Vec/etc., something like

impl From<TryReserveError> for linux_kernel_module_rust::Error {
    fn from(_: TryReserveError) -> Self { Error::ENOMEM }
}

impl<T> GFPBox for Box<T> {
    fn gfp_new<T>(x: T, flags: GFPFlags) -> Result<Box<T>, TryReserveError> {
        unsafe {
            let ptr = bindings::kmalloc(Layout::for_value(&x).size(), flags.bits());
            if ptr.is_null() {
                Err(TryReserveError::AllocError { ... })
            } else {
                ptr.write(x);
                Ok(Self::from_raw(ptr))
            }
        }
    }
}

impl<T> GFPVec for Vec<T> {
    fn gfp_with_capacity<T>(capacity: usize, flags: GFPFlags) -> Result<Vec<T>, TryReserveError> { ... }
    fn gfp_reserve(&mut self, additional: usize, flags: GFPFlags) -> Result<(), TryReserveError> { ... }
    fn gfp_push(&mut self, value: T, flags: GFPFlags) -> Result<(), TryReserveError> { ... }
    ...
}

impl GFPToOwned for str {
    type Owned = String;
    fn gfp_to_owned(&self, flags: GFPFlags) -> Result<String, TryReserveError> { ... }
}

and so forth. Then you can use these methods instead of the built-in unchecked Box::new / vec.push /"foo".to_owned() / etc. methods, but the remainder of the methods on Box / Vec / String / etc., which do not allocate, can be used as normal, and all the usual traits that apply to these types still work.

(Perhaps for production code we'd want to add a lint that you're not accidentally using the default infallible allocator, i.e., that you're not calling any of the non-gfp_ versions of these methods.)

[geofft]

Based on discussions with @joshtriplett and @ojeda.

[geofft]

https://github.com/glandium/boxext is an extension trait for Box that provides, among other things, try_new(x: T) -> Option<Box<T>>, so that's an argument that the design in this ticket is probably reasonable :)

[liubogithub]

Hi @geofft ,

The idea is soundable, we do need this in filesystem implementation where GFP_NOFS is required.
Any updates on this?