projects/linux·blob master · refs

295 lines · 9.1 KiB

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}