SarasasChipWorkshop/sarasacw-omrf
1
0
Fork 0
Code Activity

feat: improve object pool

Browse Source 
- improve object pool
- add docstring for test coverage
This commit is contained in:
yyc12345
2026-08-01 17:25:38 +08:00
parent f84e1a5d9c
commit 7834741d90
7 changed files with 356 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
omrf/src/lib.rs
-14
View File
@@ -9,17 +9,3 @@ macro_rules! resolve_enum {
<$t>::try_from($v).map_err(|_| Error::EnumOutOfRange)
};
}
#[macro_export]
macro_rules! pull_reader {
($pool:expr) => {
$pool.read().map_err(|_| Error::PoisonRwLock)
};
}
#[macro_export]
macro_rules! pull_writer {
($pool:expr) => {
$pool.write().map_err(|_| Error::PoisonRwLock)
};
}
omrf/src/object_pool.rs
+109 -65
View File
@@ -1,95 +1,139 @@
//! When exporting resources for C interface, resource management and ownership are important things.
//! In this dynamic library, we hold all resources' ownership in Rust world,
//! and only expose a token for C code manipulation.
//!
//! We need to create a container for holding all resources and providing corresponding operations.
//! So we introduce [ObjectPool] in this module for this purpose.
use slotmap::{DefaultKey, Key, KeyData, SlotMap};
//! A thread-safe object pool that hands out opaque integer tokens for objects it owns.
//!
//! The pool is backed by a sharded concurrent map ([`dashmap::DashMap`]); all operations take
//! `&self`, so no external lock guard is required -- different tokens can be accessed concurrently
//! from multiple threads (up to per-shard granularity).
//!
//! Tokens are generated by a monotonic [`AtomicU64`] counter starting at 1; `0` is reserved as the
//! invalid token (see [`INVALID_TOKEN`]). Because the counter never reuses a value, a token is
//! unique for the lifetime of the process.
use dashmap::{DashMap, mapref::one::{Ref, RefMut}};
use std::sync::atomic::{AtomicU64, Ordering};
use thiserror::Error as TeError;
/// Error occurs when operating with [ObjectPool].
/// Error occurs when operating with [`ObjectPool`].
#[derive(Debug, TeError)]
pub enum Error {
#[error("given token is not presented in object pool")]
NoSuchToken,
}
/// The token for fetching object in [ObjectPool].
/// The token for fetching an object in [`ObjectPool`].
pub type Token = u64;
/// Get the invalid token.
///
/// Invalid token is always invalid for fetching object in pool,
/// And can be useful in FFI scenario.
pub fn invalid_token() -> Token {
DefaultKey::null().data().as_ffi()
/// The canonical invalid token.
///
/// Always invalid for fetching an object in the pool, and useful as a sentinel in FFI scenarios.
/// It is `0`; the internal counter starts at `1` and never produces `0`.
pub const INVALID_TOKEN: Token = 0;
/// A thread-safe pool that manages objects keyed by unique [`Token`]s.
///
/// All methods take `&self`; concurrent access to different tokens proceeds in parallel (per
/// shard). See the [`get`](Self::get) / [`get_mut`](Self::get_mut) documentation for the deadlock
/// rule on holding multiple guards at once, and the closure-based
/// [`get_with`](Self::get_with) / [`get_mut_with`](Self::get_mut_with) for a deadlock-safe
/// alternative.
pub struct ObjectPool<T: Send + Sync> {
map: DashMap<Token, T>,
next_id: AtomicU64,
}
/// A pool for managing objects with unique tokens.
///
/// It is highly suggested to use this pool with [std::sync::RwLock] guard.
pub struct ObjectPool<T> {
objs: SlotMap<DefaultKey, T>,
}
impl<T> ObjectPool<T> {
/// Create a new [ObjectPool].
impl<T: Send + Sync> ObjectPool<T> {
/// Create a new [`ObjectPool`].
pub fn new() -> Self {
Self {
objs: SlotMap::new(),
map: DashMap::new(),
next_id: AtomicU64::new(1),
}
}
/// Convert [slotmap] crate's [DefaultKey] to our [Token].
fn key_to_token(key: &DefaultKey) -> Token {
key.data().as_ffi()
/// Insert an object into the pool and return its token.
///
/// The ownership of the object is transferred to the pool.
///
/// # Panics
///
/// Panics if the internal token counter overflows `u64`. This is practically unreachable (see
/// module-level documentation) and is treated as unrecoverable.
pub fn allocate(&self, value: T) -> Token {
let token = self
.next_id
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| v.checked_add(1))
.expect("ObjectPool token counter exhausted (u64 overflow)");
self.map.insert(token, value);
token
}
/// Convert our [Token] to [slotmap] crate's [DefaultKey].
fn token_to_key(token: Token) -> DefaultKey {
DefaultKey::from(KeyData::from_ffi(token))
/// Remove and drop the object corresponding to the given token.
pub fn free(&self, token: Token) -> Result<(), Error> {
self.map.remove(&token).ok_or(Error::NoSuchToken).map(|_| ())
}
/// Put given object into the pool and return its token.
///
/// The ownership of given object is transferred to the pool.
pub fn allocate(&mut self, value: T) -> Result<Token, Error> {
let key = self.objs.insert(value);
Ok(Self::key_to_token(&key))
}
/// Free the object in the pool corresponding to the given token.
pub fn free(&mut self, token: Token) -> Result<(), Error> {
let _ = self.pop(token)?;
Ok(())
}
/// Remove the object from the pool corresponding to the given token and return it.
///
/// Remove the object corresponding to the given token and return it.
///
/// The ownership of the object is transferred to the caller.
pub fn pop(&mut self, token: Token) -> Result<T, Error> {
match self.objs.remove(Self::token_to_key(token)) {
Some(obj) => Ok(obj),
None => Err(Error::NoSuchToken),
}
pub fn pop(&self, token: Token) -> Result<T, Error> {
self.map.remove(&token).map(|(_, v)| v).ok_or(Error::NoSuchToken)
}
/// Clear all objects in the pool.
pub fn clear(&mut self) -> () {
self.objs.clear();
/// Remove all objects from the pool.
pub fn clear(&self) {
self.map.clear();
}
/// Get a reference to the object in the pool corresponding to the given token.
pub fn get(&self, token: Token) -> Result<&T, Error> {
self.objs
.get(Self::token_to_key(token))
.ok_or(Error::NoSuchToken)
/// Get a shared (read) guard for the object keyed by `token`.
///
/// The returned guard holds a read lock on the map's shard for its whole lifetime.
///
/// # Deadlock rule -- READ CAREFULLY
///
/// **Never hold two guards from the same [`ObjectPool`] at once.** If two guards hash to the
/// same shard, the second acquisition waits for the first to release, deadlocking the thread.
/// For sequential access prefer the closure-based [`get_with`](Self::get_with), which releases
/// the lock between calls and makes holding two guards structurally impossible.
pub fn get(&self, token: Token) -> Result<Ref<'_, Token, T>, Error> {
self.map.get(&token).ok_or(Error::NoSuchToken)
}
/// Get a mutable reference to the object in the pool corresponding to the given token.
pub fn get_mut(&mut self, token: Token) -> Result<&mut T, Error> {
self.objs
.get_mut(Self::token_to_key(token))
.ok_or(Error::NoSuchToken)
/// Get an exclusive (write) guard for the object keyed by `token`.
///
/// The returned guard holds a write lock on the map's shard for its whole lifetime.
///
/// # Deadlock rule -- READ CAREFULLY
///
/// **Never hold two guards from the same [`ObjectPool`] at once.** If two guards hash to the
/// same shard, the second acquisition waits for the first to release, deadlocking the thread.
/// For sequential access prefer the closure-based [`get_mut_with`](Self::get_mut_with), which
/// releases the lock between calls and makes holding two guards structurally impossible.
pub fn get_mut(&self, token: Token) -> Result<RefMut<'_, Token, T>, Error> {
self.map.get_mut(&token).ok_or(Error::NoSuchToken)
}
/// Apply a closure to the object keyed by `token` (shared/read access).
///
/// The shard lock is held only while `f` runs and is released when this function returns, so
/// the caller cannot accidentally hold two guards across statements -- this is the
/// deadlock-safe counterpart of [`get`](Self::get).
///
/// The closure must not re-enter this [`ObjectPool`] (doing so risks the same deadlock
/// described on [`get`](Self::get)).
pub fn get_with<R>(&self, token: Token, f: impl FnOnce(&T) -> R) -> Result<R, Error> {
let guard = self.map.get(&token).ok_or(Error::NoSuchToken)?;
Ok(f(&*guard))
}
/// Apply a closure to the object keyed by `token` (exclusive/write access).
///
/// The shard lock is held only while `f` runs and is released when this function returns, so
/// the caller cannot accidentally hold two guards across statements -- this is the
/// deadlock-safe counterpart of [`get_mut`](Self::get_mut).
///
/// The closure must not re-enter this [`ObjectPool`] (doing so risks the same deadlock
/// described on [`get_mut`](Self::get_mut)).
pub fn get_mut_with<R>(&self, token: Token, f: impl FnOnce(&mut T) -> R) -> Result<R, Error> {
let mut guard = self.map.get_mut(&token).ok_or(Error::NoSuchToken)?;
Ok(f(&mut *guard))
}
}