Re: [PATCH v3 00/10] Allocation APIs

[Benno Lossin]

On 25.04.24 20:42, Danilo Krummrich wrote:
> On Thu, Apr 25, 2024 at 04:09:46PM +0000, Benno Lossin wrote:
>> On 25.04.24 17:36, Danilo Krummrich wrote:
>>> (adding folks from [1])
>>>
>>> On Tue, Apr 23, 2024 at 05:43:08PM +0200, Danilo Krummrich wrote:
>>>> Hi all,
>>>>
>>>> On 3/28/24 02:35, Wedson Almeida Filho wrote:
>>>>> From: Wedson Almeida Filho <walmeida@xxxxxxxxxxxxx>
>>>>>
>>>>> Revamp how we use the `alloc` crate.
>>>>>
>>>>> We currently have a fork of the crate with changes to `Vec`; other
>>>>> changes have been upstreamed (to the Rust project). This series removes
>>>>> the fork and exposes all the functionality as extension traits.
>>>>>
>>>>> Additionally, it also introduces allocation flag parameters to all
>>>>> functions that may result in allocations (e.g., `Box::new`, `Arc::new`,
>>>>> `Vec::push`, etc.) without the `try_` prefix -- the names are available
>>>>> because we build `alloc` with `no_global_oom_handling`.
>>>>>
>>>>> Lastly, the series also removes our reliance on the `allocator_api`
>>>>> unstable feature.
>>>>>
>>>>> Long term, we still want to make such functionality available in
>>>>> upstream Rust, but this allows us to make progress now and reduces our
>>>>> maintainance burden.
>>>>>
>>>>> In summary:
>>>>> 1. Removes `alloc` fork
>>>>> 2. Removes use of `allocator_api` unstable feature
>>>>> 3. Introduces flags (e.g., GFP_KERNEL, GFP_ATOMIC) when allocating
>>>>
>>>> With that series, how do we implement alternative allocators, such as
>>>> (k)vmalloc or DMA coherent?
>>>>
>>>> For instance, I recently sketched up some firmware bindings we want to
>>>> use in Nova providing
>>>>
>>>> fn copy<A: core::alloc::Allocator>(&self, alloc: A) -> Result<Vec<u8, A>>
>>>> [1]
>>>>
>>>> making use of Vec::try_with_capacity_in(). How would I implement
>>>> something similar now?
>>>
>>> I want to follow up on this topic after also bringing it up in yesterday's
>>> weekly Rust call.
>>>
>>> In the call a few ideas were discussed, e.g. whether we could just re-enable the
>>> allocator_api feature and try getting it stabilized.
>>>
>>> With the introduction of alloc::Flags (gfp_t abstraction) allocator_api might
>>> not be a viable choice anymore.
>>
>> Bringing in some more context from the meeting: Gary suggested we create
>> a custom trait for allocators that can also handle allocation flags:
>>
>>       pub trait AllocatorWithFlags: Allocator {
>>           type Flags;
>>
>>           fn allocate_with_flags(&self, layout: Layout, flags: Self::Flags) -> Result<NonNull<[u8]>, AllocError>;
>>
>>           /* ... */
>>       }
>>
>>       impl AllocatorWithFlags for Global { /* ... */ }
>>
>>       impl<T, A> VecExt<T> for Vec<T, A> where A: AllocatorWithFlags {
>>           /* ... */
>>       }
>>
>> I think that this would work, but we would have to ensure that users are
>> only allowed to call allocating functions if they are functions that we
>> control. For example `Vec::try_reserve` [1] would still use the normal
>> `Allocator` trait that doesn't support our flags.
>> Gary noted that this could be solved by `klint` [2].
> 
> I agree, extending the Allocator trait should work.
> 
> Regarding allocating functions we don't control, isn't that the case already?
> AFAICS, we're currently always falling back to GFP_KERNEL when calling
> Vec::try_reserve().

