feat(agent): add rust agent baseline scaffold
This commit is contained in:
Generated
+15
@@ -324,6 +324,21 @@ dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "awatch-agent"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"chrono",
|
||||
"clap",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"tempfile",
|
||||
"tiny_http",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "awatch-agent-rs"
|
||||
version = "0.3.0"
|
||||
|
||||
@@ -3,6 +3,7 @@ resolver = "3"
|
||||
members = [
|
||||
"crates/aw-1c-ingest",
|
||||
"crates/aw-windows-telemetry",
|
||||
"crates/awatch-agent",
|
||||
"crates/awatch-agent-rs",
|
||||
"crates/detmir-auto",
|
||||
"crates/detmir-aw-client",
|
||||
|
||||
@@ -87,8 +87,20 @@ fn run() -> Result<i32> {
|
||||
let transport = TelemetryTransport::new(&config);
|
||||
let aw_worktime = AwWorktimePublisher::new(&config);
|
||||
if cli.flush_spool {
|
||||
let flushed = transport.flush_spool()?;
|
||||
println!("{}", serde_json::json!({"ok": true, "flushed": flushed}));
|
||||
let telemetry_flushed = transport.flush_spool()?;
|
||||
let worktime_flushed = match aw_worktime.as_ref() {
|
||||
Some(publisher) => publisher.flush_spool()?,
|
||||
None => 0,
|
||||
};
|
||||
println!(
|
||||
"{}",
|
||||
serde_json::json!({
|
||||
"ok": true,
|
||||
"flushed": telemetry_flushed + worktime_flushed,
|
||||
"telemetry_flushed": telemetry_flushed,
|
||||
"worktime_flushed": worktime_flushed,
|
||||
})
|
||||
);
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
|
||||
@@ -363,17 +363,27 @@ fn sanitize_file_part(value: &str) -> String {
|
||||
}
|
||||
|
||||
pub fn spool_health(spool_dir: &Path) -> serde_json::Value {
|
||||
let queued = fs::read_dir(spool_dir)
|
||||
let telemetry_queued = count_spool_json_files(spool_dir);
|
||||
let worktime_spool_dir = spool_dir.join("aw-worktime");
|
||||
let worktime_queued = count_spool_json_files(&worktime_spool_dir);
|
||||
serde_json::json!({
|
||||
"generated_at_utc": Utc::now(),
|
||||
"spool_dir": spool_dir.display().to_string(),
|
||||
"worktime_spool_dir": worktime_spool_dir.display().to_string(),
|
||||
"queued": telemetry_queued,
|
||||
"telemetry_queued": telemetry_queued,
|
||||
"worktime_queued": worktime_queued,
|
||||
"total_queued": telemetry_queued + worktime_queued,
|
||||
})
|
||||
}
|
||||
|
||||
fn count_spool_json_files(spool_dir: &Path) -> usize {
|
||||
fs::read_dir(spool_dir)
|
||||
.ok()
|
||||
.into_iter()
|
||||
.flat_map(|entries| entries.filter_map(|entry| entry.ok()))
|
||||
.filter(|entry| entry.path().extension().is_some_and(|ext| ext == "json"))
|
||||
.count();
|
||||
serde_json::json!({
|
||||
"generated_at_utc": Utc::now(),
|
||||
"spool_dir": spool_dir.display().to_string(),
|
||||
"queued": queued,
|
||||
})
|
||||
.count()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -431,6 +441,40 @@ mod tests {
|
||||
assert!(!path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn send_or_spool_preserves_record_when_server_is_unavailable() {
|
||||
let dir = tempdir().unwrap();
|
||||
let config = AgentConfig {
|
||||
server_url: "http://127.0.0.1:9/api/telemetry".to_string(),
|
||||
retry_attempts: 1,
|
||||
timeout_seconds: 1,
|
||||
spool_dir: dir.path().to_path_buf(),
|
||||
..AgentConfig::default()
|
||||
};
|
||||
let transport = TelemetryTransport::new(&config);
|
||||
assert!(transport.send_or_spool(&record()).is_err());
|
||||
let queued = fs::read_dir(dir.path())
|
||||
.unwrap()
|
||||
.filter_map(|entry| entry.ok())
|
||||
.filter(|entry| entry.path().extension().is_some_and(|ext| ext == "json"))
|
||||
.count();
|
||||
assert_eq!(queued, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_spool_keeps_record_when_sender_fails() {
|
||||
let dir = tempdir().unwrap();
|
||||
let config = AgentConfig {
|
||||
spool_dir: dir.path().to_path_buf(),
|
||||
..AgentConfig::default()
|
||||
};
|
||||
let transport = TelemetryTransport::new(&config);
|
||||
let path = transport.spool(&record()).unwrap();
|
||||
let result = flush_spool_dir(dir.path(), |_| anyhow::bail!("transport down"));
|
||||
assert!(result.is_err());
|
||||
assert!(path.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_id_number_extracts_numeric_id() {
|
||||
let session = SessionInfo {
|
||||
@@ -488,4 +532,25 @@ mod tests {
|
||||
assert!(path.starts_with(dir.path().join("aw-worktime")));
|
||||
assert!(path.is_file());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn spool_health_reports_telemetry_and_worktime_backlog() {
|
||||
let dir = tempdir().unwrap();
|
||||
let config = AgentConfig {
|
||||
aw_api_base: Some("http://127.0.0.1:9/api/0".to_string()),
|
||||
aw_worktime_enabled: true,
|
||||
spool_dir: dir.path().to_path_buf(),
|
||||
..AgentConfig::default()
|
||||
};
|
||||
let transport = TelemetryTransport::new(&config);
|
||||
let publisher = AwWorktimePublisher::new(&config).unwrap();
|
||||
transport.spool(&record()).unwrap();
|
||||
publisher.spool(&record()).unwrap();
|
||||
|
||||
let health = spool_health(dir.path());
|
||||
assert_eq!(health["queued"].as_u64(), Some(1));
|
||||
assert_eq!(health["telemetry_queued"].as_u64(), Some(1));
|
||||
assert_eq!(health["worktime_queued"].as_u64(), Some(1));
|
||||
assert_eq!(health["total_queued"].as_u64(), Some(2));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
[package]
|
||||
name = "awatch-agent"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
rust-version.workspace = true
|
||||
license.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
chrono.workspace = true
|
||||
clap.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
tiny_http.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile.workspace = true
|
||||
@@ -0,0 +1,174 @@
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AgentConfig {
|
||||
pub agent_id: String,
|
||||
pub host_id: String,
|
||||
pub platform: String,
|
||||
pub server_url: String,
|
||||
pub spool_dir: PathBuf,
|
||||
pub health_bind: String,
|
||||
pub request_timeout_seconds: u64,
|
||||
pub retry_max_attempts: u32,
|
||||
pub retry_base_backoff_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for AgentConfig {
|
||||
fn default() -> Self {
|
||||
let hostname = local_hostname();
|
||||
Self {
|
||||
agent_id: uuid_from_seed(&format!("agent:{hostname}")),
|
||||
host_id: uuid_from_seed(&format!("host:{hostname}")),
|
||||
platform: current_platform().to_string(),
|
||||
server_url: "http://127.0.0.1:9/api/agent/telemetry".to_string(),
|
||||
spool_dir: default_spool_dir(),
|
||||
health_bind: "127.0.0.1:8787".to_string(),
|
||||
request_timeout_seconds: 10,
|
||||
retry_max_attempts: 3,
|
||||
retry_base_backoff_ms: 250,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
pub fn load(path: &Path) -> Result<Self> {
|
||||
let text = fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?;
|
||||
Self::parse_toml_like(&text)
|
||||
}
|
||||
|
||||
pub fn parse_toml_like(text: &str) -> Result<Self> {
|
||||
let mut config = Self::default();
|
||||
for raw in text.lines() {
|
||||
let line = raw.split('#').next().unwrap_or("").trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let Some((key, value)) = line.split_once('=') else {
|
||||
continue;
|
||||
};
|
||||
let key = key.trim();
|
||||
let value = value.trim().trim_matches('"');
|
||||
match key {
|
||||
"agent_id" => config.agent_id = value.to_string(),
|
||||
"host_id" => config.host_id = value.to_string(),
|
||||
"platform" => config.platform = value.to_string(),
|
||||
"server_url" => config.server_url = value.to_string(),
|
||||
"spool_dir" => config.spool_dir = PathBuf::from(value),
|
||||
"health_bind" => config.health_bind = value.to_string(),
|
||||
"request_timeout_seconds" => {
|
||||
config.request_timeout_seconds = value.parse().unwrap_or(10)
|
||||
}
|
||||
"retry_max_attempts" => config.retry_max_attempts = value.parse().unwrap_or(3),
|
||||
"retry_base_backoff_ms" => {
|
||||
config.retry_base_backoff_ms = value.parse().unwrap_or(250)
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
pub fn request_timeout(&self) -> Duration {
|
||||
Duration::from_secs(self.request_timeout_seconds)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn default_config_path() -> PathBuf {
|
||||
if cfg!(windows) {
|
||||
PathBuf::from(r"C:\ProgramData\AWatch-rus\agent\awatch-agent.toml")
|
||||
} else {
|
||||
PathBuf::from("/etc/awatch-agent/awatch-agent.toml")
|
||||
}
|
||||
}
|
||||
|
||||
fn default_spool_dir() -> PathBuf {
|
||||
if cfg!(windows) {
|
||||
PathBuf::from(r"C:\ProgramData\AWatch-rus\agent\spool")
|
||||
} else {
|
||||
PathBuf::from("/var/lib/awatch-agent/spool")
|
||||
}
|
||||
}
|
||||
|
||||
fn current_platform() -> &'static str {
|
||||
if cfg!(windows) {
|
||||
"windows"
|
||||
} else if cfg!(target_os = "macos") {
|
||||
"macos"
|
||||
} else if cfg!(target_os = "freebsd") {
|
||||
"freebsd"
|
||||
} else {
|
||||
"linux"
|
||||
}
|
||||
}
|
||||
|
||||
fn local_hostname() -> String {
|
||||
std::env::var("COMPUTERNAME")
|
||||
.or_else(|_| std::env::var("HOSTNAME"))
|
||||
.unwrap_or_else(|_| "HOST-EXAMPLE".to_string())
|
||||
}
|
||||
|
||||
fn uuid_from_seed(seed: &str) -> String {
|
||||
let digest = Sha256::digest(seed.as_bytes());
|
||||
let mut bytes = [0_u8; 16];
|
||||
bytes.copy_from_slice(&digest[..16]);
|
||||
bytes[6] = (bytes[6] & 0x0f) | 0x50;
|
||||
bytes[8] = (bytes[8] & 0x3f) | 0x80;
|
||||
format!(
|
||||
"{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
|
||||
bytes[0],
|
||||
bytes[1],
|
||||
bytes[2],
|
||||
bytes[3],
|
||||
bytes[4],
|
||||
bytes[5],
|
||||
bytes[6],
|
||||
bytes[7],
|
||||
bytes[8],
|
||||
bytes[9],
|
||||
bytes[10],
|
||||
bytes[11],
|
||||
bytes[12],
|
||||
bytes[13],
|
||||
bytes[14],
|
||||
bytes[15]
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parses_config_without_external_toml_dependency() {
|
||||
let config = AgentConfig::parse_toml_like(
|
||||
r#"
|
||||
agent_id = "00000000-0000-5000-8000-000000000001"
|
||||
host_id = "00000000-0000-5000-8000-000000000002"
|
||||
platform = "windows"
|
||||
server_url = "https://awatch.example/api/agent/telemetry"
|
||||
spool_dir = "/tmp/awatch-agent-spool"
|
||||
health_bind = "127.0.0.1:8788"
|
||||
request_timeout_seconds = 2
|
||||
retry_max_attempts = 5
|
||||
retry_base_backoff_ms = 50
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(config.platform, "windows");
|
||||
assert_eq!(config.retry_max_attempts, 5);
|
||||
assert_eq!(config.spool_dir, PathBuf::from("/tmp/awatch-agent-spool"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generated_ids_are_uuid_shaped() {
|
||||
let id = uuid_from_seed("HOST-EXAMPLE");
|
||||
assert_eq!(id.len(), 36);
|
||||
assert_eq!(&id[14..15], "5");
|
||||
assert!(matches!(&id[19..20], "8" | "9" | "a" | "b"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::config::AgentConfig;
|
||||
|
||||
pub const AGENT_VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct TelemetryEnvelope {
|
||||
pub agent_id: String,
|
||||
pub host_id: String,
|
||||
pub platform: String,
|
||||
pub timestamp: DateTime<Utc>,
|
||||
pub records: Vec<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct Heartbeat {
|
||||
pub agent_version: String,
|
||||
pub platform: String,
|
||||
pub status: AgentStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum AgentStatus {
|
||||
Online,
|
||||
Degraded,
|
||||
Offline,
|
||||
}
|
||||
|
||||
impl TelemetryEnvelope {
|
||||
pub fn empty(config: &AgentConfig) -> Self {
|
||||
Self {
|
||||
agent_id: config.agent_id.clone(),
|
||||
host_id: config.host_id.clone(),
|
||||
platform: config.platform.clone(),
|
||||
timestamp: Utc::now(),
|
||||
records: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn heartbeat(config: &AgentConfig) -> Self {
|
||||
let heartbeat = Heartbeat {
|
||||
agent_version: AGENT_VERSION.to_string(),
|
||||
platform: config.platform.clone(),
|
||||
status: AgentStatus::Online,
|
||||
};
|
||||
Self {
|
||||
records: vec![serde_json::json!({
|
||||
"type": "heartbeat",
|
||||
"payload": heartbeat,
|
||||
})],
|
||||
..Self::empty(config)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn heartbeat_contract_is_stable_and_contains_no_inventory() {
|
||||
let config = AgentConfig::default();
|
||||
let envelope = TelemetryEnvelope::heartbeat(&config);
|
||||
let value = serde_json::to_value(&envelope).unwrap();
|
||||
assert_eq!(value["agent_id"], config.agent_id);
|
||||
assert_eq!(value["host_id"], config.host_id);
|
||||
assert_eq!(value["platform"], config.platform);
|
||||
assert!(value["records"].is_array());
|
||||
assert_eq!(value["records"][0]["type"], "heartbeat");
|
||||
assert_eq!(value["records"][0]["payload"]["status"], "online");
|
||||
assert!(value["records"][0]["payload"].get("hostname").is_none());
|
||||
assert!(value["records"][0]["payload"].get("processes").is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
use std::net::ToSocketAddrs;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use tiny_http::{Header, Response, Server, StatusCode};
|
||||
|
||||
use crate::envelope::AGENT_VERSION;
|
||||
use crate::metrics::AgentMetrics;
|
||||
|
||||
pub fn serve_health(bind: &str, metrics: AgentMetrics, max_requests: Option<usize>) -> Result<()> {
|
||||
bind.to_socket_addrs()
|
||||
.with_context(|| format!("parse health bind address {bind}"))?;
|
||||
let server =
|
||||
Server::http(bind).map_err(|err| anyhow::anyhow!("bind health endpoint: {err}"))?;
|
||||
let mut served = 0_usize;
|
||||
loop {
|
||||
if max_requests.is_some_and(|limit| served >= limit) {
|
||||
return Ok(());
|
||||
}
|
||||
let Some(request) = server
|
||||
.recv_timeout(Duration::from_millis(250))
|
||||
.map_err(|err| anyhow::anyhow!("receive health request: {err}"))?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
served += 1;
|
||||
let response = match (request.method().as_str(), request.url()) {
|
||||
("GET", "/healthz") => json_response(serde_json::json!({
|
||||
"ok": true,
|
||||
"status": "online",
|
||||
"agent_version": AGENT_VERSION,
|
||||
})),
|
||||
("GET", "/metrics") => text_response(metrics.render_prometheus()),
|
||||
_ => Response::from_string("not found").with_status_code(StatusCode(404)),
|
||||
};
|
||||
request
|
||||
.respond(response)
|
||||
.map_err(|err| anyhow::anyhow!("send health response: {err}"))?;
|
||||
}
|
||||
}
|
||||
|
||||
fn json_response(value: serde_json::Value) -> Response<std::io::Cursor<Vec<u8>>> {
|
||||
let mut response = Response::from_data(serde_json::to_vec(&value).unwrap_or_default());
|
||||
if let Ok(header) = Header::from_bytes("Content-Type", "application/json") {
|
||||
response.add_header(header);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
fn text_response(value: String) -> Response<std::io::Cursor<Vec<u8>>> {
|
||||
let mut response = Response::from_string(value);
|
||||
if let Ok(header) = Header::from_bytes("Content-Type", "text/plain; version=0.0.4") {
|
||||
response.add_header(header);
|
||||
}
|
||||
response
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
use chrono::Utc;
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct LogLine<'a> {
|
||||
timestamp: String,
|
||||
level: &'a str,
|
||||
agent_id: &'a str,
|
||||
component: &'a str,
|
||||
message: &'a str,
|
||||
}
|
||||
|
||||
pub fn log_json(agent_id: &str, level: &str, component: &str, message: &str) {
|
||||
let line = LogLine {
|
||||
timestamp: Utc::now().to_rfc3339(),
|
||||
level,
|
||||
agent_id,
|
||||
component,
|
||||
message,
|
||||
};
|
||||
if let Ok(json) = serde_json::to_string(&line) {
|
||||
eprintln!("{json}");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn structured_log_shape_is_serializable() {
|
||||
let line = LogLine {
|
||||
timestamp: "2026-06-07T00:00:00Z".to_string(),
|
||||
level: "INFO",
|
||||
agent_id: "agent-1",
|
||||
component: "spool",
|
||||
message: "queued",
|
||||
};
|
||||
let value = serde_json::to_value(line).unwrap();
|
||||
assert_eq!(value["level"], "INFO");
|
||||
assert_eq!(value["component"], "spool");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
mod config;
|
||||
mod envelope;
|
||||
mod health;
|
||||
mod logging;
|
||||
mod metrics;
|
||||
mod spool;
|
||||
mod transport;
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use clap::Parser;
|
||||
|
||||
use config::{AgentConfig, default_config_path};
|
||||
use envelope::TelemetryEnvelope;
|
||||
use logging::log_json;
|
||||
use spool::LocalSpool;
|
||||
|
||||
#[derive(Debug, Parser)]
|
||||
#[command(about = "AWatch-rus Rust agent baseline scaffold")]
|
||||
struct Cli {
|
||||
#[arg(long, env = "AWATCH_AGENT_CONFIG")]
|
||||
config: Option<PathBuf>,
|
||||
|
||||
#[arg(long, env = "AWATCH_AGENT_SERVER_URL")]
|
||||
server_url: Option<String>,
|
||||
|
||||
#[arg(long, env = "AWATCH_AGENT_SPOOL_DIR")]
|
||||
spool_dir: Option<PathBuf>,
|
||||
|
||||
#[arg(long)]
|
||||
enqueue_heartbeat: bool,
|
||||
|
||||
#[arg(long)]
|
||||
flush_spool: bool,
|
||||
|
||||
#[arg(long)]
|
||||
metrics: bool,
|
||||
|
||||
#[arg(long)]
|
||||
healthz: bool,
|
||||
|
||||
#[arg(long)]
|
||||
print_envelope: bool,
|
||||
|
||||
#[arg(long)]
|
||||
max_health_requests: Option<usize>,
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let code = match run() {
|
||||
Ok(code) => code,
|
||||
Err(err) => {
|
||||
eprintln!("{err:#}");
|
||||
1
|
||||
}
|
||||
};
|
||||
std::process::exit(code);
|
||||
}
|
||||
|
||||
fn run() -> Result<i32> {
|
||||
let cli = Cli::parse();
|
||||
let mut config = load_config(cli.config.as_ref())?;
|
||||
if let Some(server_url) = cli.server_url {
|
||||
config.server_url = server_url;
|
||||
}
|
||||
if let Some(spool_dir) = cli.spool_dir {
|
||||
config.spool_dir = spool_dir;
|
||||
}
|
||||
let spool = LocalSpool::new(config.spool_dir.clone());
|
||||
let mut metrics = spool.metrics().unwrap_or_default();
|
||||
|
||||
if cli.print_envelope {
|
||||
println!(
|
||||
"{}",
|
||||
serde_json::to_string_pretty(&TelemetryEnvelope::heartbeat(&config))?
|
||||
);
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
if cli.enqueue_heartbeat {
|
||||
spool.enqueue(TelemetryEnvelope::heartbeat(&config))?;
|
||||
metrics.heartbeat_sent = metrics.heartbeat_sent.saturating_add(1);
|
||||
log_json(
|
||||
&config.agent_id,
|
||||
"INFO",
|
||||
"heartbeat",
|
||||
"heartbeat envelope queued",
|
||||
);
|
||||
}
|
||||
|
||||
if cli.flush_spool {
|
||||
let summary = transport::flush_with_retry(&config, &spool, &mut metrics)?;
|
||||
println!("{}", serde_json::to_string_pretty(&summary)?);
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
if cli.metrics {
|
||||
let mut current = spool.metrics()?;
|
||||
current.heartbeat_sent = metrics.heartbeat_sent;
|
||||
current.retry_count = metrics.retry_count;
|
||||
print!("{}", current.render_prometheus());
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
if cli.healthz {
|
||||
health::serve_health(
|
||||
&config.health_bind,
|
||||
spool.metrics()?,
|
||||
cli.max_health_requests,
|
||||
)?;
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
if !cli.enqueue_heartbeat {
|
||||
log_json(&config.agent_id, "INFO", "agent", "no action requested");
|
||||
}
|
||||
Ok(0)
|
||||
}
|
||||
|
||||
fn load_config(path: Option<&PathBuf>) -> Result<AgentConfig> {
|
||||
let path = path.cloned().unwrap_or_else(default_config_path);
|
||||
if path.exists() {
|
||||
AgentConfig::load(&path)
|
||||
} else {
|
||||
AgentConfig::parse_toml_like("")
|
||||
.with_context(|| format!("load default config because {} is absent", path.display()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn default_run_has_no_monitoring_side_effect() {
|
||||
let config = AgentConfig::parse_toml_like("").unwrap();
|
||||
let envelope = TelemetryEnvelope::heartbeat(&config);
|
||||
assert_eq!(envelope.records.len(), 1);
|
||||
assert!(envelope.records[0].get("processes").is_none());
|
||||
assert!(envelope.records[0].get("screenshots").is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct AgentMetrics {
|
||||
pub queued_records: usize,
|
||||
pub retry_count: u64,
|
||||
pub heartbeat_sent: u64,
|
||||
pub spool_size: u64,
|
||||
}
|
||||
|
||||
impl AgentMetrics {
|
||||
pub fn render_prometheus(&self) -> String {
|
||||
format!(
|
||||
concat!(
|
||||
"# HELP awatch_agent_queued_records Local spool records waiting for delivery.\n",
|
||||
"# TYPE awatch_agent_queued_records gauge\n",
|
||||
"awatch_agent_queued_records {}\n",
|
||||
"# HELP awatch_agent_retry_count Total retry attempts performed by the agent.\n",
|
||||
"# TYPE awatch_agent_retry_count counter\n",
|
||||
"awatch_agent_retry_count {}\n",
|
||||
"# HELP awatch_agent_heartbeat_sent Heartbeat envelopes generated by the agent.\n",
|
||||
"# TYPE awatch_agent_heartbeat_sent counter\n",
|
||||
"awatch_agent_heartbeat_sent {}\n",
|
||||
"# HELP awatch_agent_spool_size Local spool size in bytes.\n",
|
||||
"# TYPE awatch_agent_spool_size gauge\n",
|
||||
"awatch_agent_spool_size {}\n"
|
||||
),
|
||||
self.queued_records, self.retry_count, self.heartbeat_sent, self.spool_size
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn renders_prometheus_metrics() {
|
||||
let metrics = AgentMetrics {
|
||||
queued_records: 2,
|
||||
retry_count: 3,
|
||||
heartbeat_sent: 1,
|
||||
spool_size: 512,
|
||||
};
|
||||
let rendered = metrics.render_prometheus();
|
||||
assert!(rendered.contains("awatch_agent_queued_records 2"));
|
||||
assert!(rendered.contains("awatch_agent_retry_count 3"));
|
||||
assert!(rendered.contains("awatch_agent_heartbeat_sent 1"));
|
||||
assert!(rendered.contains("awatch_agent_spool_size 512"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::envelope::TelemetryEnvelope;
|
||||
use crate::metrics::AgentMetrics;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LocalSpool {
|
||||
root: PathBuf,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct SpoolItem {
|
||||
pub envelope: TelemetryEnvelope,
|
||||
pub enqueued_at: DateTime<Utc>,
|
||||
pub retry_count: u32,
|
||||
pub last_error: Option<String>,
|
||||
}
|
||||
|
||||
impl LocalSpool {
|
||||
pub fn new(root: impl Into<PathBuf>) -> Self {
|
||||
Self { root: root.into() }
|
||||
}
|
||||
|
||||
pub fn enqueue(&self, envelope: TelemetryEnvelope) -> Result<PathBuf> {
|
||||
self.ensure_dirs()?;
|
||||
let item = SpoolItem {
|
||||
envelope,
|
||||
enqueued_at: Utc::now(),
|
||||
retry_count: 0,
|
||||
last_error: None,
|
||||
};
|
||||
let file_name = format!(
|
||||
"{}-{}.json",
|
||||
item.enqueued_at.format("%Y%m%dT%H%M%S%.3fZ"),
|
||||
sanitize_file_part(&item.envelope.agent_id)
|
||||
);
|
||||
let path = self.pending_dir().join(file_name);
|
||||
write_json_atomic(&path, &item)?;
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
pub fn pending_paths(&self) -> Result<Vec<PathBuf>> {
|
||||
read_json_paths(&self.pending_dir())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn dead_letter_paths(&self) -> Result<Vec<PathBuf>> {
|
||||
read_json_paths(&self.dead_letter_dir())
|
||||
}
|
||||
|
||||
pub fn metrics(&self) -> Result<AgentMetrics> {
|
||||
let paths = self.pending_paths()?;
|
||||
let spool_size = paths
|
||||
.iter()
|
||||
.filter_map(|path| fs::metadata(path).ok())
|
||||
.map(|metadata| metadata.len())
|
||||
.sum();
|
||||
Ok(AgentMetrics {
|
||||
queued_records: paths.len(),
|
||||
spool_size,
|
||||
..AgentMetrics::default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn process_pending<F>(&self, max_retry_count: u32, mut sender: F) -> Result<FlushSummary>
|
||||
where
|
||||
F: FnMut(&TelemetryEnvelope) -> Result<()>,
|
||||
{
|
||||
self.ensure_dirs()?;
|
||||
let mut summary = FlushSummary::default();
|
||||
for path in self.pending_paths()? {
|
||||
let bytes = fs::read(&path).with_context(|| format!("read {}", path.display()))?;
|
||||
let mut item = match serde_json::from_slice::<SpoolItem>(&bytes) {
|
||||
Ok(item) => item,
|
||||
Err(err) => {
|
||||
self.move_to_dead_letter(&path, Some(format!("corrupt json: {err}")))?;
|
||||
summary.corrupt += 1;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
match sender(&item.envelope) {
|
||||
Ok(()) => {
|
||||
fs::remove_file(&path)
|
||||
.with_context(|| format!("remove delivered {}", path.display()))?;
|
||||
summary.delivered += 1;
|
||||
}
|
||||
Err(err) => {
|
||||
item.retry_count = item.retry_count.saturating_add(1);
|
||||
item.last_error = Some(err.to_string());
|
||||
summary.retried += 1;
|
||||
if item.retry_count >= max_retry_count {
|
||||
write_json_atomic(&path, &item)?;
|
||||
self.move_to_dead_letter(&path, item.last_error.clone())?;
|
||||
summary.dead_lettered += 1;
|
||||
} else {
|
||||
write_json_atomic(&path, &item)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
fn ensure_dirs(&self) -> Result<()> {
|
||||
fs::create_dir_all(self.pending_dir())
|
||||
.with_context(|| format!("create {}", self.pending_dir().display()))?;
|
||||
fs::create_dir_all(self.dead_letter_dir())
|
||||
.with_context(|| format!("create {}", self.dead_letter_dir().display()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn pending_dir(&self) -> PathBuf {
|
||||
self.root.join("pending")
|
||||
}
|
||||
|
||||
fn dead_letter_dir(&self) -> PathBuf {
|
||||
self.root.join("dead-letter")
|
||||
}
|
||||
|
||||
fn move_to_dead_letter(&self, path: &Path, reason: Option<String>) -> Result<()> {
|
||||
self.ensure_dirs()?;
|
||||
let file_name = path
|
||||
.file_name()
|
||||
.map(|value| value.to_string_lossy().to_string())
|
||||
.unwrap_or_else(|| format!("{}.json", Utc::now().timestamp_millis()));
|
||||
let target = self.dead_letter_dir().join(file_name);
|
||||
if let Some(reason) = reason {
|
||||
let note_path = target.with_extension("reason.txt");
|
||||
fs::write(note_path, reason)?;
|
||||
}
|
||||
fs::rename(path, target).or_else(|_| {
|
||||
fs::copy(path, self.dead_letter_dir().join("recovered-corrupt.json"))?;
|
||||
fs::remove_file(path)
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct FlushSummary {
|
||||
pub delivered: usize,
|
||||
pub retried: usize,
|
||||
pub dead_lettered: usize,
|
||||
pub corrupt: usize,
|
||||
}
|
||||
|
||||
fn read_json_paths(dir: &Path) -> Result<Vec<PathBuf>> {
|
||||
if !dir.exists() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut paths = fs::read_dir(dir)
|
||||
.with_context(|| format!("read {}", dir.display()))?
|
||||
.filter_map(|entry| entry.ok())
|
||||
.map(|entry| entry.path())
|
||||
.filter(|path| path.extension().is_some_and(|ext| ext == "json"))
|
||||
.collect::<Vec<_>>();
|
||||
paths.sort();
|
||||
Ok(paths)
|
||||
}
|
||||
|
||||
fn write_json_atomic<T: Serialize>(path: &Path, value: &T) -> Result<()> {
|
||||
let tmp = path.with_extension("json.tmp");
|
||||
fs::write(&tmp, serde_json::to_vec_pretty(value)?)
|
||||
.with_context(|| format!("write {}", tmp.display()))?;
|
||||
fs::rename(&tmp, path)
|
||||
.with_context(|| format!("rename {} to {}", tmp.display(), path.display()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn sanitize_file_part(value: &str) -> String {
|
||||
value
|
||||
.chars()
|
||||
.map(|ch| {
|
||||
if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
|
||||
ch
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use anyhow::anyhow;
|
||||
use tempfile::tempdir;
|
||||
|
||||
use super::*;
|
||||
use crate::config::AgentConfig;
|
||||
|
||||
fn envelope() -> TelemetryEnvelope {
|
||||
TelemetryEnvelope::heartbeat(&AgentConfig::default())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enqueues_and_delivers_spool_item() {
|
||||
let dir = tempdir().unwrap();
|
||||
let spool = LocalSpool::new(dir.path());
|
||||
spool.enqueue(envelope()).unwrap();
|
||||
assert_eq!(spool.pending_paths().unwrap().len(), 1);
|
||||
|
||||
let summary = spool.process_pending(3, |_| Ok(())).unwrap();
|
||||
assert_eq!(summary.delivered, 1);
|
||||
assert_eq!(spool.pending_paths().unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retry_keeps_item_until_max_retry_then_dead_letters() {
|
||||
let dir = tempdir().unwrap();
|
||||
let spool = LocalSpool::new(dir.path());
|
||||
spool.enqueue(envelope()).unwrap();
|
||||
|
||||
let first = spool
|
||||
.process_pending(2, |_| Err(anyhow!("transport down")))
|
||||
.unwrap();
|
||||
assert_eq!(first.retried, 1);
|
||||
assert_eq!(first.dead_lettered, 0);
|
||||
assert_eq!(spool.pending_paths().unwrap().len(), 1);
|
||||
|
||||
let second = spool
|
||||
.process_pending(2, |_| Err(anyhow!("transport down")))
|
||||
.unwrap();
|
||||
assert_eq!(second.dead_lettered, 1);
|
||||
assert_eq!(spool.pending_paths().unwrap().len(), 0);
|
||||
assert_eq!(spool.dead_letter_paths().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corrupt_spool_item_moves_to_dead_letter() {
|
||||
let dir = tempdir().unwrap();
|
||||
let spool = LocalSpool::new(dir.path());
|
||||
fs::create_dir_all(dir.path().join("pending")).unwrap();
|
||||
fs::write(dir.path().join("pending/bad.json"), b"{not-json").unwrap();
|
||||
|
||||
let summary = spool.process_pending(3, |_| Ok(())).unwrap();
|
||||
assert_eq!(summary.corrupt, 1);
|
||||
assert_eq!(spool.pending_paths().unwrap().len(), 0);
|
||||
assert_eq!(spool.dead_letter_paths().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn metrics_report_queue_and_size() {
|
||||
let dir = tempdir().unwrap();
|
||||
let spool = LocalSpool::new(dir.path());
|
||||
spool.enqueue(envelope()).unwrap();
|
||||
let metrics = spool.metrics().unwrap();
|
||||
assert_eq!(metrics.queued_records, 1);
|
||||
assert!(metrics.spool_size > 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use reqwest::blocking::Client;
|
||||
|
||||
use crate::config::AgentConfig;
|
||||
use crate::envelope::TelemetryEnvelope;
|
||||
use crate::metrics::AgentMetrics;
|
||||
use crate::spool::{FlushSummary, LocalSpool};
|
||||
|
||||
pub fn send_envelope(config: &AgentConfig, envelope: &TelemetryEnvelope) -> Result<()> {
|
||||
let client = Client::builder()
|
||||
.timeout(config.request_timeout())
|
||||
.build()
|
||||
.context("build agent telemetry HTTP client")?;
|
||||
client
|
||||
.post(&config.server_url)
|
||||
.json(envelope)
|
||||
.send()
|
||||
.and_then(|response| response.error_for_status())
|
||||
.map(|_| ())
|
||||
.map_err(|err| anyhow!("agent telemetry POST failed: {err}"))
|
||||
}
|
||||
|
||||
pub fn flush_with_retry(
|
||||
config: &AgentConfig,
|
||||
spool: &LocalSpool,
|
||||
metrics: &mut AgentMetrics,
|
||||
) -> Result<FlushSummary> {
|
||||
let mut attempt = 0_u32;
|
||||
loop {
|
||||
let summary = spool.process_pending(config.retry_max_attempts, |envelope| {
|
||||
send_envelope(config, envelope)
|
||||
})?;
|
||||
metrics.retry_count = metrics
|
||||
.retry_count
|
||||
.saturating_add(u64::try_from(summary.retried).unwrap_or(u64::MAX));
|
||||
if summary.retried == 0 || attempt + 1 >= config.retry_max_attempts {
|
||||
return Ok(summary);
|
||||
}
|
||||
let backoff = exponential_backoff(config.retry_base_backoff_ms, attempt);
|
||||
thread::sleep(backoff);
|
||||
attempt += 1;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn exponential_backoff(base_ms: u64, attempt: u32) -> Duration {
|
||||
let factor = 1_u64.checked_shl(attempt.min(10)).unwrap_or(1024);
|
||||
Duration::from_millis(base_ms.saturating_mul(factor))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn backoff_is_exponential_and_bounded() {
|
||||
assert_eq!(exponential_backoff(100, 0), Duration::from_millis(100));
|
||||
assert_eq!(exponential_backoff(100, 3), Duration::from_millis(800));
|
||||
assert_eq!(exponential_backoff(100, 99), Duration::from_millis(102400));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user