Learning Pingora 02 - A Simple HTTP Server

30
/// This trait defines how to map a request to a response
31
#[cfg_attr(not(doc_async_trait), async_trait)]
32
pub trait ServeHttp {
33
    /// Define the mapping from a request to a response.
34
    /// Note that the request header is already read, but the implementation needs to read the
35
    /// request body if any.
36
    ///
37
    /// # Limitation
38
    /// In this API, the entire response has to be generated before the end of this call.
39
    /// So it is not suitable for streaming response or interactive communications.
40
    /// Users need to implement their own [`super::HttpServerApp`] for those use cases.
41
    async fn response(&self, http_session: &mut ServerSession) -> Response<Vec<u8>>;
42
}

110
/// A helper struct for HTTP server with http modules embedded
111
pub struct HttpServer<SV> {
112
    app: SV,
113
    modules: HttpModules,
114
}

131
#[cfg_attr(not(doc_async_trait), async_trait)]
132
impl<SV> HttpServerApp for HttpServer<SV>
133
where
134
    SV: ServeHttp + Send + Sync,
135
{
136
    async fn process_new_http(
137
        self: &Arc<Self>,
138
        mut http: ServerSession,
139
        shutdown: &ShutdownWatch,
140
    ) -> Option<Stream> {
141
        // ...

33
pub trait HttpModule {
34
    fn request_header_filter(&mut self, _req: &mut RequestHeader) -> Result<()> {
35
        Ok(())
36
    }
37

38
    fn request_body_filter(&mut self, body: Option<Bytes>) -> Result<Option<Bytes>> {
39
        Ok(body)
40
    }
41

42
    fn response_filter(&mut self, _t: &mut HttpTask) -> Result<()> {
43
        Ok(())
44
    }
45

46
    fn as_any(&self) -> &dyn Any;
47
    fn as_any_mut(&mut self) -> &mut dyn Any;
48
}
49

50
type Module = Box<dyn HttpModule + 'static + Send + Sync>

31
/// An enum to hold all possible HTTP response events.
32
#[derive(Debug)]
33
pub enum HttpTask {
34
    /// the response header and the boolean end of response flag
35
    Header(Box<pingora_http::ResponseHeader>, bool),
36
    /// A piece of response header and the end of response boolean flag
37
    Body(Option<bytes::Bytes>, bool),
38
    /// HTTP response trailer
39
    Trailer(Option<Box<http::HeaderMap>>),
40
    /// Signal that the response is already finished
41
    Done,
42
    /// Signal that the reading of the response encounters errors.
43
    Failed(pingora_error::BError),
44
}

54
/// Trait to init the http module ctx for each request
55
pub trait HttpModuleBuilder {
56
    /// The order the module will run
57
    ///
58
    /// The lower the value, the later it runs relative to other filters.
59
    /// If the order of the filter is not important, leave it to the default 0.
60
    fn order(&self) -> i16 {
61
        0
62
    }
63

64
    /// Initialize and return the per request module context
65
    fn init(&self) -> Module;
66
}
67

68
pub type ModuleBuilder = Box<dyn HttpModuleBuilder + 'static + Send + Sync>

70
/// The object to hold multiple http modules
71
pub struct HttpModules {
72
    modules: Vec<ModuleBuilder>,
73
    module_index: OnceCell<Arc<HashMap<TypeId, usize>>>,
74
}
75

76
impl HttpModules {
77
    /// Create a new [HttpModules]
78
    pub fn new() -> Self {
79
        HttpModules {
80
            modules: vec![],
81
            module_index: OnceCell::new(),
82
        }
83
    }
84

85
    /// Add a new [ModuleBuilder] to [HttpModules]
86
    ///
87
    /// Each type of [HttpModule] can be only added once.
88
    /// # Panic
89
    /// Panic if any [HttpModule] is added more than once.
90
    pub fn add_module(&mut self, builder: ModuleBuilder) {
91
        if self.module_index.get().is_some() {
92
            // We use a shared module_index the index would be out of sync if we
93
            // add more modules.
94
            panic!("cannot add module after ctx is already built")
95
        }
96
        self.modules.push(builder);
97
        // not the most efficient way but should be fine
98
        // largest order first
99
        self.modules.sort_by_key(|m| -m.order());
100
    }
101

102
    /// Build the contexts of all the modules added to this [HttpModules]
103
    pub fn build_ctx(&self) -> HttpModuleCtx {
104
        let module_ctx: Vec<_> = self.modules.iter().map(|b| b.init()).collect();
105
        let module_index = self
106
            .module_index
107
            .get_or_init(|| {
108
                let mut module_index = HashMap::with_capacity(self.modules.len());
109
                for (i, c) in module_ctx.iter().enumerate() {
110
                    let exist = module_index.insert(c.as_any().type_id(), i);
111
                    if exist.is_some() {
112
                        panic!("duplicated filters found")
113
                    }
114
                }
115
                Arc::new(module_index)
116
            })
117
            .clone();
118

119
        HttpModuleCtx {
120
            module_ctx,
121
            module_index,
122
        }
123
    }
124
}

126
/// The Contexts of multiple modules
127
///
128
/// This is the object that will apply all the included modules to a certain HTTP request.
129
/// The modules are ordered according to their `order()`.
130
pub struct HttpModuleCtx {
131
    // the modules in the order of execution
132
    module_ctx: Vec<Module>,
133
    // find the module in the vec with its type ID
134
    module_index: Arc<HashMap<TypeId, usize>>,
135
}
136

137
impl HttpModuleCtx {
138
    /// Create a placeholder empty [HttpModuleCtx].
139
    ///
140
    /// [HttpModules] should be used to create nonempty [HttpModuleCtx].
141
    pub fn empty() -> Self {
142
        HttpModuleCtx {
143
            module_ctx: vec![],
144
            module_index: Arc::new(HashMap::new()),
145
        }
146
    }
147

148
    /// Get a ref to [HttpModule] if any.
149
    pub fn get<T: 'static>(&self) -> Option<&T> {
150
        let idx = self.module_index.get(&TypeId::of::<T>())?;
151
        let ctx = &self.module_ctx[*idx];
152
        Some(
153
            ctx.as_any()
154
                .downcast_ref::<T>()
155
                .expect("type should always match"),
156
        )
157
    }
158

159
    /// Get a mut ref to [HttpModule] if any.
160
    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
161
        let idx = self.module_index.get(&TypeId::of::<T>())?;
162
        let ctx = &mut self.module_ctx[*idx];
163
        Some(
164
            ctx.as_any_mut()
165
                .downcast_mut::<T>()
166
                .expect("type should always match"),
167
        )
168
    }
169

170
    /// Run the `request_header_filter` for all the modules according to their orders.
171
    pub fn request_header_filter(&mut self, req: &mut RequestHeader) -> Result<()> {
172
        for filter in self.module_ctx.iter_mut() {
173
            filter.request_header_filter(req)?;
174
        }
175
        Ok(())
176
    }
177

178
    /// Run the `request_body_filter` for all the modules according to their orders.
179
    pub fn request_body_filter(&mut self, mut body: Option<Bytes>) -> Result<Option<Bytes>> {
180
        for filter in self.module_ctx.iter_mut() {
181
            body = filter.request_body_filter(body)?;
182
        }
183
        Ok(body)
184
    }
185

186
    /// Run the `response_filter` for all the modules according to their orders.
187
    pub fn response_filter(&mut self, t: &mut HttpTask) -> Result<()> {
188
        for filter in self.module_ctx.iter_mut() {
189
            filter.response_filter(t)?;
190
        }
191
        Ok(())
192
    }
193
}

