axum_prometheus/lifecycle/
service.rs

1use std::task::{Context, Poll};
2
3use http::{Request, Response};
4use http_body::Body;
5use tower::Service;
6use tower_http::classify::MakeClassifier;
7
8use super::{
9    body::ResponseBody, future::ResponseFuture, layer::LifeCycleLayer, Callbacks, OnBodyChunk,
10};
11
12#[derive(Clone, Debug)]
13pub struct LifeCycle<S, MC, Callbacks, OnBodyChunk> {
14    pub(super) inner: S,
15    pub(super) make_classifier: MC,
16    pub(super) callbacks: Callbacks,
17    pub(super) on_body_chunk: OnBodyChunk,
18}
19
20impl<S, MC, Callbacks, OnBodyChunk> LifeCycle<S, MC, Callbacks, OnBodyChunk> {
21    pub fn new(
22        inner: S,
23        make_classifier: MC,
24        callbacks: Callbacks,
25        on_body_chunk: OnBodyChunk,
26    ) -> Self {
27        Self {
28            inner,
29            make_classifier,
30            callbacks,
31            on_body_chunk,
32        }
33    }
34
35    pub fn layer(
36        make_classifier: MC,
37        callbacks: Callbacks,
38        on_body_chunk: OnBodyChunk,
39    ) -> LifeCycleLayer<MC, Callbacks, OnBodyChunk> {
40        LifeCycleLayer::new(make_classifier, callbacks, on_body_chunk)
41    }
42
43    /// Gets a reference to the underlying service.
44    pub fn get_ref(&self) -> &S {
45        &self.inner
46    }
47
48    /// Gets a mutable reference to the underlying service.
49    pub fn get_mut(&mut self) -> &mut S {
50        &mut self.inner
51    }
52
53    /// Consumes `self`, returning the underlying service.
54    pub fn into_inner(self) -> S {
55        self.inner
56    }
57}
58
59impl<S, MC, ReqBody, ResBody, CallbacksT, OnBodyChunkT> Service<Request<ReqBody>>
60    for LifeCycle<S, MC, CallbacksT, OnBodyChunkT>
61where
62    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
63    ResBody: Body,
64    MC: MakeClassifier,
65    CallbacksT: Callbacks<MC::FailureClass> + Clone,
66    S::Error: std::fmt::Display + 'static,
67    OnBodyChunkT: OnBodyChunk<ResBody::Data, Data = CallbacksT::Data> + Clone,
68    CallbacksT::Data: Clone,
69{
70    type Response = Response<
71        ResponseBody<ResBody, MC::ClassifyEos, CallbacksT, OnBodyChunkT, CallbacksT::Data>,
72    >;
73    type Error = S::Error;
74    type Future =
75        ResponseFuture<S::Future, MC::Classifier, CallbacksT, OnBodyChunkT, CallbacksT::Data>;
76
77    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
78        self.inner.poll_ready(cx)
79    }
80
81    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
82        let callbacks_data = self.callbacks.prepare(&req);
83
84        let classifier = self.make_classifier.make_classifier(&req);
85
86        ResponseFuture {
87            inner: self.inner.call(req),
88            classifier: Some(classifier),
89            callbacks: Some(self.callbacks.clone()),
90            callbacks_data: Some(callbacks_data),
91            on_body_chunk: Some(self.on_body_chunk.clone()),
92        }
93    }
94}
axum_prometheus/lifecycle/service.rs

axum_prometheus/lifecycle/
service.rs
