1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

//! Utilities for compute-heavy tasks which need to be run in parallel.

use std::ops::Deref;
use std::sync::Mutex;

/// A factory which produces a resource for use with [`ResourcePool`].
pub trait Resource {
    /// The type of the resource to be produced.
    type Output;

    /// An error type.
    type Error;

    /// Constructor for the resource.
    fn try_create(&self) -> Result<Self::Output, Self::Error>;
}

/// An unbounded pool of on-demand generated resources. This is useful when
/// distributing work across threads which needs access to an expensive-to-build
/// context.
///
/// A new resource is created when there isn't one available in the pool. When
/// it's dropped, it's returned to the pool. Old resources remain in the pool
/// until the pool itself is dropped.
///
/// ```
/// # use std::cell::RefCell;
/// # use branchless::core::task::{Resource, ResourceHandle, ResourcePool};
/// struct MyResource {
///     num_instantiations: RefCell<usize>,
/// }
///
/// impl Resource for MyResource {
///     type Output = String;
///     type Error = std::convert::Infallible;
///     fn try_create(&self) -> Result<Self::Output, Self::Error> {
///         let mut r = self.num_instantiations.borrow_mut();
///         *r += 1;
///         Ok(format!("This is resource #{}", *r))
///     }
/// }
///
/// # fn main() {
/// let resource = MyResource { num_instantiations: Default::default() };
/// let pool = ResourcePool::new(resource);
///
/// // Any number of the resource can be created.
/// let r1: ResourceHandle<'_, MyResource> = pool.try_create().unwrap();
/// assert_eq!(&*r1, "This is resource #1");
/// let r2 = pool.try_create().unwrap();
/// assert_eq!(&*r2, "This is resource #2");
/// drop(r2);
/// drop(r1);
///
/// // After releasing a resource, an attempt to get a resource returns an
/// // existing one from the pool.
/// let r1_again = pool.try_create().unwrap();
/// assert_eq!(&*r1_again, "This is resource #1");
/// # }
/// ```
pub struct ResourcePool<R: Resource> {
    factory: R,
    resources: Mutex<Vec<R::Output>>,
}

impl<R: Resource> std::fmt::Debug for ResourcePool<R> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ResourcePool")
            .field("factory", &"<not shown>")
            .field(
                "resources.len()",
                &match self.resources.try_lock() {
                    Ok(resources) => resources.len().to_string(),
                    Err(_) => "<could not determine>".to_string(),
                },
            )
            .finish()
    }
}

/// A handle to an instance of a resource created by [`Resource::try_create`].
/// When this value is dropped, the underlying resource returns to the owning
/// [`ResourcePool`].
pub struct ResourceHandle<'pool, R: Resource> {
    parent: &'pool ResourcePool<R>,
    inner: Option<R::Output>,
}

impl<'a, R: Resource> Drop for ResourceHandle<'a, R> {
    fn drop(&mut self) {
        let mut resources = self
            .parent
            .resources
            .lock()
            .expect("Poisoned mutex for ResourceHandle");
        resources.push(self.inner.take().unwrap());
    }
}

impl<'a, R: Resource> Deref for ResourceHandle<'a, R> {
    type Target = R::Output;

    fn deref(&self) -> &Self::Target {
        self.inner.as_ref().unwrap()
    }
}

impl<R: Resource> ResourcePool<R> {
    /// Constructor.
    pub fn new(factory: R) -> Self {
        ResourcePool {
            factory,
            resources: Default::default(),
        }
    }

    /// If there are any resources available in the pool, return an arbitrary
    /// one. Otherwise, invoke the constructor function of the associated
    /// [`Resource`] and return a [`ResourceHandle`] for it.
    pub fn try_create(&self) -> Result<ResourceHandle<R>, R::Error> {
        let resource = {
            let mut resources = self
                .resources
                .lock()
                .expect("Poisoned mutex for ResourcePool");
            match resources.pop() {
                Some(resource) => resource,
                None => self.factory.try_create()?,
            }
        };
        Ok(ResourceHandle {
            parent: self,
            inner: Some(resource),
        })
    }
}