kernel/
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
10const BITS_PER_LONG: usize = bindings::BITS_PER_LONG as usize;
11
12/// Represents a dynamic ID pool backed by a [`BitmapVec`].
13///
14/// Clients acquire and release IDs from unset bits in a bitmap.
15///
16/// The capacity of the ID pool may be adjusted by users as
17/// needed. The API supports the scenario where users need precise control
18/// over the time of allocation of a new backing bitmap, which may require
19/// release of spinlock.
20/// Due to concurrent updates, all operations are re-verified to determine
21/// if the grow or shrink is sill valid.
22///
23/// # Examples
24///
25/// Basic usage
26///
27/// ```
28/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
29/// use kernel::id_pool::IdPool;
30///
31/// let mut pool = IdPool::new(64, GFP_KERNEL)?;
32/// for i in 0..64 {
33///     assert_eq!(i, pool.acquire_next_id(i).ok_or(ENOSPC)?);
34/// }
35///
36/// pool.release_id(23);
37/// assert_eq!(23, pool.acquire_next_id(0).ok_or(ENOSPC)?);
38///
39/// assert_eq!(None, pool.acquire_next_id(0));  // time to realloc.
40/// let resizer = pool.grow_request().ok_or(ENOSPC)?.realloc(GFP_KERNEL)?;
41/// pool.grow(resizer);
42///
43/// assert_eq!(pool.acquire_next_id(0), Some(64));
44/// # Ok::<(), Error>(())
45/// ```
46///
47/// Releasing spinlock to grow the pool
48///
49/// ```no_run
50/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
51/// use kernel::sync::{new_spinlock, SpinLock};
52/// use kernel::id_pool::IdPool;
53///
54/// fn get_id_maybe_realloc(guarded_pool: &SpinLock<IdPool>) -> Result<usize, AllocError> {
55///     let mut pool = guarded_pool.lock();
56///     loop {
57///         match pool.acquire_next_id(0) {
58///             Some(index) => return Ok(index),
59///             None => {
60///                 let alloc_request = pool.grow_request();
61///                 drop(pool);
62///                 let resizer = alloc_request.ok_or(AllocError)?.realloc(GFP_KERNEL)?;
63///                 pool = guarded_pool.lock();
64///                 pool.grow(resizer)
65///             }
66///         }
67///     }
68/// }
69/// ```
70pub struct IdPool {
71    map: BitmapVec,
72}
73
74/// Indicates that an [`IdPool`] should change to a new target size.
75pub struct ReallocRequest {
76    num_ids: usize,
77}
78
79/// Contains a [`BitmapVec`] of a size suitable for reallocating [`IdPool`].
80pub struct PoolResizer {
81    new: BitmapVec,
82}
83
84impl ReallocRequest {
85    /// Allocates a new backing [`BitmapVec`] for [`IdPool`].
86    ///
87    /// This method only prepares reallocation and does not complete it.
88    /// Reallocation will complete after passing the [`PoolResizer`] to the
89    /// [`IdPool::grow`] or [`IdPool::shrink`] operation, which will check
90    /// that reallocation still makes sense.
91    pub fn realloc(&self, flags: Flags) -> Result<PoolResizer, AllocError> {
92        let new = BitmapVec::new(self.num_ids, flags)?;
93        Ok(PoolResizer { new })
94    }
95}
96
97impl IdPool {
98    /// Constructs a new [`IdPool`].
99    ///
100    /// A capacity below [`BITS_PER_LONG`] is adjusted to
101    /// [`BITS_PER_LONG`].
102    ///
103    /// [`BITS_PER_LONG`]: srctree/include/asm-generic/bitsperlong.h
104    #[inline]
105    pub fn new(num_ids: usize, flags: Flags) -> Result<Self, AllocError> {
106        let num_ids = core::cmp::max(num_ids, BITS_PER_LONG);
107        let map = BitmapVec::new(num_ids, flags)?;
108        Ok(Self { map })
109    }
110
111    /// Returns how many IDs this pool can currently have.
112    #[inline]
113    pub fn capacity(&self) -> usize {
114        self.map.len()
115    }
116
117    /// Returns a [`ReallocRequest`] if the [`IdPool`] can be shrunk, [`None`] otherwise.
118    ///
119    /// The capacity of an [`IdPool`] cannot be shrunk below [`BITS_PER_LONG`].
120    ///
121    /// [`BITS_PER_LONG`]: srctree/include/asm-generic/bitsperlong.h
122    ///
123    /// # Examples
124    ///
125    /// ```
126    /// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
127    /// use kernel::id_pool::{ReallocRequest, IdPool};
128    ///
129    /// let mut pool = IdPool::new(1024, GFP_KERNEL)?;
130    /// let alloc_request = pool.shrink_request().ok_or(AllocError)?;
131    /// let resizer = alloc_request.realloc(GFP_KERNEL)?;
132    /// pool.shrink(resizer);
133    /// assert_eq!(pool.capacity(), kernel::bindings::BITS_PER_LONG as usize);
134    /// # Ok::<(), AllocError>(())
135    /// ```
136    #[inline]
137    pub fn shrink_request(&self) -> Option<ReallocRequest> {
138        let cap = self.capacity();
139        // Shrinking below [`BITS_PER_LONG`] is never possible.
140        if cap <= BITS_PER_LONG {
141            return None;
142        }
143        // Determine if the bitmap can shrink based on the position of
144        // its last set bit. If the bit is within the first quarter of
145        // the bitmap then shrinking is possible. In this case, the
146        // bitmap should shrink to half its current size.
147        let Some(bit) = self.map.last_bit() else {
148            return Some(ReallocRequest {
149                num_ids: BITS_PER_LONG,
150            });
151        };
152        if bit >= (cap / 4) {
153            return None;
154        }
155        let num_ids = usize::max(BITS_PER_LONG, cap / 2);
156        Some(ReallocRequest { num_ids })
157    }
158
159    /// Shrinks pool by using a new [`BitmapVec`], if still possible.
160    #[inline]
161    pub fn shrink(&mut self, mut resizer: PoolResizer) {
162        // Between request to shrink that led to allocation of `resizer` and now,
163        // bits may have changed.
164        // Verify that shrinking is still possible. In case shrinking to
165        // the size of `resizer` is no longer possible, do nothing,
166        // drop `resizer` and move on.
167        let Some(updated) = self.shrink_request() else {
168            return;
169        };
170        if updated.num_ids > resizer.new.len() {
171            return;
172        }
173
174        resizer.new.copy_and_extend(&self.map);
175        self.map = resizer.new;
176    }
177
178    /// Returns a [`ReallocRequest`] for growing this [`IdPool`], if possible.
179    ///
180    /// The capacity of an [`IdPool`] cannot be grown above [`i32::MAX`].
181    #[inline]
182    pub fn grow_request(&self) -> Option<ReallocRequest> {
183        let num_ids = self.capacity() * 2;
184        if num_ids > i32::MAX.try_into().unwrap() {
185            return None;
186        }
187        Some(ReallocRequest { num_ids })
188    }
189
190    /// Grows pool by using a new [`BitmapVec`], if still necessary.
191    ///
192    /// The `resizer` arguments has to be obtained by calling [`Self::grow_request`]
193    /// on this object and performing a [`ReallocRequest::realloc`].
194    #[inline]
195    pub fn grow(&mut self, mut resizer: PoolResizer) {
196        // Between request to grow that led to allocation of `resizer` and now,
197        // another thread may have already grown the capacity.
198        // In this case, do nothing, drop `resizer` and move on.
199        if resizer.new.len() <= self.capacity() {
200            return;
201        }
202
203        resizer.new.copy_and_extend(&self.map);
204        self.map = resizer.new;
205    }
206
207    /// Acquires a new ID by finding and setting the next zero bit in the
208    /// bitmap.
209    ///
210    /// Upon success, returns its index. Otherwise, returns [`None`]
211    /// to indicate that a [`Self::grow_request`] is needed.
212    #[inline]
213    pub fn acquire_next_id(&mut self, offset: usize) -> Option<usize> {
214        let next_zero_bit = self.map.next_zero_bit(offset);
215        if let Some(nr) = next_zero_bit {
216            self.map.set_bit(nr);
217        }
218        next_zero_bit
219    }
220
221    /// Releases an ID.
222    #[inline]
223    pub fn release_id(&mut self, id: usize) {
224        self.map.clear_bit(id);
225    }
226}