id_pool.rs

  1// SPDX-License-Identifier: GPL-2.0
  2
  3// Copyright (C) 2025 Google LLC.
  4
  5//! Rust API for an ID pool backed by a [`BitmapVec`].
  6
  7use crate::alloc::{AllocError, Flags};
  8use crate::bitmap::BitmapVec;
  9
 10/// Represents a dynamic ID pool backed by a [`BitmapVec`].
 11///
 12/// Clients acquire and release IDs from unset bits in a bitmap.
 13///
 14/// The capacity of the ID pool may be adjusted by users as
 15/// needed. The API supports the scenario where users need precise control
 16/// over the time of allocation of a new backing bitmap, which may require
 17/// release of spinlock.
 18/// Due to concurrent updates, all operations are re-verified to determine
 19/// if the grow or shrink is sill valid.
 20///
 21/// # Examples
 22///
 23/// Basic usage
 24///
 25/// ```
 26/// use kernel::alloc::AllocError;
 27/// use kernel::id_pool::{IdPool, UnusedId};
 28///
 29/// let mut pool = IdPool::with_capacity(64, GFP_KERNEL)?;
 30/// for i in 0..64 {
 31///     assert_eq!(i, pool.find_unused_id(i).ok_or(ENOSPC)?.acquire());
 32/// }
 33///
 34/// pool.release_id(23);
 35/// assert_eq!(23, pool.find_unused_id(0).ok_or(ENOSPC)?.acquire());
 36///
 37/// assert!(pool.find_unused_id(0).is_none());  // time to realloc.
 38/// let resizer = pool.grow_request().ok_or(ENOSPC)?.realloc(GFP_KERNEL)?;
 39/// pool.grow(resizer);
 40///
 41/// assert_eq!(pool.find_unused_id(0).ok_or(ENOSPC)?.acquire(), 64);
 42/// # Ok::<(), Error>(())
 43/// ```
 44///
 45/// Releasing spinlock to grow the pool
 46///
 47/// ```no_run
 48/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
 49/// use kernel::sync::{new_spinlock, SpinLock};
 50/// use kernel::id_pool::IdPool;
 51///
 52/// fn get_id_maybe_realloc(guarded_pool: &SpinLock<IdPool>) -> Result<usize, AllocError> {
 53///     let mut pool = guarded_pool.lock();
 54///     loop {
 55///         match pool.find_unused_id(0) {
 56///             Some(index) => return Ok(index.acquire()),
 57///             None => {
 58///                 let alloc_request = pool.grow_request();
 59///                 drop(pool);
 60///                 let resizer = alloc_request.ok_or(AllocError)?.realloc(GFP_KERNEL)?;
 61///                 pool = guarded_pool.lock();
 62///                 pool.grow(resizer)
 63///             }
 64///         }
 65///     }
 66/// }
 67/// ```
 68pub struct IdPool {
 69    map: BitmapVec,
 70}
 71
 72/// Indicates that an [`IdPool`] should change to a new target size.
 73pub struct ReallocRequest {
 74    num_ids: usize,
 75}
 76
 77/// Contains a [`BitmapVec`] of a size suitable for reallocating [`IdPool`].
 78pub struct PoolResizer {
 79    new: BitmapVec,
 80}
 81
 82impl ReallocRequest {
 83    /// Allocates a new backing [`BitmapVec`] for [`IdPool`].
 84    ///
 85    /// This method only prepares reallocation and does not complete it.
 86    /// Reallocation will complete after passing the [`PoolResizer`] to the
 87    /// [`IdPool::grow`] or [`IdPool::shrink`] operation, which will check
 88    /// that reallocation still makes sense.
 89    pub fn realloc(&self, flags: Flags) -> Result<PoolResizer, AllocError> {
 90        let new = BitmapVec::new(self.num_ids, flags)?;
 91        Ok(PoolResizer { new })
 92    }
 93}
 94
 95impl IdPool {
 96    /// Constructs a new [`IdPool`].
 97    ///
 98    /// The pool will have a capacity of [`MAX_INLINE_LEN`].
 99    ///
100    /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN
101    #[inline]
102    pub fn new() -> Self {
103        Self {
104            map: BitmapVec::new_inline(),
105        }
106    }
107
108    /// Constructs a new [`IdPool`] with space for a specific number of bits.
109    ///
110    /// A capacity below [`MAX_INLINE_LEN`] is adjusted to [`MAX_INLINE_LEN`].
111    ///
112    /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN
113    #[inline]
114    pub fn with_capacity(num_ids: usize, flags: Flags) -> Result<Self, AllocError> {
115        let num_ids = usize::max(num_ids, BitmapVec::MAX_INLINE_LEN);
116        let map = BitmapVec::new(num_ids, flags)?;
117        Ok(Self { map })
118    }
119
120    /// Returns how many IDs this pool can currently have.
121    #[inline]
122    pub fn capacity(&self) -> usize {
123        self.map.len()
124    }
125
126    /// Returns a [`ReallocRequest`] if the [`IdPool`] can be shrunk, [`None`] otherwise.
127    ///
128    /// The capacity of an [`IdPool`] cannot be shrunk below [`MAX_INLINE_LEN`].
129    ///
130    /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN
131    ///
132    /// # Examples
133    ///
134    /// ```
135    /// use kernel::{
136    ///     alloc::AllocError,
137    ///     bitmap::BitmapVec,
138    ///     id_pool::{
139    ///         IdPool,
140    ///         ReallocRequest,
141    ///     },
142    /// };
143    ///
144    /// let mut pool = IdPool::with_capacity(1024, GFP_KERNEL)?;
145    /// let alloc_request = pool.shrink_request().ok_or(AllocError)?;
146    /// let resizer = alloc_request.realloc(GFP_KERNEL)?;
147    /// pool.shrink(resizer);
148    /// assert_eq!(pool.capacity(), BitmapVec::MAX_INLINE_LEN);
149    /// # Ok::<(), AllocError>(())
150    /// ```
151    #[inline]
152    pub fn shrink_request(&self) -> Option<ReallocRequest> {
153        let cap = self.capacity();
154        // Shrinking below `MAX_INLINE_LEN` is never possible.
155        if cap <= BitmapVec::MAX_INLINE_LEN {
156            return None;
157        }
158        // Determine if the bitmap can shrink based on the position of
159        // its last set bit. If the bit is within the first quarter of
160        // the bitmap then shrinking is possible. In this case, the
161        // bitmap should shrink to half its current size.
162        let Some(bit) = self.map.last_bit() else {
163            return Some(ReallocRequest {
164                num_ids: BitmapVec::MAX_INLINE_LEN,
165            });
166        };
167        if bit >= (cap / 4) {
168            return None;
169        }
170        let num_ids = usize::max(BitmapVec::MAX_INLINE_LEN, cap / 2);
171        Some(ReallocRequest { num_ids })
172    }
173
174    /// Shrinks pool by using a new [`BitmapVec`], if still possible.
175    #[inline]
176    pub fn shrink(&mut self, mut resizer: PoolResizer) {
177        // Between request to shrink that led to allocation of `resizer` and now,
178        // bits may have changed.
179        // Verify that shrinking is still possible. In case shrinking to
180        // the size of `resizer` is no longer possible, do nothing,
181        // drop `resizer` and move on.
182        let Some(updated) = self.shrink_request() else {
183            return;
184        };
185        if updated.num_ids > resizer.new.len() {
186            return;
187        }
188
189        resizer.new.copy_and_extend(&self.map);
190        self.map = resizer.new;
191    }
192
193    /// Returns a [`ReallocRequest`] for growing this [`IdPool`], if possible.
194    ///
195    /// The capacity of an [`IdPool`] cannot be grown above [`MAX_LEN`].
196    ///
197    /// [`MAX_LEN`]: BitmapVec::MAX_LEN
198    #[inline]
199    pub fn grow_request(&self) -> Option<ReallocRequest> {
200        let num_ids = self.capacity() * 2;
201        if num_ids > BitmapVec::MAX_LEN {
202            return None;
203        }
204        Some(ReallocRequest { num_ids })
205    }
206
207    /// Grows pool by using a new [`BitmapVec`], if still necessary.
208    ///
209    /// The `resizer` arguments has to be obtained by calling [`Self::grow_request`]
210    /// on this object and performing a [`ReallocRequest::realloc`].
211    #[inline]
212    pub fn grow(&mut self, mut resizer: PoolResizer) {
213        // Between request to grow that led to allocation of `resizer` and now,
214        // another thread may have already grown the capacity.
215        // In this case, do nothing, drop `resizer` and move on.
216        if resizer.new.len() <= self.capacity() {
217            return;
218        }
219
220        resizer.new.copy_and_extend(&self.map);
221        self.map = resizer.new;
222    }
223
224    /// Finds an unused ID in the bitmap.
225    ///
226    /// Upon success, returns its index. Otherwise, returns [`None`]
227    /// to indicate that a [`Self::grow_request`] is needed.
228    #[inline]
229    #[must_use]
230    pub fn find_unused_id(&mut self, offset: usize) -> Option<UnusedId<'_>> {
231        // INVARIANT: `next_zero_bit()` returns None or an integer less than `map.len()`
232        Some(UnusedId {
233            id: self.map.next_zero_bit(offset)?,
234            pool: self,
235        })
236    }
237
238    /// Releases an ID.
239    #[inline]
240    pub fn release_id(&mut self, id: usize) {
241        self.map.clear_bit(id);
242    }
243}
244
245/// Represents an unused id in an [`IdPool`].
246///
247/// # Invariants
248///
249/// The value of `id` is less than `pool.map.len()`.
250pub struct UnusedId<'pool> {
251    id: usize,
252    pool: &'pool mut IdPool,
253}
254
255impl<'pool> UnusedId<'pool> {
256    /// Get the unused id as an usize.
257    ///
258    /// Be aware that the id has not yet been acquired in the pool. The
259    /// [`acquire`] method must be called to prevent others from taking the id.
260    ///
261    /// [`acquire`]: UnusedId::acquire()
262    #[inline]
263    #[must_use]
264    pub fn as_usize(&self) -> usize {
265        self.id
266    }
267
268    /// Get the unused id as an u32.
269    ///
270    /// Be aware that the id has not yet been acquired in the pool. The
271    /// [`acquire`] method must be called to prevent others from taking the id.
272    ///
273    /// [`acquire`]: UnusedId::acquire()
274    #[inline]
275    #[must_use]
276    pub fn as_u32(&self) -> u32 {
277        // CAST: By the type invariants:
278        // `self.id < pool.map.len() <= BitmapVec::MAX_LEN = i32::MAX`.
279        self.id as u32
280    }
281
282    /// Acquire the unused id.
283    #[inline]
284    pub fn acquire(self) -> usize {
285        self.pool.map.set_bit(self.id);
286        self.id
287    }
288}
289
290impl Default for IdPool {
291    #[inline]
292    fn default() -> Self {
293        Self::new()
294    }
295}