diff --git a/src/cloud.rs b/src/cloud.rs index 4c2745a..0f31c6b 100644 --- a/src/cloud.rs +++ b/src/cloud.rs @@ -28,7 +28,7 @@ use super::{ traffic::TrafficStats, types::{Error, HeaderMagic, NodeId, Protocol, Range, Table}, udpmessage::{decode, encode, Message}, - util::{addr_nice, resolve, CtrlC, Duration, Time, TimeSource} + util::{addr_nice, resolve, CtrlC, Duration, StatsdMsg, Time, TimeSource} }; pub type Hash = BuildHasherDefault; @@ -600,35 +600,43 @@ impl GenericCloud, pub stats_file: Option, pub statsd_server: Option, + pub statsd_prefix: Option, pub user: Option, pub group: Option } @@ -89,6 +90,7 @@ impl Default for Config { pid_file: None, stats_file: None, statsd_server: None, + statsd_prefix: None, user: None, group: None } @@ -168,6 +170,9 @@ impl Config { if let Some(val) = file.statsd_server { self.statsd_server = Some(val); } + if let Some(val) = file.statsd_prefix { + self.statsd_prefix = Some(val); + } if let Some(val) = file.user { self.user = Some(val); } @@ -246,6 +251,9 @@ impl Config { if let Some(val) = args.statsd_server { self.statsd_server = Some(val); } + if let Some(val) = args.statsd_prefix { + self.statsd_prefix = Some(val); + } if let Some(val) = args.user { self.user = Some(val); } @@ -307,6 +315,7 @@ pub struct ConfigFile { pub pid_file: Option, pub stats_file: Option, pub statsd_server: Option, + pub statsd_prefix: Option, pub user: Option, pub group: Option } @@ -342,6 +351,7 @@ group: nogroup pid_file: /run/vpncloud.run stats_file: /var/log/vpncloud.stats statsd_server: example.com:1234 +statsd_prefix: prefix "; assert_eq!(serde_yaml::from_str::(config_file).unwrap(), ConfigFile { device_type: Some(Type::Tun), @@ -368,7 +378,8 @@ statsd_server: example.com:1234 group: Some("nogroup".to_string()), pid_file: Some("/run/vpncloud.run".to_string()), stats_file: Some("/var/log/vpncloud.stats".to_string()), - statsd_server: Some("example.com:1234".to_string()) + statsd_server: Some("example.com:1234".to_string()), + statsd_prefix: Some("prefix".to_string()) }) } @@ -400,7 +411,8 @@ fn config_merge() { group: Some("nogroup".to_string()), pid_file: Some("/run/vpncloud.run".to_string()), stats_file: Some("/var/log/vpncloud.stats".to_string()), - statsd_server: Some("example.com:1234".to_string()) + statsd_server: Some("example.com:1234".to_string()), + statsd_prefix: Some("prefix".to_string()) }); assert_eq!(config, Config { device_type: Type::Tun, @@ -427,6 +439,7 @@ fn config_merge() { pid_file: Some("/run/vpncloud.run".to_string()), stats_file: Some("/var/log/vpncloud.stats".to_string()), statsd_server: Some("example.com:1234".to_string()), + statsd_prefix: Some("prefix".to_string()), ..Default::default() }); config.merge_args(Args { @@ -453,6 +466,7 @@ fn config_merge() { pid_file: Some("/run/vpncloud-mynet.run".to_string()), stats_file: Some("/var/log/vpncloud-mynet.stats".to_string()), statsd_server: Some("example.com:2345".to_string()), + statsd_prefix: Some("prefix2".to_string()), user: Some("root".to_string()), group: Some("root".to_string()), ..Default::default() @@ -486,6 +500,7 @@ fn config_merge() { pid_file: Some("/run/vpncloud-mynet.run".to_string()), stats_file: Some("/var/log/vpncloud-mynet.stats".to_string()), statsd_server: Some("example.com:2345".to_string()), + statsd_prefix: Some("prefix2".to_string()), daemonize: true }); } diff --git a/src/main.rs b/src/main.rs index f5c062e..ef51a91 100644 --- a/src/main.rs +++ b/src/main.rs @@ -173,6 +173,10 @@ pub struct Args { #[structopt(long)] statsd_server: Option, + /// Use the given prefix for statsd records + #[structopt(long)] + statsd_prefix: Option, + /// Run as other user #[structopt(long)] user: Option, diff --git a/src/traffic.rs b/src/traffic.rs index f446f43..ad868fd 100644 --- a/src/traffic.rs +++ b/src/traffic.rs @@ -10,7 +10,7 @@ use std::{ }; use super::{ - cloud::Hash, + cloud::{Hash, STATS_INTERVAL}, types::Address, util::{addr_nice, Bytes} }; @@ -157,10 +157,10 @@ impl TrafficStats { out, " - peer: \"{}\"\n in: {{ display: \"{}/s\", bytes: {}, packets: {} }}\n out: {{ display: \"{}/s\", bytes: {}, packets: {} }}", addr_nice(**addr), - Bytes(data.in_bytes / 60), + Bytes(data.in_bytes / STATS_INTERVAL as u64), data.in_bytes, data.in_packets, - Bytes(data.out_bytes / 60), + Bytes(data.out_bytes / STATS_INTERVAL as u64), data.out_bytes, data.out_packets )?; @@ -175,10 +175,10 @@ impl TrafficStats { " - addrs: [\"{}\", \"{}\"]\n in: {{ display: \"{}/s\", bytes: {}, packets: {} }}\n out: {{ display: \"{}/s\", bytes: {}, packets: {} }}", remote, local, - Bytes(data.in_bytes / 60), + Bytes(data.in_bytes / STATS_INTERVAL as u64), data.in_bytes, data.in_packets, - Bytes(data.out_bytes / 60), + Bytes(data.out_bytes / STATS_INTERVAL as u64), data.out_bytes, data.out_packets )?; @@ -187,14 +187,14 @@ impl TrafficStats { writeln!( out, "invalid_protocol_traffic: {{ display: \"{}/s\", bytes: {}, packets: {} }}", - Bytes(self.dropped.in_bytes / 60), + Bytes(self.dropped.in_bytes / STATS_INTERVAL as u64), self.dropped.in_bytes, self.dropped.in_packets )?; writeln!( out, "dropped_payload_traffic: {{ display: \"{}/s\", bytes: {}, packets: {} }}", - Bytes(self.dropped.out_bytes / 60), + Bytes(self.dropped.out_bytes / STATS_INTERVAL as u64), self.dropped.out_bytes, self.dropped.out_packets )?; diff --git a/src/types.rs b/src/types.rs index 6526cf0..51dcb21 100644 --- a/src/types.rs +++ b/src/types.rs @@ -228,6 +228,9 @@ pub trait Table { fn remove(&mut self, _: &Address) -> bool; fn remove_all(&mut self, _: &SocketAddr); fn len(&self) -> usize; + fn is_empty(&self) -> bool { + self.len() == 0 + } } pub trait Protocol: Sized { diff --git a/src/util.rs b/src/util.rs index 7444a1d..8e54911 100644 --- a/src/util.rs +++ b/src/util.rs @@ -319,6 +319,35 @@ pub fn from_base62(data: &str) -> Result, char> { } +#[derive(Default)] +pub struct StatsdMsg { + entries: Vec, + key: Vec +} + +impl StatsdMsg { + pub fn new() -> Self { + Default::default() + } + + pub fn add(&mut self, key: &str, val: T, type_: &str) -> &mut Self { + self.entries.push(format!("{}.{}:{}|{}", self.key.join("."), key, val, type_)); + self + } + + pub fn with_ns(&mut self, ns: &str, f: F) -> &mut Self { + self.key.push(ns.to_string()); + f(self); + self.key.pop(); + self + } + + pub fn build(&self) -> String { + self.entries.join("\n") + } +} + + #[test] fn base62() { assert_eq!("", to_base62(&[0]));