Merge pull request #152 from semiotic-ai/aasseman/issue146

feat(aggregator): add prometheus metrics

Author: aasseman

tap_aggregator/Cargo.toml

@@ -23,6 +23,10 @@ serde_json = { version = "1.0.96", features = ["raw_value"] }
 strum = { version = "0.24.1", features = ["strum_macros", "derive"] }
 tracing-subscriber = { version = "0.3.17" }
 log = "0.4.19"
+prometheus = "0.13.3"
+axum = "0.6.18"
+futures-util = "0.3.28"
+lazy_static = "1.4.0"
 
 [dev-dependencies]
 jsonrpsee = { version = "0.18.0", features = ["http-client", "jsonrpsee-core"] }

tap_aggregator/src/lib.rs

@@ -5,4 +5,5 @@ pub mod aggregator;
 pub mod api_versioning;
 pub mod error_codes;
 pub mod jsonrpsee_helpers;
+pub mod metrics;
 pub mod server;

tap_aggregator/src/main.rs

@@ -9,6 +9,7 @@ use ethers_signers::{coins_bip39::English, MnemonicBuilder};
 use tokio::signal::unix::{signal, SignalKind};
 
 use log::{debug, info, warn};
+use tap_aggregator::metrics;
 use tap_aggregator::server;
 
 #[derive(Parser, Debug)]
@@ -36,6 +37,11 @@ struct Args {
     /// Defaults to 32.
     #[arg(long, default_value_t = 32, env = "TAP_MAX_CONNECTIONS")]
     max_connections: u32,
+
+    /// Metrics server port.
+    /// Defaults to 5000.
+    #[arg(long, default_value_t = 5000, env = "TAP_METRICS_PORT")]
+    metrics_port: u16,
 }
 
 #[tokio::main]
@@ -50,12 +56,17 @@ async fn main() -> Result<()> {
     let args = Args::parse();
     debug!("Settings: {:?}", args);
 
+    // Start the metrics server.
+    // We just let it gracelessly get killed at the end of main()
+    tokio::spawn(metrics::run_server(args.metrics_port));
+
     // Create a wallet from the mnemonic.
     let wallet = MnemonicBuilder::<English>::default()
         .phrase(args.mnemonic.as_str())
         .build()?;
 
     // Start the JSON-RPC server.
+    // This await is non-blocking
     let (handle, _) = server::run_server(
         args.port,
         wallet,

tap_aggregator/src/metrics.rs

@@ -0,0 +1,59 @@
+// Copyright 2023-, Semiotic AI, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::{net::SocketAddr, panic};
+
+use axum::{http::StatusCode, response::IntoResponse, routing::get, Router, Server};
+use futures_util::FutureExt;
+use jsonrpsee::tracing::error;
+use log::{debug, info};
+use prometheus::TextEncoder;
+
+async fn handler_metrics() -> (StatusCode, String) {
+    let metric_families = prometheus::gather();
+    let encoder = TextEncoder::new();
+
+    match encoder.encode_to_string(&metric_families) {
+        Ok(s) => (StatusCode::OK, s),
+        Err(e) => {
+            error!("Error encoding metrics: {}", e);
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Error encoding metrics: {}", e),
+            )
+        }
+    }
+}
+
+async fn handler_404() -> impl IntoResponse {
+    (StatusCode::NOT_FOUND, "404 Not Found")
+}
+
+async fn _run_server(port: u16) {
+    let app = Router::new()
+        .route("/metrics", get(handler_metrics))
+        .fallback(handler_404);
+    let addr = SocketAddr::from(([0, 0, 0, 0], port));
+    let server = Server::bind(&addr).serve(app.into_make_service());
+
+    info!("Metrics server listening on {}", addr);
+
+    let res = server.await;
+
+    debug!("Metrics server stopped");
+
+    if let Err(err) = res {
+        panic!("Metrics server error: {:#?}", err);
+    };
+}
+
+pub async fn run_server(port: u16) {
+    // Code here is to abort program if there is a panic in _run_server
+    // Otherwise, when spawning the task, the panic will be silently ignored
+    let res = panic::AssertUnwindSafe(_run_server(port))
+        .catch_unwind()
+        .await;
+    if res.is_err() {
+        std::process::abort();
+    }
+}

tap_aggregator/src/server.rs

@@ -10,6 +10,8 @@ use jsonrpsee::{
     server::ServerBuilder,
     {core::async_trait, server::ServerHandle},
 };
+use lazy_static::lazy_static;
+use prometheus::{register_counter, register_int_counter, Counter, IntCounter};
 
 use crate::aggregator::check_and_aggregate_receipts;
 use crate::api_versioning::{
@@ -23,6 +25,51 @@ use tap_core::{
     receipt_aggregate_voucher::ReceiptAggregateVoucher, tap_receipt::Receipt,
 };
 
+// Register the metrics into the global metrics registry.
+lazy_static! {
+    static ref AGGREGATION_SUCCESS_COUNTER: IntCounter = register_int_counter!(
+        "aggregation_success_count",
+        "Number of successful receipt aggregation requests."
+    )
+    .unwrap();
+}
+lazy_static! {
+    static ref AGGREGATION_FAILURE_COUNTER: IntCounter = register_int_counter!(
+        "aggregation_failure_count",
+        "Number of failed receipt aggregation requests (for any reason)."
+    )
+    .unwrap();
+}
+lazy_static! {
+    static ref DEPRECATION_WARNING_COUNT: IntCounter = register_int_counter!(
+        "deprecation_warning_count",
+        "Number of deprecation warnings sent to clients."
+    )
+    .unwrap();
+}
+lazy_static! {
+    static ref VERSION_ERROR_COUNT: IntCounter = register_int_counter!(
+        "version_error_count",
+        "Number of API version errors sent to clients."
+    )
+    .unwrap();
+}
+lazy_static! {
+    static ref TOTAL_AGGREGATED_RECEIPTS: IntCounter = register_int_counter!(
+        "total_aggregated_receipts",
+        "Total number of receipts successfully aggregated."
+    )
+    .unwrap();
+}
+// Using float for the GRT value because it can somewhat easily exceed the maximum value of int64.
+lazy_static! {
+    static ref TOTAL_GRT_AGGREGATED: Counter = register_counter!(
+        "total_aggregated_grt",
+        "Total successfully aggregated GRT value (wei)."
+    )
+    .unwrap();
+}
+
 /// Generates the `RpcServer` trait that is used to define the JSON-RPC API.
 ///
 /// Note that because of the way the `rpc` macro works, we cannot document the RpcServer trait here.
@@ -80,6 +127,45 @@ fn check_api_version_deprecation(api_version: &TapRpcApiVersion) -> Option<JsonR
     }
 }
 
