chore: move task to MetricsRunner

Author: sa-idx-admin

src/core.rs

@@ -134,7 +134,7 @@ pub async fn main(cli: &Cli) -> Result<(), Error> {
         custom_domain_providers,
         &mut tasks,
         http_client.clone(),
-        route_provider,
+        Arc::clone(&route_provider),
         &registry,
         clickhouse.clone(),
         vector.clone(),
@@ -183,7 +183,7 @@ pub async fn main(cli: &Cli) -> Result<(), Error> {
 
     // Setup metrics
     if let Some(addr) = cli.metrics.metrics_listen {
-        let router = metrics::setup(&registry, &mut tasks);
+        let router = metrics::setup(&registry, &mut tasks, route_provider);
 
         let srv = Arc::new(http::Server::new(
             http::server::Addr::Tcp(addr),

src/metrics/mod.rs

@@ -16,6 +16,7 @@ use axum::{
     Router,
 };
 use http::header::{CONTENT_TYPE, ORIGIN, REFERER, USER_AGENT};
+use ic_agent::agent::route_provider::RouteProvider;
 use ic_bn_lib::{
     http::{
         body::CountingBody,
@@ -54,9 +55,9 @@ pub const HTTP_DURATION_BUCKETS: &[f64] = &[0.05, 0.2, 1.0, 2.0];
 pub const HTTP_REQUEST_SIZE_BUCKETS: &[f64] = &[128.0, KB, 2.0 * KB, 4.0 * KB, 8.0 * KB];
 pub const HTTP_RESPONSE_SIZE_BUCKETS: &[f64] = &[1.0 * KB, 8.0 * KB, 64.0 * KB, 256.0 * KB];
 
-pub fn setup(registry: &Registry, tasks: &mut TaskManager) -> Router {
+pub fn setup(registry: &Registry, tasks: &mut TaskManager, route_provider: Arc<dyn RouteProvider>) -> Router {
     let cache = Arc::new(runner::MetricsCache::new());
-    let runner = Arc::new(runner::MetricsRunner::new(cache.clone(), registry));
+    let runner = Arc::new(runner::MetricsRunner::new(cache.clone(), registry, route_provider));
     tasks.add_interval("metrics_runner", runner, Duration::from_secs(5));
 
     Router::new()

src/metrics/runner.rs

@@ -6,8 +6,9 @@ use async_trait::async_trait;
 use axum::{extract::State, response::IntoResponse};
 use bytes::{BufMut, Bytes, BytesMut};
 use http::header::CONTENT_TYPE;
+use ic_agent::agent::route_provider::RouteProvider;
 use ic_bn_lib::tasks::Run;
-use prometheus::{register_int_gauge_with_registry, Encoder, IntGauge, Registry, TextEncoder};
+use prometheus::{Encoder, IntGauge, Registry, TextEncoder, register_int_gauge_with_registry};
 use tikv_jemalloc_ctl::{epoch, stats};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, warn};
@@ -31,15 +32,22 @@ pub struct MetricsRunner {
     metrics_cache: Arc<MetricsCache>,
     registry: Registry,
     encoder: TextEncoder,
-
+    route_provider: Arc<dyn RouteProvider>,
     // Metrics
     mem_allocated: IntGauge,
     mem_resident: IntGauge,
+    // API boundary nodes metrics
+    total_api_boundary_nodes: IntGauge,
+    healthy_api_boundary_nodes: IntGauge,
 }
 
 // Snapshots & encodes the metrics for the handler to export
 impl MetricsRunner {
-    pub fn new(metrics_cache: Arc<MetricsCache>, registry: &Registry) -> Self {
+    pub fn new(
+        metrics_cache: Arc<MetricsCache>,
+        registry: &Registry,
+        route_provider: Arc<dyn RouteProvider>,
+    ) -> Self {
         let mem_allocated = register_int_gauge_with_registry!(
             format!("memory_allocated"),
             format!("Allocated memory in bytes"),
@@ -54,12 +62,29 @@ impl MetricsRunner {
         )
         .unwrap();
 
+        let total_api_boundary_nodes = register_int_gauge_with_registry!(
+            format!("total_api_boundary_nodes"),
+            format!("Total number of existing API boundary nodes (both healthy and unhealthy)."),
+            registry
+        )
+        .unwrap();
+
+        let healthy_api_boundary_nodes = register_int_gauge_with_registry!(
+            format!("healthy_api_boundary_nodes"),
+            format!("Number of currently healthy API boundary nodes"),
+            registry,
+        )
+        .unwrap();
+
         Self {
             metrics_cache,
             registry: registry.clone(),
             encoder: TextEncoder::new(),
             mem_allocated,
             mem_resident,
+            route_provider,
+            healthy_api_boundary_nodes,
+            total_api_boundary_nodes,
         }
     }
 }
@@ -85,6 +110,12 @@ impl MetricsRunner {
             .buffer
             .store(Arc::new(buffer.into_inner().freeze()));
 
+        // Update API boundary nodes stats
+        let stats = self.route_provider.routes_stats();
+        self.total_api_boundary_nodes.set(stats.total as i64);
+        self.healthy_api_boundary_nodes
+            .set(stats.healthy.unwrap_or(0) as i64);
+
         Ok(())
     }
 }

src/routing/ic/route_provider.rs

@@ -1,28 +1,23 @@
-use std::{sync::Arc, time::Duration};
+use std::sync::Arc;
+
 use anyhow::anyhow;
-use async_trait::async_trait;
 use candid::Principal;
 use ic_agent::agent::http_transport::reqwest_transport::reqwest::Client as AgentClient;
 use ic_agent::agent::route_provider::{
-    RoundRobinRouteProvider, RouteProvider,
     dynamic_routing::{
         dynamic_route_provider::DynamicRouteProviderBuilder, node::Node,
         snapshot::latency_based_routing::LatencyRoutingSnapshot,
     },
+    RoundRobinRouteProvider, RouteProvider,
 };
-use ic_bn_lib::tasks::Run;
-use prometheus::{IntGauge, Registry, register_int_gauge_with_registry};
-use tokio_util::sync::CancellationToken;
 use tracing::info;
 use url::Url;
 
 use crate::routing::ic::{
-    health_check::{CHECK_TIMEOUT, HealthChecker},
-    nodes_fetcher::{MAINNET_ROOT_SUBNET_ID, NodesFetcher},
+    health_check::{HealthChecker, CHECK_TIMEOUT},
+    nodes_fetcher::{NodesFetcher, MAINNET_ROOT_SUBNET_ID},
 };
 
-pub const API_BOUNDARY_NODES_STATS_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
-
 pub async fn setup_route_provider(
     urls: &[Url],
     ic_use_discovery: bool,
@@ -70,56 +65,3 @@ pub async fn setup_route_provider(
 
     Ok(route_provider)
 }
-
-struct ApiBoundaryNodesMetrics {
-    total_nodes: IntGauge,
-    healthy_nodes: IntGauge,
-}
-
-pub struct ApiBoundaryNodesStats {
-    route_provider: Arc<dyn RouteProvider>,
-    metrics: ApiBoundaryNodesMetrics,
-}
-
-impl ApiBoundaryNodesMetrics {
-    fn new(registry: &Registry) -> Self {
-        Self {
-            total_nodes: register_int_gauge_with_registry!(
-                format!("total_api_boundary_nodes"),
-                format!(
-                    "Total number of existing API boundary nodes (both healthy and unhealthy)."
-                ),
-                registry
-            )
-            .unwrap(),
-
-            healthy_nodes: register_int_gauge_with_registry!(
-                format!("healthy_api_boundary_nodes"),
-                format!("Number of currently healthy API boundary nodes"),
-                registry,
-            )
-            .unwrap(),
-        }
-    }
-}
-
-impl ApiBoundaryNodesStats {
-    pub fn new(route_provider: Arc<dyn RouteProvider>, registry: &Registry) -> Self {
-        Self {
-            route_provider,
-            metrics: ApiBoundaryNodesMetrics::new(registry),
-        }
-    }
-}
-
-#[async_trait]
-impl Run for ApiBoundaryNodesStats {
-    async fn run(&self, _: CancellationToken) -> Result<(), anyhow::Error> {
-        let stats = self.route_provider.routes_stats();
-        self.metrics.total_nodes.set(stats.total as i64);
-        self.metrics
-            .healthy_nodes
-            .set(stats.healthy.unwrap_or(0) as i64);
-        Ok(())
-    }
-}

src/routing/mod.rs

@@ -21,7 +21,6 @@ use domain::{CustomDomainStorage, DomainResolver, ProvidesCustomDomains};
 use fqdn::FQDN;
 use http::{method::Method, uri::PathAndQuery, StatusCode, Uri};
 use ic_agent::agent::route_provider::RouteProvider;
-use ic::route_provider::{ApiBoundaryNodesStats, API_BOUNDARY_NODES_STATS_REFRESH_INTERVAL};
 use ic_bn_lib::{
     http::{
         cache::{Cache, KeyExtractorUriRange, Opts},
@@ -319,12 +318,6 @@ pub fn setup_router(
     // Prepare the HTTP->IC library
     let client = ic::setup(cli, http_client.clone(), route_provider.clone())?;
 
-    tasks.add_interval(
-        "api_boundary_nodes_stats",
-        Arc::new(ApiBoundaryNodesStats::new(route_provider.clone(), registry)),
-        API_BOUNDARY_NODES_STATS_REFRESH_INTERVAL,
-    );
-
     // Prepare the states
     let state_handler = Arc::new(handler::HandlerState::new(
         client,