src/prometheus_metrics.rs

   1 // TODO: should really not unwrap everywhere here....
   2 #![allow(clippy::unwrap_used)]
   3 use actix_web::{rt::System, web, App, HttpResponse, HttpServer, Responder};
   4 use lemmy_api_common::context::LemmyContext;
   5 use lemmy_utils::settings::structs::PrometheusConfig;
   6 use prometheus::{default_registry, Encoder, Gauge, Opts, TextEncoder};
   7 use std::{
   8   net::{IpAddr, Ipv4Addr},
   9   sync::Arc,
  10   thread,
  11 };
  12
  13 struct PromContext {
  14   lemmy: LemmyContext,
  15   db_pool_metrics: DbPoolMetrics,
  16 }
  17
  18 struct DbPoolMetrics {
  19   max_size: Gauge,
  20   size: Gauge,
  21   available: Gauge,
  22 }
  23
  24 static DEFAULT_BIND: IpAddr = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
  25 static DEFAULT_PORT: i32 = 10002;
  26
  27 pub fn serve_prometheus(config: Option<&PrometheusConfig>, lemmy_context: LemmyContext) {
  28   let context = Arc::new(PromContext {
  29     lemmy: lemmy_context,
  30     db_pool_metrics: create_db_pool_metrics(),
  31   });
  32
  33   let (bind, port) = match config {
  34     Some(config) => (
  35       config.bind.unwrap_or(DEFAULT_BIND),
  36       config.port.unwrap_or(DEFAULT_PORT),
  37     ),
  38     None => (DEFAULT_BIND, DEFAULT_PORT),
  39   };
  40
  41   // spawn thread that blocks on handling requests
  42   // only mapping /metrics to a handler
  43   thread::spawn(move || {
  44     let sys = System::new();
  45     sys.block_on(async {
  46       let server = HttpServer::new(move || {
  47         App::new()
  48           .app_data(web::Data::new(Arc::clone(&context)))
  49           .route("/metrics", web::get().to(metrics))
  50       })
  51       .bind((bind, port as u16))
  52       .unwrap_or_else(|_| panic!("Cannot bind to {}:{}", bind, port))
  53       .run();
  54
  55       if let Err(err) = server.await {
  56         eprintln!("Prometheus server error: {}", err);
  57       }
  58     })
  59   });
  60 }
  61
  62 // handler for the /metrics path
  63 async fn metrics(context: web::Data<Arc<PromContext>>) -> impl Responder {
  64   // collect metrics
  65   collect_db_pool_metrics(&context).await;
  66
  67   let mut buffer = Vec::new();
  68   let encoder = TextEncoder::new();
  69
  70   // gather metrics from registry and encode in prometheus format
  71   let metric_families = prometheus::gather();
  72   encoder.encode(&metric_families, &mut buffer).unwrap();
  73   let output = String::from_utf8(buffer).unwrap();
  74
  75   HttpResponse::Ok().body(output)
  76 }
  77
  78 // create lemmy_db_pool_* metrics and register them with the default registry
  79 fn create_db_pool_metrics() -> DbPoolMetrics {
  80   let metrics = DbPoolMetrics {
  81     max_size: Gauge::with_opts(Opts::new(
  82       "lemmy_db_pool_max_connections",
  83       "Maximum number of connections in the pool",
  84     ))
  85     .unwrap(),
  86     size: Gauge::with_opts(Opts::new(
  87       "lemmy_db_pool_connections",
  88       "Current number of connections in the pool",
  89     ))
  90     .unwrap(),
  91     available: Gauge::with_opts(Opts::new(
  92       "lemmy_db_pool_available_connections",
  93       "Number of available connections in the pool",
  94     ))
  95     .unwrap(),
  96   };
  97
  98   default_registry()
  99     .register(Box::new(metrics.max_size.clone()))
 100     .unwrap();
 101   default_registry()
 102     .register(Box::new(metrics.size.clone()))
 103     .unwrap();
 104   default_registry()
 105     .register(Box::new(metrics.available.clone()))
 106     .unwrap();
 107
 108   metrics
 109 }
 110
 111 async fn collect_db_pool_metrics(context: &PromContext) {
 112   let pool_status = context.lemmy.inner_pool().status();
 113   context
 114     .db_pool_metrics
 115     .max_size
 116     .set(pool_status.max_size as f64);
 117   context.db_pool_metrics.size.set(pool_status.size as f64);
 118   context
 119     .db_pool_metrics
 120     .available
 121     .set(pool_status.available as f64);
 122 }