db_testkit/
wrapper.rs

1use std::{
2    future::Future,
3    ops::{Deref, DerefMut},
4    pin::Pin,
5};
6
7use parking_lot::Mutex;
8
9type Stack<T> = Vec<T>;
10type Init<T> =
11    Box<dyn Fn() -> Pin<Box<dyn Future<Output = T> + Send + 'static>> + Send + Sync + 'static>;
12type Reset<T> =
13    Box<dyn Fn(T) -> Pin<Box<dyn Future<Output = T> + Send + 'static>> + Send + Sync + 'static>;
14
15#[allow(dead_code)]
16pub(crate) struct ObjectPool<T> {
17    objects: Mutex<Stack<T>>,
18    init: Init<T>,
19    reset: Reset<T>,
20}
21
22#[allow(dead_code)]
23impl<T> ObjectPool<T> {
24    pub(crate) fn new(init: Init<T>, reset: Reset<T>) -> Self {
25        Self {
26            objects: Mutex::new(Stack::new()),
27            init,
28            reset,
29        }
30    }
31
32    pub(crate) async fn pull(&self) -> Reusable<T> {
33        let object = self.objects.lock().pop();
34        let object = if let Some(object) = object {
35            (self.reset)(object).await
36        } else {
37            (self.init)().await
38        };
39        Reusable::new(self, object)
40    }
41
42    fn attach(&self, t: T) {
43        self.objects.lock().push(t);
44    }
45}
46
47/// Reusable object wrapper
48pub struct Reusable<'a, T> {
49    pool: &'a ObjectPool<T>,
50    data: Option<T>,
51}
52
53#[allow(dead_code)]
54impl<'a, T> Reusable<'a, T> {
55    fn new(pool: &'a ObjectPool<T>, t: T) -> Self {
56        Self {
57            pool,
58            data: Some(t),
59        }
60    }
61}
62
63const DATA_MUST_CONTAIN_SOME: &str = "data must always contain a [Some] value";
64
65impl<'a, T> Deref for Reusable<'a, T> {
66    type Target = T;
67
68    #[inline]
69    fn deref(&self) -> &Self::Target {
70        self.data.as_ref().expect(DATA_MUST_CONTAIN_SOME)
71    }
72}
73
74impl<'a, T> DerefMut for Reusable<'a, T> {
75    #[inline]
76    fn deref_mut(&mut self) -> &mut Self::Target {
77        self.data.as_mut().expect(DATA_MUST_CONTAIN_SOME)
78    }
79}
80
81impl<'a, T> Drop for Reusable<'a, T> {
82    #[inline]
83    fn drop(&mut self) {
84        self.pool
85            .attach(self.data.take().expect(DATA_MUST_CONTAIN_SOME));
86    }
87}