metrics_example/
http_service.rs

1// Copyright 2022 Protocol Labs.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21use std::{
22    net::SocketAddr,
23    sync::{Arc, Mutex},
24};
25
26use axum::{extract::State, http::StatusCode, response::IntoResponse, routing::get, Router};
27use prometheus_client::{encoding::text::encode, registry::Registry};
28use tokio::net::TcpListener;
29
30const METRICS_CONTENT_TYPE: &str = "application/openmetrics-text;charset=utf-8;version=1.0.0";
31
32pub(crate) async fn metrics_server(registry: Registry) -> Result<(), std::io::Error> {
33    // Serve on localhost.
34    let addr: SocketAddr = ([127, 0, 0, 1], 0).into();
35    let service = MetricService::new(registry);
36    let server = Router::new()
37        .route("/metrics", get(respond_with_metrics))
38        .with_state(service);
39    let tcp_listener = TcpListener::bind(addr).await?;
40    let local_addr = tcp_listener.local_addr()?;
41    tracing::info!(metrics_server=%format!("http://{}/metrics", local_addr));
42    axum::serve(tcp_listener, server.into_make_service()).await?;
43    Ok(())
44}
45
46#[derive(Clone)]
47pub(crate) struct MetricService {
48    reg: Arc<Mutex<Registry>>,
49}
50
51async fn respond_with_metrics(state: State<MetricService>) -> impl IntoResponse {
52    let mut sink = String::new();
53    let reg = state.get_reg();
54    encode(&mut sink, &reg.lock().unwrap()).unwrap();
55
56    (
57        StatusCode::OK,
58        [(axum::http::header::CONTENT_TYPE, METRICS_CONTENT_TYPE)],
59        sink,
60    )
61}
62
63type SharedRegistry = Arc<Mutex<Registry>>;
64
65impl MetricService {
66    fn new(registry: Registry) -> Self {
67        Self {
68            reg: Arc::new(Mutex::new(registry)),
69        }
70    }
71
72    fn get_reg(&self) -> SharedRegistry {
73        Arc::clone(&self.reg)
74    }
75}