refactor: add inject labels middleware

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

Author: gusinacio

crates/service/src/middleware.rs

@@ -3,5 +3,7 @@
 
 mod inject_allocation;
 mod inject_deployment;
+mod inject_labels;
 mod inject_receipt;
 mod inject_sender;
+mod metrics;

crates/service/src/middleware/inject_labels.rs

@@ -0,0 +1,137 @@
+// Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Injects Metric Labels
+//!
+//! Require Sender, Allocation and Deployment extensions
+
+use std::sync::Arc;
+
+use axum::{extract::Request, middleware::Next, response::Response};
+use thegraph_core::DeploymentId;
+
+use super::{
+    inject_allocation::Allocation,
+    inject_sender::Sender,
+    metrics::{MetricLabelProvider, MetricLabels},
+};
+
+const NO_DEPLOYMENT_ID: &str = "no-deployment";
+const NO_ALLOCATION: &str = "no-allocation";
+const NO_SENDER: &str = "no-sender";
+
+#[derive(Clone, Default)]
+pub struct SenderAllocationDeploymentLabels {
+    sender: Option<String>,
+    allocation: Option<String>,
+    deployment_id: Option<String>,
+}
+
+impl MetricLabelProvider for SenderAllocationDeploymentLabels {
+    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);
+        }
+        if let Some(allocation) = &self.allocation {
+            list.push(allocation.as_str());
+        } else {
+            list.push(NO_ALLOCATION);
+        }
+        if let Some(sender) = &self.sender {
+            list.push(sender.as_str());
+        } else {
+            list.push(NO_SENDER);
+        }
+        list
+    }
+}
+
+pub async fn labels_middleware(mut request: Request, next: Next) -> Response {
+    let sender: Option<String> = request
+        .extensions()
+        .get::<Sender>()
+        .map(|s| s.clone().into());
+
+    let allocation: Option<String> = request
+        .extensions()
+        .get::<Allocation>()
+        .map(|s| s.clone().into());
+
+    let deployment_id: Option<String> = request
+        .extensions()
+        .get::<DeploymentId>()
+        .map(|s| s.clone().to_string());
+
+    let labels: MetricLabels = Arc::new(SenderAllocationDeploymentLabels {
+        sender,
+        allocation,
+        deployment_id,
+    });
+    request.extensions_mut().insert(labels);
+
+    next.run(request).await
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::middleware::{
+        inject_allocation::Allocation, inject_sender::Sender, metrics::MetricLabels,
+    };
+
+    use super::labels_middleware;
+
+    use alloy::primitives::Address;
+    use axum::{
+        body::Body,
+        http::{Extensions, Request},
+        middleware::from_fn,
+        routing::get,
+        Router,
+    };
+    use reqwest::StatusCode;
+    use test_assets::ESCROW_SUBGRAPH_DEPLOYMENT;
+    use tower::ServiceExt;
+
+    #[tokio::test]
+    async fn test_receipt_middleware() {
+        let middleware = from_fn(labels_middleware);
+
+        let deployment = *ESCROW_SUBGRAPH_DEPLOYMENT;
+        let sender = Address::ZERO;
+        let allocation = Address::ZERO;
+
+        let handle = move |extensions: Extensions| async move {
+            let metrics = extensions
+                .get::<MetricLabels>()
+                .expect("Should decode tap receipt");
+            assert_eq!(
+                metrics.get_labels(),
+                vec![
+                    &deployment.to_string(),
+                    &allocation.to_string(),
+                    &sender.to_string(),
+                ]
+            );
+            Body::empty()
+        };
+
+        let app = Router::new().route("/", get(handle)).layer(middleware);
+
+        let res = app
+            .oneshot(
+                Request::builder()
+                    .uri("/")
+                    .extension(Sender(sender))
+                    .extension(deployment)
+                    .extension(Allocation(allocation))
+                    .body(Body::empty())
+                    .unwrap(),
+            )
+            .await
+            .unwrap();
+        assert_eq!(res.status(), StatusCode::OK);
+    }
+}

crates/service/src/middleware/inject_sender.rs

@@ -1,7 +1,7 @@
 // Copyright 2023-, Edge & Node, GraphOps, and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use alloy::dyn_abi::Eip712Domain;
+use alloy::{dyn_abi::Eip712Domain, primitives::Address};
 use axum::{
     extract::{Request, State},
     middleware::Next,
@@ -20,11 +20,11 @@ pub struct SenderState {
 }
 
 #[derive(Clone)]
-pub struct Sender(String);
+pub struct Sender(pub Address);
 
 impl From<Sender> for String {
     fn from(value: Sender) -> Self {
-        value.0
+        value.0.to_string()
     }
 }
 
@@ -39,7 +39,7 @@ pub async fn sender_middleware(
             .escrow_accounts
             .borrow()
             .get_sender_for_signer(&signer)?;
-        request.extensions_mut().insert(Sender(sender.to_string()));
+        request.extensions_mut().insert(Sender(sender));
     }
 
     Ok(next.run(request).await)
@@ -60,7 +60,9 @@ mod tests {
     };
     use indexer_monitor::EscrowAccounts;
     use reqwest::StatusCode;
-    use test_assets::{create_signed_receipt, ESCROW_ACCOUNTS_BALANCES, ESCROW_ACCOUNTS_SENDERS_TO_SIGNERS};
+    use test_assets::{
+        create_signed_receipt, ESCROW_ACCOUNTS_BALANCES, ESCROW_ACCOUNTS_SENDERS_TO_SIGNERS,
+    };
     use tokio::sync::watch;
     use tower::ServiceExt;
 
@@ -80,7 +82,7 @@ mod tests {
 
         async fn handle(extensions: Extensions) -> Body {
             let sender = extensions.get::<Sender>().expect("Should contain sender");
-            assert_eq!(sender.0, test_assets::TAP_SENDER.1.to_string());
+            assert_eq!(sender.0, test_assets::TAP_SENDER.1);
             Body::empty()
         }
 

crates/service/src/middleware/metrics.rs

@@ -0,0 +1,8 @@
+use std::sync::Arc;
+
+
+pub type MetricLabels = Arc<dyn MetricLabelProvider + 'static + Send + Sync>;
+
+pub trait MetricLabelProvider {
+    fn get_labels(&self) -> Vec<&str>;
+}