diff --git a/CHANGELOG.md b/CHANGELOG.md index d097589..5b7662f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,12 +11,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `nettrack` BPF module for network I/O monitoring Config options: `frequency`: sampling frequency in Hz - `scaled`: report rates in bits/s instead of cumulative bytes + `scaled`: report rates in bits/s instead of interval bytes +- Configurable `nettrack` protocols, directions, metrics, and host/process/peer scopes +- Bounded per-process and per-peer network aggregation with independent peer sampling frequency ### Fixed +- Count successful TCP and UDP payload bytes instead of requested send sizes +- Replace `nettrack` perf-event and ring-buffer boundaries with cumulative map snapshots - Spawn a separate thread for timer interrupt to fix starvation issues ([189c785](https://github.com/category-labs/manytrace/commit/189c785)) ## [0.1.1] - 2025-07-11 ### Added -- Initial release with core functionality \ No newline at end of file +- Initial release with core functionality diff --git a/README.md b/README.md index ebe25a7..d7267df 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,28 @@ counters = ["cpu-cycles", "instructions", "cache-misses", "ipc"] # page-faults, context-switches, cpu-migrations, ipc ``` +#### network tracking + +collects successful TCP and UDP application payload bytes. counters are aggregated in BPF maps +and sampled from userspace, so network tracking does not use periodic per-CPU interrupts or a +ring buffer. host byte throughput is the low-overhead default; process and peer aggregation are +independent opt-in scopes. + +```toml +[bpf.nettrack] +frequency = 5 +peer_frequency = 1 +scaled = true +protocols = ["tcp", "udp"] +directions = ["send", "receive"] +scopes = ["host"] # add "process" or "peer" for more detail +metrics = ["bytes"] # optional: "operations", "errors" +``` + +process-name filters require `bpf.thread_tracker`. peer tracks contain remote IP addresses and +ports and may make traces sensitive; their kernel-map and Perfetto cardinality are bounded by +`max_peer_entries` and `max_peer_tracks`. + #### user tracing ![spans](_assets/manytrace_spans.png) @@ -103,4 +125,4 @@ random_process_id = true ## License -Licensed under the GNU General Public License ([LICENSE](LICENSE) or https://www.gnu.org/licenses/gpl-3.0.html). \ No newline at end of file +Licensed under the GNU General Public License ([LICENSE](LICENSE) or https://www.gnu.org/licenses/gpl-3.0.html). diff --git a/bpf/src/bpf/nettrack.bpf.c b/bpf/src/bpf/nettrack.bpf.c index 2fdbec1..af8e2c5 100644 --- a/bpf/src/bpf/nettrack.bpf.c +++ b/bpf/src/bpf/nettrack.bpf.c @@ -14,149 +14,369 @@ // along with this program. If not, see . #include "vmlinux.h" +#include +#include #include #include -#include char LICENSE[] SEC("license") = "GPL"; +#define AF_INET 2 +#define AF_INET6 10 +#define EAGAIN 11 +#define MSG_TRUNC 0x20 + +enum nettrack_protocol { + NET_PROTOCOL_TCP = 0, + NET_PROTOCOL_UDP = 1, +}; + struct net_stats { - u64 tcp_send_bytes; - u64 tcp_recv_bytes; - u64 udp_send_bytes; - u64 udp_recv_bytes; - u64 last_update_ns; + u64 send_bytes; + u64 recv_bytes; + u64 send_operations; + u64 recv_operations; + u64 send_errors; + u64 recv_errors; + u64 send_would_block; + u64 recv_would_block; + u64 tcp_recv_eof; + u64 udp_recv_truncated; + u64 first_seen_ns; + u64 last_seen_ns; }; -struct net_event { - u64 timestamp; - u64 tcp_send_bytes; - u64 tcp_recv_bytes; - u64 udp_send_bytes; - u64 udp_recv_bytes; - u64 total_send_bytes; - u64 total_recv_bytes; - u32 cpu; +struct process_key { + u32 tgid; + u8 protocol; + u8 padding[3]; +}; + +struct peer_key { + u32 tgid; + u8 protocol; + u8 family; + u16 local_port; + u16 remote_port; + u16 padding; + u8 remote_addr[16]; }; struct { __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY); - __uint(max_entries, 1); + __uint(max_entries, 2); __type(key, u32); __type(value, struct net_stats); -} net_stats SEC(".maps"); +} global_stats SEC(".maps"); struct { - __uint(type, BPF_MAP_TYPE_RINGBUF); - __uint(max_entries, 1024 * 1024); -} events SEC(".maps"); + __uint(type, BPF_MAP_TYPE_LRU_HASH); + __uint(max_entries, 1024); + __type(key, struct process_key); + __type(value, struct net_stats); +} process_stats SEC(".maps"); + +struct { + __uint(type, BPF_MAP_TYPE_LRU_HASH); + __uint(max_entries, 4096); + __type(key, struct peer_key); + __type(value, struct net_stats); +} peer_stats SEC(".maps"); + +struct { + __uint(type, BPF_MAP_TYPE_HASH); + __uint(max_entries, 1024); + __type(key, u32); + __type(value, u32); +} tracked_tgids SEC(".maps"); + +const volatile struct { + bool tcp_enabled; + bool udp_enabled; + bool send_enabled; + bool receive_enabled; + bool host_enabled; + bool process_enabled; + bool peer_enabled; + bool bytes_enabled; + bool operations_enabled; + bool errors_enabled; + bool filter_enabled; +} cfg = { + .tcp_enabled = true, + .udp_enabled = true, + .send_enabled = true, + .receive_enabled = true, + .host_enabled = true, + .process_enabled = false, + .peer_enabled = false, + .bytes_enabled = true, + .operations_enabled = false, + .errors_enabled = false, + .filter_enabled = false, +}; + +static __always_inline bool should_track_tgid(u32 tgid) +{ + if (tgid == 0) + return false; + if (!cfg.filter_enabled) + return true; + return bpf_map_lookup_elem(&tracked_tgids, &tgid) != NULL; +} -static __always_inline u64 get_time() { - return bpf_ktime_get_ns(); +static __always_inline void initialize_stats(struct net_stats *stats, u64 now) +{ + __builtin_memset(stats, 0, sizeof(*stats)); + stats->first_seen_ns = now; + stats->last_seen_ns = now; } -static __always_inline void submit_net_event(struct net_stats *stats, u64 now) { - struct net_event *e; - - e = bpf_ringbuf_reserve(&events, sizeof(*e), 0); - if (!e) { +static __always_inline void update_stats(struct net_stats *stats, bool send, int ret, + int protocol, u32 msg_flags, u64 now, + bool shared) +{ + if (shared) + __sync_lock_test_and_set(&stats->last_seen_ns, now); + else + stats->last_seen_ns = now; + + if (ret < 0) { + if (!cfg.errors_enabled) + return; + if (send) { + if (shared) + __sync_fetch_and_add(&stats->send_errors, 1); + else + stats->send_errors++; + if (ret == -EAGAIN) { + if (shared) + __sync_fetch_and_add(&stats->send_would_block, 1); + else + stats->send_would_block++; + } + } else { + if (shared) + __sync_fetch_and_add(&stats->recv_errors, 1); + else + stats->recv_errors++; + if (ret == -EAGAIN) { + if (shared) + __sync_fetch_and_add(&stats->recv_would_block, 1); + else + stats->recv_would_block++; + } + } return; } - - e->timestamp = now; - e->tcp_send_bytes = stats->tcp_send_bytes; - e->tcp_recv_bytes = stats->tcp_recv_bytes; - e->udp_send_bytes = stats->udp_send_bytes; - e->udp_recv_bytes = stats->udp_recv_bytes; - e->total_send_bytes = stats->tcp_send_bytes + stats->udp_send_bytes; - e->total_recv_bytes = stats->tcp_recv_bytes + stats->udp_recv_bytes; - e->cpu = bpf_get_smp_processor_id(); - - bpf_ringbuf_submit(e, 0); - - stats->tcp_send_bytes = 0; - stats->tcp_recv_bytes = 0; - stats->udp_send_bytes = 0; - stats->udp_recv_bytes = 0; - stats->last_update_ns = now; + + if (send) { + if (cfg.bytes_enabled && ret > 0) { + if (shared) + __sync_fetch_and_add(&stats->send_bytes, (u64)ret); + else + stats->send_bytes += (u64)ret; + } + if (cfg.operations_enabled) { + if (shared) + __sync_fetch_and_add(&stats->send_operations, 1); + else + stats->send_operations++; + } + } else { + if (cfg.bytes_enabled && ret > 0) { + if (shared) + __sync_fetch_and_add(&stats->recv_bytes, (u64)ret); + else + stats->recv_bytes += (u64)ret; + } + if (cfg.operations_enabled) { + if (shared) + __sync_fetch_and_add(&stats->recv_operations, 1); + else + stats->recv_operations++; + + if (protocol == NET_PROTOCOL_TCP && ret == 0) { + if (shared) + __sync_fetch_and_add(&stats->tcp_recv_eof, 1); + else + stats->tcp_recv_eof++; + } + if (protocol == NET_PROTOCOL_UDP && (msg_flags & MSG_TRUNC)) { + if (shared) + __sync_fetch_and_add(&stats->udp_recv_truncated, 1); + else + stats->udp_recv_truncated++; + } + } + } } -SEC("kprobe/udp_sendmsg") -int BPF_KPROBE(trace_udp_sendmsg, struct sock *sk, struct msghdr *msg, size_t size) +static __always_inline void update_process_stats(u32 tgid, int protocol, bool send, + int ret, u32 msg_flags, u64 now) { - u32 zero = 0; - struct net_stats *stats = bpf_map_lookup_elem(&net_stats, &zero); - if (!stats) { - return 0; + struct process_key key = { + .tgid = tgid, + .protocol = protocol, + }; + struct net_stats *stats = bpf_map_lookup_elem(&process_stats, &key); + if (stats) { + update_stats(stats, send, ret, protocol, msg_flags, now, true); + return; } - - stats->udp_send_bytes += size; - - return 0; + + struct net_stats initial; + initialize_stats(&initial, now); + update_stats(&initial, send, ret, protocol, msg_flags, now, false); + if (bpf_map_update_elem(&process_stats, &key, &initial, BPF_NOEXIST) == 0) + return; + + stats = bpf_map_lookup_elem(&process_stats, &key); + if (stats) + update_stats(stats, send, ret, protocol, msg_flags, now, true); } -SEC("kretprobe/udp_recvmsg") -int BPF_KRETPROBE(trace_udp_recvmsg_ret, int ret) +static __always_inline bool peer_from_sock(struct sock *sk, struct peer_key *key) { - u32 zero = 0; - struct net_stats *stats = bpf_map_lookup_elem(&net_stats, &zero); - if (!stats) { - return 0; + u16 family = BPF_CORE_READ(sk, __sk_common.skc_family); + key->family = family; + key->local_port = BPF_CORE_READ(sk, __sk_common.skc_num); + key->remote_port = bpf_ntohs(BPF_CORE_READ(sk, __sk_common.skc_dport)); + + if (family == AF_INET) { + u32 address = BPF_CORE_READ(sk, __sk_common.skc_daddr); + __builtin_memcpy(key->remote_addr, &address, sizeof(address)); + return address != 0 && key->remote_port != 0; } - - if (ret > 0) { - stats->udp_recv_bytes += (u64)ret; + if (family == AF_INET6) { + BPF_CORE_READ_INTO(key->remote_addr, sk, __sk_common.skc_v6_daddr); + return key->remote_port != 0; } - - return 0; + return false; } -SEC("kprobe/tcp_sendmsg") -int BPF_KPROBE(trace_tcp_sendmsg, struct sock *sk, struct msghdr *msg, size_t size) +static __always_inline bool peer_from_message(struct msghdr *msg, struct peer_key *key) { - u32 zero = 0; - struct net_stats *stats = bpf_map_lookup_elem(&net_stats, &zero); - if (!stats) { - return 0; + void *name; + u16 family; + + if (!msg) + return false; + name = BPF_CORE_READ(msg, msg_name); + if (!name) + return false; + if (bpf_probe_read_kernel(&family, sizeof(family), name) != 0) + return false; + + if (family == AF_INET) { + struct sockaddr_in address; + if (bpf_probe_read_kernel(&address, sizeof(address), name) != 0) + return false; + key->family = AF_INET; + key->remote_port = bpf_ntohs(address.sin_port); + __builtin_memcpy(key->remote_addr, &address.sin_addr, sizeof(address.sin_addr)); + return key->remote_port != 0; } - - stats->tcp_send_bytes += size; - - return 0; + if (family == AF_INET6) { + struct sockaddr_in6 address; + if (bpf_probe_read_kernel(&address, sizeof(address), name) != 0) + return false; + key->family = AF_INET6; + key->remote_port = bpf_ntohs(address.sin6_port); + __builtin_memcpy(key->remote_addr, &address.sin6_addr, sizeof(address.sin6_addr)); + return key->remote_port != 0; + } + return false; } -SEC("kretprobe/tcp_recvmsg") -int BPF_KRETPROBE(trace_tcp_recvmsg_ret, int ret) +static __always_inline void update_peer_stats(struct sock *sk, struct msghdr *msg, + u32 tgid, int protocol, bool send, + int ret, u32 msg_flags, u64 now) { - u32 zero = 0; - struct net_stats *stats = bpf_map_lookup_elem(&net_stats, &zero); - if (!stats) { - return 0; - } - - if (ret > 0) { - stats->tcp_recv_bytes += (u64)ret; + struct peer_key key = { + .tgid = tgid, + .protocol = protocol, + }; + + key.local_port = BPF_CORE_READ(sk, __sk_common.skc_num); + bool found = false; + if (protocol == NET_PROTOCOL_UDP) + found = peer_from_message(msg, &key); + if (!found) + found = peer_from_sock(sk, &key); + if (!found) + return; + + struct net_stats *stats = bpf_map_lookup_elem(&peer_stats, &key); + if (stats) { + update_stats(stats, send, ret, protocol, msg_flags, now, true); + return; } - - return 0; + + struct net_stats initial; + initialize_stats(&initial, now); + update_stats(&initial, send, ret, protocol, msg_flags, now, false); + if (bpf_map_update_elem(&peer_stats, &key, &initial, BPF_NOEXIST) == 0) + return; + + stats = bpf_map_lookup_elem(&peer_stats, &key); + if (stats) + update_stats(stats, send, ret, protocol, msg_flags, now, true); } -SEC("perf_event") -int handle_boundary_event(void *ctx) +static __always_inline int account_socket_io(struct sock *sk, struct msghdr *msg, + int protocol, bool send, int ret) { - u32 zero = 0; - struct net_stats *stats = bpf_map_lookup_elem(&net_stats, &zero); - if (!stats) { + if ((send && !cfg.send_enabled) || (!send && !cfg.receive_enabled)) return 0; - } - - u64 now = get_time(); - if (stats->last_update_ns == 0) { - stats->last_update_ns = now; + + u64 pid_tgid = bpf_get_current_pid_tgid(); + u32 tgid = pid_tgid >> 32; + if (!should_track_tgid(tgid)) return 0; + + if (!sk || (protocol == NET_PROTOCOL_TCP && !cfg.tcp_enabled) + || (protocol == NET_PROTOCOL_UDP && !cfg.udp_enabled)) + return 0; + + u64 now = bpf_ktime_get_ns(); + u32 msg_flags = msg ? BPF_CORE_READ(msg, msg_flags) : 0; + + if (cfg.host_enabled) { + u32 key = protocol; + struct net_stats *stats = bpf_map_lookup_elem(&global_stats, &key); + if (stats) + update_stats(stats, send, ret, protocol, msg_flags, now, false); } - - submit_net_event(stats, now); - + if (cfg.process_enabled) + update_process_stats(tgid, protocol, send, ret, msg_flags, now); + if (cfg.peer_enabled) + update_peer_stats(sk, msg, tgid, protocol, send, ret, msg_flags, now); return 0; } + +SEC("fexit/tcp_sendmsg") +int BPF_PROG(trace_tcp_sendmsg, struct sock *sk, struct msghdr *msg, size_t size, int ret) +{ + return account_socket_io(sk, msg, NET_PROTOCOL_TCP, true, ret); +} + +SEC("fexit/tcp_recvmsg") +int BPF_PROG(trace_tcp_recvmsg, struct sock *sk, struct msghdr *msg, size_t len, + int flags, int *addr_len, int ret) +{ + return account_socket_io(sk, msg, NET_PROTOCOL_TCP, false, ret); +} + +SEC("fexit/udp_sendmsg") +int BPF_PROG(trace_udp_sendmsg, struct sock *sk, struct msghdr *msg, size_t len, int ret) +{ + return account_socket_io(sk, msg, NET_PROTOCOL_UDP, true, ret); +} + +SEC("fexit/udp_recvmsg") +int BPF_PROG(trace_udp_recvmsg, struct sock *sk, struct msghdr *msg, size_t len, + int flags, int *addr_len, int ret) +{ + return account_socket_io(sk, msg, NET_PROTOCOL_UDP, false, ret); +} diff --git a/bpf/src/lib.rs b/bpf/src/lib.rs index 8f8fa13..83e39f3 100644 --- a/bpf/src/lib.rs +++ b/bpf/src/lib.rs @@ -46,7 +46,7 @@ fn get_monotonic_timestamp() -> u64 { } pub use cpuutil::CpuUtilConfig; -pub use nettrack::NetTrackConfig; +pub use nettrack::{NetDirection, NetMetric, NetProtocol, NetScope, NetTrackConfig}; pub use perfcounter::PerfCounterConfig; pub use profiler::ProfilerConfig; pub use schedtrace::SchedTraceConfig; @@ -124,7 +124,12 @@ impl BpfConfig { .as_ref() .map(|cfg| !cfg.filter_process.is_empty() || global_filter.is_some()) .unwrap_or(false); - cpu_needs || profiler_needs || schedtrace_needs + let nettrack_needs = self + .nettrack + .as_ref() + .map(|cfg| !cfg.filter_process.is_empty() || global_filter.is_some()) + .unwrap_or(false); + cpu_needs || profiler_needs || schedtrace_needs || nettrack_needs }; if needs_process_filtering && self.thread_tracker.is_none() { @@ -195,15 +200,26 @@ impl BpfConfig { (None, HashSet::new()) }; - let nettrack = if let Some(cfg) = self.nettrack { + let (nettrack, nettrack_filters) = if let Some(mut cfg) = self.nettrack { + if let Some(ref global) = global_filter { + cfg.filter_process = global.clone(); + } debug!( module = "nettrack", frequency = cfg.frequency, + peer_frequency = cfg.peer_frequency, + protocols = ?cfg.protocols, + directions = ?cfg.directions, + scopes = ?cfg.scopes, + metrics = ?cfg.metrics, + pid_filters = ?cfg.pid_filters, + filter_process = ?cfg.filter_process, "initializing network tracking" ); - Some(nettrack::Object::new(cfg)) + let filters: HashSet = cfg.filter_process.iter().cloned().collect(); + (Some(nettrack::Object::new(cfg)), filters) } else { - None + (None, HashSet::new()) }; let perfcounter = if let Some(cfg) = self.perfcounter { @@ -229,6 +245,7 @@ impl BpfConfig { cpuutil_filters, profiler_filters, schedtrace_filters, + nettrack_filters, watched_tgids: HashSet::new(), }) } @@ -245,6 +262,7 @@ pub struct BpfObject { cpuutil_filters: HashSet, profiler_filters: HashSet, schedtrace_filters: HashSet, + nettrack_filters: HashSet, watched_tgids: HashSet, } @@ -308,15 +326,18 @@ impl BpfObject { let has_cpuutil = cpuutil_rc.is_some(); let has_profiler = profiler_rc.is_some(); let has_schedtrace = schedtrace_rc.is_some(); + let has_nettrack = nettrack_rc.is_some(); - if has_cpuutil || has_profiler || has_schedtrace { + if has_cpuutil || has_profiler || has_schedtrace || has_nettrack { let mut user_callback = callback.clone(); let cpuutil_filters = self.cpuutil_filters.clone(); let profiler_filters = self.profiler_filters.clone(); let schedtrace_filters = self.schedtrace_filters.clone(); + let nettrack_filters = self.nettrack_filters.clone(); let cpuutil_ref = cpuutil_rc.clone(); let profiler_ref = profiler_rc.clone(); let schedtrace_ref = schedtrace_rc.clone(); + let nettrack_ref = nettrack_rc.clone(); let mut watched_tgids = std::mem::take(&mut self.watched_tgids); let wrapper_callback = move |message: Message<'_>| -> i32 { @@ -380,6 +401,23 @@ impl BpfObject { } } } + if let Some(ref net) = nettrack_ref { + let base_name = name.split('/').next_back().unwrap_or(name); + if nettrack_filters.contains(name) + || nettrack_filters.contains(base_name) + { + if let Err(e) = net.borrow_mut().filter(pid) { + tracing::warn!( + "Failed to add process {} (pid {}) to nettrack filter: {}", + name, + pid, + e + ); + } else { + watched_tgids.insert(pid); + } + } + } } protocol::TrackType::Thread { tid, pid } => { if watched_tgids.contains(&pid) { @@ -702,6 +740,54 @@ nettrack = {} let nettrack_config = config.nettrack.as_ref().unwrap(); assert_eq!(nettrack_config.frequency, 9); // default_frequency assert_eq!(nettrack_config.ringbuf, 1024 * 1024); // default_ringbuf_size + assert_eq!(nettrack_config.scopes, vec![NetScope::Host]); + assert_eq!(nettrack_config.metrics, vec![NetMetric::Bytes]); + assert_eq!(nettrack_config.peer_frequency, 1); + } + + #[test] + fn test_detailed_nettrack_config() { + let config = BpfConfig::from_toml_str( + r#" +[nettrack] +frequency = 5 +peer_frequency = 2 +protocols = ["tcp"] +directions = ["send"] +scopes = ["host", "process", "peer"] +metrics = ["bytes", "operations", "errors"] +pid_filters = [1234] +max_process_entries = 64 +max_peer_entries = 128 +max_peer_tracks = 32 +"#, + ) + .unwrap(); + let config = config.nettrack.unwrap(); + assert_eq!(config.protocols, vec![NetProtocol::Tcp]); + assert_eq!(config.directions, vec![NetDirection::Send]); + assert_eq!( + config.scopes, + vec![NetScope::Host, NetScope::Process, NetScope::Peer] + ); + assert_eq!( + config.metrics, + vec![NetMetric::Bytes, NetMetric::Operations, NetMetric::Errors] + ); + assert_eq!(config.max_peer_entries, 128); + assert_eq!(config.max_peer_tracks, 32); + } + + #[test] + fn test_nettrack_process_filter_requires_threadtrack() { + let config = BpfConfig::from_toml_str( + r#" +[nettrack] +filter_process = ["monad-node"] +"#, + ) + .unwrap(); + assert!(config.build().is_err()); } #[test] diff --git a/bpf/src/nettrack.rs b/bpf/src/nettrack.rs index ee988db..e62d4e7 100644 --- a/bpf/src/nettrack.rs +++ b/bpf/src/nettrack.rs @@ -19,37 +19,223 @@ mod nettrack_skel { use nettrack_skel::*; -use crate::{perf_event, BpfError, Filterable}; +use crate::{get_monotonic_timestamp, BpfError, Filterable}; use libbpf_rs::skel::{OpenSkel, Skel, SkelBuilder}; -use libbpf_rs::{OpenObject, RingBufferBuilder}; -use libbpf_sys::{PERF_COUNT_SW_CPU_CLOCK, PERF_TYPE_SOFTWARE}; +use libbpf_rs::{MapCore, MapFlags, OpenObject}; use protocol::{Counter, Event, Labels, Message, Track, TrackId, TrackType}; use rand::Rng; use serde::{Deserialize, Serialize}; use std::borrow::Cow; -use std::collections::HashMap; -use std::convert::TryFrom; -use std::marker::PhantomData; +use std::collections::{HashMap, HashSet}; +use std::fmt; +use std::hash::Hash; use std::mem::MaybeUninit; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; use std::time::Duration; use tracing::debug; +const AF_INET: u8 = 2; +const AF_INET6: u8 = 10; +const TCP_INDEX: usize = 0; +const UDP_INDEX: usize = 1; +const MAP_BATCH_SIZE: u32 = 256; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum NetProtocol { + Tcp, + Udp, +} + +impl NetProtocol { + fn index(self) -> usize { + match self { + Self::Tcp => TCP_INDEX, + Self::Udp => UDP_INDEX, + } + } + + fn from_index(index: u8) -> Option { + match index as usize { + TCP_INDEX => Some(Self::Tcp), + UDP_INDEX => Some(Self::Udp), + _ => None, + } + } + + fn name(self) -> &'static str { + match self { + Self::Tcp => "tcp", + Self::Udp => "udp", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum NetDirection { + Send, + Receive, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum NetScope { + Host, + Process, + Peer, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum NetMetric { + Bytes, + Operations, + Errors, +} + +fn default_frequency() -> u64 { + 9 +} + +fn default_peer_frequency() -> u64 { + 1 +} + +fn default_ringbuf_size() -> usize { + 1024 * 1024 +} + +fn default_protocols() -> Vec { + vec![NetProtocol::Tcp, NetProtocol::Udp] +} + +fn default_directions() -> Vec { + vec![NetDirection::Send, NetDirection::Receive] +} + +fn default_scopes() -> Vec { + vec![NetScope::Host] +} + +fn default_metrics() -> Vec { + vec![NetMetric::Bytes] +} + +fn default_process_entries() -> u32 { + 1024 +} + +fn default_peer_entries() -> u32 { + 4096 +} + +fn default_peer_tracks() -> usize { + 512 +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct NetTrackConfig { #[serde(default = "default_frequency")] pub frequency: u64, + #[serde(default = "default_peer_frequency")] + pub peer_frequency: u64, + /// Kept for configuration compatibility. Map-based sampling has no ring buffer. #[serde(default = "default_ringbuf_size")] pub ringbuf: usize, #[serde(default)] pub scaled: bool, + #[serde(default = "default_protocols")] + pub protocols: Vec, + #[serde(default = "default_directions")] + pub directions: Vec, + #[serde(default = "default_scopes")] + pub scopes: Vec, + #[serde(default = "default_metrics")] + pub metrics: Vec, + #[serde(default)] + pub pid_filters: Vec, + #[serde(default)] + pub filter_process: Vec, + #[serde(default = "default_process_entries")] + pub max_process_entries: u32, + #[serde(default = "default_peer_entries")] + pub max_peer_entries: u32, + #[serde(default = "default_peer_tracks")] + pub max_peer_tracks: usize, } -fn default_frequency() -> u64 { - 9 +impl Default for NetTrackConfig { + fn default() -> Self { + Self { + frequency: default_frequency(), + peer_frequency: default_peer_frequency(), + ringbuf: default_ringbuf_size(), + scaled: false, + protocols: default_protocols(), + directions: default_directions(), + scopes: default_scopes(), + metrics: default_metrics(), + pid_filters: Vec::new(), + filter_process: Vec::new(), + max_process_entries: default_process_entries(), + max_peer_entries: default_peer_entries(), + max_peer_tracks: default_peer_tracks(), + } + } } -fn default_ringbuf_size() -> usize { - 1024 * 1024 +impl NetTrackConfig { + fn has_protocol(&self, protocol: NetProtocol) -> bool { + self.protocols.contains(&protocol) + } + + fn has_direction(&self, direction: NetDirection) -> bool { + self.directions.contains(&direction) + } + + fn has_scope(&self, scope: NetScope) -> bool { + self.scopes.contains(&scope) + } + + fn has_metric(&self, metric: NetMetric) -> bool { + self.metrics.contains(&metric) + } + + fn validate(&self) -> Result<(), BpfError> { + if self.frequency == 0 { + return Err(BpfError::LoadError( + "nettrack frequency must be greater than zero".to_string(), + )); + } + if self.has_scope(NetScope::Peer) && self.peer_frequency == 0 { + return Err(BpfError::LoadError( + "nettrack peer_frequency must be greater than zero".to_string(), + )); + } + if self.protocols.is_empty() + || self.directions.is_empty() + || self.scopes.is_empty() + || self.metrics.is_empty() + { + return Err(BpfError::LoadError( + "nettrack protocols, directions, scopes, and metrics must not be empty".to_string(), + )); + } + if self.has_scope(NetScope::Process) && self.max_process_entries == 0 { + return Err(BpfError::LoadError( + "nettrack max_process_entries must be greater than zero".to_string(), + )); + } + if self.has_scope(NetScope::Peer) + && (self.max_peer_entries == 0 || self.max_peer_tracks == 0) + { + return Err(BpfError::LoadError( + "nettrack peer entry and track limits must be greater than zero".to_string(), + )); + } + Ok(()) + } } pub struct Object { @@ -69,351 +255,715 @@ impl Object { where F: for<'a> FnMut(Message<'a>) -> i32 + 'bd, { - let nettrack = NetTrack::new( - &mut self.object, - self.config.clone(), - callback, - self.config.frequency, - )?; - Ok(nettrack) + self.config.validate()?; + NetTrack::new(&mut self.object, self.config.clone(), callback) } } #[repr(C)] -#[derive(Debug)] -pub struct NetEvent { - pub timestamp: u64, - pub tcp_send_bytes: u64, - pub tcp_recv_bytes: u64, - pub udp_send_bytes: u64, - pub udp_recv_bytes: u64, - pub total_send_bytes: u64, - pub total_recv_bytes: u64, - pub cpu: u32, -} - -unsafe impl plain::Plain for NetEvent {} - -#[derive(Debug)] -struct CounterIds { - tcp_send: u64, - tcp_recv: u64, - udp_send: u64, - udp_recv: u64, - total_send: u64, - total_recv: u64, -} - -impl CounterIds { - fn new() -> Self { - let mut rng = rand::thread_rng(); +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +struct StatsValue { + send_bytes: u64, + recv_bytes: u64, + send_operations: u64, + recv_operations: u64, + send_errors: u64, + recv_errors: u64, + send_would_block: u64, + recv_would_block: u64, + tcp_recv_eof: u64, + udp_recv_truncated: u64, + first_seen_ns: u64, + last_seen_ns: u64, +} + +unsafe impl plain::Plain for StatsValue {} + +impl StatsValue { + fn accumulate(&mut self, other: Self) { + self.send_bytes = self.send_bytes.saturating_add(other.send_bytes); + self.recv_bytes = self.recv_bytes.saturating_add(other.recv_bytes); + self.send_operations = self.send_operations.saturating_add(other.send_operations); + self.recv_operations = self.recv_operations.saturating_add(other.recv_operations); + self.send_errors = self.send_errors.saturating_add(other.send_errors); + self.recv_errors = self.recv_errors.saturating_add(other.recv_errors); + self.send_would_block = self.send_would_block.saturating_add(other.send_would_block); + self.recv_would_block = self.recv_would_block.saturating_add(other.recv_would_block); + self.tcp_recv_eof = self.tcp_recv_eof.saturating_add(other.tcp_recv_eof); + self.udp_recv_truncated = self + .udp_recv_truncated + .saturating_add(other.udp_recv_truncated); + if self.first_seen_ns == 0 + || (other.first_seen_ns != 0 && other.first_seen_ns < self.first_seen_ns) + { + self.first_seen_ns = other.first_seen_ns; + } + self.last_seen_ns = self.last_seen_ns.max(other.last_seen_ns); + } + + fn delta(self, previous: Option) -> Self { + let Some(previous) = previous.filter(|p| p.first_seen_ns == self.first_seen_ns) else { + return self; + }; Self { - tcp_send: rng.gen(), - tcp_recv: rng.gen(), - udp_send: rng.gen(), - udp_recv: rng.gen(), - total_send: rng.gen(), - total_recv: rng.gen(), + send_bytes: self.send_bytes.saturating_sub(previous.send_bytes), + recv_bytes: self.recv_bytes.saturating_sub(previous.recv_bytes), + send_operations: self + .send_operations + .saturating_sub(previous.send_operations), + recv_operations: self + .recv_operations + .saturating_sub(previous.recv_operations), + send_errors: self.send_errors.saturating_sub(previous.send_errors), + recv_errors: self.recv_errors.saturating_sub(previous.recv_errors), + send_would_block: self + .send_would_block + .saturating_sub(previous.send_would_block), + recv_would_block: self + .recv_would_block + .saturating_sub(previous.recv_would_block), + tcp_recv_eof: self.tcp_recv_eof.saturating_sub(previous.tcp_recv_eof), + udp_recv_truncated: self + .udp_recv_truncated + .saturating_sub(previous.udp_recv_truncated), + first_seen_ns: self.first_seen_ns, + last_seen_ns: self.last_seen_ns, } } } -#[derive(Default, Clone)] -struct NetStats { - tcp_send_total: u64, - tcp_recv_total: u64, - udp_send_total: u64, - udp_recv_total: u64, - last_timestamp: u64, +#[repr(C)] +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)] +struct ProcessKey { + tgid: u32, + protocol: u8, + padding: [u8; 3], } -impl NetStats { - fn accumulate(&mut self, event: &NetEvent) { - self.tcp_send_total += event.tcp_send_bytes; - self.tcp_recv_total += event.tcp_recv_bytes; - self.udp_send_total += event.udp_send_bytes; - self.udp_recv_total += event.udp_recv_bytes; - self.last_timestamp = event.timestamp; - } +unsafe impl plain::Plain for ProcessKey {} - fn reset(&mut self) { - self.tcp_send_total = 0; - self.tcp_recv_total = 0; - self.udp_send_total = 0; - self.udp_recv_total = 0; +#[repr(C)] +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)] +struct PeerKey { + tgid: u32, + protocol: u8, + family: u8, + local_port: u16, + remote_port: u16, + padding: u16, + remote_addr: [u8; 16], +} + +unsafe impl plain::Plain for PeerKey {} + +impl PeerKey { + fn protocol(self) -> Option { + NetProtocol::from_index(self.protocol) } - fn total_send(&self) -> u64 { - self.tcp_send_total + self.udp_send_total + fn remote_ip(self) -> Option { + match self.family { + AF_INET => Some(IpAddr::V4(Ipv4Addr::new( + self.remote_addr[0], + self.remote_addr[1], + self.remote_addr[2], + self.remote_addr[3], + ))), + AF_INET6 => Some(IpAddr::V6(Ipv6Addr::from(self.remote_addr))), + _ => None, + } } - fn total_recv(&self) -> u64 { - self.tcp_recv_total + self.udp_recv_total + fn track_name(self) -> String { + let protocol = self.protocol().map(NetProtocol::name).unwrap_or("network"); + match self.remote_ip() { + Some(IpAddr::V4(ip)) => format!("{protocol} {ip}:{}", self.remote_port), + Some(IpAddr::V6(ip)) => format!("{protocol} [{ip}]:{}", self.remote_port), + None => format!("{protocol} unknown-peer"), + } } } -impl<'a> TryFrom<&'a [u8]> for &'a NetEvent { - type Error = BpfError; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum ScopeKey { + Host, + Process(u32), + Peer(PeerKey), +} - fn try_from(data: &'a [u8]) -> Result { - plain::from_bytes(data) - .map_err(|e| BpfError::MapError(format!("failed to parse net event: {:?}", e))) - } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum CounterKind { + SendBytes, + RecvBytes, + SendOperations, + RecvOperations, + SendErrors, + RecvErrors, + SendWouldBlock, + RecvWouldBlock, + TcpRecvEof, + UdpRecvTruncated, + TotalSendBytes, + TotalRecvBytes, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct SeriesKey { + scope: ScopeKey, + protocol: NetProtocol, + kind: CounterKind, +} + +struct CounterSample { + kind: CounterKind, + name: String, + value: f64, + unit: &'static str, } pub struct NetTrack<'this, F> { - #[allow(dead_code)] skel: NettrackSkel<'this>, - ringbuf: libbpf_rs::RingBuffer<'this>, - _perf_links: Vec, - _phantom: PhantomData, + callback: F, + config: NetTrackConfig, + last_sample_ns: u64, + last_peer_sample_ns: u64, + global_previous: [StatsValue; 2], + process_previous: HashMap, + peer_previous: HashMap, + series_ids: HashMap, + submitted_tracks: HashSet, + peer_track_ids: HashMap, } impl<'this, F> NetTrack<'this, F> where F: for<'a> FnMut(Message<'a>) -> i32 + 'this, { - fn load_and_attach_skel( + fn new( open_object: &'this mut MaybeUninit, - config: &NetTrackConfig, - ) -> Result, BpfError> { + config: NetTrackConfig, + callback: F, + ) -> Result { let skel_builder = NettrackSkelBuilder::default(); - let mut open_skel = skel_builder .open(open_object) - .map_err(|e| BpfError::LoadError(format!("failed to open bpf skeleton: {}", e)))?; + .map_err(|e| BpfError::LoadError(format!("failed to open bpf skeleton: {e}")))?; + let process_entries = if config.has_scope(NetScope::Process) { + config.max_process_entries + } else { + 1 + }; + let peer_entries = if config.has_scope(NetScope::Peer) { + config.max_peer_entries + } else { + 1 + }; + open_skel + .maps + .process_stats + .set_max_entries(process_entries) + .map_err(|e| BpfError::LoadError(format!("failed to size process map: {e}")))?; open_skel .maps - .events - .set_max_entries(config.ringbuf as u32) - .map_err(|e| BpfError::LoadError(format!("failed to set ring buffer size: {}", e)))?; + .peer_stats + .set_max_entries(peer_entries) + .map_err(|e| BpfError::LoadError(format!("failed to size peer map: {e}")))?; + + let filter_enabled = !config.pid_filters.is_empty() || !config.filter_process.is_empty(); + let rodata = open_skel.maps.rodata_data.as_mut().unwrap(); + rodata + .cfg + .tcp_enabled + .write(config.has_protocol(NetProtocol::Tcp)); + rodata + .cfg + .udp_enabled + .write(config.has_protocol(NetProtocol::Udp)); + rodata + .cfg + .send_enabled + .write(config.has_direction(NetDirection::Send)); + rodata + .cfg + .receive_enabled + .write(config.has_direction(NetDirection::Receive)); + rodata + .cfg + .host_enabled + .write(config.has_scope(NetScope::Host)); + rodata + .cfg + .process_enabled + .write(config.has_scope(NetScope::Process)); + rodata + .cfg + .peer_enabled + .write(config.has_scope(NetScope::Peer)); + rodata + .cfg + .bytes_enabled + .write(config.has_metric(NetMetric::Bytes)); + rodata + .cfg + .operations_enabled + .write(config.has_metric(NetMetric::Operations)); + rodata + .cfg + .errors_enabled + .write(config.has_metric(NetMetric::Errors)); + rodata.cfg.filter_enabled.write(filter_enabled); + + let tcp_enabled = config.has_protocol(NetProtocol::Tcp); + let udp_enabled = config.has_protocol(NetProtocol::Udp); + let send_enabled = config.has_direction(NetDirection::Send); + let receive_enabled = config.has_direction(NetDirection::Receive); + open_skel + .progs + .trace_tcp_sendmsg + .set_autoload(tcp_enabled && send_enabled); + open_skel + .progs + .trace_tcp_recvmsg + .set_autoload(tcp_enabled && receive_enabled); + open_skel + .progs + .trace_udp_sendmsg + .set_autoload(udp_enabled && send_enabled); + open_skel + .progs + .trace_udp_recvmsg + .set_autoload(udp_enabled && receive_enabled); let mut skel = open_skel .load() - .map_err(|e| BpfError::LoadError(format!("failed to load bpf program: {}", e)))?; + .map_err(|e| BpfError::LoadError(format!("failed to load bpf program: {e}")))?; + + for &pid in &config.pid_filters { + let pid = u32::try_from(pid).map_err(|_| { + BpfError::LoadError(format!("nettrack PID filter must be positive: {pid}")) + })?; + skel.maps + .tracked_tgids + .update(&pid.to_ne_bytes(), &1u32.to_ne_bytes(), MapFlags::ANY) + .map_err(|e| BpfError::MapError(format!("failed to add PID filter: {e}")))?; + } skel.attach() - .map_err(|e| BpfError::AttachError(format!("failed to attach bpf programs: {}", e)))?; + .map_err(|e| BpfError::AttachError(format!("failed to attach bpf programs: {e}")))?; + + debug!( + protocols = ?config.protocols, + directions = ?config.directions, + scopes = ?config.scopes, + metrics = ?config.metrics, + frequency = config.frequency, + peer_frequency = config.peer_frequency, + "initialized map-based network tracking" + ); - Ok(skel) + let now = get_monotonic_timestamp(); + Ok(Self { + skel, + callback, + config, + last_sample_ns: now, + last_peer_sample_ns: now, + global_previous: [StatsValue::default(); 2], + process_previous: HashMap::new(), + peer_previous: HashMap::new(), + series_ids: HashMap::new(), + submitted_tracks: HashSet::new(), + peer_track_ids: HashMap::new(), + }) } - fn setup_perf_events( - skel: &mut NettrackSkel<'this>, - freq: u64, - ) -> Result, BpfError> { - let perf_fds = - perf_event::perf_event_per_cpu(PERF_TYPE_SOFTWARE, PERF_COUNT_SW_CPU_CLOCK, freq) - .map_err(|e| BpfError::LoadError(format!("failed to open perf events: {}", e)))?; + fn parse_value(data: &[u8], kind: &str) -> Result { + plain::from_bytes::(data) + .copied() + .map_err(|e| BpfError::MapError(format!("failed to parse {kind}: {e:?}"))) + } - let prog = &mut skel.progs.handle_boundary_event; - let links = perf_event::attach_perf_event(&perf_fds, prog) - .map_err(|e| BpfError::AttachError(format!("failed to attach perf events: {}", e)))?; + fn parse_key(data: &[u8], kind: &str) -> Result { + plain::from_bytes::(data) + .copied() + .map_err(|e| BpfError::MapError(format!("failed to parse {kind}: {e:?}"))) + } - Ok(links) + fn read_global(&self, protocol: NetProtocol) -> Result { + let key = (protocol.index() as u32).to_ne_bytes(); + let values = self + .skel + .maps + .global_stats + .lookup_percpu(&key, MapFlags::ANY) + .map_err(|e| BpfError::MapError(format!("failed to read host network map: {e}")))? + .unwrap_or_default(); + let mut total = StatsValue::default(); + for value in values { + total.accumulate(Self::parse_value(&value, "host network statistics")?); + } + Ok(total) } - fn new( - open_object: &'this mut MaybeUninit, - config: NetTrackConfig, - callback: F, - freq: u64, - ) -> Result { - let mut skel = Self::load_and_attach_skel(open_object, &config)?; - let counter_ids = CounterIds::new(); - let ringbuf = Self::create_ring_buffer(&skel, callback, counter_ids, config.scaled)?; - let perf_links = Self::setup_perf_events(&mut skel, freq)?; + fn read_processes(&self) -> Result, BpfError> { + self.skel + .maps + .process_stats + .lookup_batch(MAP_BATCH_SIZE, MapFlags::ANY, MapFlags::ANY) + .map_err(|e| BpfError::MapError(format!("failed to read process network map: {e}")))? + .map(|(key, value)| { + Ok(( + Self::parse_key(&key, "process network key")?, + Self::parse_value(&value, "process network statistics")?, + )) + }) + .collect() + } - Ok(NetTrack { - skel, - ringbuf, - _perf_links: perf_links, - _phantom: PhantomData, - }) + fn read_peers(&self) -> Result, BpfError> { + self.skel + .maps + .peer_stats + .lookup_batch(MAP_BATCH_SIZE, MapFlags::ANY, MapFlags::ANY) + .map_err(|e| BpfError::MapError(format!("failed to read peer network map: {e}")))? + .map(|(key, value)| { + Ok(( + Self::parse_key(&key, "peer network key")?, + Self::parse_value(&value, "peer network statistics")?, + )) + }) + .collect() } - fn create_counters<'a>( - stats: &NetStats, - ids: &CounterIds, - scaled: bool, - time_diff_s: Option, - ) -> Vec<(&'a str, f64, u64)> { - if scaled { - let time_diff = time_diff_s.unwrap_or(1.0); - vec![ - ( - "tcp_send_rate", - (stats.tcp_send_total as f64 * 8.0) / time_diff, - ids.tcp_send, - ), + fn samples_for( + &self, + protocol: NetProtocol, + stats: StatsValue, + elapsed_seconds: f64, + ) -> Vec { + let prefix = protocol.name(); + let mut samples = Vec::new(); + if self.config.has_metric(NetMetric::Bytes) { + let (send_name, recv_name, send_value, recv_value, unit) = if self.config.scaled { ( - "tcp_recv_rate", - (stats.tcp_recv_total as f64 * 8.0) / time_diff, - ids.tcp_recv, - ), - ( - "udp_send_rate", - (stats.udp_send_total as f64 * 8.0) / time_diff, - ids.udp_send, - ), + format!("{prefix}_send_rate"), + format!("{prefix}_recv_rate"), + stats.send_bytes as f64 * 8.0 / elapsed_seconds, + stats.recv_bytes as f64 * 8.0 / elapsed_seconds, + "bits/s", + ) + } else { ( - "udp_recv_rate", - (stats.udp_recv_total as f64 * 8.0) / time_diff, - ids.udp_recv, - ), + format!("{prefix}_send"), + format!("{prefix}_recv"), + stats.send_bytes as f64, + stats.recv_bytes as f64, + "bytes", + ) + }; + samples.push(CounterSample { + kind: CounterKind::SendBytes, + name: send_name, + value: send_value, + unit, + }); + samples.push(CounterSample { + kind: CounterKind::RecvBytes, + name: recv_name, + value: recv_value, + unit, + }); + } + if self.config.has_metric(NetMetric::Operations) { + let operation = if protocol == NetProtocol::Udp { + "datagrams" + } else { + "calls" + }; + samples.push(CounterSample { + kind: CounterKind::SendOperations, + name: format!("{prefix}_send_{operation}"), + value: stats.send_operations as f64, + unit: "count", + }); + samples.push(CounterSample { + kind: CounterKind::RecvOperations, + name: format!("{prefix}_recv_{operation}"), + value: stats.recv_operations as f64, + unit: "count", + }); + if protocol == NetProtocol::Tcp { + samples.push(CounterSample { + kind: CounterKind::TcpRecvEof, + name: "tcp_recv_eof".to_string(), + value: stats.tcp_recv_eof as f64, + unit: "count", + }); + } else { + samples.push(CounterSample { + kind: CounterKind::UdpRecvTruncated, + name: "udp_recv_truncated".to_string(), + value: stats.udp_recv_truncated as f64, + unit: "count", + }); + } + } + if self.config.has_metric(NetMetric::Errors) { + for (kind, suffix, value) in [ + (CounterKind::SendErrors, "send_errors", stats.send_errors), + (CounterKind::RecvErrors, "recv_errors", stats.recv_errors), ( - "total_send_rate", - (stats.total_send() as f64 * 8.0) / time_diff, - ids.total_send, + CounterKind::SendWouldBlock, + "send_would_block", + stats.send_would_block, ), ( - "total_recv_rate", - (stats.total_recv() as f64 * 8.0) / time_diff, - ids.total_recv, + CounterKind::RecvWouldBlock, + "recv_would_block", + stats.recv_would_block, ), - ] - } else { - vec![ - ("tcp_send", stats.tcp_send_total as f64, ids.tcp_send), - ("tcp_recv", stats.tcp_recv_total as f64, ids.tcp_recv), - ("udp_send", stats.udp_send_total as f64, ids.udp_send), - ("udp_recv", stats.udp_recv_total as f64, ids.udp_recv), - ("total_send", stats.total_send() as f64, ids.total_send), - ("total_recv", stats.total_recv() as f64, ids.total_recv), - ] + ] { + samples.push(CounterSample { + kind, + name: format!("{prefix}_{suffix}"), + value: value as f64, + unit: "count", + }); + } + } + samples + } + + fn ensure_peer_track(&mut self, key: PeerKey) -> Option { + if let Some(id) = self.peer_track_ids.get(&key) { + return Some(*id); } + if self.peer_track_ids.len() >= self.config.max_peer_tracks { + return None; + } + let id = rand::thread_rng().gen::(); + let name = key.track_name(); + let result = (self.callback)(Message::Event(Event::Track(Track { + name: &name, + track_type: TrackType::Custom { id }, + parent: Some(TrackType::Process { + pid: key.tgid as i32, + }), + }))); + if result != 0 { + return None; + } + self.peer_track_ids.insert(key, id); + Some(id) } - fn submit_counter( - name: &str, - value: f64, - id: u64, + fn emit_sample( + &mut self, + scope: ScopeKey, + protocol: NetProtocol, + sample: CounterSample, timestamp: u64, - submitted_tracks: &mut HashMap, - callback: &mut G, - unit: &str, - ) -> i32 - where - G: FnMut(Message) -> i32, - { + ) -> i32 { + let parent = match scope { + ScopeKey::Host => None, + ScopeKey::Process(tgid) => Some(TrackType::Process { pid: tgid as i32 }), + ScopeKey::Peer(key) => { + let Some(id) = self.ensure_peer_track(key) else { + return 0; + }; + Some(TrackType::Custom { id }) + } + }; + let series_key = SeriesKey { + scope, + protocol, + kind: sample.kind, + }; + let id = *self + .series_ids + .entry(series_key) + .or_insert_with(|| rand::thread_rng().gen::()); let track_id = TrackId::Counter { id }; - - if submitted_tracks.insert(track_id, ()).is_none() { - let track = Message::Event(Event::Track(Track { - name, + if self.submitted_tracks.insert(track_id) { + let result = (self.callback)(Message::Event(Event::Track(Track { + name: &sample.name, track_type: TrackType::Counter { id, - unit: Some(unit), + unit: Some(sample.unit), }, - parent: None, - })); - let result = callback(track); + parent, + }))); if result != 0 { return result; } } debug!( - name = name, - value = value, - timestamp = timestamp, + name = sample.name, + value = sample.value, + timestamp, "emitting network counter" ); - - callback(Message::Event(Event::Counter(Counter { - name, - value, + (self.callback)(Message::Event(Event::Counter(Counter { + name: &sample.name, + value: sample.value, timestamp, track_id, labels: Cow::Owned(Labels::new()), - unit: Some(unit), + unit: Some(sample.unit), }))) } - fn create_ring_buffer( - skel: &NettrackSkel<'this>, - mut callback: F, - counter_ids: CounterIds, - scaled: bool, - ) -> Result, BpfError> { - let mut submitted_tracks: HashMap = HashMap::new(); - - let mut builder = RingBufferBuilder::new(); - let nprocs = libbpf_rs::num_possible_cpus().unwrap(); - let mut boundaries_reported = 0; - let mut stats = NetStats::default(); - let mut prev_timestamp = 0u64; - - builder - .add(&skel.maps.events, move |data| { - let net_event: &NetEvent = data.try_into().unwrap(); - stats.accumulate(net_event); - boundaries_reported += 1; - - if boundaries_reported == nprocs { - boundaries_reported = 0; - - if scaled && prev_timestamp == 0 { - prev_timestamp = stats.last_timestamp; - stats.reset(); - return 0; - } + fn emit_stats( + &mut self, + scope: ScopeKey, + protocol: NetProtocol, + stats: StatsValue, + elapsed_seconds: f64, + timestamp: u64, + ) -> i32 { + for sample in self.samples_for(protocol, stats, elapsed_seconds) { + let result = self.emit_sample(scope, protocol, sample, timestamp); + if result != 0 { + return result; + } + } + 0 + } - let time_diff_s = if scaled && prev_timestamp > 0 { - let time_diff_ns = stats.last_timestamp.saturating_sub(prev_timestamp); - if time_diff_ns > 0 { - Some(time_diff_ns as f64 / 1_000_000_000.0) - } else { - None - } - } else { - None - }; - - let counters = Self::create_counters(&stats, &counter_ids, scaled, time_diff_s); - let unit = if scaled { "bits/s" } else { "bytes" }; - - for (name, value, id) in counters { - let result = Self::submit_counter( - name, - value, - id, - stats.last_timestamp, - &mut submitted_tracks, - &mut callback, - unit, - ); - if result != 0 { - return result; - } + fn sample_host_and_processes(&mut self, now: u64) -> Result<(), BpfError> { + let elapsed_ns = now.saturating_sub(self.last_sample_ns); + let elapsed_seconds = (elapsed_ns.max(1) as f64) / 1_000_000_000.0; + + if self.config.has_scope(NetScope::Host) { + let mut host_delta = StatsValue::default(); + for protocol in [NetProtocol::Tcp, NetProtocol::Udp] { + if !self.config.has_protocol(protocol) { + continue; + } + let current = self.read_global(protocol)?; + let previous = self.global_previous[protocol.index()]; + let delta = current.delta(Some(previous)); + host_delta.accumulate(delta); + self.global_previous[protocol.index()] = current; + if self.emit_stats(ScopeKey::Host, protocol, delta, elapsed_seconds, now) != 0 { + break; + } + } + if self.config.has_metric(NetMetric::Bytes) { + let (send_name, recv_name, send_value, recv_value, unit) = if self.config.scaled { + ( + "total_send_rate", + "total_recv_rate", + host_delta.send_bytes as f64 * 8.0 / elapsed_seconds, + host_delta.recv_bytes as f64 * 8.0 / elapsed_seconds, + "bits/s", + ) + } else { + ( + "total_send", + "total_recv", + host_delta.send_bytes as f64, + host_delta.recv_bytes as f64, + "bytes", + ) + }; + for sample in [ + CounterSample { + kind: CounterKind::TotalSendBytes, + name: send_name.to_string(), + value: send_value, + unit, + }, + CounterSample { + kind: CounterKind::TotalRecvBytes, + name: recv_name.to_string(), + value: recv_value, + unit, + }, + ] { + if self.emit_sample(ScopeKey::Host, NetProtocol::Tcp, sample, now) != 0 { + break; } + } + } + } - prev_timestamp = stats.last_timestamp; - stats.reset(); + if self.config.has_scope(NetScope::Process) { + for (key, current) in self.read_processes()? { + let Some(protocol) = NetProtocol::from_index(key.protocol) else { + continue; + }; + let previous = self.process_previous.insert(key, current); + let delta = current.delta(previous); + if self.emit_stats( + ScopeKey::Process(key.tgid), + protocol, + delta, + elapsed_seconds, + now, + ) != 0 + { + break; } + } + } - 0 - }) - .map_err(|e| BpfError::MapError(format!("failed to add ring buffer: {}", e)))?; + self.last_sample_ns = now; + Ok(()) + } - builder - .build() - .map_err(|e| BpfError::MapError(format!("failed to build ring buffer: {}", e))) + fn sample_peers(&mut self, now: u64) -> Result<(), BpfError> { + let elapsed_ns = now.saturating_sub(self.last_peer_sample_ns); + let elapsed_seconds = (elapsed_ns.max(1) as f64) / 1_000_000_000.0; + for (key, current) in self.read_peers()? { + let Some(protocol) = key.protocol() else { + continue; + }; + let previous = self.peer_previous.insert(key, current); + let delta = current.delta(previous); + if self.emit_stats(ScopeKey::Peer(key), protocol, delta, elapsed_seconds, now) != 0 { + break; + } + } + self.last_peer_sample_ns = now; + Ok(()) } - pub fn poll(&mut self, timeout: Duration) -> Result<(), BpfError> { - self.ringbuf - .poll(timeout) - .map_err(|e| BpfError::MapError(format!("failed to poll ring buffer: {}", e)))?; + fn sample_if_due(&mut self, force: bool) -> Result<(), BpfError> { + let now = get_monotonic_timestamp(); + let sample_interval = 1_000_000_000u64 / self.config.frequency; + if force || now.saturating_sub(self.last_sample_ns) >= sample_interval { + self.sample_host_and_processes(now)?; + } + if self.config.has_scope(NetScope::Peer) { + let peer_interval = 1_000_000_000u64 / self.config.peer_frequency; + if force || now.saturating_sub(self.last_peer_sample_ns) >= peer_interval { + self.sample_peers(now)?; + } + } Ok(()) } + pub fn poll(&mut self, _timeout: Duration) -> Result<(), BpfError> { + self.sample_if_due(false) + } + pub fn consume(&mut self) -> Result<(), BpfError> { - self.ringbuf - .consume() - .map_err(|e| BpfError::MapError(format!("failed to consume ring buffer: {}", e)))?; - Ok(()) + self.sample_if_due(false) + } + + pub fn flush(&mut self) -> Result<(), BpfError> { + self.sample_if_due(true) + } + + pub fn add_pid_filter(&mut self, pid: u32) -> Result<(), BpfError> { + self.skel + .maps + .tracked_tgids + .update(&pid.to_ne_bytes(), &1u32.to_ne_bytes(), MapFlags::ANY) + .map_err(|e| BpfError::MapError(format!("failed to add PID filter: {e}"))) } } @@ -421,179 +971,247 @@ impl<'this, F> Filterable for NetTrack<'this, F> where F: for<'a> FnMut(Message<'a>) -> i32 + 'this, { - fn filter(&mut self, _pid: i32) -> Result<(), BpfError> { - Ok(()) + fn filter(&mut self, pid: i32) -> Result<(), BpfError> { + let pid = u32::try_from(pid) + .map_err(|_| BpfError::MapError(format!("invalid PID filter: {pid}")))?; + self.add_pid_filter(pid) + } +} + +impl fmt::Display for PeerKey { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.track_name()) } } #[cfg(test)] -mod root_tests { +mod tests { use super::*; - use rstest::*; + use std::cell::RefCell; use std::io::{Read, Write}; use std::net::{TcpListener, TcpStream, UdpSocket}; - use std::{ - cell::RefCell, - rc::Rc, - thread, - time::{Duration, Instant}, - }; - - fn is_root() -> bool { - unsafe { libc::geteuid() == 0 } - } - - #[derive(Debug, Clone)] - struct TestCounter { - name: String, - value: f64, - #[allow(dead_code)] - timestamp: u64, + use std::rc::Rc; + use std::thread; + + fn generate_test_traffic() -> (u16, u16) { + let listener = TcpListener::bind("127.0.0.1:0").expect("failed to bind TCP listener"); + let tcp_address = listener.local_addr().unwrap(); + let tcp_server = thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut request = vec![0u8; 4096]; + stream.read_exact(&mut request).unwrap(); + stream.write_all(&vec![1u8; 2048]).unwrap(); + }); + let mut tcp_client = TcpStream::connect(tcp_address).unwrap(); + tcp_client.write_all(&vec![2u8; 4096]).unwrap(); + let mut response = vec![0u8; 2048]; + tcp_client.read_exact(&mut response).unwrap(); + drop(tcp_client); + tcp_server.join().unwrap(); + + let udp_server_socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + let udp_address = udp_server_socket.local_addr().unwrap(); + let udp_server = thread::spawn(move || { + let mut request = vec![0u8; 1000]; + let (size, peer) = udp_server_socket.recv_from(&mut request).unwrap(); + assert_eq!(size, request.len()); + assert_eq!( + udp_server_socket.send_to(&vec![3u8; 500], peer).unwrap(), + 500 + ); + }); + let udp_client = UdpSocket::bind("127.0.0.1:0").unwrap(); + assert_eq!( + udp_client.send_to(&vec![4u8; 1000], udp_address).unwrap(), + 1000 + ); + let mut udp_response = vec![0u8; 500]; + assert_eq!(udp_client.recv_from(&mut udp_response).unwrap().0, 500); + udp_server.join().unwrap(); + (tcp_address.port(), udp_address.port()) } - struct NetTrackFixture { - events: Rc>>, + #[test] + fn defaults_are_low_overhead() { + let config = NetTrackConfig::default(); + assert_eq!(config.scopes, vec![NetScope::Host]); + assert_eq!(config.metrics, vec![NetMetric::Bytes]); + assert_eq!(config.peer_frequency, 1); } - #[fixture] - fn nettrack_setup() -> NetTrackFixture { - let events = Rc::new(RefCell::new(Vec::new())); - - NetTrackFixture { events } + #[test] + fn cumulative_values_produce_interval_deltas() { + let previous = StatsValue { + send_bytes: 100, + recv_bytes: 40, + first_seen_ns: 10, + ..StatsValue::default() + }; + let current = StatsValue { + send_bytes: 160, + recv_bytes: 75, + first_seen_ns: 10, + ..StatsValue::default() + }; + let delta = current.delta(Some(previous)); + assert_eq!(delta.send_bytes, 60); + assert_eq!(delta.recv_bytes, 35); } - fn create_test_callback(events: Rc>>) -> impl FnMut(Message) -> i32 { - move |message| { - if let Message::Event(Event::Counter(c)) = message { - let test_counter = TestCounter { - name: c.name.to_string(), - value: c.value, - timestamp: c.timestamp, - }; - events.borrow_mut().push(test_counter); - } - 0 - } + #[test] + fn changed_generation_uses_current_value() { + let previous = StatsValue { + send_bytes: 100, + first_seen_ns: 10, + ..StatsValue::default() + }; + let current = StatsValue { + send_bytes: 20, + first_seen_ns: 20, + ..StatsValue::default() + }; + assert_eq!(current.delta(Some(previous)).send_bytes, 20); } - fn start_tcp_server(listener: TcpListener) -> thread::JoinHandle<()> { - thread::spawn(move || { - let (mut stream, _) = listener.accept().expect("failed to accept connection"); - let mut buf = [0; 1024]; - while stream.read(&mut buf).unwrap_or(0) > 0 { - stream - .write_all(b"response") - .expect("failed to write response"); - } - }) + #[test] + fn peer_names_format_ipv4_and_ipv6() { + let ipv4 = PeerKey { + protocol: TCP_INDEX as u8, + family: AF_INET, + remote_port: 443, + remote_addr: [127, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + ..PeerKey::default() + }; + assert_eq!(ipv4.track_name(), "tcp 127.0.0.1:443"); + + let ipv6 = PeerKey { + protocol: UDP_INDEX as u8, + family: AF_INET6, + remote_port: 53, + remote_addr: Ipv6Addr::LOCALHOST.octets(), + ..PeerKey::default() + }; + assert_eq!(ipv6.track_name(), "udp [::1]:53"); } - fn start_udp_server(socket: UdpSocket) -> thread::JoinHandle<()> { - thread::spawn(move || { - let mut buf = [0; 1024]; - loop { - if let Ok((size, src)) = socket.recv_from(&mut buf) { - if size == 0 { - break; - } - socket - .send_to(&buf[..size], src) - .expect("failed to send UDP response"); - } - } - }) + #[test] + fn invalid_configuration_is_rejected() { + let config = NetTrackConfig { + frequency: 0, + ..NetTrackConfig::default() + }; + assert!(config.validate().is_err()); } - fn generate_network_traffic<'b, F>( - tcp_addr: std::net::SocketAddr, - udp_addr: std::net::SocketAddr, - duration: Duration, - nettrack: &mut NetTrack<'b, F>, - ) where - F: for<'a> FnMut(Message<'a>) -> i32 + 'b, - { - let start = Instant::now(); - let data = vec![b'x'; 1024]; - - while start.elapsed() < duration { - if let Ok(mut tcp_stream) = TcpStream::connect(tcp_addr) { - let _ = tcp_stream.write_all(&data); - let mut response = [0; 1024]; - let _ = tcp_stream.read(&mut response); + #[test] + #[ignore = "requires root"] + fn root_tests_exact_tcp_and_udp_payload_bytes() { + assert_eq!(unsafe { libc::geteuid() }, 0); + + let counters = Rc::new(RefCell::new(HashMap::::new())); + let callback_counters = counters.clone(); + let callback = move |message: Message<'_>| { + if let Message::Event(Event::Counter(counter)) = message { + callback_counters + .borrow_mut() + .insert(counter.name.to_string(), counter.value); } - - let udp_client = UdpSocket::bind("127.0.0.1:0").expect("failed to bind UDP client"); - udp_client - .send_to(&data, udp_addr) - .expect("failed to send UDP data"); - let mut response = [0; 1024]; - let _ = udp_client.recv(&mut response); - - let _ = nettrack.consume(); - } + 0 + }; + let config = NetTrackConfig { + scopes: vec![NetScope::Host], + metrics: vec![NetMetric::Bytes], + pid_filters: vec![unsafe { libc::getpid() }], + ..NetTrackConfig::default() + }; + let mut object = Object::new(config); + let mut tracker = object.build(callback).expect("failed to load nettrack"); + + generate_test_traffic(); + + tracker.flush().expect("failed to sample network maps"); + let counters = counters.borrow(); + assert_eq!(counters.get("tcp_send"), Some(&6144.0)); + assert_eq!(counters.get("tcp_recv"), Some(&6144.0)); + assert_eq!(counters.get("udp_send"), Some(&1500.0)); + assert_eq!(counters.get("udp_recv"), Some(&1500.0)); + assert_eq!(counters.get("total_send"), Some(&7644.0)); + assert_eq!(counters.get("total_recv"), Some(&7644.0)); } - fn verify_events(events: &[TestCounter]) { - let event_types = [ - "tcp_send", - "tcp_recv", - "udp_send", - "udp_recv", - "total_send", - "total_recv", - ]; - - for event_type in &event_types { - let type_events: Vec<_> = events.iter().filter(|c| c.name == *event_type).collect(); - assert!(!type_events.is_empty(), "no {} events", event_type); - } - - for event in events { - assert!(event.value >= 0.0, "bytes should not be negative"); - } + #[derive(Default)] + struct ScopedCapture { + process_counter_ids: HashSet, + process_counters: HashMap, + peer_tracks: Vec, } - #[rstest] + #[test] #[ignore = "requires root"] - fn test_nettrack_tcp_udp_rates(nettrack_setup: NetTrackFixture) { - assert!(is_root()); - + fn root_tests_process_and_peer_scopes() { + assert_eq!(unsafe { libc::geteuid() }, 0); + + let capture = Rc::new(RefCell::new(ScopedCapture::default())); + let callback_capture = capture.clone(); + let callback = move |message: Message<'_>| { + let mut capture = callback_capture.borrow_mut(); + if let Message::Event(event) = message { + match event { + Event::Track(track) => match track.track_type { + TrackType::Custom { .. } + if matches!(track.parent, Some(TrackType::Process { .. })) => + { + capture.peer_tracks.push(track.name.to_string()); + } + TrackType::Counter { id, .. } + if matches!(track.parent, Some(TrackType::Process { .. })) => + { + capture.process_counter_ids.insert(TrackId::Counter { id }); + } + _ => {} + }, + Event::Counter(counter) + if capture.process_counter_ids.contains(&counter.track_id) => + { + capture + .process_counters + .insert(counter.name.to_string(), counter.value); + } + _ => {} + } + } + 0 + }; let config = NetTrackConfig { - frequency: 100, - ringbuf: default_ringbuf_size(), - scaled: false, + scopes: vec![NetScope::Process, NetScope::Peer], + metrics: vec![NetMetric::Bytes], + pid_filters: vec![unsafe { libc::getpid() }], + max_peer_tracks: 16, + ..NetTrackConfig::default() }; - let mut object = Object::new(config); - let callback = create_test_callback(nettrack_setup.events.clone()); - let mut nettrack = object.build(callback).expect("failed to build nettrack"); - - let tcp_listener = TcpListener::bind("127.0.0.1:0").expect("failed to bind TCP listener"); - let tcp_addr = tcp_listener - .local_addr() - .expect("failed to get TCP address"); - - let udp_socket = UdpSocket::bind("127.0.0.1:0").expect("failed to bind UDP socket"); - let udp_addr = udp_socket.local_addr().expect("failed to get UDP address"); - - let tcp_thread = start_tcp_server(tcp_listener); - let udp_thread = start_udp_server(udp_socket); - - thread::sleep(Duration::from_millis(100)); - generate_network_traffic( - tcp_addr, - udp_addr, - Duration::from_millis(500), - &mut nettrack, + let mut tracker = object.build(callback).expect("failed to load nettrack"); + let (tcp_port, udp_port) = generate_test_traffic(); + tracker.flush().expect("failed to sample network maps"); + + let capture = capture.borrow(); + assert_eq!(capture.process_counters.get("tcp_send"), Some(&6144.0)); + assert_eq!(capture.process_counters.get("tcp_recv"), Some(&6144.0)); + assert_eq!(capture.process_counters.get("udp_send"), Some(&1500.0)); + assert_eq!(capture.process_counters.get("udp_recv"), Some(&1500.0)); + assert!( + capture + .peer_tracks + .contains(&format!("tcp 127.0.0.1:{tcp_port}")), + "missing TCP peer track: {:?}", + capture.peer_tracks + ); + assert!( + capture + .peer_tracks + .contains(&format!("udp 127.0.0.1:{udp_port}")), + "missing UDP peer track: {:?}", + capture.peer_tracks ); - thread::sleep(Duration::from_millis(200)); - let _ = nettrack.consume(); - - drop(tcp_thread); - drop(udp_thread); - - let collected_events = nettrack_setup.events.borrow(); - assert!(!collected_events.is_empty(), "no events were captured"); - verify_events(&collected_events); } } diff --git a/bpf/testdata/nettrack_e2e.toml b/bpf/testdata/nettrack_e2e.toml new file mode 100644 index 0000000..7b077d3 --- /dev/null +++ b/bpf/testdata/nettrack_e2e.toml @@ -0,0 +1,11 @@ +[bpf.nettrack] +frequency = 10 +peer_frequency = 5 +scaled = false +protocols = ["tcp", "udp"] +directions = ["send", "receive"] +scopes = ["host", "process", "peer"] +metrics = ["bytes", "operations", "errors"] +max_process_entries = 64 +max_peer_entries = 128 +max_peer_tracks = 32 diff --git a/reference.toml b/reference.toml index 5a8dba1..6b1abab 100644 --- a/reference.toml +++ b/reference.toml @@ -124,13 +124,13 @@ ringbuf = 262144 # Tracks network I/O rates for TCP and UDP traffic [bpf.nettrack] -# Sampling frequency in Hz (how often to collect network statistics) +# Host and process map sampling frequency in Hz # Default: 9 frequency = 9 -# Size of the ring buffer for network tracking events (in bytes) -# Default: 1048576 (1 MB) -ringbuf = 1048576 +# Peer maps are normally read less often to limit userspace and trace overhead +# Default: 1 +peer_frequency = 1 # Enable scaled mode to convert byte counts to bits/s rates # When true: reports network rates in bits per second (bits/s) @@ -138,6 +138,23 @@ ringbuf = 1048576 # Default: false scaled = false +# Protocols, directions, aggregation scopes, and metrics to collect. +# Low-overhead defaults are ["tcp", "udp"], ["send", "receive"], +# ["host"], and ["bytes"]. Process and peer aggregation are opt-in. +protocols = ["tcp", "udp"] +directions = ["send", "receive"] +scopes = ["host"] +metrics = ["bytes"] + +# Optional PID and process-name filters. Process names require thread_tracker. +pid_filters = [] +filter_process = [] + +# Bounded map and Perfetto cardinality controls +max_process_entries = 1024 +max_peer_entries = 4096 +max_peer_tracks = 512 + # --------------------------- # Performance Counter Module # --------------------------- @@ -195,4 +212,4 @@ random_process_id = true [[user]] socket = "/var/run/manytrace/debug.sock" -log_filter = "TRACE" \ No newline at end of file +log_filter = "TRACE" diff --git a/tracing-manytrace/examples/net_demo.rs b/tracing-manytrace/examples/net_demo.rs new file mode 100644 index 0000000..1b64a96 --- /dev/null +++ b/tracing-manytrace/examples/net_demo.rs @@ -0,0 +1,253 @@ +// Copyright (C) 2025 Category Labs, Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +use agent::AgentBuilder; +use clap::Parser; +use std::hint::black_box; +use std::io::{self, Read, Write}; +use std::net::{Shutdown, TcpListener, TcpStream, UdpSocket}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::thread; +use std::time::{Duration, Instant}; +use tracing::{info, info_span, instrument, warn}; +use tracing_manytrace::{ManytraceLayer, TracingExtension}; +use tracing_subscriber::prelude::*; + +const TCP_PAYLOAD_BYTES: usize = 32 * 1024; +const UDP_PAYLOAD_BYTES: usize = 1200; +const TICK: Duration = Duration::from_millis(100); + +#[derive(Parser)] +#[command(about = "Instrumented TCP/UDP workload for a manytrace demo")] +struct Args { + #[arg(help = "Path to the manytrace agent socket")] + socket_path: String, + + #[arg(default_value_t = 12, help = "Workload duration in seconds")] + seconds: u64, +} + +#[inline(never)] +fn cpu_burn(seed: u64, iterations: usize) -> u64 { + let mut value = seed | 1; + for i in 0..iterations { + value = value + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(i as u64 ^ value.rotate_left(17)); + } + black_box(value) +} + +fn phase(elapsed: Duration) -> (&'static str, usize, usize) { + match elapsed.as_secs() % 4 { + 0 => ("warmup", 1, 2), + 1 => ("medium", 4, 8), + 2 => ("peak", 12, 24), + _ => ("cooldown", 2, 5), + } +} + +#[instrument(skip(listener))] +fn tcp_server(listener: TcpListener) -> io::Result<()> { + let (mut stream, peer) = listener.accept()?; + info!(%peer, "TCP peer accepted"); + let mut buffer = vec![0u8; TCP_PAYLOAD_BYTES]; + + loop { + match stream.read_exact(&mut buffer) { + Ok(()) => stream.write_all(&buffer)?, + Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => break, + Err(error) => return Err(error), + } + } + + info!("TCP server observed client EOF"); + Ok(()) +} + +#[instrument(skip(stream))] +fn tcp_client(mut stream: TcpStream, duration: Duration) -> io::Result<()> { + let mut probe = [0u8; 1]; + stream.set_nonblocking(true)?; + if let Err(error) = stream.read(&mut probe) { + info!(kind = ?error.kind(), "expected nonblocking TCP receive error"); + } + stream.set_nonblocking(false)?; + + let payload = vec![0x5a; TCP_PAYLOAD_BYTES]; + let mut echo = vec![0u8; TCP_PAYLOAD_BYTES]; + let started = Instant::now(); + let deadline = started + duration; + + while Instant::now() < deadline { + let tick_started = Instant::now(); + let elapsed = tick_started.duration_since(started); + let (phase_name, tcp_bursts, _) = phase(elapsed); + let span = info_span!( + "tcp_burst", + phase = phase_name, + bursts = tcp_bursts, + payload_bytes = TCP_PAYLOAD_BYTES + ); + let _entered = span.enter(); + for _ in 0..tcp_bursts { + stream.write_all(&payload)?; + stream.read_exact(&mut echo)?; + } + black_box(cpu_burn(elapsed.as_nanos() as u64, 80_000)); + if let Some(remaining) = TICK.checked_sub(tick_started.elapsed()) { + thread::sleep(remaining); + } + } + + stream.shutdown(Shutdown::Write)?; + let eof_bytes = stream.read(&mut probe)?; + info!(eof_bytes, "TCP client observed server EOF"); + Ok(()) +} + +#[instrument(skip(socket, running))] +fn udp_server(socket: UdpSocket, running: Arc) -> io::Result<()> { + socket.set_read_timeout(Some(TICK))?; + let mut buffer = [0u8; UDP_PAYLOAD_BYTES]; + while running.load(Ordering::Relaxed) { + match socket.recv_from(&mut buffer) { + Ok((size, peer)) => { + socket.send_to(&buffer[..size], peer)?; + } + Err(error) + if matches!( + error.kind(), + io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(error), + } + } + Ok(()) +} + +#[instrument(skip(socket))] +fn udp_client(socket: UdpSocket, peer: std::net::SocketAddr, duration: Duration) -> io::Result<()> { + socket.connect(peer)?; + let mut probe = [0u8; UDP_PAYLOAD_BYTES]; + socket.set_nonblocking(true)?; + if let Err(error) = socket.recv(&mut probe) { + info!(kind = ?error.kind(), "expected nonblocking UDP receive error"); + } + socket.set_nonblocking(false)?; + socket.set_read_timeout(Some(Duration::from_secs(1)))?; + + if let Err(error) = socket.send(&vec![0u8; 70_000]) { + info!(kind = ?error.kind(), "expected oversized UDP send error"); + } + + let payload = [0xa5; UDP_PAYLOAD_BYTES]; + let started = Instant::now(); + let deadline = started + duration; + while Instant::now() < deadline { + let tick_started = Instant::now(); + let elapsed = tick_started.duration_since(started); + let (phase_name, _, udp_datagrams) = phase(elapsed); + let span = info_span!( + "udp_burst", + phase = phase_name, + datagrams = udp_datagrams, + payload_bytes = UDP_PAYLOAD_BYTES + ); + let _entered = span.enter(); + for _ in 0..udp_datagrams { + socket.send(&payload)?; + let received = socket.recv(&mut probe)?; + if received != payload.len() { + warn!(received, expected = payload.len(), "short UDP echo"); + } + } + if let Some(remaining) = TICK.checked_sub(tick_started.elapsed()) { + thread::sleep(remaining); + } + } + Ok(()) +} + +#[instrument] +fn cpu_worker(duration: Duration) { + let started = Instant::now(); + let deadline = started + duration; + let mut checksum = 1u64; + while Instant::now() < deadline { + let span = info_span!("compute_batch", iteration = checksum & 0xff); + let _entered = span.enter(); + checksum ^= cpu_burn(checksum, 1_500_000); + thread::yield_now(); + } + info!(checksum, "CPU worker complete"); +} + +fn main() -> Result<(), Box> { + let args = Args::parse(); + let extension = Arc::new(TracingExtension::new()); + let _agent = AgentBuilder::new(args.socket_path) + .register_tracing(Box::new((*extension).clone())) + .build()?; + + tracing_subscriber::registry() + .with(ManytraceLayer::new(extension.clone())) + .with(tracing_subscriber::fmt::layer()) + .init(); + + info!("waiting for manytrace to connect"); + while !extension.is_active() { + thread::sleep(Duration::from_millis(50)); + } + + let duration = Duration::from_secs(args.seconds); + info!(seconds = args.seconds, "starting network workload"); + + let tcp_listener = TcpListener::bind("127.0.0.1:0")?; + let tcp_address = tcp_listener.local_addr()?; + let udp_server_socket = UdpSocket::bind("127.0.0.1:0")?; + let udp_address = udp_server_socket.local_addr()?; + let udp_client_socket = UdpSocket::bind("127.0.0.1:0")?; + let udp_running = Arc::new(AtomicBool::new(true)); + + let tcp_server_handle = thread::Builder::new() + .name("tcp-server".to_string()) + .spawn(move || tcp_server(tcp_listener))?; + let tcp_client_handle = thread::Builder::new() + .name("tcp-client".to_string()) + .spawn(move || tcp_client(TcpStream::connect(tcp_address)?, duration))?; + let udp_running_server = udp_running.clone(); + let udp_server_handle = thread::Builder::new() + .name("udp-server".to_string()) + .spawn(move || udp_server(udp_server_socket, udp_running_server))?; + let udp_client_handle = thread::Builder::new() + .name("udp-client".to_string()) + .spawn(move || udp_client(udp_client_socket, udp_address, duration))?; + let cpu_handle = thread::Builder::new() + .name("cpu-worker".to_string()) + .spawn(move || cpu_worker(duration))?; + + tcp_client_handle.join().expect("TCP client panicked")?; + udp_client_handle.join().expect("UDP client panicked")?; + udp_running.store(false, Ordering::Relaxed); + tcp_server_handle.join().expect("TCP server panicked")?; + udp_server_handle.join().expect("UDP server panicked")?; + cpu_handle.join().expect("CPU worker panicked"); + + info!("network workload complete"); + thread::sleep(Duration::from_millis(250)); + Ok(()) +}