axum_prometheus/lifecycle/
layer.rs

1use tower::Layer;
2
3use super::service::LifeCycle;
4
5/// [`Layer`] for adding callbacks to the lifecycle of request.
6///
7/// See the [module docs](crate::lifecycle) for more details.
8///
9/// [`Layer`]: tower::Layer
10#[derive(Debug, Clone)]
11pub struct LifeCycleLayer<MC, Callbacks, OnBodyChunk> {
12    pub(super) make_classifier: MC,
13    pub(super) callbacks: Callbacks,
14    pub(super) on_body_chunk: OnBodyChunk,
15}
16
17impl<MC, Callbacks, OnBodyChunk> LifeCycleLayer<MC, Callbacks, OnBodyChunk> {
18    /// Create a new `LifeCycleLayer`.
19    pub fn new(make_classifier: MC, callbacks: Callbacks, on_body_chunk: OnBodyChunk) -> Self {
20        LifeCycleLayer {
21            make_classifier,
22            callbacks,
23            on_body_chunk,
24        }
25    }
26
27    pub(crate) fn on_body_chunk(&mut self, on_body_chunk: OnBodyChunk) {
28        self.on_body_chunk = on_body_chunk;
29    }
30}
31
32impl<S, MC, Callbacks, OnBodyChunk> Layer<S> for LifeCycleLayer<MC, Callbacks, OnBodyChunk>
33where
34    MC: Clone,
35    Callbacks: Clone,
36    OnBodyChunk: Clone,
37{
38    type Service = LifeCycle<S, MC, Callbacks, OnBodyChunk>;
39
40    fn layer(&self, inner: S) -> Self::Service {
41        LifeCycle {
42            inner,
43            make_classifier: self.make_classifier.clone(),
44            callbacks: self.callbacks.clone(),
45            on_body_chunk: self.on_body_chunk.clone(),
46        }
47    }
48}