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