axum_prometheus/lifecycle/
mod.rs

1//! Request lifecycle hooks that can be used to further customize how and what callbacks to run
2//! on events.
3//!
4//! `axum-prometheus` is built on top of lifecycle hooks. Using this module allows you to customize
5//! behavior even more.
6use bytes::Buf;
7use http::{HeaderMap, Request, Response};
8use tower_http::classify::ClassifiedResponse;
9
10mod body;
11mod future;
12pub mod layer;
13pub mod service;
14
15/// Trait that defines callbacks for [`LifeCycle`] to call.
16///
17/// [`LifeCycle`]: service::LifeCycle
18pub trait Callbacks<FailureClass>: Sized {
19    /// Additional data to attach to callbacks.
20    type Data;
21
22    /// Create an instance of `Self::Data` from the request.
23    ///
24    /// This method is called immediately after the request is received by [`Service::call`].
25    ///
26    /// The value returned here will be passed to the other methods in this trait.
27    ///
28    /// [`Service::call`]: tower::Service::call
29    fn prepare<B>(&mut self, request: &Request<B>) -> Self::Data;
30
31    /// Perform some action when a response has been generated.
32    ///
33    /// This method is called when the inner [`Service`]'s response future
34    /// completes with `Ok(response)`, regardless if the response is classified
35    /// as a success or a failure.
36    ///
37    /// If the response is the start of a stream (as determined by the
38    /// classifier passed to [`LifeCycle::new`] or [`LifeCycleLayer::new`]) then
39    /// `classification` will be [`ClassifiedResponse::RequiresEos(())`],
40    /// otherwise it will be [`ClassifiedResponse::Ready`].
41    ///
42    /// The default implementation does nothing and returns immediately.
43    ///
44    /// [`ClassifiedResponse::RequiresEos(())`]: tower_http::classify::ClassifiedResponse::RequiresEos
45    /// [`Service`]: tower::Service
46    /// [`LifeCycle::new`]: service::LifeCycle::new
47    /// [`LifeCycleLayer::new`]: layer::LifeCycleLayer::new
48    #[inline]
49    fn on_response<B>(
50        &mut self,
51        _response: &Response<B>,
52        _classification: ClassifiedResponse<FailureClass, ()>,
53        _data: &mut Self::Data,
54    ) {
55    }
56
57    /// Perform some action when a stream has ended.
58    ///
59    /// This is called when [`Body::poll_frame`] produces `Ok(trailers)` with the [`Frame::into_trailers`] method,
60    /// regardless if the trailers are classified as a failure.
61    ///
62    /// A stream that ends successfully will trigger two callbacks.
63    /// [`on_response`] will be called once the response has been generated and
64    /// the stream has started and [`on_eos`] will be called once the stream has
65    /// ended.
66    ///
67    /// If the trailers were classified as a success then `classification` will
68    /// be `Ok(())` otherwise `Err(failure_class)`.
69    ///
70    /// The default implementation does nothing and returns immediately.
71    ///
72    /// [`on_response`]: Callbacks::on_response
73    /// [`on_eos`]: Callbacks::on_eos
74    /// [`Body::poll_frame`]: http_body::Body::poll_frame
75    /// [`Frame::into_trailers`]: http_body::Frame::into_trailers
76    #[inline]
77    fn on_eos(
78        self,
79        _trailers: Option<&HeaderMap>,
80        _classification: Result<(), FailureClass>,
81        _data: Self::Data,
82    ) {
83    }
84
85    /// Perform some action when an error has been encountered.
86    ///
87    /// This method is only called in the following scenarios:
88    ///
89    /// - The inner [`Service`]'s response future resolves to an error.
90    /// - [`Body::poll_frame`] returns an error.
91    ///
92    /// That means this method is _not_ called if a response is classified as a
93    /// failure (then [`on_response`] is called) or an end-of-stream is
94    /// classified as a failure (then [`on_eos`] is called).
95    ///
96    /// `failed_at` specifies where the error happened.
97    ///
98    /// The default implementation does nothing and returns immediately.
99    ///
100    /// [`Service`]: tower::Service
101    /// [`on_response`]: Callbacks::on_response
102    /// [`on_eos`]: Callbacks::on_eos
103    /// [`Service::call`]: tower::Service::call
104    /// [`Body::poll_frame`]: http_body::Body::poll_frame
105    fn on_failure(
106        self,
107        _failed_at: FailedAt,
108        _failure_classification: FailureClass,
109        _data: &mut Self::Data,
110    ) {
111    }
112}
113
114/// A trait that allows to hook into [`http_body::Body::poll_frame`]'s lifecycle.
115pub trait OnBodyChunk<B: Buf> {
116    type Data;
117
118    /// Perform some action when a response body chunk has been generated.
119    ///
120    /// This is called when [`Body::poll_frame`] returns `Some(Ok(frame))`, and [`Frame::into_data`] returns with `Some(chunk)`,
121    /// regardless if the chunk is empty or not.
122    ///
123    /// The default implementation does nothing and returns immediately.
124    ///
125    /// [`Body::poll_frame`]: http_body::Body::poll_frame
126    /// [`Frame::into_data`]: http_body::Frame::into_data
127    #[inline]
128    fn call(&mut self, _body: &B, _exact_body_size: Option<u64>, _data: &mut Self::Data) {}
129}
130
131/// Enum used to specify where an error was encountered.
132#[derive(Debug)]
133pub enum FailedAt {
134    /// Generating the response failed.
135    Response,
136    /// Generating the response body failed.
137    Body,
138    /// Generating the response trailers failed.
139    Trailers,
140}
axum_prometheus/lifecycle/mod.rs

axum_prometheus/lifecycle/
mod.rs
