src/prometheus_metrics.rs

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