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}