feat: add status code for failure metric

Signed-off-by: Gustavo Inacio <[email protected]>

Author: gusinacio

crates/service/src/metrics.rs

@@ -30,7 +30,7 @@ lazy_static! {
     pub static ref HANDLER_FAILURE: CounterVec = register_counter_vec!(
         "indexer_query_handler_failed_total",
         "Failed queries to handler",
-        &["deployment"]
+        &["deployment", "status_code"]
     ).unwrap();
 
     /// Metric registered in global registry for

crates/service/src/middleware/labels.rs

@@ -8,7 +8,7 @@ use thegraph_core::DeploymentId;
 
 use super::{
     allocation::Allocation,
-    prometheus_metrics::{MetricLabelProvider, MetricLabels},
+    prometheus_metrics::{FailureLabels, HandlerLabels, MetricLabelProvider},
     sender::Sender,
 };
 
@@ -27,6 +27,23 @@ pub struct SenderAllocationDeploymentLabels {
     deployment_id: Option<String>,
 }
 
+#[derive(Clone, Default)]
+pub struct DeploymentLabels {
+    deployment_id: Option<String>,
+}
+
+impl MetricLabelProvider for DeploymentLabels {
+    fn get_labels(&self) -> Vec<&str> {
+        let mut list = vec![];
+        if let Some(deployment_id) = &self.deployment_id {
+            list.push(deployment_id.as_str());
+        } else {
+            list.push(NO_DEPLOYMENT_ID);
+        }
+        list
+    }
+}
+
 impl MetricLabelProvider for SenderAllocationDeploymentLabels {
     fn get_labels(&self) -> Vec<&str> {
         let mut list = vec![];
@@ -68,12 +85,15 @@ pub async fn labels_middleware(mut request: Request, next: Next) -> Response {
         .get::<DeploymentId>()
         .map(|s| s.clone().to_string());
 
-    let labels: MetricLabels = Arc::new(SenderAllocationDeploymentLabels {
+    let handler_labels = HandlerLabels(Arc::new(SenderAllocationDeploymentLabels {
         sender,
         allocation,
-        deployment_id,
-    });
-    request.extensions_mut().insert(labels);
+        deployment_id: deployment_id.clone(),
+    }));
+
+    let failure_labels = FailureLabels(Arc::new(DeploymentLabels { deployment_id }));
+    request.extensions_mut().insert(handler_labels);
+    request.extensions_mut().insert(failure_labels);
 
     next.run(request).await
 }

crates/service/src/middleware/prometheus_metrics.rs

@@ -14,6 +14,13 @@ use std::{
 };
 use tower::{Layer, Service};
 
+use crate::error::StatusCodeExt;
+
+#[derive(Clone)]
+pub struct HandlerLabels(pub MetricLabels);
+#[derive(Clone)]
+pub struct FailureLabels(pub MetricLabels);
+
 pub type MetricLabels = Arc<dyn MetricLabelProvider + 'static + Send + Sync>;
 
 pub trait MetricLabelProvider {
@@ -61,6 +68,7 @@ impl<S, ReqBody> Service<Request<ReqBody>> for PrometheusMetricsMiddleware<S>
 where
     S: Service<Request<ReqBody>> + Clone + 'static,
     ReqBody: 'static,
+    Result<S::Response, S::Error>: StatusCodeExt,
 {
     type Response = S::Response;
     type Error = S::Error;
@@ -71,12 +79,14 @@ where
     }
 
     fn call(&mut self, request: Request<ReqBody>) -> PrometheusMetricsFuture<S::Future> {
-        let labels = request.extensions().get::<MetricLabels>().cloned();
+        let handler_labels = request.extensions().get::<HandlerLabels>().cloned();
+        let failure_labels = request.extensions().get::<FailureLabels>().cloned();
         PrometheusMetricsFuture {
             timer: None,
             histogram: self.histogram.clone(),
             failure: self.failure.clone(),
-            labels,
+            handler_labels,
+            failure_labels,
             fut: self.inner.call(request),
         }
     }
@@ -90,7 +100,8 @@ pub struct PrometheusMetricsFuture<F> {
     histogram: prometheus::HistogramVec,
     failure: prometheus::CounterVec,
 
-    labels: Option<MetricLabels>,
+    handler_labels: Option<HandlerLabels>,
+    failure_labels: Option<FailureLabels>,
 
     #[pin]
     fut: F,
@@ -99,32 +110,41 @@ pub struct PrometheusMetricsFuture<F> {
 impl<F, R, E> Future for PrometheusMetricsFuture<F>
 where
     F: Future<Output = Result<R, E>>,
+    Result<R, E>: StatusCodeExt,
 {
     type Output = F::Output;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = self.project();
-        let Some(labels) = &this.labels else {
+        if this.handler_labels.is_none() && this.failure_labels.is_none() {
             return this.fut.poll(cx);
-        };
+        }
 
         if this.timer.is_none() {
-            // Start timer so we can track duration of request.
-            let duration_metric = this.histogram.with_label_values(&labels.get_labels());
-            *this.timer = Some(duration_metric.start_timer());
+            if let Some(HandlerLabels(labels)) = &this.handler_labels {
+                // Start timer so we can track duration of request.
+                let duration_metric = this.histogram.with_label_values(&labels.get_labels());
+                *this.timer = Some(duration_metric.start_timer());
+            }
         }
 
         match this.fut.poll(cx) {
             Poll::Ready(result) => {
-                if result.is_err() {
-                    let _ = this
-                        .failure
-                        .get_metric_with_label_values(&labels.get_labels());
+                let status_code = result.status_code();
+                if !status_code.is_success() {
+                    if let Some(FailureLabels(labels)) = &this.failure_labels {
+                        let mut labels = labels.get_labels();
+                        labels.push(status_code.as_str());
+                        this.failure
+                            .get_metric_with_label_values(&labels)
+                            .expect("Couldn't register metric");
+                    }
                 }
                 // Record the duration of this request.
                 if let Some(timer) = this.timer.take() {
                     timer.observe_duration();
                 }
+
                 Poll::Ready(result)
             }
             Poll::Pending => Poll::Pending,
@@ -141,9 +161,13 @@ mod tests {
         http::{Request, Response},
     };
     use prometheus::core::Collector;
+    use reqwest::StatusCode;
     use tower::{Service, ServiceBuilder, ServiceExt};
 
-    use crate::middleware::prometheus_metrics::{MetricLabels, PrometheusMetricsMiddlewareLayer};
+    use crate::{
+        error::StatusCodeExt,
+        middleware::prometheus_metrics::{MetricLabels, PrometheusMetricsMiddlewareLayer},
+    };
 
     use super::MetricLabelProvider;
 
@@ -153,12 +177,22 @@ mod tests {
             vec!["label1,", "label2", "label3"]
         }
     }
-    async fn handle(_: Request<Body>) -> anyhow::Result<Response<Body>> {
+
+    #[derive(Debug)]
+    struct ErrorResponse;
+
+    impl StatusCodeExt for ErrorResponse {
+        fn status_code(&self) -> StatusCode {
+            StatusCode::INTERNAL_SERVER_ERROR
+        }
+    }
+
+    async fn handle(_: Request<Body>) -> Result<Response<Body>, ErrorResponse> {
         Ok(Response::new(Body::default()))
     }
 
-    async fn handle_err(_: Request<Body>) -> anyhow::Result<Response<Body>> {
-        Err(anyhow::anyhow!("Error"))
+    async fn handle_err(_: Request<Body>) -> Result<Response<Body>, ErrorResponse> {
+        Err(ErrorResponse)
     }
 
     #[tokio::test]
@@ -175,7 +209,7 @@ mod tests {
         let failure_metric = prometheus::register_counter_vec_with_registry!(
             "failure_metric",
             "Test",
-            &["deployment", "sender", "allocation"],
+            &["deployment", "sender", "allocation", "status_code"],
             registry,
         )
         .unwrap();