26
/// A HTTP application that reports Prometheus metrics.
27
///
28
/// This application will report all the [static metrics](https://docs.rs/prometheus/latest/prometheus/index.html#static-metrics)
29
/// collected via the [Prometheus](https://docs.rs/prometheus/) crate;
30
pub struct PrometheusHttpApp;
31

32
#[cfg_attr(not(doc_async_trait), async_trait)]
33
impl ServeHttp for PrometheusHttpApp {
34
    async fn response(&self, _http_session: &mut ServerSession) -> Response<Vec<u8>> {
35
        let encoder = TextEncoder::new();
36
        let metric_families = prometheus::gather();
37
        let mut buffer = vec![];
38
        encoder.encode(&metric_families, &mut buffer).unwrap();
39
        Response::builder()
40
            .status(200)
41
            .header(http::header::CONTENT_TYPE, encoder.format_type())
42
            .header(http::header::CONTENT_LENGTH, buffer.len())
43
            .body(buffer)
44
            .unwrap()
45
    }
46
}
47

48
/// The [HttpServer] for [PrometheusHttpApp]
49
///
50
/// This type provides the functionality of [PrometheusHttpApp] with compression enabled
51
pub type PrometheusServer = HttpServer<PrometheusHttpApp>;
52

53
impl PrometheusServer {
54
    pub fn new() -> Self {
55
        let mut server = Self::new_app(PrometheusHttpApp);
56
        // enable gzip level 7 compression
57
        server.add_module(ResponseCompressionBuilder::enable(7));
58
        server
59
    }
60
}

20
/// HTTP response compression module
21
pub struct ResponseCompression(ResponseCompressionCtx);
22

23
impl HttpModule for ResponseCompression {
24
    fn as_any(&self) -> &dyn std::any::Any {
25
        self
26
    }
27
    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
28
        self
29
    }
30

31
    fn request_header_filter(&mut self, req: &mut RequestHeader) -> Result<()> {
32
        self.0.request_filter(req);
33
        Ok(())
34
    }
35

36
    fn response_filter(&mut self, t: &mut HttpTask) -> Result<()> {
37
        self.0.response_filter(t);
38
        Ok(())
39
    }
40
}
41

42
/// The builder for HTTP response compression module
43
pub struct ResponseCompressionBuilder {
44
    level: u32,
45
}
46

47
impl ResponseCompressionBuilder {
48
    /// Return a [ModuleBuilder] for [ResponseCompression] with the given compression level
49
    pub fn enable(level: u32) -> ModuleBuilder {
50
        Box::new(ResponseCompressionBuilder { level })
51
    }
52
}
53

54
impl HttpModuleBuilder for ResponseCompressionBuilder {
55
    fn init(&self) -> Module {
56
        Box::new(ResponseCompression(ResponseCompressionCtx::new(
57
            self.level, false,
58
        )))
59
    }
60

61
    fn order(&self) -> i16 {
62
        // run the response filter later than most others filters
63
        i16::MIN / 2
64
    }
65
}