diff --git a/client/src/main.rs b/client/src/main.rs index c718733d..b42f4a62 100644 --- a/client/src/main.rs +++ b/client/src/main.rs @@ -84,6 +84,38 @@ pub struct Args { location: String, // #[clap(long = "debug", help = "debug mode, default:false")] // debug: bool, + #[clap( + short = 'i', + long = "iface", + value_parser, + default_value = "", + value_delimiter = ',', + require_delimiter = true, + help = "iface list, eg: eth0,eth1" + )] + iface: Vec, + #[clap( + short = 'e', + long = "exclude-iface", + value_parser, + default_value = "lo,docker,vnet,veth,vmbr,kube,br-", + value_delimiter = ',', + help = "exclude iface" + )] + exclude_iface: Vec, +} + +pub fn skip_iface(name: &str, args: &Args) -> bool { + if !args.iface.is_empty() { + if args.iface.iter().any(|fa| name.eq(fa)) { + return false; + } + return true; + } + if args.exclude_iface.iter().any(|sk| name.contains(sk)) { + return true; + } + false } fn sample_all(args: &Args, stat_base: &StatRequest) -> StatRequest { @@ -215,6 +247,8 @@ async fn refresh_ip_info(args: &Args) { async fn main() -> Result<()> { pretty_env_logger::init(); let mut args = Args::parse(); + args.iface.retain(|e| !e.trim().is_empty()); + args.exclude_iface.retain(|e| !e.trim().is_empty()); dbg!(&args); if args.ip_info { @@ -243,7 +277,7 @@ async fn main() -> Result<()> { { eprintln!("enable feature native"); status::start_cpu_percent_collect_t(); - status::start_net_speed_collect_t(); + status::start_net_speed_collect_t(&args); } // use sysinfo diff --git a/client/src/status.rs b/client/src/status.rs index bd5ba5ae..890fdc98 100644 --- a/client/src/status.rs +++ b/client/src/status.rs @@ -19,6 +19,7 @@ use std::sync::Mutex; use std::thread; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use crate::skip_iface; use crate::Args; use stat_common::server_status::StatRequest; @@ -101,8 +102,7 @@ pub fn tupd() -> (u32, u32, u32, u32) { (t, u, p, d) } -static IFACE_IGNORE_VEC: &[&str] = &["lo", "docker", "vnet", "veth", "vmbr", "kube", "br-"]; -pub fn get_vnstat_traffic() -> (u64, u64, u64, u64) { +pub fn get_vnstat_traffic(args: &Args) -> (u64, u64, u64, u64) { let local_now = Local::now(); let (mut network_in, mut network_out, mut m_network_in, mut m_network_out) = (0, 0, 0, 0); let a = Command::new("/usr/bin/vnstat") @@ -114,9 +114,12 @@ pub fn get_vnstat_traffic() -> (u64, u64, u64, u64) { let j: HashMap<&str, serde_json::Value> = serde_json::from_str(b).unwrap(); for iface in j["interfaces"].as_array().unwrap() { let name = iface["name"].as_str().unwrap(); - if IFACE_IGNORE_VEC.iter().any(|sk| name.contains(*sk)) { + + // spec iface + if skip_iface(name, args) { continue; } + let total_o = iface["traffic"]["total"].as_object().unwrap(); let month_v = iface["traffic"]["month"].as_array().unwrap(); network_in += total_o["rx"].as_u64().unwrap(); @@ -142,7 +145,7 @@ static TRAFFIC_REGEX: &str = lazy_static! { static ref TRAFFIC_REGEX_RE: Regex = Regex::new(TRAFFIC_REGEX).unwrap(); } -pub fn get_sys_traffic() -> (u64, u64) { +pub fn get_sys_traffic(args: &Args) -> (u64, u64) { let (mut network_in, mut network_out) = (0, 0); let file = File::open("/proc/net/dev").unwrap(); let buf_reader = BufReader::new(file); @@ -152,9 +155,12 @@ pub fn get_sys_traffic() -> (u64, u64) { TRAFFIC_REGEX_RE.captures(&l).and_then(|caps| { // println!("caps[0]=>{:?}", caps.get(0).unwrap().as_str()); let name = caps.get(1).unwrap().as_str(); - if IFACE_IGNORE_VEC.iter().any(|sk| name.contains(*sk)) { + + // spec iface + if skip_iface(name, args) { return None; } + let net_in = caps.get(2).unwrap().as_str().parse::().unwrap(); let net_out = caps.get(10).unwrap().as_str().parse::().unwrap(); @@ -203,8 +209,9 @@ lazy_static! { } #[allow(unused)] -pub fn start_net_speed_collect_t() { - thread::spawn(|| loop { +pub fn start_net_speed_collect_t(args: &Args) { + let args_1 = args.clone(); + thread::spawn(move || loop { let _ = File::open("/proc/net/dev").map(|file| { let buf_reader = BufReader::new(file); let (mut avgrx, mut avgtx) = (0, 0); @@ -215,9 +222,11 @@ pub fn start_net_speed_collect_t() { continue; } - if IFACE_IGNORE_VEC.iter().any(|sk| v[0].contains(*sk)) { + // spec iface + if skip_iface(v[0], &args_1) { continue; } + let v1: Vec<&str> = v[1].split_whitespace().collect(); avgrx += v1[0].parse::().unwrap(); avgtx += v1[8].parse::().unwrap(); @@ -420,13 +429,13 @@ pub fn sample(args: &Args, stat: &mut StatRequest) { stat.thread = d; if args.vnstat { - let (network_in, network_out, m_network_in, m_network_out) = get_vnstat_traffic(); + let (network_in, network_out, m_network_in, m_network_out) = get_vnstat_traffic(args); stat.network_in = network_in; stat.network_out = network_out; stat.last_network_in = network_in - m_network_in; stat.last_network_out = network_out - m_network_out; } else { - let (network_in, network_out) = get_sys_traffic(); + let (network_in, network_out) = get_sys_traffic(args); stat.network_in = network_in; stat.network_out = network_out; } diff --git a/client/src/sys_info.rs b/client/src/sys_info.rs index bda2c27f..b71161dc 100644 --- a/client/src/sys_info.rs +++ b/client/src/sys_info.rs @@ -7,13 +7,13 @@ use std::thread; use std::time::Duration; use sysinfo::{CpuExt, DiskExt, NetworkExt, RefreshKind, System, SystemExt}; +use crate::skip_iface; use crate::status; use crate::status::get_vnstat_traffic; use crate::Args; use stat_common::server_status::{StatRequest, SysInfo}; const SAMPLE_PERIOD: u64 = 1000; //ms -static IFACE_IGNORE_VEC: &[&str] = &["lo", "docker", "vnet", "veth", "vmbr", "kube", "br-"]; lazy_static! { pub static ref G_EXPECT_FS: Vec<&'static str> = [ @@ -61,13 +61,15 @@ lazy_static! { pub static ref G_NET_SPEED: Arc> = Arc::new(Default::default()); } -pub fn start_net_speed_collect_t() { +pub fn start_net_speed_collect_t(args: &Args) { let mut sys = System::new_all(); sys.refresh_all(); + let args_1 = args.clone(); thread::spawn(move || loop { let (mut net_rx, mut net_tx) = (0_u64, 0_u64); for (name, data) in sys.networks() { - if IFACE_IGNORE_VEC.iter().any(|sk| name.contains(*sk)) { + // spec iface + if skip_iface(name, &args_1) { continue; } net_rx += data.received(); @@ -144,7 +146,7 @@ pub fn sample(args: &Args, stat: &mut StatRequest) { // traffic if args.vnstat { - let (network_in, network_out, m_network_in, m_network_out) = get_vnstat_traffic(); + let (network_in, network_out, m_network_in, m_network_out) = get_vnstat_traffic(args); stat.network_in = network_in; stat.network_out = network_out; stat.last_network_in = network_in - m_network_in; @@ -153,7 +155,8 @@ pub fn sample(args: &Args, stat: &mut StatRequest) { sys.refresh_networks(); let (mut network_in, mut network_out) = (0_u64, 0_u64); for (name, data) in sys.networks() { - if IFACE_IGNORE_VEC.iter().any(|sk| name.contains(*sk)) { + // spec iface + if skip_iface(name, args) { continue; } network_in += data.total_received(); diff --git a/client/stat_client.py b/client/stat_client.py index caec1f04..eb3bf5fa 100644 --- a/client/stat_client.py +++ b/client/stat_client.py @@ -71,11 +71,11 @@ def get_cpu(): return psutil.cpu_percent(interval=INTERVAL) -def get_sys_traffic(): +def get_sys_traffic(options): net_in, net_out = 0, 0 net = psutil.net_io_counters(pernic=True) for k, v in net.items(): - if any([e in k for e in IFACE_IGNORE_LIST]): + if skip_iface(k, options): continue else: net_in += v[1] @@ -83,7 +83,7 @@ def get_sys_traffic(): return net_in, net_out -def get_vnstat_traffic(): +def get_vnstat_traffic(options): now = datetime.now() vnstat_res = subprocess.check_output( "/usr/bin/vnstat --json m", shell=True) @@ -92,7 +92,7 @@ def get_vnstat_traffic(): for iface in json_dict.get("interfaces", []): name = iface["name"] - if any([e in name for e in IFACE_IGNORE_LIST]): + if skip_iface(name, options): continue traffic = iface["traffic"] @@ -200,11 +200,11 @@ def _ping_thread(target, mark): time.sleep(INTERVAL) -def _net_speed(): +def _net_speed(options): while True: avgrx, avgtx = 0, 0 for name, stats in psutil.net_io_counters(pernic=True).items(): - if any([e in name for e in IFACE_IGNORE_LIST]): + if skip_iface(name, options): continue avgrx += stats.bytes_recv avgtx += stats.bytes_sent @@ -255,6 +255,9 @@ def start_rt_collect_t(options): # net speed t_list.append(threading.Thread( target=_net_speed, + kwargs={ + 'options': options, + } )) for t in t_list: @@ -308,13 +311,13 @@ def sample(options, stat_base): if options.vnstat: (network_in, network_out, m_network_in, - m_network_out) = get_vnstat_traffic() + m_network_out) = get_vnstat_traffic(options) stat_data['network_in'] = network_in stat_data['network_out'] = network_out stat_data['last_network_in'] = network_in - m_network_in stat_data['last_network_out'] = network_out - m_network_out else: - net_in, net_out = get_sys_traffic() + net_in, net_out = get_sys_traffic(options) stat_data['network_in'] = net_in stat_data['network_out'] = net_out @@ -462,6 +465,14 @@ def gen_sys_id(sys_info): ) return hashlib.md5(s.encode("utf-8")).hexdigest() +def skip_iface(name, options): + if len(options.iface) > 0: + if any([name == e for e in options.iface]): + return False + return True + if any([e in name for e in options.exclude_iface]): + return True + return False def main(): usage = """usage: python3 %prog [options] arg @@ -503,8 +514,16 @@ def main(): default="", help="host type [default: %default]") parser.add_option("--location", dest="location", default="", help="location [default: %default]") + parser.add_option("-i", "--iface", dest="iface", + default="", help="iface list, eg: eth0,eth1 [default: %default]") + parser.add_option("-e", "--exclude-iface", dest="exclude_iface", + default="lo,docker,vnet,veth,vmbr,kube,br-", + help="exclude iface [default: %default]") (options, args) = parser.parse_args() + parse_iface_list = lambda ifaces: list(filter(lambda s: len(s), map(str.strip, ifaces.split(",")))) + options.iface = parse_iface_list(options.iface) + options.exclude_iface = parse_iface_list(options.exclude_iface) print(json.dumps(options.__dict__, indent=2)) if options.vnstat: