diff options
Diffstat (limited to 'rust/kernel/id_pool.rs')
| -rw-r--r-- | rust/kernel/id_pool.rs | 295 |
1 files changed, 295 insertions, 0 deletions
diff --git a/rust/kernel/id_pool.rs b/rust/kernel/id_pool.rs new file mode 100644 index 000000000000..384753fe0e44 --- /dev/null +++ b/rust/kernel/id_pool.rs @@ -0,0 +1,295 @@ +// SPDX-License-Identifier: GPL-2.0 + +// Copyright (C) 2025 Google LLC. + +//! Rust API for an ID pool backed by a [`BitmapVec`]. + +use crate::alloc::{AllocError, Flags}; +use crate::bitmap::BitmapVec; + +/// Represents a dynamic ID pool backed by a [`BitmapVec`]. +/// +/// Clients acquire and release IDs from unset bits in a bitmap. +/// +/// The capacity of the ID pool may be adjusted by users as +/// needed. The API supports the scenario where users need precise control +/// over the time of allocation of a new backing bitmap, which may require +/// release of spinlock. +/// Due to concurrent updates, all operations are re-verified to determine +/// if the grow or shrink is sill valid. +/// +/// # Examples +/// +/// Basic usage +/// +/// ``` +/// use kernel::alloc::AllocError; +/// use kernel::id_pool::{IdPool, UnusedId}; +/// +/// let mut pool = IdPool::with_capacity(64, GFP_KERNEL)?; +/// for i in 0..64 { +/// assert_eq!(i, pool.find_unused_id(i).ok_or(ENOSPC)?.acquire()); +/// } +/// +/// pool.release_id(23); +/// assert_eq!(23, pool.find_unused_id(0).ok_or(ENOSPC)?.acquire()); +/// +/// assert!(pool.find_unused_id(0).is_none()); // time to realloc. +/// let resizer = pool.grow_request().ok_or(ENOSPC)?.realloc(GFP_KERNEL)?; +/// pool.grow(resizer); +/// +/// assert_eq!(pool.find_unused_id(0).ok_or(ENOSPC)?.acquire(), 64); +/// # Ok::<(), Error>(()) +/// ``` +/// +/// Releasing spinlock to grow the pool +/// +/// ```no_run +/// use kernel::alloc::{AllocError, flags::GFP_KERNEL}; +/// use kernel::sync::{new_spinlock, SpinLock}; +/// use kernel::id_pool::IdPool; +/// +/// fn get_id_maybe_realloc(guarded_pool: &SpinLock<IdPool>) -> Result<usize, AllocError> { +/// let mut pool = guarded_pool.lock(); +/// loop { +/// match pool.find_unused_id(0) { +/// Some(index) => return Ok(index.acquire()), +/// None => { +/// let alloc_request = pool.grow_request(); +/// drop(pool); +/// let resizer = alloc_request.ok_or(AllocError)?.realloc(GFP_KERNEL)?; +/// pool = guarded_pool.lock(); +/// pool.grow(resizer) +/// } +/// } +/// } +/// } +/// ``` +pub struct IdPool { + map: BitmapVec, +} + +/// Indicates that an [`IdPool`] should change to a new target size. +pub struct ReallocRequest { + num_ids: usize, +} + +/// Contains a [`BitmapVec`] of a size suitable for reallocating [`IdPool`]. +pub struct PoolResizer { + new: BitmapVec, +} + +impl ReallocRequest { + /// Allocates a new backing [`BitmapVec`] for [`IdPool`]. + /// + /// This method only prepares reallocation and does not complete it. + /// Reallocation will complete after passing the [`PoolResizer`] to the + /// [`IdPool::grow`] or [`IdPool::shrink`] operation, which will check + /// that reallocation still makes sense. + pub fn realloc(&self, flags: Flags) -> Result<PoolResizer, AllocError> { + let new = BitmapVec::new(self.num_ids, flags)?; + Ok(PoolResizer { new }) + } +} + +impl IdPool { + /// Constructs a new [`IdPool`]. + /// + /// The pool will have a capacity of [`MAX_INLINE_LEN`]. + /// + /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN + #[inline] + pub fn new() -> Self { + Self { + map: BitmapVec::new_inline(), + } + } + + /// Constructs a new [`IdPool`] with space for a specific number of bits. + /// + /// A capacity below [`MAX_INLINE_LEN`] is adjusted to [`MAX_INLINE_LEN`]. + /// + /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN + #[inline] + pub fn with_capacity(num_ids: usize, flags: Flags) -> Result<Self, AllocError> { + let num_ids = usize::max(num_ids, BitmapVec::MAX_INLINE_LEN); + let map = BitmapVec::new(num_ids, flags)?; + Ok(Self { map }) + } + + /// Returns how many IDs this pool can currently have. + #[inline] + pub fn capacity(&self) -> usize { + self.map.len() + } + + /// Returns a [`ReallocRequest`] if the [`IdPool`] can be shrunk, [`None`] otherwise. + /// + /// The capacity of an [`IdPool`] cannot be shrunk below [`MAX_INLINE_LEN`]. + /// + /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN + /// + /// # Examples + /// + /// ``` + /// use kernel::{ + /// alloc::AllocError, + /// bitmap::BitmapVec, + /// id_pool::{ + /// IdPool, + /// ReallocRequest, + /// }, + /// }; + /// + /// let mut pool = IdPool::with_capacity(1024, GFP_KERNEL)?; + /// let alloc_request = pool.shrink_request().ok_or(AllocError)?; + /// let resizer = alloc_request.realloc(GFP_KERNEL)?; + /// pool.shrink(resizer); + /// assert_eq!(pool.capacity(), BitmapVec::MAX_INLINE_LEN); + /// # Ok::<(), AllocError>(()) + /// ``` + #[inline] + pub fn shrink_request(&self) -> Option<ReallocRequest> { + let cap = self.capacity(); + // Shrinking below `MAX_INLINE_LEN` is never possible. + if cap <= BitmapVec::MAX_INLINE_LEN { + return None; + } + // Determine if the bitmap can shrink based on the position of + // its last set bit. If the bit is within the first quarter of + // the bitmap then shrinking is possible. In this case, the + // bitmap should shrink to half its current size. + let Some(bit) = self.map.last_bit() else { + return Some(ReallocRequest { + num_ids: BitmapVec::MAX_INLINE_LEN, + }); + }; + if bit >= (cap / 4) { + return None; + } + let num_ids = usize::max(BitmapVec::MAX_INLINE_LEN, cap / 2); + Some(ReallocRequest { num_ids }) + } + + /// Shrinks pool by using a new [`BitmapVec`], if still possible. + #[inline] + pub fn shrink(&mut self, mut resizer: PoolResizer) { + // Between request to shrink that led to allocation of `resizer` and now, + // bits may have changed. + // Verify that shrinking is still possible. In case shrinking to + // the size of `resizer` is no longer possible, do nothing, + // drop `resizer` and move on. + let Some(updated) = self.shrink_request() else { + return; + }; + if updated.num_ids > resizer.new.len() { + return; + } + + resizer.new.copy_and_extend(&self.map); + self.map = resizer.new; + } + + /// Returns a [`ReallocRequest`] for growing this [`IdPool`], if possible. + /// + /// The capacity of an [`IdPool`] cannot be grown above [`MAX_LEN`]. + /// + /// [`MAX_LEN`]: BitmapVec::MAX_LEN + #[inline] + pub fn grow_request(&self) -> Option<ReallocRequest> { + let num_ids = self.capacity() * 2; + if num_ids > BitmapVec::MAX_LEN { + return None; + } + Some(ReallocRequest { num_ids }) + } + + /// Grows pool by using a new [`BitmapVec`], if still necessary. + /// + /// The `resizer` arguments has to be obtained by calling [`Self::grow_request`] + /// on this object and performing a [`ReallocRequest::realloc`]. + #[inline] + pub fn grow(&mut self, mut resizer: PoolResizer) { + // Between request to grow that led to allocation of `resizer` and now, + // another thread may have already grown the capacity. + // In this case, do nothing, drop `resizer` and move on. + if resizer.new.len() <= self.capacity() { + return; + } + + resizer.new.copy_and_extend(&self.map); + self.map = resizer.new; + } + + /// Finds an unused ID in the bitmap. + /// + /// Upon success, returns its index. Otherwise, returns [`None`] + /// to indicate that a [`Self::grow_request`] is needed. + #[inline] + #[must_use] + pub fn find_unused_id(&mut self, offset: usize) -> Option<UnusedId<'_>> { + // INVARIANT: `next_zero_bit()` returns None or an integer less than `map.len()` + Some(UnusedId { + id: self.map.next_zero_bit(offset)?, + pool: self, + }) + } + + /// Releases an ID. + #[inline] + pub fn release_id(&mut self, id: usize) { + self.map.clear_bit(id); + } +} + +/// Represents an unused id in an [`IdPool`]. +/// +/// # Invariants +/// +/// The value of `id` is less than `pool.map.len()`. +pub struct UnusedId<'pool> { + id: usize, + pool: &'pool mut IdPool, +} + +impl<'pool> UnusedId<'pool> { + /// Get the unused id as an usize. + /// + /// Be aware that the id has not yet been acquired in the pool. The + /// [`acquire`] method must be called to prevent others from taking the id. + /// + /// [`acquire`]: UnusedId::acquire() + #[inline] + #[must_use] + pub fn as_usize(&self) -> usize { + self.id + } + + /// Get the unused id as an u32. + /// + /// Be aware that the id has not yet been acquired in the pool. The + /// [`acquire`] method must be called to prevent others from taking the id. + /// + /// [`acquire`]: UnusedId::acquire() + #[inline] + #[must_use] + pub fn as_u32(&self) -> u32 { + // CAST: By the type invariants: + // `self.id < pool.map.len() <= BitmapVec::MAX_LEN = i32::MAX`. + self.id as u32 + } + + /// Acquire the unused id. + #[inline] + pub fn acquire(self) -> usize { + self.pool.map.set_bit(self.id); + self.id + } +} + +impl Default for IdPool { + #[inline] + fn default() -> Self { + Self::new() + } +} |