refactor(proxy): improve tunnel throughput and observability

This commit is contained in:
fawney19
2026-05-19 23:49:36 +08:00
parent 57655bdb25
commit f5deed8709
30 changed files with 1306 additions and 173 deletions

View File

@@ -23,5 +23,7 @@ reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
libc = "0.2"
sysinfo = "0.32"
tokio.workspace = true
tokio-tungstenite = { version = "0.28", features = ["rustls-tls-webpki-roots"] }

View File

@@ -7,9 +7,9 @@ use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_gateway::tunnel_protocol as protocol;
use aether_testkit::{
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for, run_http_load_probe,
ExecutionRuntimeHarness, ExecutionRuntimeHarnessConfig, GatewayHarness, GatewayHarnessConfig,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult, SpawnedServer,
TunnelHarness, TunnelHarnessConfig,
BenchmarkRuntimeSnapshot, ExecutionRuntimeHarness, ExecutionRuntimeHarnessConfig,
GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
HttpLoadProbeResult, SpawnedServer, TunnelHarness, TunnelHarnessConfig,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::StatusCode;
@@ -84,9 +84,11 @@ struct CapacityCurvePointResult {
throughput_rps: u64,
p50_ms: u64,
p95_ms: u64,
p99_ms: u64,
max_ms: u64,
mean_ms: u64,
metrics: GateMetricSnapshot,
runtime: BenchmarkRuntimeSnapshot,
}
#[derive(Debug, Serialize)]
@@ -387,9 +389,11 @@ fn capacity_point(
throughput_rps,
p50_ms: result.p50_ms,
p95_ms: result.p95_ms,
p99_ms: result.p99_ms,
max_ms: result.max_ms,
mean_ms: result.mean_ms,
metrics,
runtime: result.runtime,
}
}
@@ -636,6 +640,12 @@ async fn connect_protocol_peer(
request
.headers_mut()
.insert("x-node-id", http::HeaderValue::from_static("node-baseline"));
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
http::HeaderValue::from_static(
aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION_STR,
),
);
request.headers_mut().insert(
"x-node-name",
http::HeaderValue::from_static("proxy-baseline"),

View File

@@ -12,8 +12,8 @@ use aether_runtime_state::{
RedisClientConfig, RedisClientFactory, RedisLockRunner, RedisLockRunnerConfig,
};
use aether_testkit::{
init_test_runtime_for, reserve_local_port, ManagedPostgresServer, ManagedRedisServer,
TunnelHarness, TunnelHarnessConfig,
init_test_runtime_for, reserve_local_port, BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot,
ManagedPostgresServer, ManagedRedisServer, TunnelHarness, TunnelHarnessConfig,
};
use futures_util::{FutureExt, StreamExt};
use serde::Serialize;
@@ -89,9 +89,11 @@ struct RecoverySummary {
recovered_after_restart_ms: Option<u64>,
p50_ms: u64,
p95_ms: u64,
p99_ms: u64,
max_ms: u64,
mean_ms: u64,
phase_counts: PhaseCounts,
runtime: BenchmarkRuntimeSnapshot,
}
#[derive(Debug, Clone, Serialize)]
@@ -101,6 +103,7 @@ struct PostgresSlowQueryRecoveryReport {
recovery_claim_succeeded: bool,
recovery_claim_latency_ms: u64,
recovery_claimed_items: usize,
runtime: BenchmarkRuntimeSnapshot,
}
#[derive(Debug, Clone, Serialize)]
@@ -143,10 +146,14 @@ impl RecoveryCollector {
}
}
async fn summarize(&self, recovered_after_restart_ms: Option<u64>) -> RecoverySummary {
async fn summarize(
&self,
recovered_after_restart_ms: Option<u64>,
runtime: BenchmarkRuntimeSnapshot,
) -> RecoverySummary {
let mut latencies = self.latencies_ms.lock().await.clone();
latencies.sort_unstable();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
let (p50_ms, p95_ms, p99_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
let phase_counts = self.phase_counts.lock().await.clone();
RecoverySummary {
total_attempts: self.successful_attempts.load(Ordering::Acquire)
@@ -156,9 +163,11 @@ impl RecoveryCollector {
recovered_after_restart_ms,
p50_ms,
p95_ms,
p99_ms,
max_ms,
mean_ms,
phase_counts,
runtime,
}
}
}
@@ -236,6 +245,7 @@ async fn benchmark_redis_restart_recovery(
redis_server: Arc<Mutex<ManagedRedisServer>>,
config: &FailureRecoveryBaselineConfig,
) -> Result<RecoverySummary, Box<dyn std::error::Error>> {
let mut runtime_sampler = BenchmarkRuntimeSampler::new();
let redis_url = redis_server.lock().await.redis_url().to_string();
let factory = RedisClientFactory::new(RedisClientConfig {
url: redis_url,
@@ -334,7 +344,10 @@ async fn benchmark_redis_restart_recovery(
.map_err(std::io::Error::other)?;
Ok(collector
.summarize(load_optional_atomic_u64(&recovered_after_restart_ms))
.summarize(
load_optional_atomic_u64(&recovered_after_restart_ms),
runtime_sampler.snapshot(),
)
.await)
}
@@ -342,6 +355,7 @@ async fn benchmark_postgres_slow_query_recovery(
postgres_url: &str,
config: &FailureRecoveryBaselineConfig,
) -> Result<PostgresSlowQueryRecoveryReport, Box<dyn std::error::Error>> {
let mut runtime_sampler = BenchmarkRuntimeSampler::new();
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: postgres_url.to_string(),
min_connections: 1,
@@ -410,6 +424,7 @@ async fn benchmark_postgres_slow_query_recovery(
recovery_claim_succeeded: !claimed_ids.is_empty(),
recovery_claim_latency_ms,
recovery_claimed_items: claimed_ids.len(),
runtime: runtime_sampler.snapshot(),
})
}
@@ -444,6 +459,7 @@ async fn bootstrap_failure_recovery_lease_table(
async fn benchmark_tunnel_restart_recovery(
config: &FailureRecoveryBaselineConfig,
) -> Result<RecoverySummary, Box<dyn std::error::Error>> {
let mut runtime_sampler = BenchmarkRuntimeSampler::new();
let port = reserve_local_port()?;
let tunnel_config = TunnelHarnessConfig::default();
let initial_tunnel = TunnelHarness::start_on_port(tunnel_config.clone(), port).await?;
@@ -500,6 +516,12 @@ async fn benchmark_tunnel_restart_recovery(
.parse()
.map_err(|err| format!("failed to build x-node-id header: {err}"))?,
);
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION_STR
.parse()
.expect("protocol version header value should be valid"),
);
request.headers_mut().insert(
"x-node-name",
format!("recovery-node-{worker_index}-{current}")
@@ -556,7 +578,10 @@ async fn benchmark_tunnel_restart_recovery(
.map_err(std::io::Error::other)?;
Ok(collector
.summarize(load_optional_atomic_u64(&recovered_after_restart_ms))
.summarize(
load_optional_atomic_u64(&recovered_after_restart_ms),
runtime_sampler.snapshot(),
)
.await)
}
@@ -596,15 +621,16 @@ fn load_optional_atomic_u64(value: &AtomicU64) -> Option<u64> {
}
}
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64) {
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
return (0, 0, 0, 0, 0);
}
let max_ms = *latencies.last().unwrap_or(&0);
let mean_ms = latencies.iter().sum::<u64>() / latencies.len() as u64;
let p50_ms = percentile(latencies, 50);
let p95_ms = percentile(latencies, 95);
(p50_ms, p95_ms, max_ms, mean_ms)
let p99_ms = percentile(latencies, 99);
(p50_ms, p95_ms, p99_ms, max_ms, mean_ms)
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {

View File

@@ -3,8 +3,9 @@ use std::time::Duration;
use aether_gateway::tunnel_protocol as protocol;
use aether_testkit::{
init_test_runtime_for, run_http_load_probe, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
HttpLoadProbeResult, TunnelHarness, TunnelHarnessConfig,
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for, run_http_load_probe,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult, TunnelHarness,
TunnelHarnessConfig,
};
use futures_util::{SinkExt, StreamExt};
use reqwest::Method;
@@ -38,6 +39,23 @@ impl Default for GatewayTunnelBaselineConfig {
struct GatewayTunnelBaselineReport {
suite: &'static str,
scenario: HttpLoadProbeResult,
tunnel_metrics: TunnelMetricsSnapshot,
}
#[derive(Debug, Serialize)]
struct TunnelMetricsSnapshot {
proxy_connections: u64,
active_streams: u64,
outbound_queue_depth_total: u64,
outbound_queue_depth_max: u64,
outbound_queue_capacity_total: u64,
outbound_queue_rejected_full_total: u64,
outbound_queue_rejected_closed_total: u64,
proxy_connection_congested_total: u64,
proxy_connection_write_latency_last_us_max: u64,
proxy_connection_write_latency_ewma_us_max: u64,
proxy_connections_protocol_v1: u64,
proxy_connections_protocol_v2: u64,
}
#[tokio::main]
@@ -81,11 +99,13 @@ async fn run_suite(
.await
.map_err(std::io::Error::other)?;
let tunnel_metrics = capture_tunnel_metrics(tunnel.base_url()).await?;
drop(peer);
Ok(GatewayTunnelBaselineReport {
suite: "gateway_tunnel_stream_baseline",
scenario: result,
tunnel_metrics,
})
}
@@ -127,6 +147,12 @@ async fn connect_protocol_peer(
request
.headers_mut()
.insert("x-node-id", http::HeaderValue::from_static("node-baseline"));
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
http::HeaderValue::from_static(
aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION_STR,
),
);
request.headers_mut().insert(
"x-node-name",
http::HeaderValue::from_static("proxy-baseline"),
@@ -162,6 +188,80 @@ async fn connect_protocol_peer(
}))
}
async fn capture_tunnel_metrics(
base_url: &str,
) -> Result<TunnelMetricsSnapshot, Box<dyn std::error::Error>> {
let samples = fetch_prometheus_samples(&format!("{base_url}/metrics"))
.await
.map_err(std::io::Error::other)?;
Ok(TunnelMetricsSnapshot {
proxy_connections: find_metric_value_u64(&samples, "tunnel_proxy_connections", &[])
.unwrap_or_default(),
active_streams: find_metric_value_u64(&samples, "tunnel_active_streams", &[])
.unwrap_or_default(),
outbound_queue_depth_total: find_metric_value_u64(
&samples,
"tunnel_proxy_outbound_queue_depth_total",
&[],
)
.unwrap_or_default(),
outbound_queue_depth_max: find_metric_value_u64(
&samples,
"tunnel_proxy_outbound_queue_depth_max",
&[],
)
.unwrap_or_default(),
outbound_queue_capacity_total: find_metric_value_u64(
&samples,
"tunnel_proxy_outbound_queue_capacity_total",
&[],
)
.unwrap_or_default(),
outbound_queue_rejected_full_total: find_metric_value_u64(
&samples,
"tunnel_proxy_outbound_queue_rejected_full_total",
&[],
)
.unwrap_or_default(),
outbound_queue_rejected_closed_total: find_metric_value_u64(
&samples,
"tunnel_proxy_outbound_queue_rejected_closed_total",
&[],
)
.unwrap_or_default(),
proxy_connection_congested_total: find_metric_value_u64(
&samples,
"tunnel_proxy_connection_congested_total",
&[],
)
.unwrap_or_default(),
proxy_connection_write_latency_last_us_max: find_metric_value_u64(
&samples,
"tunnel_proxy_connection_write_latency_last_us_max",
&[],
)
.unwrap_or_default(),
proxy_connection_write_latency_ewma_us_max: find_metric_value_u64(
&samples,
"tunnel_proxy_connection_write_latency_ewma_us_max",
&[],
)
.unwrap_or_default(),
proxy_connections_protocol_v1: find_metric_value_u64(
&samples,
"tunnel_proxy_connections_protocol_v1",
&[],
)
.unwrap_or_default(),
proxy_connections_protocol_v2: find_metric_value_u64(
&samples,
"tunnel_proxy_connections_protocol_v2",
&[],
)
.unwrap_or_default(),
})
}
async fn handle_binary_frame<S>(
sink: &mut S,
data: Vec<u8>,

View File

@@ -8,10 +8,11 @@ use aether_runtime_state::{
RedisClientConfig, RuntimeSemaphore, RuntimeSemaphoreConfig, RuntimeState,
};
use aether_testkit::{
init_test_runtime_for, run_multi_url_http_load_probe, ExecutionRuntimeHarness,
ExecutionRuntimeHarnessConfig, GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig,
HttpLoadProbeResponseMode, ManagedRedisServer, MultiUrlHttpLoadProbeResult, SpawnedServer,
TunnelHarness, TunnelHarnessConfig,
init_test_runtime_for, run_multi_url_http_load_probe, BenchmarkRuntimeSampler,
BenchmarkRuntimeSnapshot, ExecutionRuntimeHarness, ExecutionRuntimeHarnessConfig,
GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
ManagedRedisServer, MultiUrlHttpLoadProbeResult, SpawnedServer, TunnelHarness,
TunnelHarnessConfig,
};
use axum::body::to_bytes;
use axum::extract::Request;
@@ -85,9 +86,11 @@ struct WebSocketAdmissionProbeResult {
successful_attempts: usize,
p50_ms: u64,
p95_ms: u64,
p99_ms: u64,
max_ms: u64,
mean_ms: u64,
status_counts: BTreeMap<u16, usize>,
runtime: BenchmarkRuntimeSnapshot,
}
#[tokio::main]
@@ -443,6 +446,7 @@ async fn run_tunnel_proxy_connection_probe(
if urls.is_empty() {
return Err("tunnel proxy connection probe requires at least one target url".to_string());
}
let mut runtime_sampler = BenchmarkRuntimeSampler::new();
let next_attempt = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let latencies_ms = Arc::new(tokio::sync::Mutex::new(Vec::with_capacity(
config.tunnel_attempts,
@@ -489,6 +493,12 @@ async fn run_tunnel_proxy_connection_probe(
.parse()
.map_err(|err| format!("failed to build x-node-id header: {err}"))?,
);
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION_STR
.parse()
.expect("protocol version header value should be valid"),
);
request.headers_mut().insert(
"x-node-name",
format!("baseline-node-{worker_index}-{current}")
@@ -557,7 +567,7 @@ async fn run_tunnel_proxy_connection_probe(
let mut latencies = latencies_ms.lock().await.clone();
latencies.sort_unstable();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
let (p50_ms, p95_ms, p99_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
let target_attempt_counts = target_attempt_counts.lock().await.clone();
let status_counts = status_counts.lock().await.clone();
@@ -572,21 +582,24 @@ async fn run_tunnel_proxy_connection_probe(
successful_attempts: successful_attempts.load(std::sync::atomic::Ordering::Acquire),
p50_ms,
p95_ms,
p99_ms,
max_ms,
mean_ms,
status_counts,
runtime: runtime_sampler.snapshot(),
})
}
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64) {
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
return (0, 0, 0, 0, 0);
}
let max_ms = *latencies.last().unwrap_or(&0);
let mean_ms = latencies.iter().sum::<u64>() / latencies.len() as u64;
let p50_ms = percentile(latencies, 50);
let p95_ms = percentile(latencies, 95);
(p50_ms, p95_ms, max_ms, mean_ms)
let p99_ms = percentile(latencies, 99);
(p50_ms, p95_ms, p99_ms, max_ms, mean_ms)
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {

View File

@@ -284,6 +284,12 @@ async fn connect_protocol_peer(
request
.headers_mut()
.insert("x-node-id", http::HeaderValue::from_static(NODE_ID));
request.headers_mut().insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
http::HeaderValue::from_static(
aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION_STR,
),
);
request.headers_mut().insert(
"x-node-name",
http::HeaderValue::from_static("proxy-owner-relay-baseline"),

View File

@@ -6,6 +6,7 @@ mod load;
mod metrics;
mod postgres;
mod redis;
mod runtime;
mod server;
mod tracing;
mod tunnel;
@@ -24,6 +25,7 @@ pub use metrics::{
};
pub use postgres::{prepare_aether_postgres_schema, ManagedPostgresServer};
pub use redis::ManagedRedisServer;
pub use runtime::{BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot};
pub use server::{reserve_local_port, SpawnedServer};
pub use tracing::{init_test_runtime, init_test_runtime_for, test_runtime_config};
pub use tunnel::{TunnelHarness, TunnelHarnessConfig};

View File

@@ -7,6 +7,8 @@ use http::{HeaderMap, HeaderName, HeaderValue};
use reqwest::{Client, Method};
use tokio::sync::Mutex;
use crate::runtime::{BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot};
#[derive(Debug, Clone, Copy, Default, serde::Serialize, PartialEq, Eq)]
pub enum HttpLoadProbeResponseMode {
#[default]
@@ -66,12 +68,16 @@ pub struct HttpLoadProbeResult {
pub response_mode: HttpLoadProbeResponseMode,
pub total_requests: usize,
pub concurrency: usize,
pub duration_ms: u64,
pub throughput_rps: u64,
pub p99_ms: u64,
pub completed_requests: usize,
pub failed_requests: usize,
pub p50_ms: u64,
pub p95_ms: u64,
pub max_ms: u64,
pub mean_ms: u64,
pub runtime: BenchmarkRuntimeSnapshot,
pub status_counts: BTreeMap<u16, usize>,
}
@@ -83,12 +89,16 @@ pub struct MultiUrlHttpLoadProbeResult {
pub response_mode: HttpLoadProbeResponseMode,
pub total_requests: usize,
pub concurrency: usize,
pub duration_ms: u64,
pub throughput_rps: u64,
pub p99_ms: u64,
pub completed_requests: usize,
pub failed_requests: usize,
pub p50_ms: u64,
pub p95_ms: u64,
pub max_ms: u64,
pub mean_ms: u64,
pub runtime: BenchmarkRuntimeSnapshot,
pub status_counts: BTreeMap<u16, usize>,
}
@@ -108,12 +118,16 @@ pub async fn run_http_load_probe(
response_mode: result.response_mode,
total_requests: result.total_requests,
concurrency: result.concurrency,
duration_ms: result.duration_ms,
throughput_rps: result.throughput_rps,
p99_ms: result.p99_ms,
completed_requests: result.completed_requests,
failed_requests: result.failed_requests,
p50_ms: result.p50_ms,
p95_ms: result.p95_ms,
max_ms: result.max_ms,
mean_ms: result.mean_ms,
runtime: result.runtime,
status_counts: result.status_counts,
})
}
@@ -141,6 +155,8 @@ async fn run_http_load_probe_against_urls(
let request_headers = build_headers(&config.headers)?;
let request_body = config.body.clone().map(Arc::new);
let response_mode = config.response_mode;
let mut runtime_sampler = BenchmarkRuntimeSampler::new();
let started_at = Instant::now();
let next_request = Arc::new(AtomicUsize::new(0));
let latencies_ms = Arc::new(Mutex::new(Vec::with_capacity(config.total_requests)));
@@ -220,7 +236,13 @@ async fn run_http_load_probe_against_urls(
let target_request_counts = target_request_counts.lock().await.clone();
let mut latencies = latencies_ms.lock().await.clone();
latencies.sort_unstable();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
let (p50_ms, p95_ms, p99_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
let duration_ms = started_at.elapsed().as_millis() as u64;
let throughput_rps = if duration_ms == 0 {
completed_requests.load(Ordering::Acquire) as u64
} else {
((completed_requests.load(Ordering::Acquire) as u64) * 1_000) / duration_ms.max(1)
};
Ok(MultiUrlHttpLoadProbeResult {
target_urls: urls.to_vec(),
@@ -229,12 +251,16 @@ async fn run_http_load_probe_against_urls(
response_mode: config.response_mode,
total_requests: config.total_requests,
concurrency: config.concurrency,
duration_ms,
throughput_rps,
p99_ms,
completed_requests: completed_requests.load(Ordering::Acquire),
failed_requests: failed_requests.load(Ordering::Acquire),
p50_ms,
p95_ms,
max_ms,
mean_ms,
runtime: runtime_sampler.snapshot(),
status_counts,
})
}
@@ -251,16 +277,17 @@ fn build_headers(headers: &BTreeMap<String, String>) -> Result<HeaderMap, String
Ok(result)
}
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64) {
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
return (0, 0, 0, 0, 0);
}
let max_ms = *latencies.last().unwrap_or(&0);
let mean_ms = latencies.iter().sum::<u64>() / latencies.len() as u64;
let p50_ms = percentile(latencies, 50);
let p95_ms = percentile(latencies, 95);
(p50_ms, p95_ms, max_ms, mean_ms)
let p99_ms = percentile(latencies, 99);
(p50_ms, p95_ms, p99_ms, max_ms, mean_ms)
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {
@@ -311,10 +338,11 @@ mod tests {
#[test]
fn summarizes_latency_distribution() {
let (p50_ms, p95_ms, max_ms, mean_ms) =
let (p50_ms, p95_ms, p99_ms, max_ms, mean_ms) =
summarize_latencies(&[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]);
assert_eq!(p50_ms, 60);
assert_eq!(p95_ms, 100);
assert_eq!(p99_ms, 100);
assert_eq!(max_ms, 100);
assert_eq!(mean_ms, 55);
}

View File

@@ -0,0 +1,141 @@
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use serde::Serialize;
use sysinfo::{get_current_pid, Pid, ProcessesToUpdate, System};
#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
pub struct BenchmarkRuntimeSnapshot {
pub sampled_at_unix_secs: u64,
pub elapsed_ms: u64,
pub system_cpu_usage_basis_points: u64,
pub process_cpu_usage_basis_points: u64,
pub memory_total_bytes: u64,
pub memory_used_bytes: u64,
pub memory_available_bytes: u64,
pub memory_used_basis_points: u64,
pub process_memory_bytes: u64,
pub process_virtual_memory_bytes: u64,
pub process_memory_basis_points: u64,
pub fd_open_count: u64,
pub fd_limit: u64,
}
pub struct BenchmarkRuntimeSampler {
started_at: Instant,
system: System,
current_pid: Option<Pid>,
}
impl BenchmarkRuntimeSampler {
pub fn new() -> Self {
let mut system = System::new_all();
let current_pid = get_current_pid().ok();
if let Some(pid) = current_pid {
system.refresh_processes(ProcessesToUpdate::Some(&[pid]), true);
}
system.refresh_cpu_usage();
system.refresh_memory();
Self {
started_at: Instant::now(),
system,
current_pid,
}
}
pub fn snapshot(&mut self) -> BenchmarkRuntimeSnapshot {
self.system.refresh_cpu_usage();
self.system.refresh_memory();
if let Some(pid) = self.current_pid {
self.system
.refresh_processes(ProcessesToUpdate::Some(&[pid]), true);
}
let memory_total_bytes = self.system.total_memory();
let memory_used_bytes = self.system.used_memory();
let memory_available_bytes = self.system.available_memory();
let (process_cpu_usage_basis_points, process_memory_bytes, process_virtual_memory_bytes) =
self.current_pid
.and_then(|pid| self.system.process(pid))
.map(|process| {
(
percent_to_basis_points(process.cpu_usage() as f64),
process.memory(),
process.virtual_memory(),
)
})
.unwrap_or((0, 0, 0));
BenchmarkRuntimeSnapshot {
sampled_at_unix_secs: current_unix_secs(),
elapsed_ms: self.started_at.elapsed().as_millis() as u64,
system_cpu_usage_basis_points: percent_to_basis_points(
self.system.global_cpu_usage() as f64
),
process_cpu_usage_basis_points,
memory_total_bytes,
memory_used_bytes,
memory_available_bytes,
memory_used_basis_points: ratio_to_basis_points(memory_used_bytes, memory_total_bytes),
process_memory_bytes,
process_virtual_memory_bytes,
process_memory_basis_points: ratio_to_basis_points(
process_memory_bytes,
memory_total_bytes,
),
fd_open_count: open_file_descriptors().unwrap_or(0),
fd_limit: file_descriptor_limit(),
}
}
}
impl Default for BenchmarkRuntimeSampler {
fn default() -> Self {
Self::new()
}
}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or(0)
}
fn percent_to_basis_points(value: f64) -> u64 {
if !value.is_finite() || value.is_sign_negative() {
0
} else {
(value * 100.0).round().clamp(0.0, u64::MAX as f64) as u64
}
}
fn ratio_to_basis_points(value: u64, total: u64) -> u64 {
value.saturating_mul(10_000).checked_div(total).unwrap_or(0)
}
fn open_file_descriptors() -> Option<u64> {
#[cfg(unix)]
{
for dir in ["/proc/self/fd", "/dev/fd"] {
if let Ok(entries) = std::fs::read_dir(dir) {
return Some(entries.count() as u64);
}
}
}
None
}
fn file_descriptor_limit() -> u64 {
#[cfg(unix)]
{
let mut limit = libc::rlimit {
rlim_cur: 0,
rlim_max: 0,
};
let result = unsafe { libc::getrlimit(libc::RLIMIT_NOFILE, &mut limit) };
if result == 0 {
return limit.rlim_cur;
}
}
0
}