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