Yes we're falling back to that, but
1. there are currently no calls to `try_reserve` in tree,
2. if you use eg a `vmalloc` allocator, then I don't know if it would be
    fine to reallocate that pointer using `krealloc`. I assumed that that
    would not be OK (hence my extra care with functions outside of our
    control).

> But yes, I also think it would be better to enforce being explicit.
> 
> Given that, is there any value extending the existing Allocator trait at all?

This is what I meant in the meeting by "do you really need the allocator
trait?". What you lose is the ability to use `Vec` and `Box`, instead
you have to use your own wrapper types (see below). But what you gain is
freedom to experiment. In the end we should still try to upstream our
findings to Rust or at least share our knowledge, but doing that from
the get-go is not ideal for productivity.

>> But we only need to extend the allocator API, if you want to use the std
>> library types that allocate. If you would also be happy with a custom
>> newtype wrapper, then we could also do that.
> 
> What do you mean with "custom newtype wrapper"?

You create a newtype struct ("newtype" means that it wraps an inner type
and adds/removes/changes features from the inner type):

     pub struct BigVec<T>(Vec<T>);

And then you implement the common operations on it:

     impl<T> BigVec<T> {
         pub fn push(&mut self, item: T) -> Result {
             self.reserve(1)?;

             self.0.spare_capacity_mut()[0].write(item);

             // SAFETY: <omitted for brevity>
             unsafe { self.0.set_len(self.0.len() + 1) };
             Ok(())
         }

         pub fn reserve(&mut self, additional: usize) -> Result {
             /*
              * implemented like `VecExt::reserve` from this patchset,
              * except that it uses `vmalloc` instead of `krealloc`.
              */
         }
     }

If we need several of these, we can also create a general API that
makes it easier to create them and avoids the duplication.

>> I think that we probably want a more general solution (ie `Allocator`
>> enriched with flags), but we would have to design that before you can
>> use it.
>>
>>
>> [1]: https://doc.rust-lang.org/alloc/vec/struct.Vec.html#method.try_reserve
>> [2]: https://github.com/Rust-for-Linux/klint
>>
>>>
>>> I think it would work for (k)vmalloc, where we could pass the page flags through
>>> const generics for instance.
>>>
>>> But I don't see how it could work with kmem_cache, where we can't just create a
>>> new allocator instance when we want to change the page flags, but need to
>>> support allocations with different page flags on the same allocator (same
>>> kmem_cache) instance.
>>
>> I think that you can write the `kmem_cache` abstraction without using
>> the allocator api. You just give the function that allocates a `flags`
>> argument like in C.
> 
> Guess you mean letting the kmem_cache implementation construct the corresponding
> container? Something like:
> 
> KmemCache<T>::alloc_box(flags: alloc::Flags) -> Box<T>
> 
> I think that'd make a lot of sense, since the size of an allocation is fixed
> anyways.

Yes, but you would probably need a newtype return type, since you need
to control how it is being freed. Also in the example above there is no
`kmem_cache` object that stores the state.

I think that the API would look more like this:

     impl<T> KMemCache<T> {
         pub fn alloc(&self, value: T, flags: Flags) -> Result<KMemObj<'_, T>>;
     }

Here are a couple of interesting elements of this API:
- I don't know if `kmem_cache` uses the same flags as the alloc module,
   if not you will need a custom `Flags` type.
- I assume that an object allocated by a `kmem_cache` is only valid
   while the cache still exists (ie if you drop the cache, all allocated
   objects are also invalidated). That is why the return type contains a
   lifetime (`'_` refers to the elided lifetime on the `&self` parameter.
   It means that the `KMemObj` is only valid while the `KMemCache` is
   also valid).
- The return type is its own kind of smart pointer that allows you to
   modify the inner value like `Box`, but it takes care of all the
   `kmem_cache` specifics (eg ensuring that the associated cache is still
   valid, freeing the object etc).r

Since I assumed several things, in the end the API might look different,
but I think that this could be a more fruitful starting point.

-- 
Cheers,
Benno