use parking_lot::RwLock;
use prometheus::{Encoder, Opts, Registry, TextEncoder};
use serde::Serialize;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
pub struct NetworkMetrics {
connections: ConnectionMetrics,
bandwidth: BandwidthMetrics,
dht: DhtMetrics,
protocols: ProtocolMetrics,
start_time: Instant,
}
impl NetworkMetrics {
pub fn new() -> Self {
Self {
connections: ConnectionMetrics::new(),
bandwidth: BandwidthMetrics::new(),
dht: DhtMetrics::new(),
protocols: ProtocolMetrics::new(),
start_time: Instant::now(),
}
}
pub fn connections(&self) -> &ConnectionMetrics {
&self.connections
}
pub fn bandwidth(&self) -> &BandwidthMetrics {
&self.bandwidth
}
pub fn dht(&self) -> &DhtMetrics {
&self.dht
}
pub fn protocols(&self) -> &ProtocolMetrics {
&self.protocols
}
pub fn uptime(&self) -> Duration {
self.start_time.elapsed()
}
pub fn snapshot(&self) -> MetricsSnapshot {
MetricsSnapshot {
uptime_secs: self.start_time.elapsed().as_secs(),
connections: self.connections.snapshot(),
bandwidth: self.bandwidth.snapshot(),
dht: self.dht.snapshot(),
}
}
pub fn create_prometheus_registry(&self) -> Result<Registry, prometheus::Error> {
let registry = Registry::new();
let connections_established = prometheus::IntCounterVec::new(
Opts::new(
"ipfrs_connections_established_total",
"Total number of connections established",
),
&["direction"],
)?;
connections_established
.with_label_values(&["inbound"])
.inc_by(self.connections.total_inbound());
connections_established
.with_label_values(&["outbound"])
.inc_by(self.connections.total_outbound());
registry.register(Box::new(connections_established))?;
let connections_failed = prometheus::IntCounter::new(
"ipfrs_connections_failed_total",
"Total number of failed connection attempts",
)?;
connections_failed.inc_by(self.connections.total_failed());
registry.register(Box::new(connections_failed))?;
let connections_active = prometheus::IntGauge::new(
"ipfrs_connections_active",
"Number of currently active connections",
)?;
connections_active.set(self.connections.active() as i64);
registry.register(Box::new(connections_active))?;
let bytes_sent = prometheus::IntCounter::new(
"ipfrs_bytes_sent_total",
"Total bytes sent over the network",
)?;
bytes_sent.inc_by(self.bandwidth.total_sent());
registry.register(Box::new(bytes_sent))?;
let bytes_received = prometheus::IntCounter::new(
"ipfrs_bytes_received_total",
"Total bytes received from the network",
)?;
bytes_received.inc_by(self.bandwidth.total_received());
registry.register(Box::new(bytes_received))?;
let dht_queries = prometheus::IntCounterVec::new(
Opts::new("ipfrs_dht_queries_total", "Total DHT queries by status"),
&["status"],
)?;
let dht_snapshot = self.dht.snapshot();
dht_queries
.with_label_values(&["success"])
.inc_by(dht_snapshot.queries_successful);
dht_queries
.with_label_values(&["failed"])
.inc_by(dht_snapshot.queries_failed);
registry.register(Box::new(dht_queries))?;
let providers_published = prometheus::IntCounter::new(
"ipfrs_dht_providers_published_total",
"Total provider records published to DHT",
)?;
providers_published.inc_by(dht_snapshot.providers_published);
registry.register(Box::new(providers_published))?;
let providers_found = prometheus::IntCounter::new(
"ipfrs_dht_providers_found_total",
"Total providers found via DHT queries",
)?;
providers_found.inc_by(dht_snapshot.providers_found);
registry.register(Box::new(providers_found))?;
let routing_table_size = prometheus::IntGauge::new(
"ipfrs_dht_routing_table_size",
"Current DHT routing table size",
)?;
routing_table_size.set(dht_snapshot.routing_table_size as i64);
registry.register(Box::new(routing_table_size))?;
let uptime = prometheus::IntGauge::new("ipfrs_uptime_seconds", "Node uptime in seconds")?;
uptime.set(self.uptime().as_secs() as i64);
registry.register(Box::new(uptime))?;
Ok(registry)
}
pub fn export_prometheus(&self) -> Result<String, prometheus::Error> {
let registry = self.create_prometheus_registry()?;
let encoder = TextEncoder::new();
let metric_families = registry.gather();
let mut buffer = Vec::new();
encoder
.encode(&metric_families, &mut buffer)
.map_err(|e| prometheus::Error::Msg(e.to_string()))?;
String::from_utf8(buffer).map_err(|e| prometheus::Error::Msg(e.to_string()))
}
}
impl Default for NetworkMetrics {
fn default() -> Self {
Self::new()
}
}
pub struct ConnectionMetrics {
connections_established: AtomicU64,
connections_failed: AtomicU64,
active_connections: AtomicU64,
inbound_connections: AtomicU64,
outbound_connections: AtomicU64,
connection_durations: RwLock<Vec<Duration>>,
}
impl ConnectionMetrics {
fn new() -> Self {
Self {
connections_established: AtomicU64::new(0),
connections_failed: AtomicU64::new(0),
active_connections: AtomicU64::new(0),
inbound_connections: AtomicU64::new(0),
outbound_connections: AtomicU64::new(0),
connection_durations: RwLock::new(Vec::new()),
}
}
pub fn connection_established(&self, inbound: bool) {
self.connections_established.fetch_add(1, Ordering::Relaxed);
self.active_connections.fetch_add(1, Ordering::Relaxed);
if inbound {
self.inbound_connections.fetch_add(1, Ordering::Relaxed);
} else {
self.outbound_connections.fetch_add(1, Ordering::Relaxed);
}
}
pub fn connection_closed(&self, duration: Duration) {
self.active_connections.fetch_sub(1, Ordering::Relaxed);
let mut durations = self.connection_durations.write();
if durations.len() >= 1000 {
durations.remove(0);
}
durations.push(duration);
}
pub fn connection_failed(&self) {
self.connections_failed.fetch_add(1, Ordering::Relaxed);
}
pub fn total_established(&self) -> u64 {
self.connections_established.load(Ordering::Relaxed)
}
pub fn total_failed(&self) -> u64 {
self.connections_failed.load(Ordering::Relaxed)
}
pub fn active(&self) -> u64 {
self.active_connections.load(Ordering::Relaxed)
}
pub fn total_inbound(&self) -> u64 {
self.inbound_connections.load(Ordering::Relaxed)
}
pub fn total_outbound(&self) -> u64 {
self.outbound_connections.load(Ordering::Relaxed)
}
pub fn avg_duration(&self) -> Option<Duration> {
let durations = self.connection_durations.read();
if durations.is_empty() {
None
} else {
let total: Duration = durations.iter().sum();
Some(total / durations.len() as u32)
}
}
pub fn snapshot(&self) -> ConnectionMetricsSnapshot {
ConnectionMetricsSnapshot {
total_established: self.total_established(),
total_failed: self.total_failed(),
active: self.active(),
total_inbound: self.total_inbound(),
total_outbound: self.total_outbound(),
avg_duration_ms: self.avg_duration().map(|d| d.as_millis() as u64),
}
}
}
pub struct BandwidthMetrics {
bytes_sent: AtomicU64,
bytes_received: AtomicU64,
protocol_bandwidth: RwLock<HashMap<String, (u64, u64)>>,
}
impl BandwidthMetrics {
fn new() -> Self {
Self {
bytes_sent: AtomicU64::new(0),
bytes_received: AtomicU64::new(0),
protocol_bandwidth: RwLock::new(HashMap::new()),
}
}
pub fn record_sent(&self, bytes: u64) {
self.bytes_sent.fetch_add(bytes, Ordering::Relaxed);
}
pub fn record_received(&self, bytes: u64) {
self.bytes_received.fetch_add(bytes, Ordering::Relaxed);
}
pub fn record_protocol_traffic(&self, protocol: &str, sent: u64, received: u64) {
let mut bandwidth = self.protocol_bandwidth.write();
let entry = bandwidth.entry(protocol.to_string()).or_insert((0, 0));
entry.0 += sent;
entry.1 += received;
}
pub fn total_sent(&self) -> u64 {
self.bytes_sent.load(Ordering::Relaxed)
}
pub fn total_received(&self) -> u64 {
self.bytes_received.load(Ordering::Relaxed)
}
pub fn snapshot(&self) -> BandwidthMetricsSnapshot {
BandwidthMetricsSnapshot {
total_sent: self.total_sent(),
total_received: self.total_received(),
}
}
}
pub struct DhtMetrics {
queries_made: AtomicU64,
queries_successful: AtomicU64,
queries_failed: AtomicU64,
providers_published: AtomicU64,
provider_queries: AtomicU64,
providers_found: AtomicU64,
routing_table_size: AtomicU64,
}
impl DhtMetrics {
fn new() -> Self {
Self {
queries_made: AtomicU64::new(0),
queries_successful: AtomicU64::new(0),
queries_failed: AtomicU64::new(0),
providers_published: AtomicU64::new(0),
provider_queries: AtomicU64::new(0),
providers_found: AtomicU64::new(0),
routing_table_size: AtomicU64::new(0),
}
}
pub fn query_made(&self) {
self.queries_made.fetch_add(1, Ordering::Relaxed);
}
pub fn query_successful(&self) {
self.queries_successful.fetch_add(1, Ordering::Relaxed);
}
pub fn query_failed(&self) {
self.queries_failed.fetch_add(1, Ordering::Relaxed);
}
pub fn provider_published(&self) {
self.providers_published.fetch_add(1, Ordering::Relaxed);
}
pub fn provider_query(&self) {
self.provider_queries.fetch_add(1, Ordering::Relaxed);
}
pub fn providers_found(&self, count: u64) {
self.providers_found.fetch_add(count, Ordering::Relaxed);
}
pub fn set_routing_table_size(&self, size: u64) {
self.routing_table_size.store(size, Ordering::Relaxed);
}
pub fn snapshot(&self) -> DhtMetricsSnapshot {
DhtMetricsSnapshot {
queries_made: self.queries_made.load(Ordering::Relaxed),
queries_successful: self.queries_successful.load(Ordering::Relaxed),
queries_failed: self.queries_failed.load(Ordering::Relaxed),
providers_published: self.providers_published.load(Ordering::Relaxed),
provider_queries: self.provider_queries.load(Ordering::Relaxed),
providers_found: self.providers_found.load(Ordering::Relaxed),
routing_table_size: self.routing_table_size.load(Ordering::Relaxed),
}
}
}
pub struct ProtocolMetrics {
messages: RwLock<HashMap<String, ProtocolStats>>,
}
#[derive(Default, Clone)]
struct ProtocolStats {
messages_sent: u64,
messages_received: u64,
bytes_sent: u64,
bytes_received: u64,
}
impl ProtocolMetrics {
fn new() -> Self {
Self {
messages: RwLock::new(HashMap::new()),
}
}
pub fn message_sent(&self, protocol: &str, bytes: u64) {
let mut messages = self.messages.write();
let stats = messages.entry(protocol.to_string()).or_default();
stats.messages_sent += 1;
stats.bytes_sent += bytes;
}
pub fn message_received(&self, protocol: &str, bytes: u64) {
let mut messages = self.messages.write();
let stats = messages.entry(protocol.to_string()).or_default();
stats.messages_received += 1;
stats.bytes_received += bytes;
}
pub fn get_stats(&self, protocol: &str) -> Option<(u64, u64, u64, u64)> {
let messages = self.messages.read();
messages.get(protocol).map(|s| {
(
s.messages_sent,
s.messages_received,
s.bytes_sent,
s.bytes_received,
)
})
}
}
#[derive(Debug, Clone, Serialize)]
pub struct MetricsSnapshot {
pub uptime_secs: u64,
pub connections: ConnectionMetricsSnapshot,
pub bandwidth: BandwidthMetricsSnapshot,
pub dht: DhtMetricsSnapshot,
}
#[derive(Debug, Clone, Serialize)]
pub struct ConnectionMetricsSnapshot {
pub total_established: u64,
pub total_failed: u64,
pub active: u64,
pub total_inbound: u64,
pub total_outbound: u64,
pub avg_duration_ms: Option<u64>,
}
#[derive(Debug, Clone, Serialize)]
pub struct BandwidthMetricsSnapshot {
pub total_sent: u64,
pub total_received: u64,
}
#[derive(Debug, Clone, Serialize)]
pub struct DhtMetricsSnapshot {
pub queries_made: u64,
pub queries_successful: u64,
pub queries_failed: u64,
pub providers_published: u64,
pub provider_queries: u64,
pub providers_found: u64,
pub routing_table_size: u64,
}
pub type SharedMetrics = Arc<NetworkMetrics>;
pub fn new_shared_metrics() -> SharedMetrics {
Arc::new(NetworkMetrics::new())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_connection_metrics() {
let metrics = NetworkMetrics::new();
metrics.connections.connection_established(true);
metrics.connections.connection_established(false);
assert_eq!(metrics.connections.active(), 2);
assert_eq!(metrics.connections.total_established(), 2);
assert_eq!(metrics.connections.total_inbound(), 1);
assert_eq!(metrics.connections.total_outbound(), 1);
metrics
.connections
.connection_closed(Duration::from_secs(10));
assert_eq!(metrics.connections.active(), 1);
metrics.connections.connection_failed();
assert_eq!(metrics.connections.total_failed(), 1);
}
#[test]
fn test_bandwidth_metrics() {
let metrics = NetworkMetrics::new();
metrics.bandwidth.record_sent(1000);
metrics.bandwidth.record_received(2000);
assert_eq!(metrics.bandwidth.total_sent(), 1000);
assert_eq!(metrics.bandwidth.total_received(), 2000);
}
#[test]
fn test_dht_metrics() {
let metrics = NetworkMetrics::new();
metrics.dht.query_made();
metrics.dht.query_successful();
metrics.dht.query_made();
metrics.dht.query_failed();
metrics.dht.providers_found(5);
let snapshot = metrics.dht.snapshot();
assert_eq!(snapshot.queries_made, 2);
assert_eq!(snapshot.queries_successful, 1);
assert_eq!(snapshot.queries_failed, 1);
assert_eq!(snapshot.providers_found, 5);
}
#[test]
fn test_protocol_metrics() {
let metrics = NetworkMetrics::new();
metrics.protocols.message_sent("/ipfs/kad/1.0.0", 100);
metrics.protocols.message_received("/ipfs/kad/1.0.0", 200);
let stats = metrics.protocols.get_stats("/ipfs/kad/1.0.0");
assert!(stats.is_some());
let (sent, received, bytes_sent, bytes_received) =
stats.expect("test: protocol stats should be present after recording messages");
assert_eq!(sent, 1);
assert_eq!(received, 1);
assert_eq!(bytes_sent, 100);
assert_eq!(bytes_received, 200);
}
#[test]
fn test_metrics_snapshot() {
let metrics = NetworkMetrics::new();
metrics.connections.connection_established(true);
metrics.bandwidth.record_sent(100);
metrics.dht.query_made();
let snapshot = metrics.snapshot();
assert_eq!(snapshot.connections.active, 1);
assert_eq!(snapshot.bandwidth.total_sent, 100);
assert_eq!(snapshot.dht.queries_made, 1);
}
}