[PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec

Shivam Kalra posted 3 patches 1 week, 2 days ago
There is a newer version of this series
[PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec
Posted by Shivam Kalra 1 week, 2 days ago
Add methods to shrink the capacity of a Vec to free excess memory.
This is useful for drivers that experience variable workloads and want
to reclaim memory after spikes.

- `shrink_to(min_capacity, flags)`: Shrinks the capacity of the vector
  with a lower bound. The capacity will remain at least as large as
  both the length and the supplied value.

- `shrink_to_fit(flags)`: Shrinks the capacity of the vector as much
  as possible.

This implementation guarantees shrinking (unless already optimal),
because the kernel allocators don't support in-place shrinking,
a new allocation is always made.

Suggested-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Shivam Kalra <shivamklr@cock.li>
---
 rust/kernel/alloc/kvec.rs | 111 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 111 insertions(+)

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index ac8d6f763ae81..9c02734ced88f 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -733,6 +733,117 @@ pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
         }
         self.truncate(num_kept);
     }
+
+    /// Shrinks the capacity of the vector with a lower bound.
+    ///
+    /// The capacity will remain at least as large as both the length
+    /// and the supplied value.
+    ///
+    /// If the current capacity is less than the lower limit, this is a no-op.
+    ///
+    /// This requires allocating a new buffer and copying the elements, then freeing
+    /// the old buffer.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = KVec::with_capacity(100, GFP_KERNEL)?;
+    /// v.push(1, GFP_KERNEL)?;
+    /// v.push(2, GFP_KERNEL)?;
+    /// assert!(v.capacity() >= 100);
+    ///
+    /// v.shrink_to(50, GFP_KERNEL)?;
+    /// assert!(v.capacity() >= 50);
+    /// assert!(v.capacity() < 100);
+    ///
+    /// v.shrink_to(0, GFP_KERNEL)?;
+    /// assert!(v.capacity() >= 2);
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn shrink_to(&mut self, min_capacity: usize, flags: Flags) -> Result<(), AllocError> {
+        // Calculate the target capacity: max(len, min_capacity).
+        let target_cap = core::cmp::max(self.len(), min_capacity);
+
+        // If we're already within limits, nothing to do.
+        if self.capacity() <= target_cap {
+            return Ok(());
+        }
+
+        // ZSTs have no allocation to shrink.
+        if Self::is_zst() {
+            return Ok(());
+        }
+
+        // Handle empty vector or target capacity 0: just free the allocation and reset.
+        if target_cap == 0 {
+            // Only free if we actually have an allocation.
+            if !self.layout.is_empty() {
+                // SAFETY:
+                // - `self.ptr` was previously allocated with `A`.
+                // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
+                unsafe { A::free(self.ptr.cast(), self.layout.into()) };
+            }
+            self.ptr = NonNull::dangling();
+            self.layout = ArrayLayout::empty();
+            return Ok(());
+        }
+
+        // Create a new layout that exactly fits the target capacity.
+        // SAFETY: `target_cap` is guaranteed to be <= `self.capacity()`, and the original
+        // capacity was validated, so `target_cap * size_of::<T>() <= isize::MAX`.
+        let new_layout = unsafe { ArrayLayout::<T>::new_unchecked(target_cap) };
+
+        // Allocate a new, smaller buffer.
+        let new_ptr = A::alloc(new_layout.into(), flags, NumaNode::NO_NODE)?;
+
+        // SAFETY:
+        // - `self.as_ptr()` is valid for reads of `self.len` elements by the type invariant.
+        // - `new_ptr` is valid for writes of `self.len` elements (we just allocated it).
+        // - The regions do not overlap (different allocations).
+        // - Both pointers are properly aligned.
+        unsafe {
+            ptr::copy_nonoverlapping(self.as_ptr(), new_ptr.as_ptr().cast::<T>(), self.len);
+        }
+
+        // Free the old buffer.
+        // SAFETY:
+        // - `self.ptr` was previously allocated with `A`.
+        // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
+        unsafe { A::free(self.ptr.cast(), self.layout.into()) };
+
+        // SAFETY: `new_ptr.as_ptr()` is non-null because `A::alloc` succeeded.
+        self.ptr = unsafe { NonNull::new_unchecked(new_ptr.as_ptr().cast::<T>()) };
+        self.layout = new_layout;
+
+        Ok(())
+    }
+
+    /// Shrinks the capacity of the vector as much as possible.
+    ///
+    /// The capacity will be reduced to match the length, freeing any excess memory.
+    /// This requires allocating a new buffer and copying the elements, then freeing
+    /// the old buffer.
+    ///
+    /// If the allocation of the new buffer fails, the vector is left unchanged and
+    /// an error is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = KVec::with_capacity(100, GFP_KERNEL)?;
+    /// v.push(1, GFP_KERNEL)?;
+    /// v.push(2, GFP_KERNEL)?;
+    /// v.push(3, GFP_KERNEL)?;
+    /// assert!(v.capacity() >= 100);
+    ///
+    /// v.shrink_to_fit(GFP_KERNEL)?;
+    /// assert_eq!(v.capacity(), 3);
+    /// assert_eq!(&v, &[1, 2, 3]);
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn shrink_to_fit(&mut self, flags: Flags) -> Result<(), AllocError> {
+        self.shrink_to(0, flags)
+    }
 }
 
 impl<T: Clone, A: Allocator> Vec<T, A> {
-- 
2.43.0
Re: [PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec
Posted by Danilo Krummrich 6 days, 8 hours ago
On Sat Jan 31, 2026 at 4:40 PM CET, Shivam Kalra wrote:
> This implementation guarantees shrinking (unless already optimal),
> because the kernel allocators don't support in-place shrinking,
> a new allocation is always made.

I'm not sure we should go in this direction. There is a reason why krealloc()
does not migrate memory between kmalloc buckets, i.e. the cost of migration vs.
memory saving.

For Vmalloc buffers the story is a bit different though. When I wrote vrealloc()
I left some TODO comments [1][2].

  (1) If a smaller buffer is requested we can shrink the vm_area, i.e. unmap and
      free unused pages.

  (2) If a bigger buffer is requested we can grow the vm_area, i.e. allocate and
      map additional pages. (At least as long as we have enough space in the
      virtual address space.)

So, I think we should just use A::realloc(), leave the rest to the underlying
specific realloc() implementations and address the TODOs in vrealloc() if
necessary.

[1] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4162
[2] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4192
Re: [PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec
Posted by Alice Ryhl 6 days, 8 hours ago
On Tue, Feb 03, 2026 at 04:19:14PM +0100, Danilo Krummrich wrote:
> On Sat Jan 31, 2026 at 4:40 PM CET, Shivam Kalra wrote:
> > This implementation guarantees shrinking (unless already optimal),
> > because the kernel allocators don't support in-place shrinking,
> > a new allocation is always made.
> 
> I'm not sure we should go in this direction. There is a reason why krealloc()
> does not migrate memory between kmalloc buckets, i.e. the cost of migration vs.
> memory saving.
> 
> For Vmalloc buffers the story is a bit different though. When I wrote vrealloc()
> I left some TODO comments [1][2].
> 
>   (1) If a smaller buffer is requested we can shrink the vm_area, i.e. unmap and
>       free unused pages.
> 
>   (2) If a bigger buffer is requested we can grow the vm_area, i.e. allocate and
>       map additional pages. (At least as long as we have enough space in the
>       virtual address space.)
> 
> So, I think we should just use A::realloc(), leave the rest to the underlying
> specific realloc() implementations and address the TODOs in vrealloc() if
> necessary.
> 
> [1] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4162
> [2] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4192

If kvrealloc() does the right thing, then let's use it.

Alice
Re: [PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec
Posted by Danilo Krummrich 6 days, 8 hours ago
On Tue Feb 3, 2026 at 4:38 PM CET, Alice Ryhl wrote:
> On Tue, Feb 03, 2026 at 04:19:14PM +0100, Danilo Krummrich wrote:
>> On Sat Jan 31, 2026 at 4:40 PM CET, Shivam Kalra wrote:
>> > This implementation guarantees shrinking (unless already optimal),
>> > because the kernel allocators don't support in-place shrinking,
>> > a new allocation is always made.
>> 
>> I'm not sure we should go in this direction. There is a reason why krealloc()
>> does not migrate memory between kmalloc buckets, i.e. the cost of migration vs.
>> memory saving.
>> 
>> For Vmalloc buffers the story is a bit different though. When I wrote vrealloc()
>> I left some TODO comments [1][2].
>> 
>>   (1) If a smaller buffer is requested we can shrink the vm_area, i.e. unmap and
>>       free unused pages.
>> 
>>   (2) If a bigger buffer is requested we can grow the vm_area, i.e. allocate and
>>       map additional pages. (At least as long as we have enough space in the
>>       virtual address space.)
>> 
>> So, I think we should just use A::realloc(), leave the rest to the underlying
>> specific realloc() implementations and address the TODOs in vrealloc() if
>> necessary.
>> 
>> [1] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4162
>> [2] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4192
>
> If kvrealloc() does the right thing, then let's use it.

It should once the TODOs of vrealloc() are addressed. The reason I left them as
TODOs was that I didn't want to implement all the shrink and grow logic for
struct vm_area without having a user that actually needs it.

If binder needs it, I think we should do it.
Re: [PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec
Posted by Shivam Kalra 6 days, 6 hours ago
On 03/02/26 21:13, Danilo Krummrich wrote:
> On Tue Feb 3, 2026 at 4:38 PM CET, Alice Ryhl wrote:
>> On Tue, Feb 03, 2026 at 04:19:14PM +0100, Danilo Krummrich wrote:
>>> On Sat Jan 31, 2026 at 4:40 PM CET, Shivam Kalra wrote:
>>>> This implementation guarantees shrinking (unless already optimal),
>>>> because the kernel allocators don't support in-place shrinking,
>>>> a new allocation is always made.
>>>
>>> I'm not sure we should go in this direction. There is a reason why krealloc()
>>> does not migrate memory between kmalloc buckets, i.e. the cost of migration vs.
>>> memory saving.
>>>
>>> For Vmalloc buffers the story is a bit different though. When I wrote vrealloc()
>>> I left some TODO comments [1][2].
>>>
>>>   (1) If a smaller buffer is requested we can shrink the vm_area, i.e. unmap and
>>>       free unused pages.
>>>
>>>   (2) If a bigger buffer is requested we can grow the vm_area, i.e. allocate and
>>>       map additional pages. (At least as long as we have enough space in the
>>>       virtual address space.)
>>>
>>> So, I think we should just use A::realloc(), leave the rest to the underlying
>>> specific realloc() implementations and address the TODOs in vrealloc() if
>>> necessary.
>>>
>>> [1] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4162
>>> [2] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4192
>>
>> If kvrealloc() does the right thing, then let's use it.
> 
> It should once the TODOs of vrealloc() are addressed. The reason I left them as
> TODOs was that I didn't want to implement all the shrink and grow logic for
> struct vm_area without having a user that actually needs it.
> 
> If binder needs it, I think we should do it.
Hi Danilo, Alice,

Thanks for the detailed feedback - I hadn't considered the kmalloc bucket
migration costs.

Given that:
- krealloc() intentionally avoids migrating data to smaller buckets when
shrinking
- vrealloc() has TODOs for in-place shrinking
- The immediate need is binder, which uses KVec (could use either allocator)

I'm thinking the pragmatic path is:

1. For v3: Simplify shrink_to() to use A::realloc() instead of
   alloc+copy+free. This ensures we get whatever optimization
   the allocator provides (including the bucket preservation for kmalloc).

2. The vrealloc() in-place shrinking could be a separate follow-up
   series, as it's a larger change to the allocator itself.

Does this approach make sense, or would you prefer I tackle the
vrealloc TODOs first?

Thanks,
Shivam
Re: [PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec
Posted by Alice Ryhl 5 days, 13 hours ago
On Tue, Feb 03, 2026 at 11:18:17PM +0530, Shivam Kalra wrote:
> On 03/02/26 21:13, Danilo Krummrich wrote:
> > On Tue Feb 3, 2026 at 4:38 PM CET, Alice Ryhl wrote:
> >> On Tue, Feb 03, 2026 at 04:19:14PM +0100, Danilo Krummrich wrote:
> >>> On Sat Jan 31, 2026 at 4:40 PM CET, Shivam Kalra wrote:
> >>>> This implementation guarantees shrinking (unless already optimal),
> >>>> because the kernel allocators don't support in-place shrinking,
> >>>> a new allocation is always made.
> >>>
> >>> I'm not sure we should go in this direction. There is a reason why krealloc()
> >>> does not migrate memory between kmalloc buckets, i.e. the cost of migration vs.
> >>> memory saving.
> >>>
> >>> For Vmalloc buffers the story is a bit different though. When I wrote vrealloc()
> >>> I left some TODO comments [1][2].
> >>>
> >>>   (1) If a smaller buffer is requested we can shrink the vm_area, i.e. unmap and
> >>>       free unused pages.
> >>>
> >>>   (2) If a bigger buffer is requested we can grow the vm_area, i.e. allocate and
> >>>       map additional pages. (At least as long as we have enough space in the
> >>>       virtual address space.)
> >>>
> >>> So, I think we should just use A::realloc(), leave the rest to the underlying
> >>> specific realloc() implementations and address the TODOs in vrealloc() if
> >>> necessary.
> >>>
> >>> [1] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4162
> >>> [2] https://elixir.bootlin.com/linux/v6.18.6/source/mm/vmalloc.c#L4192
> >>
> >> If kvrealloc() does the right thing, then let's use it.
> > 
> > It should once the TODOs of vrealloc() are addressed. The reason I left them as
> > TODOs was that I didn't want to implement all the shrink and grow logic for
> > struct vm_area without having a user that actually needs it.
> > 
> > If binder needs it, I think we should do it.
> Hi Danilo, Alice,
> 
> Thanks for the detailed feedback - I hadn't considered the kmalloc bucket
> migration costs.
> 
> Given that:
> - krealloc() intentionally avoids migrating data to smaller buckets when
> shrinking
> - vrealloc() has TODOs for in-place shrinking
> - The immediate need is binder, which uses KVec (could use either allocator)

Binder uses KVVec not KVec, which is the one that could use either allocator.

> I'm thinking the pragmatic path is:
> 
> 1. For v3: Simplify shrink_to() to use A::realloc() instead of
>    alloc+copy+free. This ensures we get whatever optimization
>    the allocator provides (including the bucket preservation for kmalloc).
> 
> 2. The vrealloc() in-place shrinking could be a separate follow-up
>    series, as it's a larger change to the allocator itself.
> 
> Does this approach make sense, or would you prefer I tackle the
> vrealloc TODOs first?

I would kind of prefer that we do this in two steps. First have
shrink_to() use the implementation it does right now. Then a follow-up
patch fix the TODOs in vrealloc().

Alice
Re: [PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec
Posted by Danilo Krummrich 5 days, 12 hours ago
On Wed Feb 4, 2026 at 11:32 AM CET, Alice Ryhl wrote:
> I would kind of prefer that we do this in two steps. First have
> shrink_to() use the implementation it does right now. Then a follow-up
> patch fix the TODOs in vrealloc().

This is fine with me as a short term workaround, but we should only do the full
copy under certain conditions only:

  (1) is_vmalloc_addr() returns true.

  (2) The new size of the allocation requires at least one page less in total.

I.e. if it is a kmalloc() buffer we don't do anything. And if it's a vmalloc()
buffer, we only shrink if we can get rid of at least one page, since otherwise
there are no savings effectively.

Shivam do you plan to follow up on the vrealloc() TODOs subsequently?
Re: [PATCH v2 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec
Posted by Shivam Kalra 5 days, 2 hours ago
On 04/02/26 17:20, Danilo Krummrich wrote:
> This is fine with me as a short term workaround, but we should only do the full
> copy under certain conditions only:
> 
>   (1) is_vmalloc_addr() returns true.
> 
>   (2) The new size of the allocation requires at least one page less in total.
> 
> I.e. if it is a kmalloc() buffer we don't do anything. And if it's a vmalloc()
> buffer, we only shrink if we can get rid of at least one page, since otherwise
> there are no savings effectively.
> 
> Shivam do you plan to follow up on the vrealloc() TODOs subsequently?
I'll work on v3 that addresses your conditions
(only shrink for vmalloc buffers when freeing >= 1 page).
I'll check if the necessary primitives (like is_vmalloc_addr) are
exposed to Rust - if not, I may need to add helper functions or bindings.

As for the vrealloc() TODOs - yes, I'm happy to follow up on that as a
separate series. Since it would be my first time touching the mm
subsystem C code, I'd appreciate any pointers on the preferred approach
when I get there.