+async fn aggregate_receipts_(
+    api_version: String,
+    wallet: &LocalWallet,
+    receipts: Vec<EIP712SignedMessage<Receipt>>,
+    previous_rav: Option<EIP712SignedMessage<ReceiptAggregateVoucher>>,
+) -> JsonRpcResult<EIP712SignedMessage<ReceiptAggregateVoucher>> {
+    // Return an error if the API version is not supported.
+    let api_version = match parse_api_version(api_version.as_str()) {
+        Ok(v) => v,
+        Err(e) => {
+            VERSION_ERROR_COUNT.inc();
+            return Err(e);
+        }
+    };
+
+    // Add a warning if the API version is to be deprecated.
+    let mut warnings: Vec<JsonRpcWarning> = Vec::new();
+    if let Some(w) = check_api_version_deprecation(&api_version) {
+        warnings.push(w);
+        DEPRECATION_WARNING_COUNT.inc();
+    }
+
+    let res = match api_version {
+        TapRpcApiVersion::V0_0 => {
+            check_and_aggregate_receipts(&receipts, previous_rav, wallet).await
+        }
+    };
+
+    // Handle aggregation error
+    match res {
+        Ok(res) => Ok(JsonRpcResponse::warn(res, warnings)),
+        Err(e) => Err(jsonrpsee::types::ErrorObject::owned(
+            JsonRpcErrorCode::Aggregation as i32,
+            e.to_string(),
+            None::<()>,
+        )),
+    }
+}
+
 #[async_trait]
 impl RpcServer for RpcImpl {
     async fn api_versions(&self) -> JsonRpcResult<TapRpcApiVersionsInfo> {
@@ -92,29 +178,21 @@ impl RpcServer for RpcImpl {
         receipts: Vec<EIP712SignedMessage<Receipt>>,
         previous_rav: Option<EIP712SignedMessage<ReceiptAggregateVoucher>>,
     ) -> JsonRpcResult<EIP712SignedMessage<ReceiptAggregateVoucher>> {
-        // Return an error if the API version is not supported.
-        let api_version = parse_api_version(api_version.as_str())?;
-
-        // Add a warning if the API version is to be deprecated.
-        let mut warnings: Vec<JsonRpcWarning> = Vec::new();
-        if let Some(w) = check_api_version_deprecation(&api_version) {
-            warnings.push(w);
-        }
-
-        let res = match api_version {
-            TapRpcApiVersion::V0_0 => {
-                check_and_aggregate_receipts(&receipts, previous_rav, &self.wallet).await
+        // Values for Prometheus metrics
+        let receipts_grt: u128 = receipts.iter().map(|r| r.message.value).sum();
+        let receipts_count: u64 = receipts.len() as u64;
+
+        match aggregate_receipts_(api_version, &self.wallet, receipts, previous_rav).await {
+            Ok(res) => {
+                TOTAL_GRT_AGGREGATED.inc_by(receipts_grt as f64);
+                TOTAL_AGGREGATED_RECEIPTS.inc_by(receipts_count);
+                AGGREGATION_SUCCESS_COUNTER.inc();
+                Ok(res)
+            }
+            Err(e) => {
+                AGGREGATION_FAILURE_COUNTER.inc();
+                Err(e)
             }
-        };
-
-        // Handle aggregation error
-        match res {
-            Ok(res) => Ok(JsonRpcResponse::warn(res, warnings)),
-            Err(e) => Err(jsonrpsee::types::ErrorObject::owned(
-                JsonRpcErrorCode::Aggregation as i32,
-                e.to_string(),
-                None::<()>,
-            )),
         }
     }
 }

tap_integration_tests/Cargo.toml

@@ -18,6 +18,7 @@ rand = "0.8.5"
 futures = "0.3.28"
 anyhow = "1.0.71"
 tokio = "1.28.2"
+prometheus = "0.13.3"
 
 [[test]]
 name = "integration_tests"