Support specified NIC interface statistics

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<String>,
+    #[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<String>,
+}
+
+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

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::<u64>().unwrap();
             let net_out = caps.get(10).unwrap().as_str().parse::<u64>().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::<u64>().unwrap();
                 avgtx += v1[8].parse::<u64>().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;
     }

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<Mutex<NetSpeed>> = 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();

client/stat_client.py

@@ -71,19 +71,19 @@ 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]
             net_out += v[0]
     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: