From cb01d445f12f1de847771a20e89287910949aecc Mon Sep 17 00:00:00 2001
From: eternal-flame-AD <yume@yumechi.jp>
Date: Sat, 23 Nov 2024 12:30:31 -0600
Subject: [PATCH] add histogram metrics

Signed-off-by: eternal-flame-AD <yume@yumechi.jp>
---
 src/main.rs | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 4c00e6c..5656285 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,9 +6,11 @@ use axum::{
 };
 use clap::Parser;
 use futures::TryFutureExt;
-use prometheus::core::{GenericCounterVec, GenericGauge};
 #[cfg(feature = "metrics")]
-use prometheus::{core::AtomicU64, Opts};
+use prometheus::{
+    core::{AtomicU64, GenericCounterVec, GenericGauge},
+    Opts,
+};
 use tokio::sync::mpsc;
 use tower_service::Service;
 use yumechi_no_kuni_proxy_worker::{
@@ -107,16 +109,28 @@ fn main() {
         )
         .expect("Failed to create prometheus counter");
 
+        let histo_request_timing = prometheus::HistogramVec::new(
+            prometheus::HistogramOpts::new("misskey_media_proxy_request_timing", "Request timing")
+                .buckets(vec![0.01, 0.02, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0]),
+            &["status", "postprocess"],
+        )
+        .expect("Failed to create prometheus histogram");
+
         reg.register(Box::new(counter_requests_served.clone()))
             .expect("Failed to register prometheus counter");
 
         reg.register(Box::new(counter_requests_received.clone()))
             .expect("Failed to register prometheus counter");
 
+        reg.register(Box::new(histo_request_timing.clone()))
+            .expect("Failed to register prometheus histogram");
+
         router = router
             .layer(middleware::from_fn(move |req, mut next: Next| {
                 counter_requests_received.with_label_values(&[]).inc();
                 let counter_requests_served = counter_requests_served.clone();
+                let histo_request_timing = histo_request_timing.clone();
+                let begin = std::time::Instant::now();
                 next.call(req)
                     .map_ok(move |res| {
                         let status = match res.status().as_u16() {
@@ -150,6 +164,12 @@ fn main() {
                         counter_requests_served
                             .with_label_values(&[&status, post_process])
                             .inc();
+
+                        let elapsed = begin.elapsed();
+                        histo_request_timing
+                            .with_label_values(&[&status, post_process])
+                            .observe(elapsed.as_secs_f64());
+
                         res
                     })
                     .map_err(|e: Infallible| e)