[PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap

Burak Emir posted 5 patches 7 months ago
There is a newer version of this series
[PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Burak Emir 7 months ago
This is a port of the Binder data structure introduced in commit
15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
Rust.

Like drivers/android/dbitmap.h, the ID pool abstraction lets
clients acquire and release IDs. The implementation uses a bitmap to
know what IDs are in use, and gives clients fine-grained control over
the time of allocation. This fine-grained control is needed in the
Android Binder. We provide an example that release a spinlock for
allocation and unit tests (rustdoc examples).

The implementation is not aware that the underlying Bitmap abstraction
handles lengths below BITS_PER_LONG without allocation.

Suggested-by: Alice Ryhl <aliceryhl@google.com>
Suggested-by: Yury Norov <yury.norov@gmail.com>
Signed-off-by: Burak Emir <bqe@google.com>
---
 MAINTAINERS            |   1 +
 rust/kernel/id_pool.rs | 201 +++++++++++++++++++++++++++++++++++++++++
 rust/kernel/lib.rs     |   1 +
 3 files changed, 203 insertions(+)
 create mode 100644 rust/kernel/id_pool.rs

diff --git a/MAINTAINERS b/MAINTAINERS
index 943d85ed1876..bc95d98f266b 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -4134,6 +4134,7 @@ R:	Yury Norov <yury.norov@gmail.com>
 S:	Maintained
 F:	lib/find_bit_benchmark_rust.rs
 F:	rust/kernel/bitmap.rs
+F:	rust/kernel/id_pool.rs
 
 BITOPS API
 M:	Yury Norov <yury.norov@gmail.com>
diff --git a/rust/kernel/id_pool.rs b/rust/kernel/id_pool.rs
new file mode 100644
index 000000000000..8f07526bb580
--- /dev/null
+++ b/rust/kernel/id_pool.rs
@@ -0,0 +1,201 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2025 Google LLC.
+
+//! Rust API for an ID pool backed by a `Bitmap`.
+
+use crate::alloc::{AllocError, Flags};
+use crate::bitmap::Bitmap;
+
+/// Represents a dynamic ID pool backed by a `Bitmap`.
+///
+/// Clients acquire and release IDs from zero bits in a bitmap.
+///
+/// The ID pool can grow or shrink as needed. It has been designed
+/// to support the scenario where users need to control the time
+/// of allocation of a new backing bitmap, which may require release
+/// of locks.
+/// These operations then, are verified to determine if the grow or
+/// shrink is sill valid.
+///
+/// # Examples
+///
+/// Basic usage
+///
+/// ```
+/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
+/// use kernel::id_pool::IdPool;
+///
+/// let mut pool = IdPool::new(64, GFP_KERNEL)?;
+/// for i in 0..64 {
+///   assert_eq!(i, pool.acquire_next_id(i).ok_or(ENOSPC)?);
+/// }
+///
+/// pool.release_id(23);
+/// assert_eq!(23, pool.acquire_next_id(0).ok_or(ENOSPC)?);
+///
+/// assert_eq!(None, pool.acquire_next_id(0));  // time to realloc.
+/// let resizer = pool.grow_alloc().alloc(GFP_KERNEL)?;
+/// pool.grow(resizer);
+///
+/// assert_eq!(pool.acquire_next_id(0), Some(64));
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// Releasing spinlock to grow the pool
+///
+/// ```no_run
+/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
+/// use kernel::sync::{new_spinlock, SpinLock};
+/// use kernel::id_pool::IdPool;
+///
+/// fn get_id_maybe_alloc(guarded_pool: &SpinLock<IdPool>) -> Result<usize, AllocError> {
+///   let mut pool = guarded_pool.lock();
+///   loop {
+///     match pool.acquire_next_id(0) {
+///       Some(index) => return Ok(index),
+///       None => {
+///         let alloc_request = pool.grow_alloc();
+///         drop(pool);
+///         let resizer = alloc_request.alloc(GFP_KERNEL)?;
+///         pool = guarded_pool.lock();
+///         pool.grow(resizer)
+///       }
+///     }
+///   }
+/// }
+/// ```
+pub struct IdPool {
+    map: Bitmap,
+}
+
+/// Returned when the `IdPool` should change size.
+pub struct AllocRequest {
+    nbits: usize,
+}
+
+/// Contains an allocated `Bitmap` for resizing `IdPool`.
+pub struct PoolResizer {
+    new: Bitmap,
+}
+
+impl AllocRequest {
+    /// Allocates a new `Bitmap` for `IdPool`.
+    pub fn alloc(&self, flags: Flags) -> Result<PoolResizer, AllocError> {
+        let new = Bitmap::new(self.nbits, flags)?;
+        Ok(PoolResizer { new })
+    }
+}
+
+impl IdPool {
+    /// Constructs a new `[IdPool]`.
+    #[inline]
+    pub fn new(nbits: usize, flags: Flags) -> Result<Self, AllocError> {
+        let map = Bitmap::new(nbits, flags)?;
+        Ok(Self { map })
+    }
+
+    /// Returns how many IDs this pool can currently have.
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.map.len()
+    }
+
+    /// Returns an [`AllocRequest`] if the [`IdPool`] can be shrunk, [`None`] otherwise.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
+    /// use kernel::id_pool::{AllocRequest, IdPool};
+    ///
+    /// let mut pool = IdPool::new(1024, GFP_KERNEL)?;
+    /// let alloc_request = pool.shrink_alloc().ok_or(AllocError)?;
+    /// let resizer = alloc_request.alloc(GFP_KERNEL)?;
+    /// pool.shrink(resizer);
+    /// assert_eq!(pool.len(), kernel::bindings::BITS_PER_LONG as usize);
+    /// # Ok::<(), AllocError>(())
+    /// ```
+    #[inline]
+    pub fn shrink_alloc(&self) -> Option<AllocRequest> {
+        let len = self.map.len();
+        if len <= bindings::BITS_PER_LONG as usize {
+            return None;
+        }
+        /*
+         * Determine if the bitmap can shrink based on the position of
+         * its last set bit. If the bit is within the first quarter of
+         * the bitmap then shrinking is possible. In this case, the
+         * bitmap should shrink to half its current size.
+         */
+        match self.map.last_bit() {
+            Some(bit) => {
+                if bit < (len >> 2) {
+                    Some(AllocRequest { nbits: len >> 1 })
+                } else {
+                    None
+                }
+            }
+            None => Some(AllocRequest {
+                nbits: bindings::BITS_PER_LONG as usize,
+            }),
+        }
+    }
+
+    /// Shrinks pool by using a new `Bitmap`, if still possible.
+    #[inline]
+    pub fn shrink(&mut self, mut resizer: PoolResizer) {
+        // Verify that shrinking is still possible. The `resizer`
+        // bitmap might have been allocated without locks, so this call
+        // could now be outdated. In this case, drop `resizer` and move on.
+        if let Some(AllocRequest { nbits }) = self.shrink_alloc() {
+            if nbits <= resizer.new.len() {
+                resizer.new.copy_and_extend(&self.map);
+                self.map = resizer.new;
+                return;
+            }
+        }
+    }
+
+    /// Returns an `AllocRequest` for growing this `IdPool`.
+    #[inline]
+    pub fn grow_alloc(&self) -> AllocRequest {
+        AllocRequest {
+            nbits: self.map.len() << 1,
+        }
+    }
+
+    /// Grows pool by using a new `Bitmap`, if still necessary.
+    #[inline]
+    pub fn grow(&mut self, mut resizer: PoolResizer) {
+        // `resizer` bitmap might have been allocated without locks,
+        // so this call could now be outdated. In this case, drop
+        // `resizer` and move on.
+        if resizer.new.len() <= self.map.len() {
+            return;
+        }
+
+        resizer.new.copy_and_extend(&self.map);
+        self.map = resizer.new;
+    }
+
+    /// Acquires a new ID by finding and setting the next zero bit in the
+    /// bitmap. Upon success, returns its index. Otherwise, returns `None`
+    /// to indicate that a `grow_alloc` is needed.
+    #[inline]
+    pub fn acquire_next_id(&mut self, offset: usize) -> Option<usize> {
+        match self.map.next_zero_bit(offset) {
+            res @ Some(nr) => {
+                self.map.set_bit(nr);
+                res
+            }
+            None => None,
+        }
+    }
+
+    /// Releases an ID.
+    #[inline]
+    pub fn release_id(&mut self, id: usize) {
+        self.map.clear_bit(id);
+    }
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 8c4161cd82ac..d7def807900a 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -54,6 +54,7 @@
 #[cfg(CONFIG_RUST_FW_LOADER_ABSTRACTIONS)]
 pub mod firmware;
 pub mod fs;
+pub mod id_pool;
 pub mod init;
 pub mod io;
 pub mod ioctl;
-- 
2.49.0.1101.gccaa498523-goog
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Jann Horn 7 months ago
On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> +/// fn get_id_maybe_alloc(guarded_pool: &SpinLock<IdPool>) -> Result<usize, AllocError> {
> +///   let mut pool = guarded_pool.lock();
> +///   loop {
> +///     match pool.acquire_next_id(0) {
> +///       Some(index) => return Ok(index),
> +///       None => {
> +///         let alloc_request = pool.grow_alloc();
> +///         drop(pool);
> +///         let resizer = alloc_request.alloc(GFP_KERNEL)?;
> +///         pool = guarded_pool.lock();
> +///         pool.grow(resizer)
> +///       }
> +///     }
> +///   }
> +/// }
> +/// ```

Hmm, I think I just understood a bit better than before what's
actually going on in the C binder code... the reason why you have this
complicated API for reallocation here in Rust is that you want the
guarded_pool to not have its own lock, but a lock provided by the
caller? And so pool functions are unable to take/drop the lock that
protects the pool?

This is just me idly wondering, and I know I tend to come up with
overcomplicated approaches and I'm bad at Rust - please just ignore
this message if you think it's not a good idea.

I wonder if there is a way in Rust to address this kind of situation
that looks nicer. Maybe you could have a function as part of the
IdPool implementation that operates on an IdPool, but instead of an
IdPool as parameter, the parameter is something like a function that
can be called to obtain a guard that gives you a mutable reference?
Could you maybe have some lock trait whose implementations store a
pointer to something like a SpinLock, with a "lock()" method that
first locks the SpinLock (creating a Guard), then looks up a specific
member of the data contained inside the SpinLock, and returns a
mutable pointer to that? Basically a subset view of the SpinLock. Then
the IdPool implementation would be able to internally do this pattern
of alternating lock/unlock. Though this way you'd still have to not
hold the lock when calling into this.

It might be an even better fit for a 1:1 translation of the C code if
you could then combine this with some kind of SpinLockUnlockGuard that
mutably borrows a SpinLockGuard; releases the underlying lock when it
is created; and reacquires the lock when it is dropped... but that's
maybe taking things too far.
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Jann Horn 7 months ago
On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> This is a port of the Binder data structure introduced in commit
> 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> Rust.

Stupid high-level side comment:

That commit looks like it changed a simple linear rbtree scan (which
is O(n) with slow steps) into a bitmap thing. A more elegant option
might have been to use an augmented rbtree, reducing the O(n) rbtree
scan to an O(log n) rbtree lookup, just like how finding a free area
used to work in MM code... That would let you drop that ID pool bitmap
entirely. But I guess actually wiring up an augmented rbtree into Rust
would be very annoying too.
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Boqun Feng 7 months ago
On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > This is a port of the Binder data structure introduced in commit
> > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > Rust.
> 
> Stupid high-level side comment:
> 
> That commit looks like it changed a simple linear rbtree scan (which
> is O(n) with slow steps) into a bitmap thing. A more elegant option
> might have been to use an augmented rbtree, reducing the O(n) rbtree
> scan to an O(log n) rbtree lookup, just like how finding a free area

I think RBTree::cursor_lower_bound() [1] does exactly what you said

[1]: https://rust.docs.kernel.org/kernel/rbtree/struct.RBTree.html#method.cursor_lower_bound

Regards,
Boqun

> used to work in MM code... That would let you drop that ID pool bitmap
> entirely. But I guess actually wiring up an augmented rbtree into Rust
> would be very annoying too.
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Alice Ryhl 7 months ago
On Mon, May 19, 2025 at 4:56 PM Boqun Feng <boqun.feng@gmail.com> wrote:
>
> On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > This is a port of the Binder data structure introduced in commit
> > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > Rust.
> >
> > Stupid high-level side comment:
> >
> > That commit looks like it changed a simple linear rbtree scan (which
> > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > scan to an O(log n) rbtree lookup, just like how finding a free area
>
> I think RBTree::cursor_lower_bound() [1] does exactly what you said

We need the smallest ID without a value, not the smallest ID in use.

Alice
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Boqun Feng 7 months ago
On Mon, May 19, 2025 at 08:46:37PM -0700, Alice Ryhl wrote:
> On Mon, May 19, 2025 at 4:56 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> >
> > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > This is a port of the Binder data structure introduced in commit
> > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > Rust.
> > >
> > > Stupid high-level side comment:
> > >
> > > That commit looks like it changed a simple linear rbtree scan (which
> > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > scan to an O(log n) rbtree lookup, just like how finding a free area
> >
> > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> 
> We need the smallest ID without a value, not the smallest ID in use.
> 

Ok, but it shouldn't be hard to write a Rust function that search that,
right? My point was mostly the Rust rbtree binding can do O(log n)
search. I have no idea about "even so, should we try something like Jann
suggested". And I think your other reply basically says no.

Regards,
Boqun

> Alice
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Alice Ryhl 7 months ago
On Mon, May 19, 2025 at 10:21 PM Boqun Feng <boqun.feng@gmail.com> wrote:
>
> On Mon, May 19, 2025 at 08:46:37PM -0700, Alice Ryhl wrote:
> > On Mon, May 19, 2025 at 4:56 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > >
> > > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > > This is a port of the Binder data structure introduced in commit
> > > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > > Rust.
> > > >
> > > > Stupid high-level side comment:
> > > >
> > > > That commit looks like it changed a simple linear rbtree scan (which
> > > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > > scan to an O(log n) rbtree lookup, just like how finding a free area
> > >
> > > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> >
> > We need the smallest ID without a value, not the smallest ID in use.
> >
>
> Ok, but it shouldn't be hard to write a Rust function that search that,
> right? My point was mostly the Rust rbtree binding can do O(log n)
> search. I have no idea about "even so, should we try something like Jann
> suggested". And I think your other reply basically says no.

We would need to store additional data in the r/b tree to know whether
to go left or right, so it would be somewhat tricky. We don't have an
implementation of that in Rust.

Alice
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Boqun Feng 7 months ago
On Tue, May 20, 2025 at 05:42:51AM -0700, Alice Ryhl wrote:
> On Mon, May 19, 2025 at 10:21 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> >
> > On Mon, May 19, 2025 at 08:46:37PM -0700, Alice Ryhl wrote:
> > > On Mon, May 19, 2025 at 4:56 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > >
> > > > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > > > This is a port of the Binder data structure introduced in commit
> > > > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > > > Rust.
> > > > >
> > > > > Stupid high-level side comment:
> > > > >
> > > > > That commit looks like it changed a simple linear rbtree scan (which
> > > > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > > > scan to an O(log n) rbtree lookup, just like how finding a free area
> > > >
> > > > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> > >
> > > We need the smallest ID without a value, not the smallest ID in use.
> > >
> >
> > Ok, but it shouldn't be hard to write a Rust function that search that,
> > right? My point was mostly the Rust rbtree binding can do O(log n)
> > search. I have no idea about "even so, should we try something like Jann
> > suggested". And I think your other reply basically says no.
> 
> We would need to store additional data in the r/b tree to know whether
> to go left or right, so it would be somewhat tricky. We don't have an

Hmm... I'm confused, I thought you can implement a search like that by
doing what RBTree::raw_entry() does except that when Ordering::Equal you
always go left or right (depending on whether you want to get an unused
ID less or greater than a key value), i.e. you always search until you
get an Vacant entry. Why do you need store additional data for that?
Maybe I'm missing something here?

Regards,
Boqun

> implementation of that in Rust.
> 
> Alice
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Alice Ryhl 7 months ago
On Tue, May 20, 2025 at 5:56 AM Boqun Feng <boqun.feng@gmail.com> wrote:
>
> On Tue, May 20, 2025 at 05:42:51AM -0700, Alice Ryhl wrote:
> > On Mon, May 19, 2025 at 10:21 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > >
> > > On Mon, May 19, 2025 at 08:46:37PM -0700, Alice Ryhl wrote:
> > > > On Mon, May 19, 2025 at 4:56 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > > >
> > > > > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > > > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > > > > This is a port of the Binder data structure introduced in commit
> > > > > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > > > > Rust.
> > > > > >
> > > > > > Stupid high-level side comment:
> > > > > >
> > > > > > That commit looks like it changed a simple linear rbtree scan (which
> > > > > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > > > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > > > > scan to an O(log n) rbtree lookup, just like how finding a free area
> > > > >
> > > > > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> > > >
> > > > We need the smallest ID without a value, not the smallest ID in use.
> > > >
> > >
> > > Ok, but it shouldn't be hard to write a Rust function that search that,
> > > right? My point was mostly the Rust rbtree binding can do O(log n)
> > > search. I have no idea about "even so, should we try something like Jann
> > > suggested". And I think your other reply basically says no.
> >
> > We would need to store additional data in the r/b tree to know whether
> > to go left or right, so it would be somewhat tricky. We don't have an
>
> Hmm... I'm confused, I thought you can implement a search like that by
> doing what RBTree::raw_entry() does except that when Ordering::Equal you
> always go left or right (depending on whether you want to get an unused
> ID less or greater than a key value), i.e. you always search until you
> get an Vacant entry. Why do you need store additional data for that?
> Maybe I'm missing something here?

Let's say you're at the root node of an r/b tree, and you see that the
root node has id 17, the left node has id 8, and the right node has id
25. Do you go left or right?

Alice
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Boqun Feng 7 months ago
On Tue, May 20, 2025 at 06:05:52AM -0700, Alice Ryhl wrote:
> On Tue, May 20, 2025 at 5:56 AM Boqun Feng <boqun.feng@gmail.com> wrote:
> >
> > On Tue, May 20, 2025 at 05:42:51AM -0700, Alice Ryhl wrote:
> > > On Mon, May 19, 2025 at 10:21 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > >
> > > > On Mon, May 19, 2025 at 08:46:37PM -0700, Alice Ryhl wrote:
> > > > > On Mon, May 19, 2025 at 4:56 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > > > >
> > > > > > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > > > > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > > > > > This is a port of the Binder data structure introduced in commit
> > > > > > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > > > > > Rust.
> > > > > > >
> > > > > > > Stupid high-level side comment:
> > > > > > >
> > > > > > > That commit looks like it changed a simple linear rbtree scan (which
> > > > > > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > > > > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > > > > > scan to an O(log n) rbtree lookup, just like how finding a free area
> > > > > >
> > > > > > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> > > > >
> > > > > We need the smallest ID without a value, not the smallest ID in use.
> > > > >
> > > >
> > > > Ok, but it shouldn't be hard to write a Rust function that search that,
> > > > right? My point was mostly the Rust rbtree binding can do O(log n)
> > > > search. I have no idea about "even so, should we try something like Jann
> > > > suggested". And I think your other reply basically says no.
> > >
> > > We would need to store additional data in the r/b tree to know whether
> > > to go left or right, so it would be somewhat tricky. We don't have an
> >
> > Hmm... I'm confused, I thought you can implement a search like that by
> > doing what RBTree::raw_entry() does except that when Ordering::Equal you
> > always go left or right (depending on whether you want to get an unused
> > ID less or greater than a key value), i.e. you always search until you
> > get an Vacant entry. Why do you need store additional data for that?
> > Maybe I'm missing something here?
> 
> Let's say you're at the root node of an r/b tree, and you see that the
> root node has id 17, the left node has id 8, and the right node has id
> 25. Do you go left or right?
> 

I went to check what commit 15d9da3f818c actually did and I understand
what you mean now ;-) In your case, the rbtree cannot have nodes with
the same key. If Jann can provide the O(log n) search that could help in
this case, I'm happy to learn about it ;-)

Regards,
Boqun

> Alice
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Jann Horn 7 months ago
On Tue, May 20, 2025 at 3:21 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> On Tue, May 20, 2025 at 06:05:52AM -0700, Alice Ryhl wrote:
> > On Tue, May 20, 2025 at 5:56 AM Boqun Feng <boqun.feng@gmail.com> wrote:
> > >
> > > On Tue, May 20, 2025 at 05:42:51AM -0700, Alice Ryhl wrote:
> > > > On Mon, May 19, 2025 at 10:21 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > > >
> > > > > On Mon, May 19, 2025 at 08:46:37PM -0700, Alice Ryhl wrote:
> > > > > > On Mon, May 19, 2025 at 4:56 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > > > > >
> > > > > > > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > > > > > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > > > > > > This is a port of the Binder data structure introduced in commit
> > > > > > > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > > > > > > Rust.
> > > > > > > >
> > > > > > > > Stupid high-level side comment:
> > > > > > > >
> > > > > > > > That commit looks like it changed a simple linear rbtree scan (which
> > > > > > > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > > > > > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > > > > > > scan to an O(log n) rbtree lookup, just like how finding a free area
> > > > > > >
> > > > > > > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> > > > > >
> > > > > > We need the smallest ID without a value, not the smallest ID in use.
> > > > > >
> > > > >
> > > > > Ok, but it shouldn't be hard to write a Rust function that search that,
> > > > > right? My point was mostly the Rust rbtree binding can do O(log n)
> > > > > search. I have no idea about "even so, should we try something like Jann
> > > > > suggested". And I think your other reply basically says no.
> > > >
> > > > We would need to store additional data in the r/b tree to know whether
> > > > to go left or right, so it would be somewhat tricky. We don't have an
> > >
> > > Hmm... I'm confused, I thought you can implement a search like that by
> > > doing what RBTree::raw_entry() does except that when Ordering::Equal you
> > > always go left or right (depending on whether you want to get an unused
> > > ID less or greater than a key value), i.e. you always search until you
> > > get an Vacant entry. Why do you need store additional data for that?
> > > Maybe I'm missing something here?
> >
> > Let's say you're at the root node of an r/b tree, and you see that the
> > root node has id 17, the left node has id 8, and the right node has id
> > 25. Do you go left or right?
> >
>
> I went to check what commit 15d9da3f818c actually did and I understand
> what you mean now ;-) In your case, the rbtree cannot have nodes with
> the same key. If Jann can provide the O(log n) search that could help in
> this case, I'm happy to learn about it ;-)

Linux has the concept of an "augmented rbtree", where you can stuff
extra information into the rbtree to keep track of things like "how
big is the biggest gap between objects in this subtree". This is how
the MM subsystem used to find free space in the virtual address space
before the maple tree refactor, a complicated example is here:

finding a free region (by looking at vm_area_struct::rb_subtree_gap to
decide whether to go left or right; this is made complicated here
because they have more search constraints):
https://elixir.bootlin.com/linux/v4.19.325/source/mm/mmap.c#L1841

But that requires an "augmented rbtree" where the rbtree code calls
back into callbacks for updating the subtree gap; the MM code has its
gap update here:
https://elixir.bootlin.com/linux/v4.19.325/source/mm/mmap.c#L261

And associates that with VMA trees through this macro magic that would
probably be a terrible fit for Rust code:
https://elixir.bootlin.com/linux/v4.19.325/source/mm/mmap.c#L400

As Alice said, this is probably not a great fit for Rust code. As she
said, an xarray or maple tree would have this kind of gap search
built-in, which would be nicer here. But if you're trying to do
insertions while holding your own outer spinlocks, I think they would
be hard (or impossible?) to use.

If you managed to avoid broad use of spinlocks, that might make it
much easier to use xarrays or maple trees (and that would also allow
you to make the bitmap API much simpler).
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Boqun Feng 7 months ago
On Tue, May 20, 2025 at 05:55:47PM +0200, Jann Horn wrote:
> On Tue, May 20, 2025 at 3:21 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > On Tue, May 20, 2025 at 06:05:52AM -0700, Alice Ryhl wrote:
> > > On Tue, May 20, 2025 at 5:56 AM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > >
> > > > On Tue, May 20, 2025 at 05:42:51AM -0700, Alice Ryhl wrote:
> > > > > On Mon, May 19, 2025 at 10:21 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > > > >
> > > > > > On Mon, May 19, 2025 at 08:46:37PM -0700, Alice Ryhl wrote:
> > > > > > > On Mon, May 19, 2025 at 4:56 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > > > > > > >
> > > > > > > > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > > > > > > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > > > > > > > This is a port of the Binder data structure introduced in commit
> > > > > > > > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > > > > > > > Rust.
> > > > > > > > >
> > > > > > > > > Stupid high-level side comment:
> > > > > > > > >
> > > > > > > > > That commit looks like it changed a simple linear rbtree scan (which
> > > > > > > > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > > > > > > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > > > > > > > scan to an O(log n) rbtree lookup, just like how finding a free area
> > > > > > > >
> > > > > > > > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> > > > > > >
> > > > > > > We need the smallest ID without a value, not the smallest ID in use.
> > > > > > >
> > > > > >
> > > > > > Ok, but it shouldn't be hard to write a Rust function that search that,
> > > > > > right? My point was mostly the Rust rbtree binding can do O(log n)
> > > > > > search. I have no idea about "even so, should we try something like Jann
> > > > > > suggested". And I think your other reply basically says no.
> > > > >
> > > > > We would need to store additional data in the r/b tree to know whether
> > > > > to go left or right, so it would be somewhat tricky. We don't have an
> > > >
> > > > Hmm... I'm confused, I thought you can implement a search like that by
> > > > doing what RBTree::raw_entry() does except that when Ordering::Equal you
> > > > always go left or right (depending on whether you want to get an unused
> > > > ID less or greater than a key value), i.e. you always search until you
> > > > get an Vacant entry. Why do you need store additional data for that?
> > > > Maybe I'm missing something here?
> > >
> > > Let's say you're at the root node of an r/b tree, and you see that the
> > > root node has id 17, the left node has id 8, and the right node has id
> > > 25. Do you go left or right?
> > >
> >
> > I went to check what commit 15d9da3f818c actually did and I understand
> > what you mean now ;-) In your case, the rbtree cannot have nodes with
> > the same key. If Jann can provide the O(log n) search that could help in
> > this case, I'm happy to learn about it ;-)
> 
> Linux has the concept of an "augmented rbtree", where you can stuff
> extra information into the rbtree to keep track of things like "how
> big is the biggest gap between objects in this subtree". This is how
> the MM subsystem used to find free space in the virtual address space
> before the maple tree refactor, a complicated example is here:
> 
> finding a free region (by looking at vm_area_struct::rb_subtree_gap to
> decide whether to go left or right; this is made complicated here
> because they have more search constraints):
> https://elixir.bootlin.com/linux/v4.19.325/source/mm/mmap.c#L1841
> 
> But that requires an "augmented rbtree" where the rbtree code calls
> back into callbacks for updating the subtree gap; the MM code has its
> gap update here:
> https://elixir.bootlin.com/linux/v4.19.325/source/mm/mmap.c#L261
> 

I see. I was missing this part.

> And associates that with VMA trees through this macro magic that would
> probably be a terrible fit for Rust code:
> https://elixir.bootlin.com/linux/v4.19.325/source/mm/mmap.c#L400
> 

Well, not sure that's true implementation-wise, I mean it's just
function callbacks while you insert or erase nodes from rbtree, which
could probably be described by a trait like:

    pub trait RBTreeAugmented<K, V> {
        fn compute(node: &Node<K, V, Self>) -> Self;
    }

    impl<K, V> RBTreeAugmented<K, V> for () {
        fn compute(_node: &Node<K, V, Self>) -> Self {
	    ()
	}
    }
and we change the Node type into:

    pub struct Node<K, V, A: RBTreeAugmented<K, V> = ()> 
    {
        links: bindings::rb_node,
        key: K,
        value: V,
	augmented: A
    }

and _propagate() can be something like:

   unsafe fn augmented_propagate<K, V, A: RBTreeAugmented<K, V>>(
       mut node: *mut rb_node, stop: *mut rb_node
   ) {
       	while !ptr::eq(node, stop) {
            let rbnode = unsafe { container_of!(node, Node<K, V, A>, links) }.cast_mut();
	    let rbnode: &mut Node<K,V,A> = unsafe { &mut *rbnode };

	    let new_augmented = A::compute(rbnode);

	    if rbnode.aurmented == new_augmented {
	        break;
	    }
		if (node->rbaugmented == augmented)			\
			break;						\
	    rbnode.augmented = augmented;				\

	    node = rb_parent(node);
	}
   }

probably works? However I guess we don't need to do that right now given
Alice's point on xarray or maple tree.

Regards,
Boqun

> As Alice said, this is probably not a great fit for Rust code. As she
> said, an xarray or maple tree would have this kind of gap search
> built-in, which would be nicer here. But if you're trying to do
> insertions while holding your own outer spinlocks, I think they would
> be hard (or impossible?) to use.
> 
> If you managed to avoid broad use of spinlocks, that might make it
> much easier to use xarrays or maple trees (and that would also allow
> you to make the bitmap API much simpler).
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Yury Norov 7 months ago
+ Carlos Llamas

On Mon, May 19, 2025 at 04:56:21PM -0700, Boqun Feng wrote:
> On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > This is a port of the Binder data structure introduced in commit
> > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > Rust.
> > 
> > Stupid high-level side comment:
> > 
> > That commit looks like it changed a simple linear rbtree scan (which
> > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > scan to an O(log n) rbtree lookup, just like how finding a free area
> 
> I think RBTree::cursor_lower_bound() [1] does exactly what you said
> 
> [1]: https://rust.docs.kernel.org/kernel/rbtree/struct.RBTree.html#method.cursor_lower_bound

Alice mentioned before that in many cases the whole pool of IDs will
fit into a single machine word if represented as bitmap. If that holds,
bitmaps will win over any other data structure that I can imagine. 

For very large ID pools, the algorithmic complexity will take over,
for sure. On the other hand, the 15d9da3f818ca explicitly mentions
that it switches implementation to bitmaps for performance reasons. 

Anyways, Burak and Alice, before we move forward, can you tell if you
ran any experiments with data structures allowing logarithmic lookup,
like rb-tree? Can you maybe measure at which point rb-tree lookup will
win over find_bit as the size of pool growth?

Can you describe how the existing dbitmap is used now? What is the
typical size of ID pools? Which operation is the bottleneck? Looking
forward, are there any expectations about ID pools size in future?

Carlos, can you please elaborate your motivation to switch to bitmaps?
Have you considered rb-trees with O(logn) lookup?

Thanks,
Yury


Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Carlos Llamas 7 months ago
On Mon, May 19, 2025 at 08:57:04PM -0400, Yury Norov wrote:
> + Carlos Llamas
> 
> On Mon, May 19, 2025 at 04:56:21PM -0700, Boqun Feng wrote:
> > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > This is a port of the Binder data structure introduced in commit
> > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > Rust.
> > > 
> > > Stupid high-level side comment:
> > > 
> > > That commit looks like it changed a simple linear rbtree scan (which
> > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > scan to an O(log n) rbtree lookup, just like how finding a free area
> > 
> > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> > 
> > [1]: https://rust.docs.kernel.org/kernel/rbtree/struct.RBTree.html#method.cursor_lower_bound
> 
> Alice mentioned before that in many cases the whole pool of IDs will
> fit into a single machine word if represented as bitmap. If that holds,
> bitmaps will win over any other data structure that I can imagine. 
> 
> For very large ID pools, the algorithmic complexity will take over,
> for sure. On the other hand, the 15d9da3f818ca explicitly mentions
> that it switches implementation to bitmaps for performance reasons. 
> 
> Anyways, Burak and Alice, before we move forward, can you tell if you
> ran any experiments with data structures allowing logarithmic lookup,
> like rb-tree? Can you maybe measure at which point rb-tree lookup will
> win over find_bit as the size of pool growth?
> 
> Can you describe how the existing dbitmap is used now? What is the
> typical size of ID pools? Which operation is the bottleneck? Looking
> forward, are there any expectations about ID pools size in future?
> 
> Carlos, can you please elaborate your motivation to switch to bitmaps?
> Have you considered rb-trees with O(logn) lookup?

Yeah, we tried rb-trees. There was even a patch that implemented the
augmented logic. See this:
https://lore.kernel.org/all/20240917030203.286-1-ebpqwerty472123@gmail.com/
IIRC, it just didn't make sense for our use case because of the extra
memory bytes required for this solution. The performance ended up being
the same (from my local testing).

I'm not certain of this but one potential factor is that the rb nodes
are in-strucutre members allocated separately. This can lead to more
cache misses when traversing them. I don't know how applicable this
would be for the Rust implementation though. Take that with a grain of
salt as I didn't actually look super close while running the tests.

I would also note, this whole logic wouldn't be required if userspace
wasn't using these descriptor IDs as vector indexes. At some point this
practice will be fixed and we can remove the "dbitmap" implementation.

Cheers,
--
Carlos Llamas
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Yury Norov 7 months ago
On Wed, May 21, 2025 at 03:57:55AM +0000, Carlos Llamas wrote:
> On Mon, May 19, 2025 at 08:57:04PM -0400, Yury Norov wrote:
> > + Carlos Llamas

...

> > Carlos, can you please elaborate your motivation to switch to bitmaps?
> > Have you considered rb-trees with O(logn) lookup?
> 
> Yeah, we tried rb-trees. There was even a patch that implemented the
> augmented logic. See this:
> https://lore.kernel.org/all/20240917030203.286-1-ebpqwerty472123@gmail.com/
> IIRC, it just didn't make sense for our use case because of the extra
> memory bytes required for this solution. The performance ended up being
> the same (from my local testing).
> 
> I'm not certain of this but one potential factor is that the rb nodes
> are in-strucutre members allocated separately. This can lead to more
> cache misses when traversing them. I don't know how applicable this
> would be for the Rust implementation though. Take that with a grain of
> salt as I didn't actually look super close while running the tests.
> 
> I would also note, this whole logic wouldn't be required if userspace
> wasn't using these descriptor IDs as vector indexes. At some point this
> practice will be fixed and we can remove the "dbitmap" implementation.

Yeah, I expected to get this kind of feedback from real-world testing.
Your reply to the patch you mentioned above answers all my questions:

https://lore.kernel.org/all/ZwdWe_I2p3zD-v1O@google.com/

Let's stick to bitmaps unless someone shows clear benefit of using any
alternative approach, both in time and memory perspectives, supported
by testing.

Thanks,
Yury
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Burak Emir 6 months, 3 weeks ago
On Wed, May 21, 2025 at 3:50 PM Yury Norov <yury.norov@gmail.com> wrote:
>
> On Wed, May 21, 2025 at 03:57:55AM +0000, Carlos Llamas wrote:
> > On Mon, May 19, 2025 at 08:57:04PM -0400, Yury Norov wrote:
> > > + Carlos Llamas
>
> ...
>
> > > Carlos, can you please elaborate your motivation to switch to bitmaps?
> > > Have you considered rb-trees with O(logn) lookup?
> >
> > Yeah, we tried rb-trees. There was even a patch that implemented the
> > augmented logic. See this:
> > https://lore.kernel.org/all/20240917030203.286-1-ebpqwerty472123@gmail.com/
> > IIRC, it just didn't make sense for our use case because of the extra
> > memory bytes required for this solution. The performance ended up being
> > the same (from my local testing).
> >
> > I'm not certain of this but one potential factor is that the rb nodes
> > are in-strucutre members allocated separately. This can lead to more
> > cache misses when traversing them. I don't know how applicable this
> > would be for the Rust implementation though. Take that with a grain of
> > salt as I didn't actually look super close while running the tests.
> >
> > I would also note, this whole logic wouldn't be required if userspace
> > wasn't using these descriptor IDs as vector indexes. At some point this
> > practice will be fixed and we can remove the "dbitmap" implementation.
>
> Yeah, I expected to get this kind of feedback from real-world testing.
> Your reply to the patch you mentioned above answers all my questions:
>
> https://lore.kernel.org/all/ZwdWe_I2p3zD-v1O@google.com/
>
> Let's stick to bitmaps unless someone shows clear benefit of using any
> alternative approach, both in time and memory perspectives, supported
> by testing.

Thanks all for the additional context.

Yury, I've addressed most of the comments.

You also commented on the API. The weirdness of the API is all due to
the separating "request to shrink/grow" from allocation.
Since allocation can happen while other threads may mess with the id
pool, one has to double check that the request to shrink/grow still
makes sense.
Dealing with this situation is required in the Android binder context
where this ID pool is used, my understanding is that one cannot
allocate there while holding the spinlock.

In the next version, I have renamed the operations to make this a bit
clearer, and made the comments a bit more explicit.
If it's ok, let's move the discussion to v9 that I will send in a
moment, I hope it clears things up a bit.

Thanks.
Burak
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Alice Ryhl 7 months ago
On Mon, May 19, 2025 at 5:57 PM Yury Norov <yury.norov@gmail.com> wrote:
>
> + Carlos Llamas
>
> On Mon, May 19, 2025 at 04:56:21PM -0700, Boqun Feng wrote:
> > On Tue, May 20, 2025 at 12:51:07AM +0200, Jann Horn wrote:
> > > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > > This is a port of the Binder data structure introduced in commit
> > > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > > Rust.
> > >
> > > Stupid high-level side comment:
> > >
> > > That commit looks like it changed a simple linear rbtree scan (which
> > > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > > scan to an O(log n) rbtree lookup, just like how finding a free area
> >
> > I think RBTree::cursor_lower_bound() [1] does exactly what you said
> >
> > [1]: https://rust.docs.kernel.org/kernel/rbtree/struct.RBTree.html#method.cursor_lower_bound
>
> Alice mentioned before that in many cases the whole pool of IDs will
> fit into a single machine word if represented as bitmap. If that holds,
> bitmaps will win over any other data structure that I can imagine.
>
> For very large ID pools, the algorithmic complexity will take over,
> for sure. On the other hand, the 15d9da3f818ca explicitly mentions
> that it switches implementation to bitmaps for performance reasons.
>
> Anyways, Burak and Alice, before we move forward, can you tell if you
> ran any experiments with data structures allowing logarithmic lookup,
> like rb-tree? Can you maybe measure at which point rb-tree lookup will
> win over find_bit as the size of pool growth?
>
> Can you describe how the existing dbitmap is used now? What is the
> typical size of ID pools? Which operation is the bottleneck? Looking
> forward, are there any expectations about ID pools size in future?

Generally, an Android phone will have around 3 processes with a large
ID pool (thousands of IDs), and essentially all other processes have a
very small number of IDs (less than 10). The large pools are typically
the same size as the number of concurrently running processes on the
device. The bitmap was added to the C driver to deal with perf issues
that came from doing linear lookups on the rbtree for the large pools
while holding a spinlock, and it did solve those perf issues.

Alice
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Alice Ryhl 7 months ago
On Mon, May 19, 2025 at 3:51 PM Jann Horn <jannh@google.com> wrote:
>
> On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > This is a port of the Binder data structure introduced in commit
> > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > Rust.
>
> Stupid high-level side comment:
>
> That commit looks like it changed a simple linear rbtree scan (which
> is O(n) with slow steps) into a bitmap thing. A more elegant option
> might have been to use an augmented rbtree, reducing the O(n) rbtree
> scan to an O(log n) rbtree lookup, just like how finding a free area
> used to work in MM code... That would let you drop that ID pool bitmap
> entirely. But I guess actually wiring up an augmented rbtree into Rust
> would be very annoying too.

If we're talking approaches to avoid the bitmap entirely, it would
probably be easier to replace the rb tree with xarray than to use an
augmented one.

Alice
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Jann Horn 7 months ago
On Tue, May 20, 2025 at 1:12 AM Alice Ryhl <aliceryhl@google.com> wrote:
> On Mon, May 19, 2025 at 3:51 PM Jann Horn <jannh@google.com> wrote:
> > On Mon, May 19, 2025 at 6:20 PM Burak Emir <bqe@google.com> wrote:
> > > This is a port of the Binder data structure introduced in commit
> > > 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> > > Rust.
> >
> > Stupid high-level side comment:
> >
> > That commit looks like it changed a simple linear rbtree scan (which
> > is O(n) with slow steps) into a bitmap thing. A more elegant option
> > might have been to use an augmented rbtree, reducing the O(n) rbtree
> > scan to an O(log n) rbtree lookup, just like how finding a free area
> > used to work in MM code... That would let you drop that ID pool bitmap
> > entirely. But I guess actually wiring up an augmented rbtree into Rust
> > would be very annoying too.
>
> If we're talking approaches to avoid the bitmap entirely, it would
> probably be easier to replace the rb tree with xarray than to use an
> augmented one.

Ah, yeah, maybe. I'm not familiar enough with xarray to know how
annoying it would be to insert stuff in an xarray that you reach
through a spinlock, which seems to be the requirement that made the
API for the bitmap interface so complicated if I understand
correctly...
Re: [PATCH v8 5/5] rust: add dynamic ID pool abstraction for bitmap
Posted by Yury Norov 7 months ago
On Mon, May 19, 2025 at 04:17:05PM +0000, Burak Emir wrote:
> This is a port of the Binder data structure introduced in commit
> 15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
> Rust.
> 
> Like drivers/android/dbitmap.h, the ID pool abstraction lets
> clients acquire and release IDs. The implementation uses a bitmap to
> know what IDs are in use, and gives clients fine-grained control over
> the time of allocation. This fine-grained control is needed in the
> Android Binder. We provide an example that release a spinlock for
> allocation and unit tests (rustdoc examples).
> 
> The implementation is not aware that the underlying Bitmap abstraction
> handles lengths below BITS_PER_LONG without allocation.
> 
> Suggested-by: Alice Ryhl <aliceryhl@google.com>
> Suggested-by: Yury Norov <yury.norov@gmail.com>
> Signed-off-by: Burak Emir <bqe@google.com>
> ---
>  MAINTAINERS            |   1 +
>  rust/kernel/id_pool.rs | 201 +++++++++++++++++++++++++++++++++++++++++
>  rust/kernel/lib.rs     |   1 +
>  3 files changed, 203 insertions(+)
>  create mode 100644 rust/kernel/id_pool.rs
> 
> diff --git a/MAINTAINERS b/MAINTAINERS
> index 943d85ed1876..bc95d98f266b 100644
> --- a/MAINTAINERS
> +++ b/MAINTAINERS
> @@ -4134,6 +4134,7 @@ R:	Yury Norov <yury.norov@gmail.com>
>  S:	Maintained
>  F:	lib/find_bit_benchmark_rust.rs
>  F:	rust/kernel/bitmap.rs
> +F:	rust/kernel/id_pool.rs
>  
>  BITOPS API
>  M:	Yury Norov <yury.norov@gmail.com>
> diff --git a/rust/kernel/id_pool.rs b/rust/kernel/id_pool.rs
> new file mode 100644
> index 000000000000..8f07526bb580
> --- /dev/null
> +++ b/rust/kernel/id_pool.rs
> @@ -0,0 +1,201 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +// Copyright (C) 2025 Google LLC.
> +
> +//! Rust API for an ID pool backed by a `Bitmap`.
> +
> +use crate::alloc::{AllocError, Flags};
> +use crate::bitmap::Bitmap;
> +
> +/// Represents a dynamic ID pool backed by a `Bitmap`.
> +///
> +/// Clients acquire and release IDs from zero bits in a bitmap.

Maybe 'unset' or 'free' instead of 'zero'?

> +///
> +/// The ID pool can grow or shrink as needed. It has been designed

s/can/may/. You don't implement automatic adjustment of a pool
capacity, but from this comment one may conclude that the pool  
may grow or shrink by itself. Can you instead say: the Capacity
of IDs may be adjusted by user as needed. Or something like that.

> +/// to support the scenario where users need to control the time
> +/// of allocation of a new backing bitmap, which may require release
> +/// of locks.
> +/// These operations then, are verified to determine if the grow or
> +/// shrink is sill valid.
> +///
> +/// # Examples
> +///
> +/// Basic usage
> +///
> +/// ```
> +/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
> +/// use kernel::id_pool::IdPool;
> +///
> +/// let mut pool = IdPool::new(64, GFP_KERNEL)?;
> +/// for i in 0..64 {
> +///   assert_eq!(i, pool.acquire_next_id(i).ok_or(ENOSPC)?);
> +/// }
> +///
> +/// pool.release_id(23);
> +/// assert_eq!(23, pool.acquire_next_id(0).ok_or(ENOSPC)?);
> +///
> +/// assert_eq!(None, pool.acquire_next_id(0));  // time to realloc.
> +/// let resizer = pool.grow_alloc().alloc(GFP_KERNEL)?;
> +/// pool.grow(resizer);
> +///
> +/// assert_eq!(pool.acquire_next_id(0), Some(64));
> +/// # Ok::<(), Error>(())
> +/// ```
> +///
> +/// Releasing spinlock to grow the pool
> +///
> +/// ```no_run
> +/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
> +/// use kernel::sync::{new_spinlock, SpinLock};
> +/// use kernel::id_pool::IdPool;
> +///
> +/// fn get_id_maybe_alloc(guarded_pool: &SpinLock<IdPool>) -> Result<usize, AllocError> {
> +///   let mut pool = guarded_pool.lock();
> +///   loop {
> +///     match pool.acquire_next_id(0) {
> +///       Some(index) => return Ok(index),
> +///       None => {
> +///         let alloc_request = pool.grow_alloc();
> +///         drop(pool);
> +///         let resizer = alloc_request.alloc(GFP_KERNEL)?;
> +///         pool = guarded_pool.lock();
> +///         pool.grow(resizer)
> +///       }
> +///     }
> +///   }
> +/// }
> +/// ```
> +pub struct IdPool {
> +    map: Bitmap,
> +}
> +
> +/// Returned when the `IdPool` should change size.
> +pub struct AllocRequest {
> +    nbits: usize,
> +}
> +
> +/// Contains an allocated `Bitmap` for resizing `IdPool`.
> +pub struct PoolResizer {
> +    new: Bitmap,
> +}
> +
> +impl AllocRequest {
> +    /// Allocates a new `Bitmap` for `IdPool`.
> +    pub fn alloc(&self, flags: Flags) -> Result<PoolResizer, AllocError> {
> +        let new = Bitmap::new(self.nbits, flags)?;
> +        Ok(PoolResizer { new })
> +    }
> +}
> +
> +impl IdPool {
> +    /// Constructs a new `[IdPool]`.
> +    #[inline]
> +    pub fn new(nbits: usize, flags: Flags) -> Result<Self, AllocError> {

Those 'bits' are 'IDs' now. I think the 1st parameter name should
reflect that: num_ids, or something.

> +        let map = Bitmap::new(nbits, flags)?;
> +        Ok(Self { map })
> +    }

What for do you split new() and alloc()? When I call 'new()' method, I
expect it will initialize all the fields.

Or I misunderstand something?

> +
> +    /// Returns how many IDs this pool can currently have.
> +    #[inline]
> +    pub fn len(&self) -> usize {
> +        self.map.len()
> +    }
> +
> +    /// Returns an [`AllocRequest`] if the [`IdPool`] can be shrunk, [`None`] otherwise.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
> +    /// use kernel::id_pool::{AllocRequest, IdPool};
> +    ///
> +    /// let mut pool = IdPool::new(1024, GFP_KERNEL)?;
> +    /// let alloc_request = pool.shrink_alloc().ok_or(AllocError)?;
> +    /// let resizer = alloc_request.alloc(GFP_KERNEL)?;
> +    /// pool.shrink(resizer);
> +    /// assert_eq!(pool.len(), kernel::bindings::BITS_PER_LONG as usize);
> +    /// # Ok::<(), AllocError>(())


So, I don't understand what for do you have this 'pool.shrink_alloc()'
line after the 'new()' call. Thinking object-oriented, when I create
a new object, and receive OK, I can use it immediately. IIUC, I can't
request an ID immediately after creating an IdPool.

> +    /// ```
> +    #[inline]
> +    pub fn shrink_alloc(&self) -> Option<AllocRequest> {
> +        let len = self.map.len();
> +        if len <= bindings::BITS_PER_LONG as usize {
> +            return None;
> +        }

How that? If I want to shrink from 60 to say 30 IDs, I expect that
you will not allow me to allocate 31st ID. But the above code makes
the whole function a no-op...

> +        /*
> +         * Determine if the bitmap can shrink based on the position of
> +         * its last set bit. If the bit is within the first quarter of
> +         * the bitmap then shrinking is possible. In this case, the
> +         * bitmap should shrink to half its current size.
> +         */
> +        match self.map.last_bit() {
> +            Some(bit) => {
> +                if bit < (len >> 2) {
> +                    Some(AllocRequest { nbits: len >> 1 })

Can you use the 'a/2' style instead of this shifts?

> +                } else {
> +                    None
> +                }
> +            }
> +            None => Some(AllocRequest {
> +                nbits: bindings::BITS_PER_LONG as usize,
> +            }),

Can you reorder the logic such that non-error path will have no extra
indentations?

> +        }
> +    }
> +
> +    /// Shrinks pool by using a new `Bitmap`, if still possible.
> +    #[inline]

I don't understand it. You have shrink_alloc(), which indeed allocates
something, and you have just 'shrink(). And it also allocates before
using the copy_and_extend().

Can you elaborate what for do you need 2 versions of 'shrink'?

> +    pub fn shrink(&mut self, mut resizer: PoolResizer) {

I believe your users will be more thankful if you just allow them to
shrink their ID pools without all that intermediate objects, like:
        my_pool.shrink(100);

> +        // Verify that shrinking is still possible. The `resizer`
> +        // bitmap might have been allocated without locks, so this call
> +        // could now be outdated. In this case, drop `resizer` and move on.
> +        if let Some(AllocRequest { nbits }) = self.shrink_alloc() {
> +            if nbits <= resizer.new.len() {
> +                resizer.new.copy_and_extend(&self.map);
> +                self.map = resizer.new;
> +                return;

Again, can you please have non-error path as non-indented as possible?

> +            }
> +        }
> +    }
> +
> +    /// Returns an `AllocRequest` for growing this `IdPool`.
> +    #[inline]
> +    pub fn grow_alloc(&self) -> AllocRequest {
> +        AllocRequest {
> +            nbits: self.map.len() << 1,
> +        }
> +    }
> +
> +    /// Grows pool by using a new `Bitmap`, if still necessary.
> +    #[inline]
> +    pub fn grow(&mut self, mut resizer: PoolResizer) {
> +        // `resizer` bitmap might have been allocated without locks,
> +        // so this call could now be outdated. In this case, drop
> +        // `resizer` and move on.
> +        if resizer.new.len() <= self.map.len() {
> +            return;
> +        }
> +
> +        resizer.new.copy_and_extend(&self.map);
> +        self.map = resizer.new;
> +    }
> +
> +    /// Acquires a new ID by finding and setting the next zero bit in the
> +    /// bitmap. Upon success, returns its index. Otherwise, returns `None`
> +    /// to indicate that a `grow_alloc` is needed.
> +    #[inline]
> +    pub fn acquire_next_id(&mut self, offset: usize) -> Option<usize> {
> +        match self.map.next_zero_bit(offset) {
> +            res @ Some(nr) => {
> +                self.map.set_bit(nr);
> +                res
> +            }
> +            None => None,
> +        }
> +    }
> +
> +    /// Releases an ID.
> +    #[inline]
> +    pub fn release_id(&mut self, id: usize) {
> +        self.map.clear_bit(id);

What if I release an empty ID? Maybe emit a warning?

> +    }
> +}

I think I totally miss the idea of PoolResizers here. The standard
dynamic array API looks like this:

        new(capacity)     -> IdPool
        drop(IdPool)      -> void
        acquire()         -> int
        release(id)       -> bool       // True if had been allocated
        acquire_next(id)* -> int        
        get_capacity()    -> int
        set_capacity(cap) -> bool
        num_ids()         -> int

 * This one is pretty rare.  Are you sure you need it?

Your API looks weird with those AllocRequests and PoolResizers,
and TBH I don't understand how that would help your users.

Can you please consider a more standard API?

Thanks,
Yury

> diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
> index 8c4161cd82ac..d7def807900a 100644
> --- a/rust/kernel/lib.rs
> +++ b/rust/kernel/lib.rs
> @@ -54,6 +54,7 @@
>  #[cfg(CONFIG_RUST_FW_LOADER_ABSTRACTIONS)]
>  pub mod firmware;
>  pub mod fs;
> +pub mod id_pool;
>  pub mod init;
>  pub mod io;
>  pub mod ioctl;
> -- 
> 2.49.0.1101.gccaa498523-goog