kora_lib/metrics/
mod.rs

1pub mod balance;
2pub mod handler;
3pub mod middleware;
4
5pub use balance::BalanceTracker;
6pub use handler::{MetricsHandlerLayer, MetricsHandlerService};
7pub use middleware::{HttpMetricsLayer, HttpMetricsService};
8pub use prometheus;
9use solana_client::nonblocking::rpc_client::RpcClient;
10use tokio::task::JoinHandle;
11
12use crate::{config::MetricsConfig, state::get_config};
13use jsonrpsee::{
14    server::{ServerBuilder, ServerHandle},
15    RpcModule,
16};
17use prometheus::{Encoder, TextEncoder};
18use std::{net::SocketAddr, sync::Arc};
19
20pub struct MetricsLayers {
21    pub http_metrics_layer: Option<HttpMetricsLayer>,
22    pub metrics_handler_layer: Option<MetricsHandlerLayer>,
23}
24
25fn get_metrics_layers(metrics_config: &MetricsConfig) -> Option<MetricsLayers> {
26    if metrics_config.enabled {
27        Some(MetricsLayers {
28            http_metrics_layer: Some(HttpMetricsLayer::new()),
29            metrics_handler_layer: Some(MetricsHandlerLayer::new(metrics_config.endpoint.clone())),
30        })
31    } else {
32        None
33    }
34}
35
36pub async fn run_metrics_server_if_required(
37    rpc_port: u16,
38    rpc_client: Arc<RpcClient>,
39) -> Result<(Option<ServerHandle>, Option<MetricsLayers>, Option<JoinHandle<()>>), anyhow::Error> {
40    let metrics_config = get_config()?.metrics.clone();
41
42    if !metrics_config.enabled {
43        return Ok((None, None, None));
44    }
45
46    // Initialize balance tracker if metrics are enabled and start background tracking
47    let balance_tracker_handle = if let Err(e) = BalanceTracker::init().await {
48        log::warn!("Failed to initialize balance tracker: {e}");
49        // Don't fail metrics server startup if balance tracker fails to initialize
50        None
51    } else {
52        // Start background balance tracking (only if initialized worked)
53        BalanceTracker::start_background_tracking(rpc_client).await
54    };
55
56    // If running on the same port as the RPC server, we don't need to run a separate metrics server
57    if metrics_config.port == rpc_port {
58        log::info!("Metrics endpoint enabled at {} on RPC server", metrics_config.endpoint);
59        return Ok((None, get_metrics_layers(&metrics_config), balance_tracker_handle));
60    }
61
62    let addr = SocketAddr::from(([0, 0, 0, 0], metrics_config.port));
63    log::info!("Metrics server started on {addr}, port {}", metrics_config.port);
64    log::info!("Metrics endpoint: {}", metrics_config.endpoint);
65
66    // Simple middleware stack for metrics-only server
67    let middleware = tower::ServiceBuilder::new()
68        .layer(MetricsHandlerLayer::new(metrics_config.endpoint.clone()));
69
70    // Configure and build the server
71    let server =
72        ServerBuilder::default().set_middleware(middleware).http_only().build(addr).await?;
73
74    // Empty RPC module since we only serve metrics
75    let module = RpcModule::new(());
76
77    let metrics_server_handle = server
78        .start(module)
79        .map_err(|e| anyhow::anyhow!("Failed to start metrics server: {}", e))?;
80
81    // Return both the metrics server handle AND the HTTP metrics middleware for the main RPC server
82    // The HTTP middleware needs to be on the RPC server to collect metrics, even if metrics are served separately
83    let metrics_layers = MetricsLayers {
84        http_metrics_layer: Some(HttpMetricsLayer::new()), // Collect metrics on RPC server
85        metrics_handler_layer: None, // Don't serve metrics on RPC server (separate server handles this)
86    };
87
88    Ok((Some(metrics_server_handle), Some(metrics_layers), balance_tracker_handle))
89}
90
91/// Gather all Prometheus metrics and encode them in text format
92pub fn gather() -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
93    let encoder = TextEncoder::new();
94    let metric_families = prometheus::gather();
95    let mut buffer = Vec::new();
96    encoder.encode(&metric_families, &mut buffer)?;
97    String::from_utf8(buffer).map_err(Into::into)
98}
kora_lib/metrics/mod.rs

kora_lib/metrics/
mod.rs
