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 const 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 /// ``` 70 pub struct IdPool { 71 map: BitmapVec, 72 } 73 74 /// Indicates that an [`IdPool`] should change to a new target size. 75 pub struct ReallocRequest { 76 num_ids: usize, 77 } 78 79 /// Contains a [`BitmapVec`] of a size suitable for reallocating [`IdPool`]. 80 pub struct PoolResizer { 81 new: BitmapVec, 82 } 83 84 impl 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 97 impl 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 } 227