dd_sds/scanner/
shared_pool.rs

1use std::ops::DerefMut;
2
3use regex_automata_fork::util::pool::{Pool, PoolGuard};
4
5type CachePoolFn<T> = Box<dyn Fn() -> T + Send + Sync>;
6pub type SharedPoolGuard<'a, T, const MAX_POOL_STACKS: usize> =
7    PoolGuard<'a, T, CachePoolFn<T>, MAX_POOL_STACKS>;
8
9/// This is a simple generic wrapper around `Pool` to make it a bit easier to use
10pub struct SharedPool<T> {
11    pool: Box<dyn AutoStacksSizePool<T>>,
12}
13
14// AutoStacksSizePool and AutoStacksSizeGuard are used to hide away the constant generic in the Pool
15pub trait AutoStacksSizePool<T>: Sync + Send {
16    fn get(&self) -> Box<dyn AutoStacksSizeGuard<T> + '_>;
17}
18
19pub trait AutoStacksSizeGuard<T> {
20    fn get_ref(&mut self) -> &mut T;
21}
22
23impl<T: Send, const MAX_POOL_STACKS: usize> AutoStacksSizeGuard<T>
24    for PoolGuard<'_, T, CachePoolFn<T>, MAX_POOL_STACKS>
25{
26    fn get_ref(&mut self) -> &mut T {
27        self.deref_mut()
28    }
29}
30
31impl<T: Send, const MAX_POOL_STACKS: usize> AutoStacksSizePool<T>
32    for Pool<T, CachePoolFn<T>, MAX_POOL_STACKS>
33{
34    fn get(&self) -> Box<dyn AutoStacksSizeGuard<T> + '_> {
35        Box::new(Pool::get(self))
36    }
37}
38
39impl<T: Send + 'static> SharedPool<T> {
40    /// Create a new shared pool with the given MAX_POOL_STACKS count.
41    ///
42    /// Pool stacks are a way to reduce contention on the pool for non-owners.
43    /// To reduce contention as much as possible, we ideally want to have a one-to-one
44    /// ratio from the number of CPUs to the number of pool stacks.
45    /// This function will create a pool with the number of stacks set to the nearest power of 2
46    /// greater than or equal to the given count, until 64. Anything greater than 32 will use 64.
47    /// Anything less than or equal to 4 will use 4.
48    ///
49    /// Examples:
50    ///     * new(_, 7) -> 8
51    ///     * new(_, 16) -> 16
52    ///     * new(_, 42) -> 64
53    pub fn new(factory: CachePoolFn<T>, count: usize) -> Self {
54        let pool = if count <= 4 {
55            Box::new(Pool::<_, _, 4>::new(factory)) as Box<dyn AutoStacksSizePool<T>>
56        } else if count <= 8 {
57            Box::new(Pool::<_, _, 8>::new(factory)) as Box<dyn AutoStacksSizePool<T>>
58        } else if count <= 16 {
59            Box::new(Pool::<_, _, 16>::new(factory)) as Box<dyn AutoStacksSizePool<T>>
60        } else if count <= 32 {
61            Box::new(Pool::<_, _, 32>::new(factory)) as Box<dyn AutoStacksSizePool<T>>
62        } else {
63            Box::new(Pool::<_, _, 64>::new(factory)) as Box<dyn AutoStacksSizePool<T>>
64        };
65        Self { pool }
66    }
67
68    pub fn get(&self) -> Box<dyn AutoStacksSizeGuard<T> + '_> {
69        self.pool.get()
70    }
71}