diff --git a/src/rust/Cargo.lock b/src/rust/Cargo.lock index 09dc4940..fc49f606 100644 --- a/src/rust/Cargo.lock +++ b/src/rust/Cargo.lock @@ -1505,6 +1505,7 @@ dependencies = [ "lqos_utils", "nix", "once_cell", + "thiserror", ] [[package]] diff --git a/src/rust/lqos_queue_tracker/src/tracking/watched_queues.rs b/src/rust/lqos_queue_tracker/src/tracking/watched_queues.rs index 5dd65901..9904e74a 100644 --- a/src/rust/lqos_queue_tracker/src/tracking/watched_queues.rs +++ b/src/rust/lqos_queue_tracker/src/tracking/watched_queues.rs @@ -2,6 +2,7 @@ use crate::queue_structure::QUEUE_STRUCTURE; use dashmap::DashMap; use log::{info, warn}; use lqos_bus::TcHandle; +use lqos_sys::num_possible_cpus; use lqos_utils::unix_time::unix_now; use once_cell::sync::Lazy; @@ -32,7 +33,7 @@ pub fn expiration_in_the_future() -> u64 { pub fn add_watched_queue(circuit_id: &str) { //info!("Watching queue {circuit_id}"); - let max = unsafe { lqos_sys::libbpf_num_possible_cpus() } * 2; + let max = num_possible_cpus().unwrap() * 2; { if WATCHED_QUEUES.contains_key(circuit_id) { warn!("Queue {circuit_id} is already being watched. Duplicate ignored."); diff --git a/src/rust/lqos_sys/Cargo.toml b/src/rust/lqos_sys/Cargo.toml index f3b39f4e..86adf996 100644 --- a/src/rust/lqos_sys/Cargo.toml +++ b/src/rust/lqos_sys/Cargo.toml @@ -14,6 +14,7 @@ log = "0" lqos_utils = { path = "../lqos_utils" } once_cell = "1" dashmap = "5" +thiserror = "1" [build-dependencies] bindgen = "0" diff --git a/src/rust/lqos_sys/src/bpf_per_cpu_map.rs b/src/rust/lqos_sys/src/bpf_per_cpu_map.rs index 83b10cb3..f494c657 100644 --- a/src/rust/lqos_sys/src/bpf_per_cpu_map.rs +++ b/src/rust/lqos_sys/src/bpf_per_cpu_map.rs @@ -1,7 +1,6 @@ use anyhow::{Error, Result}; use libbpf_sys::{ bpf_map_get_next_key, bpf_map_lookup_elem, bpf_obj_get, - libbpf_num_possible_cpus, }; use std::fmt::Debug; use std::{ @@ -10,6 +9,8 @@ use std::{ ptr::null_mut, }; +use crate::num_possible_cpus; + /// Represents an underlying BPF map, accessed via the filesystem. /// `BpfMap` *only* talks to PER-CPU variants of maps. /// @@ -43,7 +44,7 @@ where /// and allocating, calls `callback` for each key/value slice /// with references to the data returned from C. pub(crate) fn for_each(&self, callback: &mut dyn FnMut(&K, &[V])) { - let num_cpus = unsafe { libbpf_num_possible_cpus() }; + let num_cpus = num_possible_cpus().unwrap(); let mut prev_key: *mut K = null_mut(); let mut key: K = K::default(); let key_ptr: *mut K = &mut key; diff --git a/src/rust/lqos_sys/src/cpu_map.rs b/src/rust/lqos_sys/src/cpu_map.rs index 6b75813d..e565d9a0 100644 --- a/src/rust/lqos_sys/src/cpu_map.rs +++ b/src/rust/lqos_sys/src/cpu_map.rs @@ -1,8 +1,10 @@ use anyhow::{Error, Result}; -use libbpf_sys::{bpf_map_update_elem, bpf_obj_get, libbpf_num_possible_cpus}; +use libbpf_sys::{bpf_map_update_elem, bpf_obj_get}; use log::info; use std::{ffi::CString, os::raw::c_void}; +use crate::num_possible_cpus; + //* Provides an interface for querying the number of CPUs eBPF can //* see, and marking CPUs as available. Currently marks ALL eBPF //* usable CPUs as available. @@ -33,7 +35,7 @@ impl CpuMapping { } pub(crate) fn mark_cpus_available(&self) -> Result<()> { - let cpu_count = unsafe { libbpf_num_possible_cpus() } as u32; + let cpu_count = num_possible_cpus()?; let queue_size = 2048u32; let val_ptr: *const u32 = &queue_size; diff --git a/src/rust/lqos_sys/src/lib.rs b/src/rust/lqos_sys/src/lib.rs index 495de395..31d01c8f 100644 --- a/src/rust/lqos_sys/src/lib.rs +++ b/src/rust/lqos_sys/src/lib.rs @@ -17,6 +17,7 @@ mod lqos_kernel; mod tcp_rtt; mod throughput; mod xdp_ip_address; +mod linux; pub use heimdall_map::{ heimdall_expire, heimdall_for_each, heimdall_watch_ip, set_heimdall_mode, @@ -26,7 +27,7 @@ pub use ip_mapping::{ add_ip_to_tc, clear_ips_from_tc, del_ip_from_tc, list_mapped_ips, }; pub use kernel_wrapper::LibreQoSKernels; -pub use libbpf_sys::libbpf_num_possible_cpus; +pub use linux::num_possible_cpus; pub use lqos_kernel::max_tracked_ips; pub use tcp_rtt::{rtt_for_each, RttTrackingEntry}; pub use throughput::{throughput_for_each, HostCounter}; diff --git a/src/rust/lqos_sys/src/linux/mod.rs b/src/rust/lqos_sys/src/linux/mod.rs new file mode 100644 index 00000000..ae661a08 --- /dev/null +++ b/src/rust/lqos_sys/src/linux/mod.rs @@ -0,0 +1,4 @@ +//! Ports of C code that is very Linux specific. + +mod possible_cpus; +pub use possible_cpus::num_possible_cpus; \ No newline at end of file diff --git a/src/rust/lqos_sys/src/linux/possible_cpus.rs b/src/rust/lqos_sys/src/linux/possible_cpus.rs new file mode 100644 index 00000000..b832b053 --- /dev/null +++ b/src/rust/lqos_sys/src/linux/possible_cpus.rs @@ -0,0 +1,79 @@ +use std::{fs::read_to_string, path::Path}; + +use log::error; +use thiserror::Error; + +const POSSIBLE_CPUS_PATH: &str = "/sys/devices/system/cpu/possible"; + +/// Query the number of available CPUs from `/sys/devices/system/cpu/possible`, +/// and return the last digit (it will be formatted 0-3 or similar) plus one. +/// So on a 16 CPU system, `0-15` will return `16`. +pub fn num_possible_cpus() -> Result { + let path = Path::new(POSSIBLE_CPUS_PATH); + if !path.exists() { + error!("Unable to read /sys/devices/system/cpu/possible"); + return Err(PossibleCpuError::FileNotFound); + }; + + let file_contents = read_to_string(path); + if file_contents.is_err() { + error!("Unable to read contents of /sys/devices/system/cpu/possible"); + error!("{file_contents:?}"); + return Err(PossibleCpuError::UnableToRead); + } + let file_contents = file_contents.unwrap(); + + parse_cpu_string(&file_contents) +} + +fn parse_cpu_string(possible_cpus: &str) -> Result { + if let Some(last_digit) = possible_cpus.trim().split('-').last() { + if let Ok(n) = last_digit.parse::() { + Ok(n + 1) + } else { + error!("Unable to parse /sys/devices/system/cpu/possible"); + error!("{possible_cpus}"); + Err(PossibleCpuError::ParseError) + } + } else { + error!("Unable to parse /sys/devices/system/cpu/possible"); + error!("{possible_cpus}"); + Err(PossibleCpuError::ParseError) + } +} + +#[derive(Error, Debug, Clone, PartialEq)] +pub enum PossibleCpuError { + #[error("Unable to access /sys/devices/system/cpu/possible")] + FileNotFound, + #[error("Unable to read /sys/devices/system/cpu/possible")] + UnableToRead, + #[error("Unable to parse contents of /sys/devices/system/cpu/possible")] + ParseError, +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_unable_to_parse() { + assert_eq!(parse_cpu_string("blah").err().unwrap(), PossibleCpuError::ParseError); + } + + #[test] + fn test_four_cpus() { + assert_eq!(4, parse_cpu_string("0-3").unwrap()); + } + + #[test] + fn test_sixteen_cpus() { + assert_eq!(16, parse_cpu_string("0-15").unwrap()); + } + + #[test] + fn test_againt_c() { + let cpu_count = unsafe { libbpf_sys::libbpf_num_possible_cpus() } as u32; + assert_eq!(cpu_count, num_possible_cpus().unwrap()); + } +}