use std::panic::{RefUnwindSafe, UnwindSafe};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
static COUNTER: AtomicUsize = AtomicUsize::new(1);
thread_local!(
static THREAD_ID: usize = {
let next = COUNTER.fetch_add(1, Ordering::Relaxed);
if next == 0 {
panic!("regex: thread ID allocation space exhausted");
}
next
};
);
type CreateFn<T> =
Box<dyn Fn() -> T + Send + Sync + UnwindSafe + RefUnwindSafe + 'static>;
pub struct Pool<T> {
stack: Mutex<Vec<Box<T>>>,
create: CreateFn<T>,
owner: AtomicUsize,
owner_val: T,
}
unsafe impl<T: Send> Sync for Pool<T> {}
impl<T: ::std::fmt::Debug> ::std::fmt::Debug for Pool<T> {
fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
f.debug_struct("Pool")
.field("stack", &self.stack)
.field("owner", &self.owner)
.field("owner_val", &self.owner_val)
.finish()
}
}
#[derive(Debug)]
pub struct PoolGuard<'a, T: Send> {
pool: &'a Pool<T>,
value: Option<Box<T>>,
}
impl<T: Send> Pool<T> {
pub fn new(create: CreateFn<T>) -> Pool<T> {
let owner = AtomicUsize::new(0);
let owner_val = create();
Pool { stack: Mutex::new(vec![]), create, owner, owner_val }
}
#[cfg_attr(feature = "perf-inline", inline(always))]
pub fn get(&self) -> PoolGuard<'_, T> {
let caller = THREAD_ID.with(|id| *id);
let owner = self.owner.load(Ordering::Relaxed);
if caller == owner {
return self.guard_owned();
}
self.get_slow(caller, owner)
}
#[cold]
fn get_slow(&self, caller: usize, owner: usize) -> PoolGuard<'_, T> {
use std::sync::atomic::Ordering::Relaxed;
if owner == 0 {
let res = self.owner.compare_exchange(0, caller, Relaxed, Relaxed);
if res.is_ok() {
return self.guard_owned();
}
}
let mut stack = self.stack.lock().unwrap();
let value = match stack.pop() {
None => Box::new((self.create)()),
Some(value) => value,
};
self.guard_stack(value)
}
fn put(&self, value: Box<T>) {
let mut stack = self.stack.lock().unwrap();
stack.push(value);
}
fn guard_owned(&self) -> PoolGuard<'_, T> {
PoolGuard { pool: self, value: None }
}
fn guard_stack(&self, value: Box<T>) -> PoolGuard<'_, T> {
PoolGuard { pool: self, value: Some(value) }
}
}
impl<'a, T: Send> PoolGuard<'a, T> {
pub fn value(&self) -> &T {
match self.value {
None => &self.pool.owner_val,
Some(ref v) => &**v,
}
}
}
impl<'a, T: Send> Drop for PoolGuard<'a, T> {
#[cfg_attr(feature = "perf-inline", inline(always))]
fn drop(&mut self) {
if let Some(value) = self.value.take() {
self.pool.put(value);
}
}
}
#[cfg(test)]
mod tests {
use std::panic::{RefUnwindSafe, UnwindSafe};
use super::*;
#[test]
fn oibits() {
use crate::exec::ProgramCache;
fn has_oibits<T: Send + Sync + UnwindSafe + RefUnwindSafe>() {}
has_oibits::<Pool<ProgramCache>>();
}
#[test]
fn thread_owner_optimization() {
use std::cell::RefCell;
use std::sync::Arc;
let pool: Arc<Pool<RefCell<Vec<char>>>> =
Arc::new(Pool::new(Box::new(|| RefCell::new(vec!['a']))));
pool.get().value().borrow_mut().push('x');
let pool1 = pool.clone();
let t1 = std::thread::spawn(move || {
let guard = pool1.get();
let v = guard.value();
v.borrow_mut().push('y');
});
let pool2 = pool.clone();
let t2 = std::thread::spawn(move || {
let guard = pool2.get();
let v = guard.value();
v.borrow_mut().push('z');
});
t1.join().unwrap();
t2.join().unwrap();
assert_eq!(vec!['a', 'x'], *pool.get().value().borrow());
}
}