rocket_prometheus/
lib.rs

1/*!
2Prometheus instrumentation for Rocket applications.
3
4# Usage
5
6Add this crate to your `Cargo.toml` alongside Rocket 0.5.0:
7
8```toml
9[dependencies]
10rocket = "0.5.0"
11rocket_prometheus = "0.10.1"
12```
13
14Then attach and mount a [`PrometheusMetrics`] instance to your Rocket app:
15
16```rust
17use rocket_prometheus::PrometheusMetrics;
18
19#[rocket::launch]
20fn launch() -> _ {
21    let prometheus = PrometheusMetrics::new();
22    rocket::build()
23        .attach(prometheus.clone())
24        .mount("/metrics", prometheus)
25}
26```
27
28This will expose metrics like this at the /metrics endpoint of your application:
29
30```shell
31$ curl localhost:8000/metrics
32# HELP rocket_http_requests_duration_seconds HTTP request duration in seconds for all requests
33# TYPE rocket_http_requests_duration_seconds histogram
34rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.005"} 2
35rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.01"} 2
36rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.025"} 2
37rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.05"} 2
38rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.1"} 2
39rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.25"} 2
40rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.5"} 2
41rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="1"} 2
42rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="2.5"} 2
43rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="5"} 2
44rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="10"} 2
45rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="+Inf"} 2
46rocket_http_requests_duration_seconds_sum{endpoint="/metrics",method="GET",status="200"} 0.0011045669999999999
47rocket_http_requests_duration_seconds_count{endpoint="/metrics",method="GET",status="200"} 2
48# HELP rocket_http_requests_total Total number of HTTP requests
49# TYPE rocket_http_requests_total counter
50rocket_http_requests_total{endpoint="/metrics",method="GET",status="200"} 2
51```
52
53# Metrics
54
55By default this crate tracks two metrics:
56
57- `rocket_http_requests_total` (labels: endpoint, method, status): the
58  total number of HTTP requests handled by Rocket.
59- `rocket_http_requests_duration_seconds` (labels: endpoint, method, status):
60  the request duration for all HTTP requests handled by Rocket.
61
62The 'rocket' prefix of these metrics can be changed by setting the
63`ROCKET_PROMETHEUS_NAMESPACE` environment variable.
64
65## Custom Metrics
66
67Further metrics can be tracked by registering them with the registry of the
68[`PrometheusMetrics`] instance:
69
70```rust
71use once_cell::sync::Lazy;
72use rocket::{get, launch, routes};
73use rocket_prometheus::{
74    prometheus::{opts, IntCounterVec},
75    PrometheusMetrics,
76};
77
78static NAME_COUNTER: Lazy<IntCounterVec> = Lazy::new(|| {
79    IntCounterVec::new(opts!("name_counter", "Count of names"), &["name"])
80        .expect("Could not create NAME_COUNTER")
81});
82
83#[get("/hello/<name>")]
84pub fn hello(name: &str) -> String {
85    NAME_COUNTER.with_label_values(&[name]).inc();
86    format!("Hello, {}!", name)
87}
88
89#[launch]
90fn launch() -> _ {
91    let prometheus = PrometheusMetrics::new();
92    prometheus
93        .registry()
94        .register(Box::new(NAME_COUNTER.clone()))
95        .unwrap();
96    rocket::build()
97        .attach(prometheus.clone())
98        .mount("/", routes![hello])
99        .mount("/metrics", prometheus)
100}
101```
102
103*/
104#![deny(missing_docs)]
105#![deny(unsafe_code)]
106
107use std::{env, time::Instant};
108
109use prometheus::{opts, Encoder, HistogramVec, IntCounterVec, Registry, TextEncoder};
110use rocket::{
111    fairing::{Fairing, Info, Kind},
112    http::{ContentType, Method},
113    route::{Handler, Outcome},
114    Data, Request, Response, Route,
115};
116
117/// Re-export Prometheus so users can use it without having to explicitly
118/// add a specific version to their dependencies, which can result in
119/// mysterious compiler error messages.
120pub use prometheus;
121
122/// Environment variable used to configure the namespace of metrics exposed
123/// by `PrometheusMetrics`.
124const NAMESPACE_ENV_VAR: &str = "ROCKET_PROMETHEUS_NAMESPACE";
125
126#[derive(Clone)]
127#[must_use = "must be attached and mounted to a Rocket instance"]
128/// Fairing and Handler implementing request instrumentation.
129///
130/// By default this tracks two metrics:
131///
132/// - `rocket_http_requests_total` (labels: endpoint, method, status): the
133///   total number of HTTP requests handled by Rocket.
134/// - `rocket_http_requests_duration_seconds` (labels: endpoint, method, status):
135///   the request duration for all HTTP requests handled by Rocket.
136///
137/// The `rocket` prefix of these metrics can be changed by setting the
138/// `ROCKET_PROMETHEUS_NAMESPACE` environment variable.
139///
140/// # Usage
141///
142/// Simply attach and mount a `PrometheusMetrics` instance to your Rocket
143/// app as for a normal fairing / handler:
144///
145/// ```rust
146/// use rocket_prometheus::PrometheusMetrics;
147///
148/// #[rocket::launch]
149/// fn launch() -> _ {
150///     let prometheus = PrometheusMetrics::new();
151///     rocket::build()
152///         .attach(prometheus.clone())
153///         .mount("/metrics", prometheus)
154/// }
155/// ```
156///
157/// Metrics will then be available on the "/metrics" endpoint:
158///
159/// ```shell
160/// $ curl localhost:8000/metrics
161/// # HELP rocket_http_requests_duration_seconds HTTP request duration in seconds for all requests
162/// # TYPE rocket_http_requests_duration_seconds histogram
163/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.005"} 2
164/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.01"} 2
165/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.025"} 2
166/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.05"} 2
167/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.1"} 2
168/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.25"} 2
169/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="0.5"} 2
170/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="1"} 2
171/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="2.5"} 2
172/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="5"} 2
173/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="10"} 2
174/// rocket_http_requests_duration_seconds_bucket{endpoint="/metrics",method="GET",status="200",le="+Inf"} 2
175/// rocket_http_requests_duration_seconds_sum{endpoint="/metrics",method="GET",status="200"} 0.0011045669999999999
176/// rocket_http_requests_duration_seconds_count{endpoint="/metrics",method="GET",status="200"} 2
177/// # HELP rocket_http_requests_total Total number of HTTP requests
178/// # TYPE rocket_http_requests_total counter
179/// rocket_http_requests_total{endpoint="/metrics",method="GET",status="200"} 2
180/// ```
181pub struct PrometheusMetrics {
182    // Standard metrics tracked by the fairing.
183    http_requests_total: IntCounterVec,
184    http_requests_duration_seconds: HistogramVec,
185
186    // The registry used by the fairing for Rocket metrics.
187    //
188    // This registry is created by `PrometheusMetrics::with_registry` and is
189    // private to each `PrometheusMetrics` instance, allowing multiple
190    // `PrometheusMetrics` instances to share the same `extra_registry`.
191    //
192    // Previously the fairing tried to register the internal metrics on the `extra_registry`,
193    // which caused conflicts if the same registry was passed twice. This is now avoided
194    // by using an internal registry for those metrics.
195    rocket_registry: Registry,
196
197    // The registry used by the fairing for custom metrics.
198    //
199    // See `rocket_registry` for details on why these metrics are stored on a separate registry.
200    custom_registry: Registry,
201
202    // Option for request filtering. If it contains a function, said function should return true when
203    // the current request should be considered in metrics.
204    request_filter: Option<for<'a> fn(&'a Request<'_>) -> bool>,
205}
206
207impl PrometheusMetrics {
208    /// Create a new [`PrometheusMetrics`].
209    pub fn new() -> Self {
210        Self::with_registry(Registry::new())
211    }
212
213    /// Create a new [`PrometheusMetrics`] with a custom [`Registry`].
214    // Allow `clippy::missing_panics_doc` because we know:
215    // - the two metrics can't fail to be created (their config is valid)
216    // - registering the metrics can't fail (the registry is new, so there is no chance of metric duplication)
217    #[allow(clippy::missing_panics_doc)]
218    pub fn with_registry(registry: Registry) -> Self {
219        let rocket_registry = Registry::new();
220        let namespace = env::var(NAMESPACE_ENV_VAR).unwrap_or_else(|_| "rocket".into());
221
222        let http_requests_total_opts =
223            opts!("http_requests_total", "Total number of HTTP requests")
224                .namespace(namespace.clone());
225        let http_requests_total =
226            IntCounterVec::new(http_requests_total_opts, &["endpoint", "method", "status"])
227                .unwrap();
228        let http_requests_duration_seconds_opts = opts!(
229            "http_requests_duration_seconds",
230            "HTTP request duration in seconds for all requests"
231        )
232        .namespace(namespace);
233        let http_requests_duration_seconds = HistogramVec::new(
234            http_requests_duration_seconds_opts.into(),
235            &["endpoint", "method", "status"],
236        )
237        .unwrap();
238
239        rocket_registry
240            .register(Box::new(http_requests_total.clone()))
241            .unwrap();
242        rocket_registry
243            .register(Box::new(http_requests_duration_seconds.clone()))
244            .unwrap();
245
246        Self {
247            http_requests_total,
248            http_requests_duration_seconds,
249            rocket_registry,
250            custom_registry: registry,
251            request_filter: None,
252        }
253    }
254
255    /// Create a new [`PrometheusMetrics`] using the default Prometheus [`Registry`].
256    ///
257    /// This will cause the fairing to include metrics created by the various
258    /// `prometheus` macros, e.g.  `register_int_counter`.
259    pub fn with_default_registry() -> Self {
260        Self::with_registry(prometheus::default_registry().clone())
261    }
262
263    /// Get the registry used by this fairing to track additional metrics.
264    ///
265    /// You can use this to register further metrics,
266    /// causing them to be exposed along with the default metrics
267    /// on the [`PrometheusMetrics`] handler.
268    ///
269    /// Note that the `http_requests_total` and `http_requests_duration_seconds` metrics
270    /// are _not_ included in this registry.
271    ///
272    /// ```rust
273    /// use once_cell::sync::Lazy;
274    /// use prometheus::{opts, IntCounter};
275    /// use rocket_prometheus::PrometheusMetrics;
276    ///
277    /// static MY_COUNTER: Lazy<IntCounter> = Lazy::new(|| {
278    ///     IntCounter::new("my_counter", "A counter I use a lot")
279    ///         .expect("Could not create counter")
280    /// });
281    ///
282    /// let prometheus = PrometheusMetrics::new();
283    /// prometheus.registry().register(Box::new(MY_COUNTER.clone()));
284    /// ```
285    #[must_use]
286    pub const fn registry(&self) -> &Registry {
287        &self.custom_registry
288    }
289
290    /// Get the `http_requests_total` metric.
291    pub fn http_requests_total(&self) -> &IntCounterVec {
292        &self.http_requests_total
293    }
294
295    /// Get the `http_requests_duration_seconds` metric.
296    pub fn http_requests_duration_seconds(&self) -> &HistogramVec {
297        &self.http_requests_duration_seconds
298    }
299
300    /// Set a filter for which request should be considered in metrics. The filter function should
301    /// return false when the request should be ignored.
302    ///
303    /// Example:
304    ///
305    /// ```rust
306    /// use rocket_prometheus::PrometheusMetrics;
307    ///
308    /// let prometheus = PrometheusMetrics::new()
309    ///     .with_request_filter(|request| {
310    ///         request.uri().path() != "/metrics"
311    ///     });
312    /// ```
313    pub fn with_request_filter(mut self, filter: for<'a> fn(&'a Request<'_>) -> bool) -> Self {
314        self.request_filter = Some(filter);
315        self
316    }
317}
318
319impl Default for PrometheusMetrics {
320    fn default() -> Self {
321        Self::new()
322    }
323}
324
325/// Value stored in request-local state to measure response time.
326#[derive(Copy, Clone)]
327struct TimerStart(Option<Instant>);
328
329/// A status code which tries not to allocate to produce a `&str` representation.
330enum StatusCode {
331    /// A 'standard' status code, i.e. between 100 and 999.
332    ///
333    /// Most status codes should be represented as this variant,
334    /// which doesn't allocate and provides a non-allocating `&str`
335    /// representation.
336    Standard(rocket::http::Status),
337    /// A non-standard status code.
338    ///
339    /// This is the fallback option used when a status code can't be
340    /// parsed by [`http::StatusCode`]. It requires an allocation.
341    NonStandard(String),
342}
343
344// A string of packed 3-ASCII-digit status code values for the supported range
345// of [100, 999] (900 codes, 2700 bytes).
346// Taken directly from `http::status`.
347const CODE_DIGITS: &str = "\
348100101102103104105106107108109110111112113114115116117118119\
349120121122123124125126127128129130131132133134135136137138139\
350140141142143144145146147148149150151152153154155156157158159\
351160161162163164165166167168169170171172173174175176177178179\
352180181182183184185186187188189190191192193194195196197198199\
353200201202203204205206207208209210211212213214215216217218219\
354220221222223224225226227228229230231232233234235236237238239\
355240241242243244245246247248249250251252253254255256257258259\
356260261262263264265266267268269270271272273274275276277278279\
357280281282283284285286287288289290291292293294295296297298299\
358300301302303304305306307308309310311312313314315316317318319\
359320321322323324325326327328329330331332333334335336337338339\
360340341342343344345346347348349350351352353354355356357358359\
361360361362363364365366367368369370371372373374375376377378379\
362380381382383384385386387388389390391392393394395396397398399\
363400401402403404405406407408409410411412413414415416417418419\
364420421422423424425426427428429430431432433434435436437438439\
365440441442443444445446447448449450451452453454455456457458459\
366460461462463464465466467468469470471472473474475476477478479\
367480481482483484485486487488489490491492493494495496497498499\
368500501502503504505506507508509510511512513514515516517518519\
369520521522523524525526527528529530531532533534535536537538539\
370540541542543544545546547548549550551552553554555556557558559\
371560561562563564565566567568569570571572573574575576577578579\
372580581582583584585586587588589590591592593594595596597598599\
373600601602603604605606607608609610611612613614615616617618619\
374620621622623624625626627628629630631632633634635636637638639\
375640641642643644645646647648649650651652653654655656657658659\
376660661662663664665666667668669670671672673674675676677678679\
377680681682683684685686687688689690691692693694695696697698699\
378700701702703704705706707708709710711712713714715716717718719\
379720721722723724725726727728729730731732733734735736737738739\
380740741742743744745746747748749750751752753754755756757758759\
381760761762763764765766767768769770771772773774775776777778779\
382780781782783784785786787788789790791792793794795796797798799\
383800801802803804805806807808809810811812813814815816817818819\
384820821822823824825826827828829830831832833834835836837838839\
385840841842843844845846847848849850851852853854855856857858859\
386860861862863864865866867868869870871872873874875876877878879\
387880881882883884885886887888889890891892893894895896897898899\
388900901902903904905906907908909910911912913914915916917918919\
389920921922923924925926927928929930931932933934935936937938939\
390940941942943944945946947948949950951952953954955956957958959\
391960961962963964965966967968969970971972973974975976977978979\
392980981982983984985986987988989990991992993994995996997998999";
393
394/// Returns a &str representation of the `StatusCode`
395///
396/// The return value only includes a numerical representation of the
397/// status code. The canonical reason is not included.
398///
399/// This is taken directly from `http::Status::as_str`.
400#[inline]
401fn status_as_str(s: &rocket::http::Status) -> &'static str {
402    let offset = (s.code - 100) as usize;
403    let offset = offset * 3;
404
405    // Invariant: s.code has checked range [100, 999] and CODE_DIGITS is
406    // ASCII-only, of length 900 * 3 = 2700 bytes
407
408    #[cfg(debug_assertions)]
409    {
410        &CODE_DIGITS[offset..offset + 3]
411    }
412
413    #[cfg(not(debug_assertions))]
414    #[allow(unsafe_code)]
415    unsafe {
416        CODE_DIGITS.get_unchecked(offset..offset + 3)
417    }
418}
419
420impl StatusCode {
421    fn as_str(&self) -> &str {
422        match self {
423            Self::Standard(s) => status_as_str(s),
424            Self::NonStandard(s) => s.as_str(),
425        }
426    }
427}
428
429impl From<u16> for StatusCode {
430    fn from(code: u16) -> Self {
431        rocket::http::Status::from_code(code)
432            .map_or_else(|| Self::NonStandard(code.to_string()), Self::Standard)
433    }
434}
435
436#[rocket::async_trait]
437impl Fairing for PrometheusMetrics {
438    fn info(&self) -> Info {
439        Info {
440            name: "Prometheus metric collection",
441            kind: Kind::Request | Kind::Response,
442        }
443    }
444
445    async fn on_request(&self, req: &mut Request<'_>, _: &mut Data<'_>) {
446        req.local_cache(|| TimerStart(Some(Instant::now())));
447    }
448
449    async fn on_response<'r>(&self, req: &'r Request<'_>, response: &mut Response<'r>) {
450        // Don't touch metrics if the request didn't match a route.
451        if req.route().is_none() {
452            return;
453        }
454
455        // Don't touch metrics if the request touched a route we should ignore
456        if let Some(request_filter) = self.request_filter {
457            if !request_filter(req) {
458                return;
459            }
460        }
461
462        let endpoint = req.route().unwrap().uri.as_str();
463        let method = req.method().as_str();
464        let status = StatusCode::from(response.status().code);
465        self.http_requests_total
466            .with_label_values(&[endpoint, method, status.as_str()])
467            .inc();
468
469        let start_time = req.local_cache(|| TimerStart(None));
470        if let Some(duration) = start_time.0.map(|st| st.elapsed()) {
471            let duration_secs = duration.as_secs_f64();
472            self.http_requests_duration_seconds
473                .with_label_values(&[endpoint, method, status.as_str()])
474                .observe(duration_secs);
475        }
476    }
477}
478
479#[rocket::async_trait]
480impl Handler for PrometheusMetrics {
481    async fn handle<'r>(&self, req: &'r Request<'_>, _: Data<'r>) -> Outcome<'r> {
482        // Gather the metrics.
483        let mut buffer = vec![];
484        let encoder = TextEncoder::new();
485        encoder
486            .encode(&self.custom_registry.gather(), &mut buffer)
487            .unwrap();
488        encoder
489            .encode(&self.rocket_registry.gather(), &mut buffer)
490            .unwrap();
491        let body = String::from_utf8(buffer).unwrap();
492        Outcome::from(
493            req,
494            (
495                ContentType::new("text", "plain")
496                    .with_params([("version", "0.0.4"), ("charset", "utf-8")]),
497                body,
498            ),
499        )
500    }
501}
502
503impl From<PrometheusMetrics> for Vec<Route> {
504    fn from(other: PrometheusMetrics) -> Self {
505        vec![Route::new(Method::Get, "/", other)]
506    }
507}
508
509#[cfg(test)]
510mod test {
511    use super::PrometheusMetrics;
512
513    #[test]
514    fn test_multiple_instantiations() {
515        let _pm1 = PrometheusMetrics::with_default_registry();
516        let _pm2 = PrometheusMetrics::with_default_registry();
517    }
518}