353 lines
11 KiB
Rust
353 lines
11 KiB
Rust
use flate2::{Compression, read::GzDecoder, write::GzEncoder};
|
|
use serde::{Deserialize, Serialize, de::DeserializeOwned};
|
|
use std::io::{self, Read};
|
|
|
|
pub const AGENT_PROTOCOL_VERSION: &str = "1";
|
|
pub const AGENT_PROTOCOL_HEADER: &str = "x-nodeflare-agent-protocol";
|
|
pub const AGENT_CAPABILITIES_HEADER: &str = "x-nodeflare-agent-capabilities";
|
|
pub const REQUIRED_AGENT_CAPABILITIES: [&str; 4] =
|
|
["metrics-v1", "config-v1", "remote-exec-v1", "task-ack-v1"];
|
|
pub const MIN_UPLOAD_INTERVAL: u64 = 3;
|
|
pub const MAX_BATCH_BYTES: usize = 768 * 1024;
|
|
pub const MAX_DECODED_BYTES: usize = 2 * 1024 * 1024;
|
|
pub const MAX_SAMPLES: usize = 720;
|
|
|
|
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
|
#[serde(default, deny_unknown_fields)]
|
|
pub struct Report {
|
|
pub timestamp: i64,
|
|
pub cpu: f64,
|
|
pub load1: f64,
|
|
pub load5: f64,
|
|
pub load15: f64,
|
|
pub mem_used: i64,
|
|
pub mem_total: i64,
|
|
pub swap_used: i64,
|
|
pub swap_total: i64,
|
|
pub disk_used: i64,
|
|
pub disk_total: i64,
|
|
pub net_in: f64,
|
|
pub net_out: f64,
|
|
pub net_rx_total: i64,
|
|
pub net_tx_total: i64,
|
|
pub uptime: i64,
|
|
pub processes: i64,
|
|
pub tcp_connections: i64,
|
|
pub udp_connections: i64,
|
|
pub cpu_cores: i64,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub cpu_model: String,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub os: String,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub kernel: String,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub arch: String,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub virtualization: String,
|
|
pub gpu_usage: f64,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub gpu_model: String,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub agent_version: String,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub ip_v4: String,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub ip_v6: String,
|
|
pub disk_read_bps: f64,
|
|
pub disk_write_bps: f64,
|
|
pub disk_read_iops: f64,
|
|
pub disk_write_iops: f64,
|
|
pub disk_await_ms: f64,
|
|
pub disk_utilization: f64,
|
|
pub disks: Vec<DiskMetric>,
|
|
pub gpus: Vec<GpuMetric>,
|
|
pub latency_results: Vec<LatencyResult>,
|
|
}
|
|
|
|
impl Report {
|
|
fn round_for_upload(&mut self) {
|
|
for value in [
|
|
&mut self.net_in,
|
|
&mut self.net_out,
|
|
&mut self.disk_read_bps,
|
|
&mut self.disk_write_bps,
|
|
] {
|
|
*value = value.round();
|
|
}
|
|
for value in [
|
|
&mut self.cpu,
|
|
&mut self.load1,
|
|
&mut self.load5,
|
|
&mut self.load15,
|
|
&mut self.gpu_usage,
|
|
&mut self.disk_read_iops,
|
|
&mut self.disk_write_iops,
|
|
&mut self.disk_await_ms,
|
|
&mut self.disk_utilization,
|
|
] {
|
|
*value = (*value * 100.0).round() / 100.0;
|
|
}
|
|
for disk in &mut self.disks {
|
|
disk.read_bps = disk.read_bps.round();
|
|
disk.write_bps = disk.write_bps.round();
|
|
for value in [
|
|
&mut disk.read_iops,
|
|
&mut disk.write_iops,
|
|
&mut disk.await_ms,
|
|
&mut disk.utilization,
|
|
] {
|
|
*value = (*value * 100.0).round() / 100.0;
|
|
}
|
|
}
|
|
for gpu in &mut self.gpus {
|
|
gpu.usage = gpu.usage.map(|value| (value * 100.0).round() / 100.0);
|
|
}
|
|
for latency in &mut self.latency_results {
|
|
latency.latency_ms = (latency.latency_ms * 100.0).round() / 100.0;
|
|
latency.packet_loss = (latency.packet_loss * 100.0).round() / 100.0;
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
|
#[serde(default, deny_unknown_fields)]
|
|
pub struct DiskMetric {
|
|
pub name: String,
|
|
pub mount_point: String,
|
|
pub used: i64,
|
|
pub total: i64,
|
|
pub read_bps: f64,
|
|
pub write_bps: f64,
|
|
pub read_iops: f64,
|
|
pub write_iops: f64,
|
|
pub await_ms: f64,
|
|
pub utilization: f64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
|
#[serde(default, deny_unknown_fields)]
|
|
pub struct GpuMetric {
|
|
pub model: String,
|
|
pub usage: Option<f64>,
|
|
pub memory_used: i64,
|
|
pub memory_total: i64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
|
|
#[serde(default, deny_unknown_fields)]
|
|
pub struct LatencyResult {
|
|
pub task_id: String,
|
|
pub timestamp: i64,
|
|
pub latency_ms: f64,
|
|
pub packet_loss: f64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
pub struct Info {
|
|
cpu_model: String,
|
|
os: String,
|
|
kernel: String,
|
|
arch: String,
|
|
virtualization: String,
|
|
gpu_model: String,
|
|
agent_version: String,
|
|
ip_v4: String,
|
|
ip_v6: String,
|
|
}
|
|
|
|
impl Info {
|
|
fn take(report: &mut Report) -> Self {
|
|
Self {
|
|
cpu_model: std::mem::take(&mut report.cpu_model),
|
|
os: std::mem::take(&mut report.os),
|
|
kernel: std::mem::take(&mut report.kernel),
|
|
arch: std::mem::take(&mut report.arch),
|
|
virtualization: std::mem::take(&mut report.virtualization),
|
|
gpu_model: std::mem::take(&mut report.gpu_model),
|
|
agent_version: std::mem::take(&mut report.agent_version),
|
|
ip_v4: std::mem::take(&mut report.ip_v4),
|
|
ip_v6: std::mem::take(&mut report.ip_v6),
|
|
}
|
|
}
|
|
|
|
fn apply(&self, report: &mut Report) {
|
|
report.cpu_model.clone_from(&self.cpu_model);
|
|
report.os.clone_from(&self.os);
|
|
report.kernel.clone_from(&self.kernel);
|
|
report.arch.clone_from(&self.arch);
|
|
report.virtualization.clone_from(&self.virtualization);
|
|
report.gpu_model.clone_from(&self.gpu_model);
|
|
report.agent_version.clone_from(&self.agent_version);
|
|
report.ip_v4.clone_from(&self.ip_v4);
|
|
report.ip_v6.clone_from(&self.ip_v6);
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Deserialize, Serialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
pub struct Sample {
|
|
pub metrics: Report,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub info: Option<Info>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize, Serialize)]
|
|
#[serde(deny_unknown_fields)]
|
|
pub struct Update {
|
|
pub samples: Vec<Sample>,
|
|
pub persist: bool,
|
|
}
|
|
|
|
impl Update {
|
|
pub fn from_reports(reports: Vec<Report>, persist: bool, info: &mut Option<Info>) -> Self {
|
|
let samples = reports
|
|
.into_iter()
|
|
.map(|mut metrics| {
|
|
let current = Info::take(&mut metrics);
|
|
metrics.round_for_upload();
|
|
let changed = info.as_ref() != Some(¤t);
|
|
if changed {
|
|
*info = Some(current.clone());
|
|
}
|
|
Sample {
|
|
metrics,
|
|
info: changed.then_some(current),
|
|
}
|
|
})
|
|
.collect();
|
|
Self { samples, persist }
|
|
}
|
|
|
|
pub fn into_reports(self, info: &mut Option<Info>) -> io::Result<Vec<Report>> {
|
|
if self.samples.len() > MAX_SAMPLES {
|
|
return Err(io::Error::other("too many telemetry samples"));
|
|
}
|
|
self.samples
|
|
.into_iter()
|
|
.map(|mut sample| {
|
|
if let Some(current) = sample.info {
|
|
*info = Some(current);
|
|
}
|
|
info.as_ref()
|
|
.ok_or_else(|| io::Error::other("missing telemetry info"))?
|
|
.apply(&mut sample.metrics);
|
|
Ok(sample.metrics)
|
|
})
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
pub fn encode(value: &impl Serialize) -> io::Result<Vec<u8>> {
|
|
let mut encoder = GzEncoder::new(Vec::new(), Compression::fast());
|
|
serde_json::to_writer(&mut encoder, value)?;
|
|
encoder.finish()
|
|
}
|
|
|
|
pub fn decode<T: DeserializeOwned>(bytes: &[u8]) -> io::Result<T> {
|
|
if bytes.len() > MAX_DECODED_BYTES {
|
|
return Err(io::Error::other("telemetry frame is too large"));
|
|
}
|
|
let mut decoded = Vec::new();
|
|
GzDecoder::new(bytes)
|
|
.take(MAX_DECODED_BYTES as u64 + 1)
|
|
.read_to_end(&mut decoded)?;
|
|
if decoded.len() > MAX_DECODED_BYTES {
|
|
return Err(io::Error::other("expanded telemetry frame is too large"));
|
|
}
|
|
serde_json::from_slice(&decoded).map_err(io::Error::other)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn trims_float_noise_without_changing_byte_counters() {
|
|
let report = Report {
|
|
cpu: 100.0 / 3.0,
|
|
net_in: 12345.6789,
|
|
net_rx_total: i64::MAX,
|
|
disks: vec![DiskMetric {
|
|
read_bps: 125.6789,
|
|
await_ms: 1.234567,
|
|
..DiskMetric::default()
|
|
}],
|
|
latency_results: vec![LatencyResult {
|
|
latency_ms: 28.4444,
|
|
packet_loss: 100.0 / 3.0,
|
|
..LatencyResult::default()
|
|
}],
|
|
..Report::default()
|
|
};
|
|
let encoded = encode(&Update::from_reports(vec![report], false, &mut None)).unwrap();
|
|
let reports = decode::<Update>(&encoded)
|
|
.unwrap()
|
|
.into_reports(&mut None)
|
|
.unwrap();
|
|
assert_eq!(reports[0].cpu, 33.33);
|
|
assert_eq!(reports[0].net_in, 12346.0);
|
|
assert_eq!(reports[0].net_rx_total, i64::MAX);
|
|
assert_eq!(reports[0].disks[0].read_bps, 126.0);
|
|
assert_eq!(reports[0].disks[0].await_ms, 1.23);
|
|
assert_eq!(reports[0].latency_results[0].latency_ms, 28.44);
|
|
}
|
|
|
|
#[test]
|
|
fn info_is_sent_on_connect_and_changes_including_during_replay() {
|
|
let mut encoder = None;
|
|
let mut decoder = None;
|
|
for (name, changed) in [
|
|
("first", true),
|
|
("first", false),
|
|
("second", true),
|
|
("first", true),
|
|
] {
|
|
let update = Update::from_reports(
|
|
vec![Report {
|
|
cpu_model: name.into(),
|
|
cpu: 23.5,
|
|
..Report::default()
|
|
}],
|
|
false,
|
|
&mut encoder,
|
|
);
|
|
assert_eq!(update.samples[0].info.is_some(), changed);
|
|
assert!(update.samples[0].metrics.cpu_model.is_empty());
|
|
let bytes = encode(&update).unwrap();
|
|
let decoded: Update = decode(&bytes).unwrap();
|
|
let reports = decoded.into_reports(&mut decoder).unwrap();
|
|
assert_eq!(reports[0].cpu_model, name);
|
|
assert_eq!(reports[0].cpu, 23.5);
|
|
}
|
|
let update = Update::from_reports(vec![Report::default()], true, &mut None);
|
|
assert!(update.samples[0].info.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn rejects_missing_info_corruption_and_decompression_bombs() {
|
|
let update = Update {
|
|
samples: vec![Sample {
|
|
metrics: Report::default(),
|
|
info: None,
|
|
}],
|
|
persist: false,
|
|
};
|
|
assert!(update.into_reports(&mut None).is_err());
|
|
assert!(decode::<Update>(b"not gzip").is_err());
|
|
let bytes = encode(&"x".repeat(MAX_DECODED_BYTES + 1)).unwrap();
|
|
assert!(decode::<String>(&bytes).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn compresses_batches_and_supports_empty_commit_messages() {
|
|
let update = Update::from_reports(vec![Report::default(); 20], false, &mut None);
|
|
assert!(encode(&update).unwrap().len() < serde_json::to_vec(&update).unwrap().len() / 4);
|
|
let update = Update::from_reports(Vec::new(), true, &mut None);
|
|
let decoded: Update = decode(&encode(&update).unwrap()).unwrap();
|
|
assert!(decoded.persist);
|
|
assert!(decoded.samples.is_empty());
|
|
}
|
|
}
|