135 lines
4.3 KiB
Rust
135 lines
4.3 KiB
Rust
//! socktop agent entrypoint: sets up sysinfo handles and serves a WebSocket endpoint at /ws.
|
|
|
|
mod gpu;
|
|
mod metrics;
|
|
mod proto;
|
|
// sampler module removed (metrics now purely request-driven)
|
|
mod state;
|
|
mod types;
|
|
mod ws;
|
|
|
|
use axum::{Router, http::StatusCode, routing::get};
|
|
use std::net::SocketAddr;
|
|
use std::str::FromStr;
|
|
|
|
mod tls;
|
|
|
|
use state::AppState;
|
|
|
|
fn arg_flag(name: &str) -> bool {
|
|
std::env::args().any(|a| a == name)
|
|
}
|
|
fn arg_value(name: &str) -> Option<String> {
|
|
let mut it = std::env::args();
|
|
while let Some(a) = it.next() {
|
|
if a == name {
|
|
return it.next();
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
fn main() -> anyhow::Result<()> {
|
|
// Install rustls crypto provider before any TLS operations
|
|
// This is required when using axum-server's tls-rustls feature
|
|
rustls::crypto::aws_lc_rs::default_provider()
|
|
.install_default()
|
|
.ok(); // Ignore error if already installed
|
|
|
|
#[cfg(feature = "logging")]
|
|
tracing_subscriber::fmt::init();
|
|
|
|
// Configure Tokio runtime with optimized thread pool for reduced overhead.
|
|
//
|
|
// The agent is primarily I/O-bound (WebSocket, /proc file reads, sysinfo)
|
|
// with no CPU-intensive or blocking operations, so a smaller thread pool
|
|
// is beneficial:
|
|
//
|
|
// Benefits:
|
|
// - Lower memory footprint (~1-2MB per thread saved)
|
|
// - Reduced context switching overhead
|
|
// - Fewer idle threads consuming resources
|
|
// - Better for resource-constrained systems
|
|
//
|
|
// Trade-offs:
|
|
// - Slightly reduced throughput under very high concurrent connections
|
|
// - Could introduce latency if blocking operations are added (don't do this!)
|
|
//
|
|
// Default: 2 threads (sufficient for typical workloads with 1-10 clients)
|
|
// Override: Set SOCKTOP_WORKER_THREADS=4 to use more threads if needed
|
|
//
|
|
// Note: Default Tokio uses num_cpus threads which is excessive for this workload.
|
|
|
|
let worker_threads = std::env::var("SOCKTOP_WORKER_THREADS")
|
|
.ok()
|
|
.and_then(|s| s.parse::<usize>().ok())
|
|
.unwrap_or(2)
|
|
.clamp(1, 16); // Ensure 1-16 threads
|
|
|
|
let runtime = tokio::runtime::Builder::new_multi_thread()
|
|
.worker_threads(worker_threads)
|
|
.thread_name("socktop-agent")
|
|
.enable_all()
|
|
.build()?;
|
|
|
|
runtime.block_on(async_main())
|
|
}
|
|
|
|
async fn async_main() -> anyhow::Result<()> {
|
|
// Version flag (print and exit). Keep before heavy initialization.
|
|
if arg_flag("--version") || arg_flag("-V") {
|
|
println!("socktop_agent {}", env!("CARGO_PKG_VERSION"));
|
|
return Ok(());
|
|
}
|
|
|
|
let state = AppState::new();
|
|
|
|
// No background samplers: metrics collected on-demand per websocket request.
|
|
|
|
// Web app: route /ws to the websocket handler
|
|
async fn healthz() -> StatusCode {
|
|
println!("/healthz request");
|
|
StatusCode::OK
|
|
}
|
|
let app = Router::new()
|
|
.route("/ws", get(ws::ws_handler))
|
|
.route("/healthz", get(healthz))
|
|
.with_state(state.clone());
|
|
|
|
let enable_ssl =
|
|
arg_flag("--enableSSL") || std::env::var("SOCKTOP_ENABLE_SSL").ok().as_deref() == Some("1");
|
|
if enable_ssl {
|
|
// Port can be overridden by --port or SOCKTOP_PORT; default to 8443 when SSL
|
|
let port = arg_value("--port")
|
|
.or_else(|| arg_value("-p"))
|
|
.or_else(|| std::env::var("SOCKTOP_PORT").ok())
|
|
.and_then(|s| s.parse::<u16>().ok())
|
|
.unwrap_or(8443);
|
|
|
|
let (cert_path, key_path) = tls::ensure_self_signed_cert()?;
|
|
let cfg = axum_server::tls_rustls::RustlsConfig::from_pem_file(cert_path, key_path).await?;
|
|
|
|
let addr = SocketAddr::from_str(&format!("0.0.0.0:{port}"))?;
|
|
println!("socktop_agent: TLS enabled. Listening on wss://{addr}/ws");
|
|
axum_server::bind_rustls(addr, cfg)
|
|
.serve(app.into_make_service())
|
|
.await?;
|
|
return Ok(());
|
|
}
|
|
|
|
// Non-TLS HTTP/WS path
|
|
let port = arg_value("--port")
|
|
.or_else(|| arg_value("-p"))
|
|
.or_else(|| std::env::var("SOCKTOP_PORT").ok())
|
|
.and_then(|s| s.parse::<u16>().ok())
|
|
.unwrap_or(3000);
|
|
let addr = SocketAddr::from(([0, 0, 0, 0], port));
|
|
println!("socktop_agent: Listening on ws://{addr}/ws");
|
|
axum_server::bind(addr)
|
|
.serve(app.into_make_service())
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
|
|
// Unit tests for CLI parsing moved to `tests/port_parse.rs`.
|