deadpool/managed/
hooks.rs

1//! Hooks allowing to run code when creating and/or recycling objects.
2
3use std::{borrow::Cow, fmt, future::Future, pin::Pin};
4
5use super::{Manager, Metrics, ObjectInner};
6
7/// The result returned by hooks
8pub type HookResult<E> = Result<(), HookError<E>>;
9
10/// The boxed future that should be returned by async hooks
11pub type HookFuture<'a, E> = Pin<Box<dyn Future<Output = HookResult<E>> + Send + 'a>>;
12
13/// Function signature for sync callbacks
14type SyncFn<M> =
15    dyn Fn(&mut <M as Manager>::Type, &Metrics) -> HookResult<<M as Manager>::Error> + Sync + Send;
16
17/// Function siganture for async callbacks
18type AsyncFn<M> = dyn for<'a> Fn(&'a mut <M as Manager>::Type, &'a Metrics) -> HookFuture<'a, <M as Manager>::Error>
19    + Sync
20    + Send;
21
22/// Wrapper for hook functions
23pub enum Hook<M: Manager> {
24    /// Use a plain function (non-async) as a hook
25    Fn(Box<SyncFn<M>>),
26    /// Use an async function as a hook
27    AsyncFn(Box<AsyncFn<M>>),
28}
29
30impl<M: Manager> Hook<M> {
31    /// Create Hook from sync function
32    pub fn sync_fn(
33        f: impl Fn(&mut M::Type, &Metrics) -> HookResult<M::Error> + Sync + Send + 'static,
34    ) -> Self {
35        Self::Fn(Box::new(f))
36    }
37    /// Create Hook from async function
38    pub fn async_fn(
39        f: impl for<'a> Fn(&'a mut M::Type, &'a Metrics) -> HookFuture<'a, M::Error>
40            + Sync
41            + Send
42            + 'static,
43    ) -> Self {
44        Self::AsyncFn(Box::new(f))
45    }
46}
47
48impl<M: Manager> fmt::Debug for Hook<M> {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        match self {
51            Self::Fn(_) => f
52                .debug_tuple("Fn")
53                //.field(arg0)
54                .finish(),
55            Self::AsyncFn(_) => f
56                .debug_tuple("AsyncFn")
57                //.field(arg0)
58                .finish(),
59        }
60    }
61}
62
63/// Error which is returned by `pre_create`, `pre_recycle` and
64/// `post_recycle` hooks.
65#[derive(Debug)]
66pub enum HookError<E> {
67    /// Hook failed for some other reason.
68    Message(Cow<'static, str>),
69
70    /// Error caused by the backend.
71    Backend(E),
72}
73
74impl<E> HookError<E> {
75    /// Convenience constructor function for the `HookError::Message`
76    /// variant.
77    pub fn message(msg: impl Into<Cow<'static, str>>) -> Self {
78        Self::Message(msg.into())
79    }
80}
81
82impl<E: fmt::Display> fmt::Display for HookError<E> {
83    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
84        match self {
85            Self::Message(msg) => write!(f, "{}", msg),
86            Self::Backend(e) => write!(f, "{}", e),
87        }
88    }
89}
90
91impl<E: std::error::Error + 'static> std::error::Error for HookError<E> {
92    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
93        match self {
94            Self::Message(_) => None,
95            Self::Backend(e) => Some(e),
96        }
97    }
98}
99
100pub(crate) struct HookVec<M: Manager> {
101    vec: Vec<Hook<M>>,
102}
103
104// Implemented manually to avoid unnecessary trait bound on `M` type parameter.
105impl<M: Manager> fmt::Debug for HookVec<M> {
106    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
107        f.debug_struct("HookVec")
108            //.field("fns", &self.fns)
109            .finish_non_exhaustive()
110    }
111}
112
113// Implemented manually to avoid unnecessary trait bound on `M` type parameter.
114impl<M: Manager> Default for HookVec<M> {
115    fn default() -> Self {
116        Self { vec: Vec::new() }
117    }
118}
119
120impl<M: Manager> HookVec<M> {
121    pub(crate) async fn apply(
122        &self,
123        inner: &mut ObjectInner<M>,
124    ) -> Result<(), HookError<M::Error>> {
125        for hook in &self.vec {
126            match hook {
127                Hook::Fn(f) => f(&mut inner.obj, &inner.metrics)?,
128                Hook::AsyncFn(f) => f(&mut inner.obj, &inner.metrics).await?,
129            };
130        }
131        Ok(())
132    }
133    pub(crate) fn push(&mut self, hook: Hook<M>) {
134        self.vec.push(hook);
135    }
136}
137
138/// Collection of all the hooks that can be configured for a [`Pool`].
139///
140/// [`Pool`]: super::Pool
141pub(crate) struct Hooks<M: Manager> {
142    pub(crate) post_create: HookVec<M>,
143    pub(crate) pre_recycle: HookVec<M>,
144    pub(crate) post_recycle: HookVec<M>,
145}
146
147// Implemented manually to avoid unnecessary trait bound on `M` type parameter.
148impl<M: Manager> fmt::Debug for Hooks<M> {
149    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150        f.debug_struct("Hooks")
151            .field("post_create", &self.post_create)
152            .field("pre_recycle", &self.post_recycle)
153            .field("post_recycle", &self.post_recycle)
154            .finish()
155    }
156}
157
158// Implemented manually to avoid unnecessary trait bound on `M` type parameter.
159impl<M: Manager> Default for Hooks<M> {
160    fn default() -> Self {
161        Self {
162            pre_recycle: HookVec::default(),
163            post_create: HookVec::default(),
164            post_recycle: HookVec::default(),
165        }
166    }
167}
deadpool/managed/hooks.rs

deadpool/managed/
hooks.rs
