branchless/core/
task.rs

1//! Utilities for compute-heavy tasks which need to be run in parallel.
2
3use std::ops::Deref;
4use std::sync::Mutex;
5
6/// A factory which produces a resource for use with [`ResourcePool`].
7pub trait Resource {
8    /// The type of the resource to be produced.
9    type Output;
10
11    /// An error type.
12    type Error;
13
14    /// Constructor for the resource.
15    fn try_create(&self) -> Result<Self::Output, Self::Error>;
16}
17
18/// An unbounded pool of on-demand generated resources. This is useful when
19/// distributing work across threads which needs access to an expensive-to-build
20/// context.
21///
22/// A new resource is created when there isn't one available in the pool. When
23/// it's dropped, it's returned to the pool. Old resources remain in the pool
24/// until the pool itself is dropped.
25///
26/// ```
27/// # use std::cell::RefCell;
28/// # use branchless::core::task::{Resource, ResourceHandle, ResourcePool};
29/// struct MyResource {
30///     num_instantiations: RefCell<usize>,
31/// }
32///
33/// impl Resource for MyResource {
34///     type Output = String;
35///     type Error = std::convert::Infallible;
36///     fn try_create(&self) -> Result<Self::Output, Self::Error> {
37///         let mut r = self.num_instantiations.borrow_mut();
38///         *r += 1;
39///         Ok(format!("This is resource #{}", *r))
40///     }
41/// }
42///
43/// # fn main() {
44/// let resource = MyResource { num_instantiations: Default::default() };
45/// let pool = ResourcePool::new(resource);
46///
47/// // Any number of the resource can be created.
48/// let r1: ResourceHandle<'_, MyResource> = pool.try_create().unwrap();
49/// assert_eq!(&*r1, "This is resource #1");
50/// let r2 = pool.try_create().unwrap();
51/// assert_eq!(&*r2, "This is resource #2");
52/// drop(r2);
53/// drop(r1);
54///
55/// // After releasing a resource, an attempt to get a resource returns an
56/// // existing one from the pool.
57/// let r1_again = pool.try_create().unwrap();
58/// assert_eq!(&*r1_again, "This is resource #1");
59/// # }
60/// ```
61pub struct ResourcePool<R: Resource> {
62    factory: R,
63    resources: Mutex<Vec<R::Output>>,
64}
65
66impl<R: Resource> std::fmt::Debug for ResourcePool<R> {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        f.debug_struct("ResourcePool")
69            .field("factory", &"<not shown>")
70            .field(
71                "resources.len()",
72                &match self.resources.try_lock() {
73                    Ok(resources) => resources.len().to_string(),
74                    Err(_) => "<could not determine>".to_string(),
75                },
76            )
77            .finish()
78    }
79}
80
81/// A handle to an instance of a resource created by [`Resource::try_create`].
82/// When this value is dropped, the underlying resource returns to the owning
83/// [`ResourcePool`].
84pub struct ResourceHandle<'pool, R: Resource> {
85    parent: &'pool ResourcePool<R>,
86    inner: Option<R::Output>,
87}
88
89impl<'a, R: Resource> Drop for ResourceHandle<'a, R> {
90    fn drop(&mut self) {
91        let mut resources = self
92            .parent
93            .resources
94            .lock()
95            .expect("Poisoned mutex for ResourceHandle");
96        resources.push(self.inner.take().unwrap());
97    }
98}
99
100impl<'a, R: Resource> Deref for ResourceHandle<'a, R> {
101    type Target = R::Output;
102
103    fn deref(&self) -> &Self::Target {
104        self.inner.as_ref().unwrap()
105    }
106}
107
108impl<R: Resource> ResourcePool<R> {
109    /// Constructor.
110    pub fn new(factory: R) -> Self {
111        ResourcePool {
112            factory,
113            resources: Default::default(),
114        }
115    }
116
117    /// If there are any resources available in the pool, return an arbitrary
118    /// one. Otherwise, invoke the constructor function of the associated
119    /// [`Resource`] and return a [`ResourceHandle`] for it.
120    pub fn try_create(&self) -> Result<ResourceHandle<R>, R::Error> {
121        let resource = {
122            let mut resources = self
123                .resources
124                .lock()
125                .expect("Poisoned mutex for ResourcePool");
126            let resource = resources.pop();
127            match resource {
128                Some(resource) => resource,
129                None => self.factory.try_create()?,
130            }
131        };
132        Ok(ResourceHandle {
133            parent: self,
134            inner: Some(resource),
135        })
136    }
137}
branchless/core/task.rs

branchless/core/
task.rs
