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 crate::alloc::{AllocError, Flags}; 8 use 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 /// ``` 68 pub struct IdPool { 69 map: BitmapVec, 70 } 71 72 /// Indicates that an [`IdPool`] should change to a new target size. 73 pub struct ReallocRequest { 74 num_ids: usize, 75 } 76 77 /// Contains a [`BitmapVec`] of a size suitable for reallocating [`IdPool`]. 78 pub struct PoolResizer { 79 new: BitmapVec, 80 } 81 82 impl 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. realloc(&self, flags: Flags) -> Result<PoolResizer, AllocError>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 95 impl 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] new() -> Self102 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] with_capacity(num_ids: usize, flags: Flags) -> Result<Self, AllocError>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] capacity(&self) -> usize122 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] shrink_request(&self) -> Option<ReallocRequest>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] shrink(&mut self, mut resizer: PoolResizer)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] grow_request(&self) -> Option<ReallocRequest>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] grow(&mut self, mut resizer: PoolResizer)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] find_unused_id(&mut self, offset: usize) -> Option<UnusedId<'_>>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] release_id(&mut self, id: usize)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()`. 250 pub struct UnusedId<'pool> { 251 id: usize, 252 pool: &'pool mut IdPool, 253 } 254 255 impl<'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] as_usize(&self) -> usize264 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] as_u32(&self) -> u32276 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] acquire(self) -> usize284 pub fn acquire(self) -> usize { 285 self.pool.map.set_bit(self.id); 286 self.id 287 } 288 } 289 290 impl Default for IdPool { 291 #[inline] default() -> Self292 fn default() -> Self { 293 Self::new() 294 } 295 } 296