1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

//! Request lifecycle hooks that can be used to further customize how and what callbacks to run
//! on events.
//!
//! `axum-prometheus` is built on top of lifecycle hooks. Using this module allows you to customize
//! behavior even more.
use bytes::Buf;
use http::{HeaderMap, Request, Response};
use tower_http::classify::ClassifiedResponse;

mod body;
mod future;
pub mod layer;
pub mod service;

/// Trait that defines callbacks for [`LifeCycle`] to call.
///
/// [`LifeCycle`]: service::LifeCycle
pub trait Callbacks<FailureClass>: Sized {
    /// Additional data to attach to callbacks.
    type Data;

    /// Create an instance of `Self::Data` from the request.
    ///
    /// This method is called immediately after the request is received by [`Service::call`].
    ///
    /// The value returned here will be passed to the other methods in this trait.
    ///
    /// [`Service::call`]: tower::Service::call
    fn prepare<B>(&mut self, request: &Request<B>) -> Self::Data;

    /// Perform some action when a response has been generated.
    ///
    /// This method is called when the inner [`Service`]'s response future
    /// completes with `Ok(response)`, regardless if the response is classified
    /// as a success or a failure.
    ///
    /// If the response is the start of a stream (as determined by the
    /// classifier passed to [`LifeCycle::new`] or [`LifeCycleLayer::new`]) then
    /// `classification` will be [`ClassifiedResponse::RequiresEos(())`],
    /// otherwise it will be [`ClassifiedResponse::Ready`].
    ///
    /// The default implementation does nothing and returns immediately.
    ///
    /// [`ClassifiedResponse::RequiresEos(())`]: tower_http::classify::ClassifiedResponse::RequiresEos
    /// [`Service`]: tower::Service
    /// [`LifeCycle::new`]: service::LifeCycle::new
    /// [`LifeCycleLayer::new`]: layer::LifeCycleLayer::new
    #[inline]
    fn on_response<B>(
        &mut self,
        _response: &Response<B>,
        _classification: ClassifiedResponse<FailureClass, ()>,
        _data: &mut Self::Data,
    ) {
    }

    /// Perform some action when a response body chunk has been generated.
    ///
    /// This is called when [`Body::poll_data`] completes with `Some(Ok(chunk))`
    /// regardless if the chunk is empty or not.
    ///
    /// The default implementation does nothing and returns immediately.
    ///
    /// [`Body::poll_data`]: http_body::Body::poll_data
    #[inline]
    fn on_body_chunk<B: Buf>(&self, _check: &B, _data: &Self::Data) {}

    /// Perform some action when a stream has ended.
    ///
    /// This is called when [`Body::poll_trailers`] completes with
    /// `Ok(trailers)` regardless if the trailers are classified as a failure.
    ///
    /// A stream that ends successfully will trigger two callbacks.
    /// [`on_response`] will be called once the response has been generated and
    /// the stream has started and [`on_eos`] will be called once the stream has
    /// ended.
    ///
    /// If the trailers were classified as a success then `classification` will
    /// be `Ok(())` otherwise `Err(failure_class)`.
    ///
    /// The default implementation does nothing and returns immediately.
    ///
    /// [`on_response`]: Callbacks::on_response
    /// [`on_eos`]: Callbacks::on_eos
    /// [`Body::poll_trailers`]: http_body::Body::poll_trailers
    #[inline]
    fn on_eos(
        self,
        _trailers: Option<&HeaderMap>,
        _classification: Result<(), FailureClass>,
        _data: Self::Data,
    ) {
    }

    /// Perform some action when an error has been encountered.
    ///
    /// This method is only called in the following scenarios:
    ///
    /// - The inner [`Service`]'s response future resolves to an error.
    /// - [`Body::poll_data`] returns an error.
    /// - [`Body::poll_trailers`] returns an error.
    ///
    /// That means this method is _not_ called if a response is classified as a
    /// failure (then [`on_response`] is called) or an end-of-stream is
    /// classified as a failure (then [`on_eos`] is called).
    ///
    /// `failed_at` specifies where the error happened.
    ///
    /// The default implementation does nothing and returns immediately.
    ///
    /// [`Service`]: tower::Service
    /// [`on_response`]: Callbacks::on_response
    /// [`on_eos`]: Callbacks::on_eos
    /// [`Service::call`]: tower::Service::call
    /// [`Body::poll_data`]: http_body::Body::poll_data
    /// [`Body::poll_trailers`]: http_body::Body::poll_trailers
    fn on_failure(
        self,
        _failed_at: FailedAt,
        _failure_classification: FailureClass,
        _data: Self::Data,
    ) {
    }
}

/// Enum used to specify where an error was encountered.
#[derive(Debug)]
pub enum FailedAt {
    /// Generating the response failed.
    Response,
    /// Generating the response body failed.
    Body,
    /// Generating the response trailers failed.
    Trailers,
}