diff --git a/Cargo.lock b/Cargo.lock index 07ae7140..f964f4e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -767,6 +767,20 @@ dependencies = [ "syn 2.0.101", ] +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "data-encoding" version = "2.6.0" @@ -1778,11 +1792,10 @@ checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" [[package]] name = "lock_api" -version = "0.4.12" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" dependencies = [ - "autocfg", "scopeguard", ] @@ -2235,9 +2248,9 @@ checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" [[package]] name = "parking_lot" -version = "0.12.3" +version = "0.12.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" dependencies = [ "lock_api", "parking_lot_core", @@ -2245,15 +2258,15 @@ dependencies = [ [[package]] name = "parking_lot_core" -version = "0.9.10" +version = "0.9.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" dependencies = [ "cfg-if", "libc", "redox_syscall", "smallvec", - "windows-targets 0.52.6", + "windows-link 0.2.1", ] [[package]] @@ -2398,6 +2411,7 @@ dependencies = [ "clickhouse", "config", "console-subscriber", + "dashmap", "data-encoding", "default-net", "defguard_wireguard_rs", @@ -2410,6 +2424,7 @@ dependencies = [ "log", "netstat2", "openssl", + "parking_lot", "percent-encoding", "postgres-types", "prost", diff --git a/Cargo.toml b/Cargo.toml index e869e489..2de5ca4c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ config = "0.14" console-subscriber = {version = "0.4", optional = true} clickhouse = { version = "0.13.0" } clap = { version = "4.4", features = ["derive"] } +dashmap = "6.1.0" defguard_wireguard_rs = {version = "0.7.2", features=["serde"]} default-net = "0.22" fern = "0.7" @@ -59,6 +60,7 @@ qrcode = "0.14" x25519-dalek = { version = "2", features = ["static_secrets"] } warp = "0.3" zmq = "0.10" +parking_lot = "0.12.5" sha2 = "0.10" data-encoding = "2.5" diff --git a/config-agent-example.toml b/config-agent-example.toml index d47e969e..b4be0b15 100644 --- a/config-agent-example.toml +++ b/config-agent-example.toml @@ -1,57 +1,38 @@ -[debug] -# enabled = true -# web_server = "127.0.0.1" -# web_port = 3001 - -[carbon] -# address = "localhost:2003" - -[logging] -# level = "debug" - [agent] -# local = false -# metrics_enabled = true -# metrics_interval = 60 -# stat_enabled = true -# stat_job_interval = 60 +local = false +snapshot_interval = 30 +snapshot_path = "snapshots/agent_snapshot.bin" [xray] -# xray_config_path = "dev/xray-config.json" +enabled = true +xray_config_path = "/etc/xray/config.json" -[zmq] -# sub_endpoint = "tcp://localhost:3000" -# pub_endpoint = "tcp://*:3000" - -[node] -# env = "dev" -# hostname = "" -# default_interface = "" -# ipv4 = "" -# uuid = "9658b391-01cb-4031-a3f5-6cbdd749bcff" -# label = "πŸ΄β€β˜ οΈπŸ΄β€β˜ οΈπŸ΄β€β˜ οΈ dev" +[wg] +enabled = true +port = 51820 +interface = "wg0" +# privkey ΠΈ address ΠΎΠ±Ρ‹Ρ‡Π½ΠΎ Π»ΡƒΡ‡ΡˆΠ΅ Ρ‡Π΅Ρ€Π΅Π· env, Π½ΠΎ ΠΌΠΎΠΆΠ½ΠΎ ΠΈ Ρ‚ΡƒΡ‚ [h2] -#enabled = true -#path = "path/to/h2.yaml" +enabled = false +path = "etc/frkn/h2.yaml" -[wg] -#enabled = true -#port = 51599 -#network= "10.10.0.0/16" -#address = "10.10.0.1" -#interface = "utun7" -#privkey = "6KVWBZmCBNt1t5hI6KXXf4Eo4mlGNKAxETeRgcXNK0Q=" -#pubkey = "Yhg+gj8TqQ3/10a/Ae+Tv1iPin/HigmOnhRkMj7k/hw=" -#dns = ["1.1.1.1"] +[mtproto] +enabled = true +port = 8443 +secret = "dd00112233445566778899aabbccddeeff" + +[metrics] +enabled = true +publisher = "tcp://api-server-ip:5555" +interval = 10 [api] -# address = "127.0.0.1" -# port = 3005 -# endpoint = "http://localhost:3005" -# node_health_check_timeout = 60 -# user_limit_check_interval = 60 -# user_reactivate_interval = 60 -# token = "supetsecrettoken" -# user_limit_mb = 1000 -# healthcheck_interval = 60 +endpoint = "https://api.frkn.org" +token = "your-super-secret-api-token" + +[logging] +level = "warn" + +[zmq] +endpoint = "tcp://api-server-ip:3000" diff --git a/config-api-example.toml b/config-api-example.toml index df939f72..630286e4 100644 --- a/config-api-example.toml +++ b/config-api-example.toml @@ -1,39 +1,29 @@ [api] -# address = "127.0.0.1" -# port = 3005 -# endpoint = "http://localhost:3005" -# node_health_check_timeout = 60 -# user_limit_check_interval = 60 -# user_reactivate_interval = 60 -# token = "supetsecrettoken" -# user_limit_mb = 1000 -# healthcheck_interval = 60 - -[debug] -# enabled = true -# web_server = "127.0.0.1" -# web_port = 3001 +address = "0.0.0.0" +port = 3005 +token = "your-super-secret-api-token" +metrics_enabled = true +metrics_interval = 60 +db_sync_interval_sec = 300 +subscription_restore_interval = 60 +subscription_expire_interval = 60 +max_points = 100000000 +retention_seconds = 604800 # 7 days [node] -# env = "dev" -# hostname = "" -# default_interface = "" -# ipv4 = "" -# uuid = "9658b391-01cb-4031-a3f5-6cbdd749bcff" -# label = "πŸ΄β€β˜ οΈπŸ΄β€β˜ οΈπŸ΄β€β˜ οΈ dev" +env = "experimental" +label = "πŸ΄β€β˜ οΈ DarkMachine" +max_bandwidth_bps = 1000000000 #1gbps [logging] -# level = "debug" +level = "debug" [zmq] -# endpoint = "tcp://*:3000" - -[clickhouse] -# address = "http://localhost:8123" +endpoint = "tcp://*:3000" [pg] -# host = "localhost" -# port = 5432 -# db = "postgres" -# username = "postgres" -# password = "password" +host = "localhost" +port = 5432 +db = "frkn_db" +username = "postgres" +password = "password" diff --git a/config-auth-eample.toml b/config-auth-eample.toml index e8e55014..5751e5d6 100644 --- a/config-auth-eample.toml +++ b/config-auth-eample.toml @@ -1,26 +1,28 @@ +[auth] +snapshot_interval = 30 +snapshot_path = "snapshots/auth_snapshot.bin" +web_host = "https://auth.frkn.org" +web_server = "0.0.0.0" +web_port = 3005 +email_file = "users_trials.csv" -[debug] -#enable = true -#web_server = "127.0.0.1" -#web_port = 3445 +[api] +endpoint = "http://localhost:3005" +token = "your-super-secret-api-token" +[smtp] +server = "smtp.gmail.com" +port = 587 +username = "notifications@frkn.org" +password = "app-password-here" +from = "FRKN Privacy " -[logging] -#level = "debug" - -[zmq] -#endpoint = "tcp://localhost:4000" - -[auth] -#snapshot_path = "snapshots/snapshot.bin" - -[node] -#env = "dev" -#hostname = "darkmachine-test" -#uuid = "9557b391-01cb-4031-a3f5-6cbdd74fffff" - -[api] -#endpoint = "http://127.0.0.1:5005" -#token = "token" +[metrics] +publisher = "tcp://127.0.0.1:5555" +interval = 60 +[logging] +level = "info" +[zmq] +endpoint = "tcp://localhost:3000" diff --git a/dev/user-state.html b/dev/user-state.html deleted file mode 100644 index ddd6ea13..00000000 --- a/dev/user-state.html +++ /dev/null @@ -1,330 +0,0 @@ - - - - - State Info - - - -

Connections Info

- -
- - - - - - - - - - - - - - - - -
- -
- -
- - - -
- - - - - - - - - - - - - - - diff --git a/dev/xray-config.json b/dev/xray-config.json index f2c92494..3dbb0ec4 100644 --- a/dev/xray-config.json +++ b/dev/xray-config.json @@ -1,174 +1,160 @@ { - "log": { - "loglevel": "debug" - }, - "inbounds": [ - { - "tag": "VlessXhttpReality", - "listen": "0.0.0.0", - "port": 443, - "protocol": "vless", - "settings": { - "clients": [ - { - "id": "ff9a7bc1-dd25-4217-a7f1-208cff4692ab", - "email": "wl-test-xhttp@main", - "flow": "" - } - ], - "decryption": "none" - }, - "streamSettings": { - "network": "xhttp", - "security": "reality", - "realitySettings": { - "xver": 0, - "target": "ya.ru:443", - "serverNames": ["ya.ru"], - "privateKey": "0EtqcOKG4wJsURvTfzSWRvH0dPeECVL1TXY21Ydx6Ho", - "publicKey": "jgdqVgUnunocvov5MtuGKh9fVHCpSXIWega-VsHGXAE", - "shortIds": ["0498897564df"] - }, - "xhttpSettings": { - "path": "/" - } - }, - "sniffing": { - "enabled": true, - "destOverride": ["http", "tls", "quic"] - } + "log": { + "loglevel": "debug" }, - { - "tag": "VlessGrpcReality", - "listen": "0.0.0.0", - "port": 8443, - "protocol": "vless", - "settings": { - "clients": [ - { - "id": "ff9a7bc1-dd25-4217-a7f1-208cff4692ab", - "email": "wl-test-grpc@main" - } - ], - "decryption": "none" - }, - "streamSettings": { - "network": "grpc", - "security": "reality", - - - "realitySettings": { - "xver": 0, - "target": "ya.ru:443", - "serverNames": ["ya.ru"], - "privateKey": "0EtqcOKG4wJsURvTfzSWRvH0dPeECVL1TXY21Ydx6Ho", - "publicKey": "jgdqVgUnunocvov5MtuGKh9fVHCpSXIWega-VsHGXAE", - "shortIds": ["0498897564df"] + "inbounds": [ + { + "tag": "VlessXhttpReality", + "listen": "0.0.0.0", + "port": 443, + "protocol": "vless", + "settings": { + "clients": [], + "decryption": "none" + }, + "streamSettings": { + "network": "xhttp", + "security": "reality", + "realitySettings": { + "xver": 0, + "target": "ya.ru:443", + "serverNames": [ + "ya.ru" + ], + "privateKey": "0EtqcOKG4wJsURvTfzSWRvH0dPeECVL1TXY21Ydx6Ho", + "publicKey": "jgdqVgUnunocvov5MtuGKh9fVHCpSXIWega-VsHGXAE", + "shortIds": [ + "0498897564df" + ] + }, + "xhttpSettings": { + "path": "/" + } + }, + "sniffing": { + "enabled": true, + "destOverride": [ + "http", + "tls", + "quic" + ] + } + }, + { + "tag": "VlessGrpcReality", + "listen": "0.0.0.0", + "port": 8443, + "protocol": "vless", + "settings": { + "clients": [], + "decryption": "none" + }, + "streamSettings": { + "network": "grpc", + "security": "reality", + "realitySettings": { + "xver": 0, + "target": "ya.ru:443", + "serverNames": [ + "ya.ru" + ], + "privateKey": "0EtqcOKG4wJsURvTfzSWRvH0dPeECVL1TXY21AAAAAA", + "publicKey": "jgdqVgUnunocvov5MtuGKh9fVHCpSXIWega-AAAAAAA", + "shortIds": [ + "0498897564df" + ] + }, + "grpcSettings": { + "serviceName": "api.v1" + } + } }, - "grpcSettings": { - "serviceName": "api.v1" + { + "tag": "VlessTcpReality", + "listen": "0.0.0.0", + "port": 8444, + "protocol": "vless", + "settings": { + "clients": [], + "decryption": "none" + }, + "streamSettings": { + "network": "tcp", + "security": "reality", + "realitySettings": { + "serverNames": [ + "cdn.discordapp.com", + "discordapp.com" + ], + "privateKey": "0EtqcOKG4wJsURvTfzSWRvH0dPeECVL1AAAAAAAA", + "publicKey": "jgdqVgUnunocvov5MtuGKh9fVHCpSXIWega-AAAAAA", + "shortIds": [ + "e5c4d84fb339fb92" + ], + "target": "discordapp.com:443" + } + } } - } - }, - - - { - "tag": "VlessTcpReality", - "listen": "0.0.0.0", - "port": 8444, - "protocol": "vless", - "settings": { - "clients": [], - "decryption": "none" - }, - "streamSettings": { - "network": "tcp", - "security": "reality", - "realitySettings": { - "serverNames": [ - "cdn.discordapp.com", - "discordapp.com" - ], - "privateKey": "0EtqcOKG4wJsURvTfzSWRvH0dPeECVL1TXY21Ydx6Ho", - "publicKey": "jgdqVgUnunocvov5MtuGKh9fVHCpSXIWega-VsHGXAE", - "shortIds": [ - "e5c4d84fb339fb92" - ], - "target": "discordapp.com:443" + ], + "outbounds": [ + { + "protocol": "freedom", + "tag": "DIRECT" + }, + { + "protocol": "blackhole", + "tag": "BLOCK" } - } - } - ], - "outbounds": [ - { - "protocol": "freedom", - "tag": "DIRECT" + ], + "routing": { + "rules": [ + { + "ip": [ + "geoip:private" + ], + "outboundTag": "BLOCK", + "type": "field" + }, + { + "domain": [ + "geosite:private" + ], + "outboundTag": "BLOCK", + "type": "field" + }, + { + "protocol": [ + "bittorrent" + ], + "outboundTag": "BLOCK", + "type": "field" + } + ] }, - { - "protocol": "blackhole", - "tag": "BLOCK" - } - ], - "routing": { - "rules": [ - { - "inboundTag": [ - "API_INBOUND" - ], - "source": [ - "127.0.0.1", - "194.54.156.79" - ], - "outboundTag": "API", - "type": "field" - }, - { - "ip": [ - "geoip:private" - ], - "outboundTag": "BLOCK", - "type": "field" - }, - { - "domain": [ - "geosite:private" + "api": { + "listen": "127.0.0.1:23456", + "services": [ + "HandlerService", + "StatsService", + "LoggerService", + "ReflectionService" ], - "outboundTag": "BLOCK", - "type": "field" - }, - { - "protocol": [ - "bittorrent" - ], - "outboundTag": "BLOCK", - "type": "field" - } - ] - }, - "api": { - "listen": "127.0.0.1:23456", - "services": [ - "HandlerService", - "StatsService", - "LoggerService", - "ReflectionService" - ], - "tag": "API" - }, - "stats": {}, - "policy": { - "levels": { - "0": { - "statsUserUplink": true, - "statsUserDownlink": true, - "statsUserOnline": true - } + "tag": "API" }, - "system": { - "statsInboundDownlink": false, - "statsInboundUplink": false, - "statsOutboundDownlink": true, - "statsOutboundUplink": true + "stats": {}, + "policy": { + "levels": { + "0": { + "statsUserUplink": true, + "statsUserDownlink": true, + "statsUserOnline": true + } + }, + "system": { + "statsInboundDownlink": false, + "statsInboundUplink": false, + "statsOutboundDownlink": true, + "statsOutboundUplink": true + } } - } } diff --git a/src/bin/agent/core/http.rs b/src/bin/agent/core/http.rs index 40482a7a..5cc17729 100644 --- a/src/bin/agent/core/http.rs +++ b/src/bin/agent/core/http.rs @@ -7,14 +7,14 @@ use std::collections::HashMap; use pony::config::xray::Inbound; use pony::ConnectionBaseOp; -use pony::NodeStorageOp; -use pony::SubscriptionOp; use pony::Tag; use pony::{PonyError, Result}; use std::net::Ipv4Addr; use super::Agent; +use pony::http::response::{Instance, InstanceWithId, ResponseMessage}; + #[derive(Clone, Debug, Deserialize, Serialize)] pub struct ConnTypeParam { pub proto: Tag, @@ -34,6 +34,7 @@ pub struct NodeRequest { pub interface: String, pub cores: usize, pub max_bandwidth_bps: i64, + pub country: String, } #[async_trait] @@ -49,11 +50,9 @@ pub trait ApiRequests { } #[async_trait] -impl ApiRequests for Agent +impl ApiRequests for Agent where - T: NodeStorageOp + Send + Sync + Clone, C: ConnectionBaseOp + Send + Sync + Clone + 'static, - S: SubscriptionOp + Send + Sync + Clone + 'static, { async fn get_connections( &self, @@ -62,13 +61,7 @@ where proto: Tag, last_update: Option, ) -> Result<()> { - let node = { - let mem = self.memory.read().await; - mem.nodes - .get_self() - .expect("No node available to register") - .clone() - }; + let node = self.node.clone(); let id = node.uuid; let env = node.env; @@ -97,8 +90,25 @@ where let status = res.status(); let body = res.text().await?; + if status.is_success() { - log::debug!("Connections Request Accepted for {}: {} ", proto, status); + let result: ResponseMessage> = serde_json::from_str(&body)?; + let count = match result.response.instance { + Instance::Count(count) => count, + _ => { + return Err(PonyError::Custom("Unexpected instance type".into())); + } + }; + log::debug!( + "Message: {}. Connections Request Accepted for {}: {} Count: {} ", + result.message, + proto, + result.status, + count, + ); + Ok(()) + } else if status == StatusCode::NOT_MODIFIED { + log::debug!("Connections Request Accepted for {}: {} ", proto, status,); Ok(()) } else { log::error!( @@ -115,13 +125,7 @@ where } async fn register_node(&self, endpoint: String, token: String) -> Result<()> { - let node = { - let mem = self.memory.read().await; - mem.nodes - .get_self() - .expect("No node available to register") - .clone() - }; + let node = self.node.clone(); let mut endpoint_url = Url::parse(&endpoint)?; endpoint_url @@ -145,6 +149,7 @@ where interface: node.interface.clone(), cores: node.cores, max_bandwidth_bps: node.max_bandwidth_bps, + country: node.country, }; let res = HttpClient::new() diff --git a/src/bin/agent/core/metrics.rs b/src/bin/agent/core/metrics.rs index 3f3dd150..168a5221 100644 --- a/src/bin/agent/core/metrics.rs +++ b/src/bin/agent/core/metrics.rs @@ -1,74 +1,147 @@ -use chrono::Utc; - -use pony::metrics::bandwidth::bandwidth_metrics; -use pony::metrics::cpuusage::cpu_metrics; -use pony::metrics::heartbeat::heartbeat_metrics; -use pony::metrics::loadavg::loadavg_metrics; -use pony::metrics::memory::mem_metrics; -use pony::metrics::xray::xray_conn_metrics; -use pony::metrics::xray::xray_stat_metrics; -use pony::metrics::Metric; -use pony::metrics::MetricType; -use pony::metrics::Metrics; -use pony::Connection; +use crate::core::Agent; +use pony::memory::node::Node; +use pony::metrics::storage::HasMetrics; +use pony::metrics::storage::MetricBuffer; +use pony::xray_op::stats::{Prefix, StatOp}; use pony::ConnectionBaseOp; -use pony::NodeStorageOp; -use pony::SubscriptionOp; +use pony::Tag; -use crate::core::Agent; +impl HasMetrics for Agent +where + C: ConnectionBaseOp + Send + Sync + Clone + 'static, +{ + fn metrics(&self) -> &MetricBuffer { + &self.metrics + } + + fn node_settings(&self) -> &Node { + &self.node + } +} + +#[async_trait::async_trait] +pub trait BusinessMetrics { + async fn collect_inbound_metrics(&self); + async fn collect_user_metrics(&self); + async fn collect_wg_metrics(&self); +} #[async_trait::async_trait] -impl Metrics for Agent +impl BusinessMetrics for Agent where - T: NodeStorageOp + Send + Sync + Clone + 'static, - C: ConnectionBaseOp + Send + Sync + Clone + 'static + From, - S: Send + Sync + Clone + 'static + std::cmp::PartialEq + SubscriptionOp, + C: ConnectionBaseOp + Send + Sync + Clone + 'static, { - async fn collect_metrics(&self) -> Vec - where - T: NodeStorageOp + Sync + Send + Clone + 'static, - { - let mut metrics: Vec = Vec::new(); - let mem = self.memory.read().await; - let connections = mem.connections.clone(); - - let node = mem.nodes.get_self(); - - if let Some(node) = node { - let bandwidth: Vec = - bandwidth_metrics(&node.env, &node.hostname, &node.interface); - let cpuusage: Vec = cpu_metrics(&node.env, &node.hostname); - let loadavg: Vec = loadavg_metrics(&node.env, &node.hostname); - let memory: Vec = mem_metrics(&node.env, &node.hostname); - - let xray_stat: Vec = xray_stat_metrics(node.clone()); - let connections_stat: Vec = - xray_conn_metrics(connections, &node.env, &node.hostname); - - metrics.extend(bandwidth); - metrics.extend(cpuusage); - metrics.extend(loadavg); - metrics.extend(memory); - metrics.extend(xray_stat); - metrics.extend(connections_stat); + async fn collect_inbound_metrics(&self) { + let node_uuid = self.node.uuid; + let base_tags = self.node.get_base_tags(); + + for tag in self.node.inbounds.keys() { + if matches!(tag, Tag::Hysteria2 | Tag::Mtproto) { + continue; + } + + let prefix = Prefix::InboundPrefix(*tag); + + if let Ok(stats) = self.inbound_stats(prefix).await { + let mut metric_tags = base_tags.clone(); + metric_tags.insert("inbound_tag".to_string(), tag.to_string()); + + let metric_prefix = format!("net.inbound.{}", tag); + + self.metrics.push( + node_uuid, + &format!("{}.downlink", metric_prefix), + stats.downlink as f64, + metric_tags.clone(), + ); + self.metrics.push( + node_uuid, + &format!("{}.uplink", metric_prefix), + stats.uplink as f64, + metric_tags.clone(), + ); + self.metrics.push( + node_uuid, + &format!("{}.connections", metric_prefix), + stats.conn_count as f64, + metric_tags, + ); + } } + } + + async fn collect_user_metrics(&self) { + let node_uuid = self.node.uuid; + let base_tags = self.node.get_base_tags(); - log::debug!("Total metrics collected: {}", metrics.len()); + let active_conns = { + let mem = self.memory.read().await; + mem.keys().cloned().collect::>() + }; - metrics + for conn_id in active_conns { + let res = self.conn_stats(Prefix::ConnPrefix(conn_id)).await; + match res { + Ok(stats) => { + log::debug!("Successfully fetched stats for {}", conn_id); + let mut metric_tags = base_tags.clone(); + metric_tags.insert("conn_id".to_string(), conn_id.to_string()); + + self.metrics.push( + node_uuid, + "user.traffic.downlink", + stats.downlink as f64, + metric_tags.clone(), + ); + self.metrics.push( + node_uuid, + "user.traffic.uplink", + stats.uplink as f64, + metric_tags.clone(), + ); + self.metrics + .push(node_uuid, "user.online", stats.online as f64, metric_tags); + } + Err(e) => { + log::error!("Failed to get stats for user {}: {:?}", conn_id, e); + } + } + } } - async fn collect_hb_metrics(&self) -> MetricType { - let mem = self.memory.read().await; - let node = mem.nodes.get_self(); - if let Some(node) = node { - heartbeat_metrics(&node.env, &node.uuid, &node.hostname) - } else { - MetricType::F64(Metric::new( - "hb.unknown".into(), - 0.0, - Utc::now().timestamp(), - )) + async fn collect_wg_metrics(&self) { + let wg_client = match &self.wg_client { + Some(c) => c, + None => return, + }; + + let node_uuid = self.node.uuid; + let base_tags = self.node.get_base_tags(); + + let wg_conns = { + let mem = self.memory.read().await; + mem.iter() + .filter_map(|(id, conn)| { + conn.get_wireguard().map(|wg| (*id, wg.keys.pubkey.clone())) + }) + .collect::>() + }; + + for (conn_id, pubkey) in wg_conns { + if let Ok((uplink, downlink)) = wg_client.peer_stats(&pubkey) { + let mut metric_tags = base_tags.clone(); + metric_tags.insert("user_id".to_string(), conn_id.to_string()); + metric_tags.insert("proto".to_string(), "wireguard".to_string()); + + self.metrics.push( + node_uuid, + "user.traffic.downlink", + downlink as f64, + metric_tags.clone(), + ); + self.metrics + .push(node_uuid, "user.traffic.uplink", uplink as f64, metric_tags); + } } } } diff --git a/src/bin/agent/core/mod.rs b/src/bin/agent/core/mod.rs index 32174ef7..93c59453 100644 --- a/src/bin/agent/core/mod.rs +++ b/src/bin/agent/core/mod.rs @@ -1,18 +1,15 @@ -use pony::Subscription; use std::sync::Arc; use tokio::sync::Mutex; use tokio::sync::RwLock; +use pony::memory::connection::Connections; use pony::memory::node::Node; +use pony::metrics::storage::MetricBuffer; use pony::wireguard_op::WgApi; use pony::xray_op::client::HandlerClient; use pony::xray_op::client::StatsClient; use pony::zmq::subscriber::Subscriber as ZmqSubscriber; -use pony::BaseConnection as Connection; use pony::ConnectionBaseOp; -use pony::MemoryCache; -use pony::NodeStorageOp; -use pony::SubscriptionOp; mod http; pub(crate) mod metrics; @@ -21,36 +18,36 @@ mod snapshot; mod stats; mod tasks; -pub type AgentState = MemoryCache; - -pub struct Agent +pub struct Agent where - N: NodeStorageOp + Send + Sync + Clone + 'static, C: ConnectionBaseOp + Send + Sync + Clone + 'static, - S: SubscriptionOp + Send + Sync + Clone + 'static, { - pub memory: Arc>>, + pub memory: Arc>>, + pub node: Node, + pub metrics: Arc, pub subscriber: ZmqSubscriber, pub xray_stats_client: Option>>, pub xray_handler_client: Option>>, pub wg_client: Option, } -impl Agent +impl Agent where - N: NodeStorageOp + Send + Sync + Clone + 'static, C: ConnectionBaseOp + Send + Sync + Clone + 'static, - S: SubscriptionOp + Send + Sync + Clone + 'static, { pub fn new( - memory: Arc>>, + node: Node, subscriber: ZmqSubscriber, + metrics: Arc, xray_stats_client: Option>>, xray_handler_client: Option>>, wg_client: Option, ) -> Self { + let memory = Arc::new(RwLock::new(Connections::default())); Self { memory, + node, + metrics, subscriber, xray_stats_client, xray_handler_client, diff --git a/src/bin/agent/core/service.rs b/src/bin/agent/core/service.rs index 04b349fc..670dbbac 100644 --- a/src/bin/agent/core/service.rs +++ b/src/bin/agent/core/service.rs @@ -1,12 +1,8 @@ -use qrcode::render::unicode; -use qrcode::QrCode; -use std::net::Ipv4Addr; use std::path::Path; use std::sync::Arc; use tokio::signal; use tokio::sync::broadcast; use tokio::sync::Mutex; -use tokio::sync::RwLock; use tokio::task::JoinHandle; use tokio::time::sleep; use tokio::time::Duration; @@ -17,34 +13,26 @@ use pony::config::settings::AgentSettings; use pony::config::settings::NodeConfig; use pony::config::wireguard::WireguardSettings; use pony::config::xray::Config as XrayConfig; -use pony::http::debug; -use pony::memory::connection::wireguard::IpAddrMaskSerializable; -use pony::memory::connection::wireguard::Param as WgParam; +use pony::metrics::storage::MetricBuffer; +use pony::Publisher; + use pony::memory::node::Node; -use pony::metrics::Metrics; use pony::utils::*; -use pony::wireguard_op::wireguard_conn; use pony::wireguard_op::WgApi; -use pony::xray_op::client::HandlerActions; use pony::xray_op::client::HandlerClient; use pony::xray_op::client::StatsClient; use pony::xray_op::client::XrayClient; use pony::BaseConnection as Connection; -use pony::ConnectionBaseOp; -use pony::MemoryCache; -use pony::NodeStorageOp; -use pony::Proto; + use pony::Result; use pony::SnapshotManager; use pony::Subscriber as ZmqSubscriber; -use pony::Subscription; use pony::Tag; use super::snapshot::SnapshotRestore; use super::tasks::Tasks; use super::Agent; use crate::core::http::ApiRequests; -use crate::core::AgentState; pub async fn run(settings: AgentSettings) -> Result<()> { let mut tasks: Vec> = vec![]; @@ -128,12 +116,6 @@ pub async fn run(settings: AgentSettings) -> Result<()> { None }; - let subscriber = ZmqSubscriber::new( - &settings.zmq.endpoint, - &settings.node.uuid, - &settings.node.env, - ); - let node_config = NodeConfig::from_raw(settings.node.clone()); let node = Node::new( node_config?, @@ -143,31 +125,38 @@ pub async fn run(settings: AgentSettings) -> Result<()> { mtproto_config, ); - let memory: Arc> = - Arc::new(RwLock::new(MemoryCache::with_node(node.clone()))); + let zmq_endpoint = settings.zmq.endpoint.clone(); + let subscriber = ZmqSubscriber::new(&zmq_endpoint, &node.uuid, &node.env); + + let metric_publisher = Publisher::connect(&settings.metrics.publisher).await; - let agent = Arc::new(Agent::new( - memory.clone(), + let metrics = MetricBuffer { + batch: parking_lot::Mutex::new(Vec::new()), + publisher: metric_publisher, + }; + + let agent = Arc::new(Agent::::new( + node.clone(), subscriber, + Arc::new(metrics), xray_stats_client.clone(), xray_handler_client.clone(), wg_client.clone(), )); - let snapshot_manager = - SnapshotManager::new(settings.clone().agent.snapshot_path, memory.clone()); + let snapshot_path = settings.agent.snapshot_path.clone(); + let snapshot_manager = SnapshotManager::new(snapshot_path, agent.memory.clone()); let snapshot_timestamp = if Path::new(&snapshot_manager.snapshot_path).exists() { match snapshot_manager.load_snapshot().await { Ok(Some(timestamp)) => { - let count = snapshot_manager.len().await; if let Err(e) = snapshot_manager - .restore_connections(xray_handler_client.clone(), wg_client) + .restore_connections(agent.xray_handler_client.clone(), wg_client) .await { log::error!("Couldn't restore connections from memory, {}", e); - panic!("Couldn't restore connections from memory, {}", e); } + let count = snapshot_manager.len().await; log::info!( "Loaded {} connections from snapshot with ts {}", count, @@ -186,34 +175,36 @@ pub async fn run(settings: AgentSettings) -> Result<()> { } } } else { - log::info!("No snapshot found, starting fresh"); + log::warn!("No snapshot found, starting fresh"); None }; - let snapshot_manager = snapshot_manager.clone(); tokio::spawn(async move { + log::info!( + "Running snapshot task, interval {}", + settings.agent.snapshot_interval + ); + let mut interval = tokio::time::interval(std::time::Duration::from_secs( settings.agent.snapshot_interval, )); + loop { interval.tick().await; - if let Err(e) = measure_time( - snapshot_manager.create_snapshot(), - "Snapshot took".to_string(), - ) - .await - { + if let Err(e) = measure_time(snapshot_manager.create_snapshot(), "Snapshot").await { log::error!("Failed to create snapshot: {}", e); } else { let count = snapshot_manager.len().await; - log::info!("Connections snapshot saved successfully {}", count); + log::debug!( + "Connections snapshot saved successfully; {} Connections", + count + ); } } }); { log::info!("ZMQ listener starting..."); - let zmq_task = tokio::spawn({ let agent = agent.clone(); let mut shutdown = shutdown_tx.subscribe(); @@ -231,6 +222,7 @@ pub async fn run(settings: AgentSettings) -> Result<()> { if !settings.agent.local { { let settings = settings.clone(); + let node = node.clone(); loop { match agent @@ -267,278 +259,48 @@ pub async fn run(settings: AgentSettings) -> Result<()> { } }; - if settings.agent.metrics_enabled { - log::info!("Running metrics send task"); - let carbon_addr = settings.carbon.address.clone(); - let metrics_handle: JoinHandle<()> = tokio::spawn({ - let agent = agent.clone(); - let mut shutdown = shutdown_tx.subscribe(); - async move { - loop { - tokio::select! { - _ = sleep(Duration::from_secs(settings.agent.metrics_interval)) => { - match agent.send_metrics(&carbon_addr).await { - Ok(_) => {}, - Err(e) => log::error!("❌ Failed to send metrics: {}", e), - } - }, - _ = shutdown.recv() => { - log::info!("πŸ›‘ Metrics task received shutdown"); - break; - }, - } - } - } - }); - tasks.push(metrics_handle); - - log::info!("Running HB metrics send task"); - let carbon_addr = settings.carbon.address.clone(); - let metrics_handle = tokio::spawn({ - let agent = agent.clone(); - let mut shutdown = shutdown_tx.subscribe(); - async move { - loop { - tokio::select! { - _ = sleep(Duration::from_secs(settings.agent.metrics_hb_interval)) => { - match agent.send_hb_metric(&carbon_addr).await { - Ok(_) => {}, - Err(e) => log::error!("❌ Failed to send heartbeat: {}", e), - } - }, - _ = shutdown.recv() => { - log::info!("πŸ›‘ Heartbeat task received shutdown"); - break; - }, - } - } - } - }); - tasks.push(metrics_handle); - } + log::info!("Running metrics task"); - if settings.agent.stat_enabled { - log::info!("Running Stat Task"); - let xray_stats_task = tokio::spawn({ - let agent = Arc::new(agent.clone()); - let mut shutdown = shutdown_tx.subscribe(); - async move { - loop { - tokio::select! { - _ = sleep(Duration::from_secs(settings.agent.stat_job_interval)) => { - let _ = > as Clone>::clone(&agent) - .collect_stats() - .await; - }, - _ = shutdown.recv() => break, - } - } - } - }); - tasks.push(xray_stats_task); - } - - if settings.agent.stat_enabled && settings.wg.enabled { - log::info!("Running WG Stat Task"); - let wg_stats_task = tokio::spawn({ - let agent = Arc::new(agent.clone()); - let mut shutdown = shutdown_tx.subscribe(); - async move { - loop { - tokio::select! { - _ = sleep(Duration::from_secs(settings.agent.stat_job_interval)) => { - let _ = > as Clone>::clone(&agent) - .collect_wireguard_stats() - .await; - }, - _ = shutdown.recv() => break, - } + let metrics_handle: JoinHandle<()> = tokio::spawn({ + let agent = agent.clone(); + let mut shutdown = shutdown_tx.subscribe(); + async move { + loop { + tokio::select! { + _ = sleep(Duration::from_secs(settings.metrics.interval)) => { + agent.collect_metrics().await; + + }, + _ = shutdown.recv() => { + log::info!("πŸ›‘ Metrics task received shutdown"); + break; + }, } } - }); - tasks.push(wg_stats_task); - } - - if settings.xray.enabled && settings.agent.local { - if let Some(xray_handler_client) = xray_handler_client { - let conn_task = tokio::spawn({ - let conn_id = uuid::Uuid::new_v4(); - let agent = agent.clone(); - let node = node.clone(); - let mut shutdown = shutdown_tx.subscribe(); - async move { - tokio::select! { - _ = async { - { - let mut mem = agent.memory.write().await; - for (tag, _) in node.inbounds { - let proto = Proto::new_xray(&tag); - let conn = Connection::new(proto, None, None); - let _ = mem.connections.insert(conn_id, conn); - let _ = xray_handler_client.create(&conn_id, tag, None).await; - } - }; - - let mem = agent.memory.read().await; - if let Some(node) = mem.nodes.get_self() { - println!( - r#" -╔════════════════════════════════════╗ -β•‘ Xray Connection Details β•‘ -β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β• -"# - ); - for (tag, inbound) in node.inbounds { - if let Ok(conn) = create_conn_link( - tag, - &node.uuid, - &inbound, - &node.label, - node.address, - &None - ) { - println!("->> {tag} ➜ {:?}\n", conn); - let qrcode = QrCode::new(conn).unwrap(); - - let image = qrcode - .render::() - .quiet_zone(true) - .min_dimensions(1, 1) - .build(); - - - println!("{}\n", image); - } - } - } else { - panic!("Node information is missing"); - } - } => {}, - _ = shutdown.recv() => {}, - } - } - }); - tasks.push(conn_task); } - } + }); - if settings.wg.enabled && settings.agent.local { - let conn_task = tokio::spawn({ - let agent = agent.clone(); - let settings = settings.clone(); - let node = node.clone(); - let mut shutdown = shutdown_tx.subscribe(); - async move { + log::info!("Running flush metrics task"); + let metrics_flush_handle: JoinHandle<()> = tokio::spawn({ + let agent = agent.clone(); + let mut shutdown = shutdown_tx.subscribe(); + async move { + loop { tokio::select! { - _ = async { - let mut mem = agent.memory.write().await; - for tag in node.inbounds.keys() { - if tag.is_wireguard() { - let conn_id = uuid::Uuid::new_v4(); - - if let Some(wg_api) = &agent.wg_client { - if let Some(ref wg_config) = wg_config { - let next = match wg_api.next_available_ip_mask(&wg_config.network, &wg_config.address) { - Ok(ip) => ip, - Err(e) => { - log::error!("Failed to get next available WireGuard IP: {}; Check your Wireguard interface {}", - e, settings.wg.interface); - return Err(format!( "Failed to get next available WireGuard IP: {}; Check your Wireguard interface {}", - e, settings.wg.interface)); - } - }; - let wg_params = WgParam::new(next.clone()); - let proto = Proto::new_wg(&wg_params, &node.uuid); - let conn = Connection::new(proto, None, None); - let _ = mem.connections.insert(conn_id, conn.clone()); - - if let Err(e) = wg_api.create(&wg_params.keys.pubkey, next) { - log::error!( - "Failed to register WireGuard peer (tag: {} {}): {}", - tag, - wg_params.keys.pubkey, - e - ); - return Err(format!( - "Failed to register WireGuard peer (tag: {} {}): {}", - tag, - wg_params.keys.pubkey, - e,)); - } - if let Some(inbound) = node.inbounds.get(&Tag::Wireguard) { - log::debug!("Inbound {:?}", inbound); - - - - if let Some(wg) = conn.get_wireguard() { - let wg_address = ::clone(&wg.address).into(); - - if let Ok(conn) = wireguard_conn( - &conn_id, - &node.address, - inbound.clone(), - &node.label, - &wg.keys.privkey, - &wg_address, - ) { - - println!( - r#" -╔════════════════════════════════════╗ -β•‘ Wireguard Connection Details β•‘ -β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β• -"#); - println!("{}", conn); - let qrcode = QrCode::new(conn).unwrap(); - - let image = qrcode - .render::() - .quiet_zone(true) - // .module_dimensions(100, 100) - .build(); - println!("{}\n", image); - } - } - } - } - } else { - log::warn!("WG API client is not available"); - return Err("WG API client is not available".into()); - } - } - } - Ok(()) - } => {}, - _ = shutdown.recv() => {}, + _ = sleep(Duration::from_secs(settings.metrics.interval+3)) => { + agent.metrics.flush_to_zmq().await; + + }, + _ = shutdown.recv() => { + log::info!("πŸ›‘ Metrics flush task received shutdown"); + break; + }, } } - }); - tasks.push(conn_task); - } - - let token = Arc::new(settings.api.token.clone()); - if settings.debug.enabled { - log::debug!( - "Running debug server: localhost:{}", - settings.debug.web_port - ); - let mut shutdown = shutdown_tx.subscribe(); - let memory = memory.clone(); - let addr = settings - .debug - .web_server - .unwrap_or(Ipv4Addr::new(127, 0, 0, 1)); - let port = settings.debug.web_port; - let token = token.clone(); - - let debug_handle = tokio::spawn(async move { - tokio::select! { - _ = debug::start_ws_server(memory, addr, port, token) => {}, - _ = shutdown.recv() => {}, - } - }); - tasks.push(debug_handle); - } + } + }); + tasks.push(metrics_handle); + tasks.push(metrics_flush_handle); wait_all_tasks_or_ctrlc(tasks, shutdown_tx).await; Ok(()) diff --git a/src/bin/agent/core/snapshot.rs b/src/bin/agent/core/snapshot.rs index 8ad28f8e..78bb9f57 100644 --- a/src/bin/agent/core/snapshot.rs +++ b/src/bin/agent/core/snapshot.rs @@ -3,13 +3,12 @@ use std::sync::Arc; use tokio::sync::Mutex; use pony::memory::connection::wireguard::IpAddrMaskSerializable; +use pony::memory::connection::Connections; use pony::memory::snapshot::SnapshotManager; use pony::wireguard_op::WgApi; use pony::xray_op::client::HandlerActions; use pony::xray_op::client::HandlerClient; -use pony::Connection; use pony::ConnectionBaseOp as ConnOp; -use pony::NodeStorageOp as NodeOp; use pony::Result; use pony::Tag; @@ -23,11 +22,9 @@ pub trait SnapshotRestore { } #[async_trait::async_trait] -impl SnapshotRestore for SnapshotManager +impl SnapshotRestore for SnapshotManager> where - C: Archive + Send + Sync + Clone + 'static + ConnOp + From, - N: Send + Sync + Clone + 'static + NodeOp, - S: Send + Sync + Clone + 'static, + C: Archive + Send + Sync + Clone + 'static + ConnOp, { async fn restore_connections( &self, @@ -36,10 +33,16 @@ where ) -> Result<()> { let mem = self.memory.read().await; - for (conn_id, conn) in mem.connections.iter() { - let conn_id = *conn_id; - let conn = conn.clone(); - let memory = self.memory.clone(); + if mem.is_empty() { + return Err(pony::PonyError::Custom("Empty snapshot".into())); + } + + let conns: Vec<(uuid::Uuid, C)> = + mem.iter().map(|(id, conn)| (*id, conn.clone())).collect(); + + drop(mem); + + for (conn_id, conn) in conns { let wg_client = wg_client.clone(); let xray_client = xray_client.clone(); @@ -47,24 +50,16 @@ where match conn.get_proto().proto() { Tag::Wireguard => { if let Some(wg) = conn.get_wireguard() { - let node_id = { - let mem = memory.read().await; - mem.nodes.get_self().map(|n| n.uuid) - }; - - if let Some(_node_id) = node_id { - if let Some(api) = wg_client.as_ref() { - if let Err(e) = api.create( - &wg.keys.pubkey, - ::clone(&wg.address) - .into(), - ) { - log::error!( - "Failed to restore WireGuard connection {}: {}", - conn_id, - e - ); - } + if let Some(api) = wg_client.as_ref() { + if let Err(e) = api.create( + &wg.keys.pubkey, + ::clone(&wg.address).into(), + ) { + log::error!( + "Failed to restore WireGuard connection {}: {}", + conn_id, + e + ); } } } diff --git a/src/bin/agent/core/stats.rs b/src/bin/agent/core/stats.rs index d21d108b..bc70edbf 100644 --- a/src/bin/agent/core/stats.rs +++ b/src/bin/agent/core/stats.rs @@ -1,10 +1,5 @@ use async_trait::async_trait; -use futures::future::join_all; - -use std::sync::Arc; -use tokio::task::JoinHandle; -use tonic::Code; -use tonic::{Request, Status}; +use tonic::{Code, Request, Status}; use pony::memory::connection::stat::Stat as ConnectionStat; use pony::memory::node::Stat as InboundStat; @@ -15,20 +10,14 @@ use pony::xray_op::connections::ConnOp; use pony::xray_op::stats::Prefix; use pony::xray_op::stats::StatOp; use pony::ConnectionBaseOp; -use pony::ConnectionStorageBaseOp; -use pony::NodeStorageOp; -use pony::Result as PonyResult; -use pony::SubscriptionOp; use pony::Tag; use super::Agent; #[async_trait] -impl StatOp for Agent +impl StatOp for Agent where - T: NodeStorageOp + Send + Sync + Clone, C: ConnectionBaseOp + Send + Sync + Clone + 'static, - S: SubscriptionOp + Send + Sync + Clone + 'static + std::cmp::PartialEq, { async fn stat( &self, @@ -54,15 +43,7 @@ where Stat::Conn(StatKind::Online) => { format!("{}>>>{}", base_name, stat_type) } - Stat::Inbound(StatKind::Online) => { - return Err(Status::internal("Online is not supported for inbound")); - } - Stat::Outbound(_) => { - return Err(Status::internal("Outbound stat type is not implemented")); - } - Stat::Conn(StatKind::Unknown) | Stat::Inbound(StatKind::Unknown) => { - return Err(Status::internal("Unknown stat type is not implemented")); - } + _ => return Err(Status::internal("Unsupported stat type")), }; let request = Request::new(GetStatsRequest { @@ -77,10 +58,7 @@ where match response { Ok(stat) => Ok(stat.into_inner()), - Err(e) => Err(Status::internal(format!( - "Stat request failed {:?} {}", - stat_type, e - ))), + Err(e) => Err(Status::internal(format!("Stat request failed: {}", e))), } } else { Err(Status::new( @@ -92,309 +70,80 @@ where async fn reset_stat(&self, conn_id: &uuid::Uuid) -> Result<(), Status> { let id = Prefix::ConnPrefix(*conn_id); - let (downlink_result, uplink_result, online_result) = tokio::join!( + let _ = tokio::join!( self.stat(id, Stat::Conn(StatKind::Downlink), true), self.stat(id, Stat::Conn(StatKind::Uplink), true), self.stat(id, Stat::Conn(StatKind::Online), true) ); - - match (downlink_result, uplink_result, online_result) { - (Ok(downlink), Ok(uplink), Ok(online)) => { - if let (Some(downlink), Some(uplink), Some(online)) = ( - downlink.stat.clone(), - uplink.stat.clone(), - online.stat.clone(), - ) { - log::debug!( - "Connection Stats reset successfully: downlink={:?}, uplink={:?}, online={:?}", - downlink.clone(), - uplink.clone(), - online.clone() - ); - - Ok(()) - } else { - let error_msg = format!( - "Incomplete connection stats for {:?}: downlink={:?}, uplink={:?}, online={:?}", - conn_id, - downlink.stat.clone(), - uplink.stat.clone(), - online.stat.clone() - ); - Err(Status::internal(error_msg)) - } - } - (Err(e), _, _) | (_, Err(e), _) | (_, _, Err(e)) => { - let error_msg = format!( - "Failed to fetch connection stats for {:?}: error={:?}", - conn_id, e - ); - Err(Status::internal(error_msg)) - } - } + Ok(()) } async fn conn_stats(&self, conn_id: Prefix) -> Result { - log::debug!("conn_stats {:?}", conn_id); - - let (downlink_result, uplink_result, online_result) = tokio::join!( + let (down_res, up_res, online_res) = tokio::join!( self.stat(conn_id, Stat::Conn(StatKind::Downlink), false), self.stat(conn_id, Stat::Conn(StatKind::Uplink), false), self.stat(conn_id, Stat::Conn(StatKind::Online), false) ); - match (downlink_result, uplink_result, online_result) { - (Ok(downlink), Ok(uplink), Ok(online)) => { - if let (Some(downlink), Some(uplink), Some(online)) = ( - downlink.stat.clone(), - uplink.stat.clone(), - online.stat.clone(), - ) { - log::debug!( - "Connection Stats fetched successfully: downlink={:?}, uplink={:?}, online={:?}", - downlink.clone(), - uplink.clone(), - online.clone() - ); - Ok(ConnectionStat { - downlink: downlink.value, - uplink: uplink.value, - online: online.value, - }) - } else { - let error_msg = format!( - "Incomplete connection stats for {:?}: downlink={:?}, uplink={:?}, online={:?}", - conn_id, - downlink.stat.clone(), - uplink.stat.clone(), - online.stat.clone() - ); - Err(Status::internal(error_msg)) - } - } - (Err(e), _, _) | (_, Err(e), _) | (_, _, Err(e)) => { - let error_msg = format!( - "Failed to fetch connection stats for {:?}: error={:?}", - conn_id, e - ); - Err(Status::internal(error_msg)) - } + if let Err(e) = &down_res { + log::warn!("Downlink stat missing for {:?}: {}", conn_id, e); + } + if let Err(e) = &up_res { + log::warn!("Uplink stat missing for {:?}: {}", conn_id, e); } + if let Err(e) = &online_res { + log::warn!("Online stat missing for {:?}: {}", conn_id, e); + } + + Ok(ConnectionStat { + downlink: down_res + .ok() + .and_then(|s| s.stat) + .map(|v| v.value) + .unwrap_or(0), + uplink: up_res + .ok() + .and_then(|s| s.stat) + .map(|v| v.value) + .unwrap_or(0), + online: online_res + .ok() + .and_then(|s| s.stat) + .map(|v| v.value) + .unwrap_or(0), + }) } async fn inbound_stats(&self, inbound: Prefix) -> Result { - let downlink_result = self - .stat(inbound, Stat::Inbound(StatKind::Downlink), false) - .await; - - let uplink_result = self - .stat(inbound, Stat::Inbound(StatKind::Uplink), false) - .await; - - let conn_count_result = self.conn_count(*inbound.as_tag().unwrap()).await; - - match (downlink_result, uplink_result, conn_count_result) { - (Ok(downlink), Ok(uplink), Ok(conn_count)) => { - if let (Some(downlink), Some(uplink), Some(conn_count)) = - (downlink.stat.clone(), uplink.stat.clone(), conn_count) - { - log::debug!( - "Node Stats successfully fetched: inbound={:?}, downlink={:?}, uplink={:?}, conn_count={:?} ", - inbound, downlink, uplink, conn_count - ); - Ok(InboundStat { - downlink: downlink.value, - uplink: uplink.value, - conn_count, - }) - } else { - let error_msg = format!( - "Incomplete stats for inbound {:?}: downlink={:?}, uplink={:?}, conn_count={:?}", - inbound, - downlink.stat.clone(), - uplink.stat.clone(), - conn_count.clone(), - ); - Err(Status::internal(error_msg)) - } - } + let (down, up, count) = tokio::join!( + self.stat(inbound, Stat::Inbound(StatKind::Downlink), false), + self.stat(inbound, Stat::Inbound(StatKind::Uplink), false), + self.conn_count(*inbound.as_tag().unwrap()) + ); - (Err(e1), Err(e2), Err(e3)) => { - let error_msg = format!( - "ALl requests failed for inbound {:?}: downlink error: {:?}, uplink error: {:?}, conn_count: {:?}", - inbound, e1, e2, e3 - ); - Err(Status::internal(error_msg)) - } - (Err(e), _, _) | (_, Err(e), _) | (_, _, Err(e)) => { - let error_msg = format!( - "One of the requests failed for inbound {:?}: {:?}", - inbound, e - ); - Err(Status::internal(error_msg)) - } + match (down, up, count) { + (Ok(d), Ok(u), Ok(c)) => Ok(InboundStat { + downlink: d.stat.map(|s| s.value).unwrap_or(0), + uplink: u.stat.map(|s| s.value).unwrap_or(0), + conn_count: c.unwrap_or(0), + }), + _ => Err(Status::internal("Failed to fetch inbound stats")), } } async fn conn_count(&self, inbound: Tag) -> Result, Status> { if let Some(client) = &self.xray_handler_client { let mut handler_client = client.lock().await; - match handler_client.conn_count_op(inbound).await { - Ok(count) => Ok(Some(count)), - Err(e) => Err(Status::internal(format!( - "Failed to fetch conn count for inbound {}: {}", - inbound, e - ))), - } + handler_client + .conn_count_op(inbound) + .await + .map(Some) + .map_err(|e| Status::internal(e.to_string())) } else { Err(Status::new( Code::Unavailable, - "Xray Hanler client doesn't exist", + "Xray Handler client missing", )) } } } - -impl Agent -where - T: NodeStorageOp + Send + Sync + Clone, - C: ConnectionBaseOp + Send + Sync + Clone + 'static, - S: Send + Sync + Clone + 'static + std::cmp::PartialEq + SubscriptionOp, -{ - async fn collect_conn_stats(self: Arc, conn_id: uuid::Uuid) -> PonyResult<()> { - let conn_stat = self.conn_stats(Prefix::ConnPrefix(conn_id)).await?; - let mut mem = self.memory.write().await; - let _ = mem - .connections - .update_downlink(&conn_id, conn_stat.downlink); - let _ = mem.connections.update_uplink(&conn_id, conn_stat.uplink); - let _ = mem.connections.update_online(&conn_id, conn_stat.online); - Ok(()) - } - - async fn collect_inbound_stats( - self: Arc, - tag: Tag, - env: String, - node_uuid: uuid::Uuid, - ) -> PonyResult<()> { - let inbound_stat = self.inbound_stats(Prefix::InboundPrefix(tag)).await?; - let mut mem = self.memory.write().await; - let _ = mem - .nodes - .update_node_downlink(&tag, inbound_stat.downlink, &env, &node_uuid); - let _ = mem - .nodes - .update_node_uplink(&tag, inbound_stat.uplink, &env, &node_uuid); - let _ = mem - .nodes - .update_node_conn_count(&tag, inbound_stat.conn_count, &env, &node_uuid); - Ok(()) - } - - pub async fn collect_stats(self: Arc) -> PonyResult<()> { - log::debug!("Running xray stat job"); - let mut tasks: Vec> = Vec::new(); - - let conn_ids = { - let mem = self.memory.write().await; - mem.connections.keys().cloned().collect::>() - }; - - for conn_id in conn_ids { - let agent = self.clone(); - tasks.push(tokio::spawn(async move { - if let Err(e) = agent.collect_conn_stats(conn_id).await { - log::error!("Failed to collect stats for connection {}: {}", conn_id, e); - } - })); - } - - if let Some(node) = { - let mem = self.memory.read().await; - mem.nodes.get_self() - } { - let node_tags = node.inbounds.keys().cloned().collect::>(); - for tag in node_tags { - let agent = self.clone(); - let env = node.env.clone(); - let node_uuid = node.uuid; - tasks.push(tokio::spawn(async move { - if let Err(e) = agent.collect_inbound_stats(tag, env, node_uuid).await { - log::error!("Failed to collect stats for inbound {}: {}", tag.clone(), e); - } - })); - } - } - - let results = join_all(tasks).await; - for result in results { - if let Err(e) = result { - log::error!("Task panicked: {:?}", e); - } - } - - Ok(()) - } - - pub async fn collect_wireguard_stats(self: Arc) -> PonyResult<()> { - let mut tasks: Vec> = Vec::new(); - - let wg_client = match &self.wg_client { - Some(client) => client.clone(), - None => { - log::warn!("WG API client is not available, skipping stats collection"); - return Ok(()); - } - }; - - let conns = { - let mem = self.memory.read().await; - mem.connections - .iter() - .filter_map(|(id, conn)| { - if let Tag::Wireguard = conn.get_proto().proto() { - Some((*id, conn.clone())) - } else { - None - } - }) - .collect::>() - }; - - for (conn_id, conn) in conns { - let wg_client = wg_client.clone(); - let agent = self.clone(); - tasks.push(tokio::spawn(async move { - if let Some(wg) = conn.get_wireguard() { - match wg_client.peer_stats(&wg.keys.pubkey.clone()) { - Ok((uplink, downlink)) => { - let mut mem = agent.memory.write().await; - if let Some(existing) = mem.connections.get_mut(&conn_id) { - existing.set_uplink(uplink); - existing.set_downlink(downlink); - } - } - Err(e) => { - log::error!( - "Failed to collect WG stats for {} {}: {}", - conn_id, - wg.keys.pubkey, - e - ); - } - } - } - })); - } - - let results = join_all(tasks).await; - for result in results { - if let Err(e) = result { - log::error!("WireGuard stats task panicked: {:?}", e); - } - } - - Ok(()) - } -} diff --git a/src/bin/agent/core/tasks.rs b/src/bin/agent/core/tasks.rs index 4b51ed52..d9d69661 100644 --- a/src/bin/agent/core/tasks.rs +++ b/src/bin/agent/core/tasks.rs @@ -1,12 +1,14 @@ use async_trait::async_trait; -use defguard_wireguard_rs::net::IpAddrMask; use futures::future::try_join_all; -use pony::Proto; use rkyv::AlignedVec; +use rkyv::Deserialize; use rkyv::Infallible; use tokio::time::Duration; + +use defguard_wireguard_rs::net::IpAddrMask; use tonic::Status; +use pony::metrics::Metrics; use pony::xray_op::client::HandlerActions; use pony::xray_op::stats::StatOp; use pony::Action; @@ -14,29 +16,26 @@ use pony::BaseConnection as Connection; use pony::ConnectionBaseOp; use pony::ConnectionStorageBaseOp; use pony::Message; -use pony::NodeStorageOp; -use pony::SubscriptionOp; +use pony::Proto; use pony::Tag; use pony::Topic; use pony::{PonyError, Result}; -use rkyv::Deserialize; - +use super::metrics::BusinessMetrics; use super::Agent; #[async_trait] pub trait Tasks { async fn run_subscriber(&self) -> Result<()>; - async fn handle_message(&self, msg: Message) -> Result<()>; async fn handle_messages_batch(&self, msg: Vec) -> Result<()>; + async fn handle_message(&self, msg: Message) -> Result<()>; + async fn collect_metrics(&self); } #[async_trait] -impl Tasks for Agent +impl Tasks for Agent where - T: NodeStorageOp + Send + Sync + Clone, C: ConnectionBaseOp + Send + Sync + Clone + 'static + From, - S: SubscriptionOp + Send + Sync + Clone + 'static + std::cmp::PartialEq, { async fn run_subscriber(&self) -> Result<()> { let sub = self.subscriber.clone(); @@ -106,6 +105,19 @@ where } } + async fn handle_messages_batch(&self, messages: Vec) -> Result<()> { + log::debug!("Got {} messages", messages.len()); + + let handles: Vec<_> = messages + .into_iter() + .map(|msg| self.handle_message(msg)) + .collect(); + + let _results = try_join_all(handles).await?; + + Ok(()) + } + async fn handle_message(&self, msg: Message) -> Result<()> { match msg.action { Action::Create | Action::Update => { @@ -118,13 +130,7 @@ where .clone() .ok_or_else(|| PonyError::Custom("Missing WireGuard keys".into()))?; - let node_id = { - let mem = self.memory.read().await; - let node = mem.nodes.get_self(); - node.map(|n| n.uuid).ok_or_else(|| { - PonyError::Custom("Current node UUID not found".to_string()) - })? - }; + let node_id = self.node.uuid; let proto = Proto::new_wg(&wg, &node_id); let conn = Connection::new( @@ -135,8 +141,7 @@ where { let mut mem = self.memory.write().await; - mem.connections - .add(&conn_id.clone(), conn.clone().into()) + mem.add(&conn_id.clone(), conn.clone().into()) .map_err(|err| { PonyError::Custom(format!( "Failed to add WireGuard conn {}: {}", @@ -195,14 +200,9 @@ where })?; let mut mem = self.memory.write().await; - mem.connections - .add(&conn_id.clone(), conn.into()) - .map_err(|err| { - PonyError::Custom(format!( - "Failed to add conn {}: {}", - conn_id, err - )) - })?; + mem.add(&conn_id.clone(), conn.into()).map_err(|err| { + PonyError::Custom(format!("Failed to add conn {}: {}", conn_id, err)) + })?; return Ok(()); } @@ -232,14 +232,12 @@ where })?; let mut mem = self.memory.write().await; - mem.connections - .add(&conn_id.clone(), conn.into()) - .map_err(|err| { - PonyError::Custom(format!( - "Failed to add conn {}: {}", - conn_id, err - )) - })?; + mem.add(&conn_id.clone(), conn.into()).map_err(|err| { + PonyError::Custom(format!( + "Failed to add conn {}: {}", + conn_id, err + )) + })?; return Ok(()); } else { @@ -275,7 +273,7 @@ where })?; let mut mem = self.memory.write().await; - let _ = mem.connections.remove(&conn_id); + let _ = mem.remove(&conn_id); return Ok(()); } @@ -301,7 +299,7 @@ where })?; let mut mem = self.memory.write().await; - let _ = mem.connections.remove(&conn_id); + let _ = mem.remove(&conn_id); return Ok(()); } @@ -324,16 +322,21 @@ where } } - async fn handle_messages_batch(&self, messages: Vec) -> Result<()> { - log::debug!("Got {} messages", messages.len()); - - let handles: Vec<_> = messages - .into_iter() - .map(|msg| self.handle_message(msg)) - .collect(); - - let _results = try_join_all(handles).await?; + async fn collect_metrics(&self) { + self.heartbeat().await; + self.bandwidth().await; + self.cpu_usage().await; + self.loadavg().await; + self.memory().await; + self.disk_usage().await; + + if self.xray_stats_client.is_some() { + self.collect_inbound_metrics().await; + self.collect_user_metrics().await; + } - Ok(()) + if self.wg_client.is_some() { + self.collect_wg_metrics().await; + } } } diff --git a/src/bin/api/core/clickhouse/mod.rs b/src/bin/api/core/clickhouse/mod.rs deleted file mode 100644 index 32fc0c10..00000000 --- a/src/bin/api/core/clickhouse/mod.rs +++ /dev/null @@ -1,43 +0,0 @@ -use clickhouse::Client as ChClient; -use pony::metrics::Metric; -use serde::de::DeserializeOwned; -use std::fmt::Debug; - -pub mod query; -pub mod score; - -#[derive(Clone)] -pub struct ChContext { - url: String, -} - -impl ChContext { - pub fn new(url: &str) -> Self { - Self { - url: url.to_string(), - } - } - - pub fn client(&self) -> ChClient { - ChClient::default().with_url(&self.url) - } -} - -impl ChContext { - async fn execute_metric_query(&self, query: &str) -> Option> - where - T: DeserializeOwned + Debug + Send + Clone + Sync + 'static, - { - log::debug!("Executing query:\n{}", query); - let client = self.client(); - - let result = client.query(query).fetch_all::>().await; - - match &result { - Ok(rows) => log::debug!("Fetched {} rows", rows.len()), - Err(e) => log::error!("Query failed: {:?}", e), - } - - result.ok().and_then(|mut rows| rows.pop()) - } -} diff --git a/src/bin/api/core/clickhouse/query.rs b/src/bin/api/core/clickhouse/query.rs deleted file mode 100644 index 40b8a23c..00000000 --- a/src/bin/api/core/clickhouse/query.rs +++ /dev/null @@ -1,193 +0,0 @@ -use async_trait::async_trait; -use chrono::DateTime; -use chrono::Utc; -use uuid::Uuid; - -use super::ChContext; -use pony::metrics::Metric; - -#[async_trait] -pub trait Queries { - async fn fetch_node_heartbeat( - &self, - env: &str, - uuid: &Uuid, - hostname: &str, - ) -> Option>; - - async fn fetch_conn_stats( - &self, - conn_id: uuid::Uuid, - start: DateTime, - ) -> Option>>; - - async fn fetch_node_cpu_load_avg5(&self, env: &str, hostname: &str) -> Option>; - async fn fetch_node_bandwidth( - &self, - env: &str, - hostname: &str, - interface: &str, - ) -> Option>; - async fn fetch_node_cpuusage_max(&self, env: &str, hostname: &str) -> Option>; - async fn fetch_node_memory_usage_ratio(&self, env: &str, hostname: &str) - -> Option>; -} - -#[async_trait] -impl Queries for ChContext { - async fn fetch_node_heartbeat( - &self, - env: &str, - uuid: &Uuid, - hostname: &str, - ) -> Option> { - let query = format!( - r#" -SELECT Path AS metric, - toUInt8(argMax(Value, Timestamp)) AS value, - toInt64(toUnixTimestamp(toDateTime(argMax(Timestamp, Timestamp)))) AS timestamp -FROM default.graphite_data -WHERE Path = '{}.{}.{}.heartbeat' -GROUP BY Path"#, - env, hostname, uuid - ); - self.execute_metric_query(&query).await - } - - async fn fetch_conn_stats( - &self, - conn_id: uuid::Uuid, - start: DateTime, - ) -> Option>> { - let start_str = start.format("%Y-%m-%d %H:%M:%S").to_string(); - - let query = format!( - r#" -SELECT metric, - toInt64(sum(metric_value)) AS value, - anyLast(timestamp) AS timestamp -FROM ( - SELECT - toInt64(toUnixTimestamp(toDateTime(anyLast(Timestamp)))) AS timestamp, - extract(Path, '[^.]+$') AS metric, - anyLast(Value) AS metric_value - FROM default.graphite_data - WHERE Path LIKE '%.%.{conn_id}.conn_stat.%' - AND Timestamp >= toDateTime('{start_str}') - AND Timestamp < toDateTime('{start_str}') + INTERVAL 1 DAY - GROUP BY Path - ) -GROUP BY metric"#, - ); - let client = self.client(); - - let result = client.query(&query).fetch_all::>().await; - - match &result { - Ok(rows) => { - log::debug!("Fetched {} connection stats rows", rows.len()); - Some(rows.clone()) - } - Err(e) => { - log::error!("Failed to fetch connection stats: {:?}", e); - None - } - } - } - - async fn fetch_node_cpu_load_avg5(&self, env: &str, hostname: &str) -> Option> { - let query = format!( - r#" -SELECT - any(Path) AS metric, - avg(toFloat64(Value)) AS value, - toInt64(toUnixTimestamp(now())) AS timestamp -FROM default.graphite_data -WHERE Path = '{}.{}.loadavg.5m' -AND Timestamp >= now() - INTERVAL 5 MINUTE"#, - env, hostname - ); - self.execute_metric_query(&query).await - } - - async fn fetch_node_bandwidth( - &self, - env: &str, - hostname: &str, - interface: &str, - ) -> Option> { - let query = format!( - r#" -SELECT - any(Path) AS metric, - avg(toFloat64(Value)) AS value, - toInt64(toUnixTimestamp(now())) AS timestamp -FROM default.graphite_data -WHERE Path = '{}.{}.network.{}.tx_bps' -AND Timestamp >= now() - INTERVAL 5 MINUTE"#, - env, hostname, interface - ); - - self.execute_metric_query(&query).await - } - - async fn fetch_node_cpuusage_max(&self, env: &str, hostname: &str) -> Option> { - let query = format!( - r#" -SELECT - Path AS metric, - toFloat64(argMax(Value, Timestamp)) AS value, - toInt64(toUnixTimestamp(toDateTime(argMax(Timestamp, Timestamp)))) AS timestamp -FROM default.graphite_data -WHERE Path LIKE '{}.{}.cpu.%.percentage' -GROUP BY Path -ORDER BY value DESC -LIMIT 1 -"#, - env, hostname, - ); - self.execute_metric_query(&query).await - } - - async fn fetch_node_memory_usage_ratio( - &self, - env: &str, - hostname: &str, - ) -> Option> { - let query = format!( - r#" - SELECT - used.value / total.value AS value, - toInt64(toUnixTimestamp(now())) AS timestamp, - 'mem.used/mem.total' AS metric - FROM - ( - SELECT argMax(Value, Timestamp) AS value - FROM default.graphite_data - WHERE Path = '{0}.{1}.mem.used' - ) AS used, - ( - SELECT argMax(Value, Timestamp) AS value - FROM default.graphite_data - WHERE Path = '{0}.{1}.mem.total' - ) AS total - "#, - env, hostname - ); - - let client = self.client(); - let row_result = client.query(&query).fetch_one::<(f64, i64, String)>().await; - - match row_result { - Ok((value, timestamp, metric)) => Some(Metric { - value, - timestamp, - metric, - }), - Err(err) => { - log::error!("Failed to fetch memory usage ratio: {:?}", err); - None - } - } - } -} diff --git a/src/bin/api/core/clickhouse/score.rs b/src/bin/api/core/clickhouse/score.rs deleted file mode 100644 index 04e9230f..00000000 --- a/src/bin/api/core/clickhouse/score.rs +++ /dev/null @@ -1,86 +0,0 @@ -use serde::{Deserialize, Serialize}; - -use super::query::Queries; -use super::ChContext; - -#[derive(Serialize, Deserialize)] -pub struct NodeScore { - pub score: f64, - pub details: NodeScoreBreakdown, -} - -#[derive(Serialize, Deserialize)] -pub struct NodeScoreBreakdown { - pub cpu_usage: f64, - pub loadavg: f64, - pub memory_ratio: f64, - pub bandwidth: f64, -} - -impl ChContext { - pub async fn fetch_node_score( - &self, - env: &str, - hostname: &str, - interface: &str, - num_cores: usize, - max_bandwidth_bps: i64, - ) -> Option { - let cpu_usage = match self.fetch_node_cpuusage_max(env, hostname).await { - Some(m) => m.value, - None => { - log::error!("Failed to fetch CPU usage for {}.{}", env, hostname); - return None; - } - }; - - let loadavg = match self.fetch_node_cpu_load_avg5(env, hostname).await { - Some(m) => m.value, - None => { - log::error!("Failed to fetch loadavg for {}.{}", env, hostname); - return None; - } - }; - - let memory_ratio = match self.fetch_node_memory_usage_ratio(env, hostname).await { - Some(m) => m.value, - None => { - log::error!("Failed to fetch memory ratio for {}.{}", env, hostname); - return None; - } - }; - - let tx_bps = match self.fetch_node_bandwidth(env, hostname, interface).await { - Some(m) => m.value, - None => { - log::error!("Failed to fetch tx_bps for {}.{}", env, hostname); - return None; - } - }; - - let cpu_score = (cpu_usage / 100.0).clamp(0.0, 1.0); - - let load_score = if num_cores > 0 { - (loadavg / num_cores as f64).clamp(0.0, 1.0) - } else { - 1.0 - }; - - let memory_score = memory_ratio.clamp(0.0, 1.0); - - let upload_score = (tx_bps / max_bandwidth_bps as f64).clamp(0.0, 1.0); - - let score = - 0.35 * cpu_score + 0.25 * load_score + 0.25 * memory_score + 0.15 * upload_score; - - Some(NodeScore { - score, - details: NodeScoreBreakdown { - cpu_usage, - loadavg, - memory_ratio, - bandwidth: tx_bps, - }, - }) - } -} diff --git a/src/bin/api/core/http/filters.rs b/src/bin/api/core/http/filters.rs index 5aecfadc..8b79a33e 100644 --- a/src/bin/api/core/http/filters.rs +++ b/src/bin/api/core/http/filters.rs @@ -1,16 +1,17 @@ +use std::sync::Arc; use warp::Filter; +use pony::metrics::storage::MetricStorage; use pony::Connection; use pony::ConnectionApiOp; use pony::ConnectionBaseOp; use pony::NodeStorageOp; use pony::SubscriptionOp; -use crate::core::clickhouse::ChContext; use crate::MemSync; /// Provides application state filter -pub fn with_state( +pub fn with_sync( mem_sync: MemSync, ) -> impl Filter,), Error = std::convert::Infallible> + Clone where @@ -21,13 +22,6 @@ where warp::any().map(move || mem_sync.clone()) } -/// Provides ckickhouse filter -pub fn with_ch( - ch: ChContext, -) -> impl Filter + Clone { - warp::any().map(move || ch.clone()) -} - pub fn with_param_string( param: String, ) -> impl Filter + Clone { @@ -45,3 +39,9 @@ pub fn with_param_vec_string( ) -> impl Filter,), Error = std::convert::Infallible> + Clone { warp::any().map(move || param.clone()) } + +pub fn with_metrics( + metrics: Arc, +) -> impl Filter,), Error = std::convert::Infallible> + Clone { + warp::any().map(move || metrics.clone()) +} diff --git a/src/bin/api/core/http/handlers/connection.rs b/src/bin/api/core/http/handlers/connection.rs index cd7c91dc..828274e6 100644 --- a/src/bin/api/core/http/handlers/connection.rs +++ b/src/bin/api/core/http/handlers/connection.rs @@ -2,19 +2,16 @@ use chrono::DateTime; use chrono::Utc; use defguard_wireguard_rs::net::IpAddrMask; use rkyv::to_bytes; -use warp::http::StatusCode; use pony::http::helpers as http; -use pony::http::IdResponse; +use pony::http::response::Instance; use pony::http::IpParseError; use pony::http::MyRejection; -use pony::http::ResponseMessage; use pony::memory::tag::ProtoTag; use pony::utils; use pony::Connection; use pony::ConnectionApiOp; use pony::ConnectionBaseOp; -use pony::ConnectionStat; use pony::NodeStorageOp; use pony::OperationStatus as StorageOperationStatus; use pony::Proto; @@ -72,7 +69,7 @@ where }) .collect(); - if !connections_to_send.is_empty() { + if connections_to_send.is_empty() { log::debug!( "Sending {} {:?} connections for env {:?} to topic {} ", connections_to_send.len(), @@ -81,45 +78,36 @@ where topic ); - let messages: Vec<_> = connections_to_send - .iter() - .map(|(conn_id, conn)| conn.as_create_message(conn_id)) - .collect(); + return Ok(http::not_modified("")); + } - if messages.is_empty() { - return Ok(http::not_modified("")); - } + let messages: Vec<_> = connections_to_send + .iter() + .map(|(conn_id, conn)| conn.as_create_message(conn_id)) + .collect(); - let bytes = to_bytes::<_, 1024>(&messages).map_err(|e| { - log::error!("Serialization error: {}", e); + if messages.is_empty() { + return Ok(http::not_modified("")); + } + + let bytes = to_bytes::<_, 1024>(&messages).map_err(|e| { + log::error!("Serialization error: {}", e); + warp::reject::custom(MyRejection(Box::new(e))) + })?; + + memory + .publisher + .send_binary(&topic, bytes.as_ref()) + .await + .map_err(|e| { + log::error!("Publish error: {}", e); warp::reject::custom(MyRejection(Box::new(e))) })?; - memory - .publisher - .send_binary(&topic, bytes.as_ref()) - .await - .map_err(|e| { - log::error!("Publish error: {}", e); - warp::reject::custom(MyRejection(Box::new(e))) - })?; - } else { - log::debug!( - "No message {} to send for env {:?} to topic {}", - proto, - env, - topic - ); - } - - let resp = ResponseMessage::> { - status: StatusCode::OK.as_u16(), - message: "Ok".to_string(), - response: None, - }; - Ok(warp::reply::with_status( - warp::reply::json(&resp), - StatusCode::OK, + Ok(http::success_response( + "Ok".into(), + None, + Instance::Count(messages.len()), )) } @@ -342,13 +330,8 @@ where drop(mem); - let conn: Connection = Connection::new( - &conn_req.env, - conn_req.subscription_id, - ConnectionStat::default(), - proto, - expired_at, - ); + let conn: Connection = + Connection::new(&conn_req.env, conn_req.subscription_id, proto, expired_at); log::debug!("New connection to create {}", conn); let conn_id = uuid::Uuid::new_v4(); @@ -381,7 +364,7 @@ where Ok(http::success_response( format!("Connection {} has been created", id), Some(id), - http::Instance::Connection(conn), + Instance::Connection(conn), )) } @@ -456,7 +439,7 @@ where Ok(StorageOperationStatus::Ok(id)) => Ok(http::success_response( format!("Connection {} has been deleted", id), Some(id), - http::Instance::Connection(conn.clone().into()), + Instance::Connection(conn.clone().into()), )), Ok(StorageOperationStatus::NotFound(id)) => { @@ -504,7 +487,7 @@ where Ok(http::success_response( "Connection is found".to_string(), Some(conn_id), - http::Instance::Connection(conn.clone().into()), + Instance::Connection(conn.clone().into()), )) } else { Ok(http::not_found("Connection is not found")) diff --git a/src/bin/api/core/http/handlers/key.rs b/src/bin/api/core/http/handlers/key.rs index 228be382..af6d6dd2 100644 --- a/src/bin/api/core/http/handlers/key.rs +++ b/src/bin/api/core/http/handlers/key.rs @@ -1,4 +1,5 @@ use pony::http::helpers as http; +use pony::http::response::Instance; use pony::memory::key::Distributor; use pony::memory::key::Key; use pony::Connection; @@ -46,11 +47,11 @@ where return Ok(http::success_response( "Key is valid and already activated".to_string(), Some(key.id), - http::Instance::Key(key.clone()), + Instance::Key(key.clone()), )); } - let instance = http::Instance::Key(key.clone()); + let instance = Instance::Key(key.clone()); Ok(http::success_response( "Key is valid".to_string(), Some(key.id), @@ -95,7 +96,7 @@ where Ok(http::success_response( msg, Some(key.id), - http::Instance::Key(key), + Instance::Key(key), )) } Err(e) => { @@ -153,7 +154,7 @@ where Ok(http::success_response( format!("Key {} activated", key.id), Some(key.id), - http::Instance::Key(key), + Instance::Key(key), )) } Ok(StorageOperationStatus::NotFound(_)) => Ok(http::not_found("Subscription not found")), diff --git a/src/bin/api/core/http/handlers/metrics.rs b/src/bin/api/core/http/handlers/metrics.rs new file mode 100644 index 00000000..a2be0a9b --- /dev/null +++ b/src/bin/api/core/http/handlers/metrics.rs @@ -0,0 +1,37 @@ +use std::sync::Arc; + +use pony::metrics::storage::MetricStorage; + +pub async fn handle_ws_client( + socket: warp::ws::WebSocket, + node_id: uuid::Uuid, + metric: String, + storage: Arc, +) { + use futures::{SinkExt, StreamExt}; + let (mut ws_tx, _) = socket.split(); + let mut ticker = tokio::time::interval(std::time::Duration::from_millis(1000)); + + loop { + ticker.tick().await; + + let now_ms = chrono::Utc::now().timestamp_millis(); + let ten_min_ago_ms = now_ms - (10 * 60 * 1000); + + let points = storage.get_range(&node_id, &metric, ten_min_ago_ms, now_ms); + + if !points.is_empty() { + let chart_points: Vec = points + .into_iter() + .map(|p| serde_json::json!({ "x": p.timestamp, "y": p.value })) + .collect(); + + if let Ok(msg) = serde_json::to_string(&chart_points) { + if let Err(e) = ws_tx.send(warp::ws::Message::text(msg)).await { + log::error!("WS send error: {}", e); + break; + } + } + } + } +} diff --git a/src/bin/api/core/http/handlers/mod.rs b/src/bin/api/core/http/handlers/mod.rs index 2dda67ff..b3b25c61 100644 --- a/src/bin/api/core/http/handlers/mod.rs +++ b/src/bin/api/core/http/handlers/mod.rs @@ -1,6 +1,7 @@ pub mod connection; pub mod html; pub mod key; +pub mod metrics; pub mod node; pub mod sub; diff --git a/src/bin/api/core/http/handlers/node.rs b/src/bin/api/core/http/handlers/node.rs index b342b602..efb78d4e 100644 --- a/src/bin/api/core/http/handlers/node.rs +++ b/src/bin/api/core/http/handlers/node.rs @@ -1,9 +1,11 @@ +use std::sync::Arc; use warp::http::StatusCode; use pony::http::IdResponse; use pony::http::ResponseMessage; -use pony::memory::node::NodeResponse; use pony::memory::node::Status as NodeStatus; +use pony::memory::node::{NodeMetricInfo, NodeResponse}; +use pony::metrics::storage::MetricStorage; use pony::Connection; use pony::ConnectionApiOp; use pony::ConnectionBaseOp; @@ -15,8 +17,6 @@ use super::super::param::NodeIdParam; use super::super::param::NodesQueryParams; use super::super::request::NodeRequest; -use crate::core::clickhouse::score::NodeScore; -use crate::core::clickhouse::ChContext; use crate::core::sync::tasks::SyncOp; use crate::core::sync::MemSync; @@ -102,6 +102,7 @@ where pub async fn get_nodes_handler( node_param: NodesQueryParams, memory: MemSync, + metrics: Arc, ) -> Result where N: NodeStorageOp + Sync + Send + Clone + 'static, @@ -127,11 +128,56 @@ where Some(nodes) => { let node_response: Vec = nodes .into_iter() - .map(|node| node.as_node_response()) + .map(|node| { + let mut res = node.as_node_response(); + let node_uuid = node.uuid; + + res.metrics = if let Some(node_metrics_map) = metrics.inner.get(&node_uuid) { + node_metrics_map + .iter() + .filter_map(|entry| { + let series_key = entry.key(); + let points = entry.value(); + + if points.is_empty() { + return None; + } + + if !(series_key.starts_with("sys.") + || series_key.starts_with("net.")) + { + return None; + } + + let tags = metrics + .metadata + .get(series_key) + .map(|m| m.value().clone()) + .unwrap_or_default(); + + let name = series_key + .split(':') + .next() + .unwrap_or(series_key) + .to_string(); + + Some(NodeMetricInfo { + key: series_key.clone(), + name, + tags, + }) + }) + .collect() + } else { + vec![] + }; + res + }) .collect(); - let response = ResponseMessage::>> { + + let response = ResponseMessage { status: StatusCode::OK.as_u16(), - message: "List of nodes".to_string(), + message: "List of nodes with metrics".to_string(), response: Some(node_response), }; @@ -196,68 +242,3 @@ where )) } } - -/// Get score of load a node handler -// GET /node/score?id= -pub async fn get_node_score_handler( - node_param: NodeIdParam, - memory: MemSync, - ch: ChContext, -) -> Result -where - N: NodeStorageOp + Sync + Send + Clone + 'static, - C: ConnectionApiOp - + ConnectionBaseOp - + Sync - + Send - + Clone - + 'static - + From - + PartialEq, - S: SubscriptionOp + Send + Sync + Clone + 'static, -{ - let mem = memory.memory.read().await; - - if let Some(node) = mem.nodes.get_by_id(&node_param.id) { - let interface = node.interface.clone(); - let env = node.env.clone(); - let hostname = node.hostname.clone(); - let cores = node.cores; - let max_bandwidth_bps = node.max_bandwidth_bps; - - if let Some(score) = ch - .fetch_node_score(&env, &hostname, &interface, cores, max_bandwidth_bps) - .await - { - let response = ResponseMessage::> { - status: StatusCode::OK.as_u16(), - message: format!("Node score for {}", node_param.id), - response: Some(score), - }; - Ok(warp::reply::with_status( - warp::reply::json(&response), - StatusCode::OK, - )) - } else { - let response = ResponseMessage::> { - status: StatusCode::INTERNAL_SERVER_ERROR.as_u16(), - message: "Failed to calculate node score".to_string(), - response: None, - }; - Ok(warp::reply::with_status( - warp::reply::json(&response), - StatusCode::INTERNAL_SERVER_ERROR, - )) - } - } else { - let response = ResponseMessage::> { - status: StatusCode::NOT_FOUND.as_u16(), - message: "Node not found".to_string(), - response: None, - }; - Ok(warp::reply::with_status( - warp::reply::json(&response), - StatusCode::NOT_FOUND, - )) - } -} diff --git a/src/bin/api/core/http/handlers/sub.rs b/src/bin/api/core/http/handlers/sub.rs index 758965a0..24819bbb 100644 --- a/src/bin/api/core/http/handlers/sub.rs +++ b/src/bin/api/core/http/handlers/sub.rs @@ -1,18 +1,15 @@ use base64::Engine; use chrono::DateTime; use chrono::Utc; -use pony::mtproto_op::mtproto_conn; -use url::Url; -use super::super::request::TagReq; +use url::Url; use warp::http::Response; use warp::http::StatusCode; -use super::super::param::SubIdQueryParam; -use super::super::param::SubQueryParam; -use super::super::request::Subscription as SubReq; use pony::http::helpers as http; +use pony::http::response::Instance; use pony::http::ResponseMessage; +use pony::mtproto_op::mtproto_conn; use pony::utils; use pony::utils::get_uuid_last_octet_simple; use pony::xray_op::clash::generate_clash_config; @@ -20,7 +17,6 @@ use pony::xray_op::clash::generate_proxy_config; use pony::Connection; use pony::ConnectionApiOp; use pony::ConnectionBaseOp; -use pony::ConnectionStat; use pony::ConnectionStorageApiOp; use pony::NodeStorageOp; use pony::OperationStatus as StorageOperationStatus; @@ -29,7 +25,12 @@ use pony::SubscriptionOp; use pony::SubscriptionStorageOp; use pony::Tag; +use super::super::param::SubIdQueryParam; +use super::super::param::SubQueryParam; +use super::super::request::Subscription as SubReq; +use super::super::request::TagReq; use super::html::{FOOTER, HEAD, LOGO}; + use crate::core::sync::tasks::SyncOp; use crate::core::sync::MemSync; @@ -102,7 +103,7 @@ where Ok(StorageOperationStatus::Ok(id)) => Ok(http::success_response( format!("Subscription {} has been created", id), Some(sub_id), - http::Instance::Subscription(sub), + Instance::Subscription(sub), )), Ok(StorageOperationStatus::AlreadyExist(id)) => Ok(http::not_modified(&format!( "Subscription {} already exists", @@ -152,7 +153,7 @@ where Ok(StorageOperationStatus::Updated(id)) => Ok(http::success_response( format!("Subscription {} has been updated", id), Some(sub_id), - http::Instance::None, + Instance::None, )), Ok(StorageOperationStatus::NotFound(id)) => Ok(http::not_found(&format!( "Subscription {} is not found", @@ -176,55 +177,6 @@ where } } -// GET /sub/stat?id=<> -pub async fn subscription_conn_stat_handler( - sub_param: SubIdQueryParam, - memory: MemSync, -) -> Result -where - N: NodeStorageOp + Sync + Send + Clone + 'static, - S: SubscriptionOp + Send + Sync + Clone + 'static + PartialEq, - C: ConnectionApiOp - + ConnectionBaseOp - + Sync - + Send - + Clone - + 'static - + From - + PartialEq, -{ - log::debug!("Received: {:?}", sub_param); - - let mem = memory.memory.read().await; - let mut result: Vec<(uuid::Uuid, ConnectionStat, Tag)> = Vec::new(); - - if mem.subscriptions.find_by_id(&sub_param.id).is_none() { - return Ok(http::not_found("Subscription is not found")); - } - - if let Some(connections) = mem.connections.get_by_subscription_id(&sub_param.id) { - for (conn_id, conn) in connections { - let tag = conn.get_proto().proto(); - - let stat = ConnectionStat { - online: conn.get_online(), - downlink: conn.get_downlink(), - uplink: conn.get_uplink(), - }; - - result.push((conn_id, stat, tag)); - } - - Ok(http::success_response( - "List of subscription connection statistics".to_string(), - None, - http::Instance::Stat(result), - )) - } else { - Ok(http::not_found("Connections is not found")) - } -} - /// Gets Subscription info page // GET /sub/info?id=&env= pub async fn subscription_info_handler( @@ -252,26 +204,6 @@ where dt.format("%d %B %Y Β· %H:%M UTC").to_string() } - fn format_bytes(bytes: i64) -> String { - const KB: f64 = 1024.0; - const MB: f64 = KB * 1024.0; - const GB: f64 = MB * 1024.0; - const TB: f64 = GB * 1024.0; - - let b = bytes as f64; - if b >= TB { - format!("{:.2} Π’Π‘", b / TB) - } else if b >= GB { - format!("{:.2} Π“Π‘", b / GB) - } else if b >= MB { - format!("{:.2} ΠœΠ‘", b / MB) - } else if b >= KB { - format!("{:.2} ΠšΠ‘", b / KB) - } else { - format!("{} Π±Π°ΠΉΡ‚", bytes) - } - } - let env = &sub_param.env; let is_ru = env == "ru"; let is_wl = env == "wl"; @@ -307,19 +239,6 @@ where let invited = mem.subscriptions.count_invited_by(&sub.refer_code()); - let mut downlink = 0; - let mut uplink = 0; - - if let Some(conns) = mem.connections.get_by_subscription_id(id) { - for (_, c) in conns { - downlink += c.get_downlink(); - uplink += c.get_uplink(); - } - } - - let down_str = format_bytes(downlink); - let up_str = format_bytes(uplink); - let title = if is_ru { "Подписка Π½Π° Π ΠΈΠ»Π·ΠΎΠΏΡ€ΠΎΠ²ΠΎΠ΄ (RU)" } else if is_wl { @@ -564,7 +483,7 @@ r#"{head}
Id: {subscription_id}

-
Π’Ρ€Π°Ρ„ΠΈΠΊ: ↓ {down_str}    ↑ {up_str}
+
Π’Ρ€Π°Ρ„ΠΈΠΊ: ↓    ↑

Бсылки для ΠΏΠΎΠ΄ΠΊΠ»ΡŽΡ‡Π΅Π½ΠΈΡ

@@ -742,8 +661,6 @@ window.onload = () => {{ status_text = status_text, expires = expires, days = days, - down_str = down_str, - up_str = up_str, ref = sub.refer_code(), invited = invited, subscription_id = id, @@ -828,7 +745,7 @@ where Ok(http::success_response( format!("Connections are found for {}", sub_id), None, - http::Instance::Connections(cons.clone()), + Instance::Connections(cons.clone()), )) } } diff --git a/src/bin/api/core/http/request.rs b/src/bin/api/core/http/request.rs index 036ad047..ce0e0204 100644 --- a/src/bin/api/core/http/request.rs +++ b/src/bin/api/core/http/request.rs @@ -35,6 +35,7 @@ pub struct NodeRequest { pub interface: String, pub cores: usize, pub max_bandwidth_bps: i64, + pub country: String, } impl NodeRequest { @@ -53,6 +54,7 @@ impl NodeRequest { interface: self.interface.clone(), cores: self.cores, max_bandwidth_bps: self.max_bandwidth_bps, + country: self.country.clone(), } } } diff --git a/src/bin/api/core/http/routes.rs b/src/bin/api/core/http/routes.rs index a3a421f0..d751e388 100644 --- a/src/bin/api/core/http/routes.rs +++ b/src/bin/api/core/http/routes.rs @@ -15,6 +15,7 @@ use super::super::Api; use super::filters::*; use super::handlers::connection::*; use super::handlers::key::*; +use super::handlers::metrics::*; use super::handlers::node::*; use super::handlers::sub::*; use super::param::*; @@ -55,37 +56,31 @@ where async fn run(&self, params: ApiServiceConfig) -> Result<()> { let auth = auth(Arc::new(self.settings.api.token.clone())); + let cors = warp::cors() + .allow_any_origin() + .allow_methods(vec!["GET", "POST", "DELETE"]); + let get_healthcheck_route = warp::get() .and(warp::path("healthcheck")) .and(warp::path::end()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(healthcheck_handler); // Node routes let get_nodes_route = warp::get() .and(warp::path("nodes")) .and(warp::path::end()) - .and(auth.clone()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) + .and(with_metrics(self.metrics.clone())) .and_then(get_nodes_handler); - let get_node_score_route = warp::get() - .and(warp::path("node")) - .and(warp::path("score")) - .and(warp::path::end()) - .and(auth.clone()) - .and(warp::query::()) - .and(with_state(self.sync.clone())) - .and(with_ch(self.ch.clone())) - .and_then(get_node_score_handler); - let post_node_register_route = warp::post() .and(warp::path("node")) .and(warp::path::end()) .and(auth.clone()) .and(warp::body::json::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(post_node_handler); let get_node_route = warp::get() @@ -93,34 +88,25 @@ where .and(warp::path::end()) .and(auth.clone()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(get_node_handler); // Subscription Routes - let get_subscription_stat_route = warp::get() - .and(warp::path("sub")) - .and(warp::path("stat")) - .and(warp::path::end()) - .and(auth.clone()) - .and(warp::query::()) - .and(with_state(self.sync.clone())) - .and_then(subscription_conn_stat_handler); - let get_subscription_connections_route = warp::get() .and(warp::path("sub")) .and(warp::path("connections")) .and(warp::path::end()) .and(auth.clone()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(get_subscription_connections_handler); let get_subscription_route = warp::get() .and(warp::path("sub")) .and(warp::path::end()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(subscription_link_handler); let get_subscription_info_route = warp::get() @@ -128,7 +114,7 @@ where .and(warp::path("info")) .and(warp::path::end()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and(with_param_string(params.web_host)) .and(with_param_string(params.api_web_host)) .and_then(subscription_info_handler); @@ -138,7 +124,7 @@ where .and(warp::path::end()) .and(auth.clone()) .and(warp::body::json()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and(with_i64(params.bonus_days)) .and(with_param_vec_string(params.promo_codes)) .and_then(post_subscription_handler); @@ -149,7 +135,7 @@ where .and(auth.clone()) .and(warp::query::()) .and(warp::body::json()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(put_subscription_handler); // Connections Routes @@ -159,7 +145,7 @@ where .and(warp::path::end()) .and(auth.clone()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(get_connection_handler); let get_connections_route = warp::get() @@ -167,7 +153,7 @@ where .and(warp::path::end()) .and(auth.clone()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(get_connections_handler); let post_connection_route = warp::post() @@ -175,7 +161,7 @@ where .and(warp::path::end()) .and(auth.clone()) .and(warp::body::json()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(create_connection_handler); let delete_connection_route = warp::delete() @@ -183,7 +169,7 @@ where .and(warp::path::end()) .and(auth.clone()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(delete_connection_handler); // Keys Routes @@ -193,7 +179,7 @@ where .and(warp::path::end()) .and(auth.clone()) .and(warp::query::()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and(with_param_vec(params.key_sign_token.clone())) .and_then(get_key_validate_handler); @@ -202,7 +188,7 @@ where .and(warp::path::end()) .and(auth.clone()) .and(warp::body::json()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and(with_param_vec(params.key_sign_token)) .and_then(post_key_handler); @@ -211,13 +197,28 @@ where .and(warp::path("activate")) .and(warp::path::end()) .and(warp::body::json()) - .and(with_state(self.sync.clone())) + .and(with_sync(self.sync.clone())) .and_then(post_activate_key_handler); + use uuid::Uuid; + let ws_metrics_route = warp::path!("metrics" / Uuid / String / "ws") + .and(warp::ws()) + .and(with_metrics(self.metrics.clone())) + .map( + |node_id, encoded_metric: String, ws: warp::ws::Ws, storage| { + let metric_name = urlencoding::decode(&encoded_metric) + .map(|s| s.into_owned()) + .unwrap_or(encoded_metric); + + ws.on_upgrade(move |socket| { + handle_ws_client(socket, node_id, metric_name, storage) + }) + }, + ); + let routes = get_healthcheck_route // Subscription .or(get_subscription_connections_route) - .or(get_subscription_stat_route) .or(get_subscription_route) .or(get_subscription_info_route) .or(post_subscription_route) @@ -225,7 +226,6 @@ where // Node .or(get_nodes_route) .or(get_node_route) - .or(get_node_score_route) .or(post_node_register_route) // Connection .or(get_connection_route) @@ -236,7 +236,10 @@ where .or(get_key_validation_route) .or(post_key_route) .or(post_activate_key_route) - .recover(rejection); + // Metrics + .or(ws_metrics_route) + .recover(rejection) + .with(cors); if let Some(ipv4) = self.settings.api.address { warp::serve(routes) diff --git a/src/bin/api/core/metrics.rs b/src/bin/api/core/metrics.rs index ecf98850..9b18a8c3 100644 --- a/src/bin/api/core/metrics.rs +++ b/src/bin/api/core/metrics.rs @@ -1,64 +1,51 @@ -use chrono::Utc; +use pony::metrics::storage::MetricStorage; +use pony::metrics::MetricEnvelope; +use pony::zmq::subscriber::Subscriber; -use pony::metrics::cpuusage::cpu_metrics; -use pony::metrics::heartbeat::heartbeat_metrics; -use pony::metrics::loadavg::loadavg_metrics; -use pony::metrics::memory::mem_metrics; -use pony::metrics::Metric; -use pony::metrics::MetricType; -use pony::metrics::Metrics; -use pony::Connection; -use pony::ConnectionApiOp; -use pony::ConnectionBaseOp; -use pony::NodeStorageOp; -use pony::SubscriptionOp; +use rkyv::Deserialize; +use std::sync::Arc; -use crate::Api; +pub struct MetricWorker; -#[async_trait::async_trait] -impl Metrics for Api -where - N: NodeStorageOp + Send + Sync + Clone + 'static, - C: ConnectionApiOp - + ConnectionBaseOp - + Send - + Sync - + Clone - + 'static - + From - + std::cmp::PartialEq, - S: SubscriptionOp + Send + Sync + Clone + 'static, -{ - async fn collect_metrics(&self) -> Vec - where - N: NodeStorageOp + Sync + Send + Clone + 'static, - { - let mut metrics = Vec::new(); +impl MetricWorker { + pub async fn start(metric_storage: Arc, subscriber: Subscriber) { + tokio::spawn(async move { + log::info!("MetricWorker: Monitoring pipeline initialized..."); - let env = &self.settings.node.env; - let _uuid = self.settings.node.uuid; + loop { + if let Some((_topic, payload_bytes)) = subscriber.recv().await { + let archived = + match rkyv::check_archived_root::>(&payload_bytes) { + Ok(a) => a, + Err(e) => { + log::error!("MetricWorker: Binary corruption detected: {:?}", e); + continue; + } + }; - if let Some(hostname) = &self.settings.node.hostname { - metrics.extend(cpu_metrics(env, hostname)); - metrics.extend(loadavg_metrics(env, hostname)); - metrics.extend(mem_metrics(env, hostname)); - } else { - log::warn!("Hostname is not set, skipping metrics collection"); - } + let metrics: Vec = match archived + .deserialize(&mut rkyv::Infallible) + { + Ok(m) => m, + Err(e) => { + log::error!("MetricWorker: Failed to reconstruct envelopes: {:?}", e); + continue; + } + }; - log::debug!("Total metrics collected: {}", metrics.len()); - metrics - } + for metric in metrics { + log::debug!( + "Incoming: node={} metric={} value={} tags={:?}", + metric.node_id, + metric.name, + metric.value, + metric.tags + ); - async fn collect_hb_metrics(&self) -> MetricType { - if let Some(hostname) = &self.settings.node.hostname { - heartbeat_metrics(&self.settings.node.env, &self.settings.node.uuid, hostname) - } else { - MetricType::F64(Metric::new( - "hb.unknown".into(), - 0.0, - Utc::now().timestamp(), - )) - } + metric_storage.insert_envelope(metric); + } + } + } + }); } } diff --git a/src/bin/api/core/mod.rs b/src/bin/api/core/mod.rs index b1da3b73..6339bbc7 100644 --- a/src/bin/api/core/mod.rs +++ b/src/bin/api/core/mod.rs @@ -1,26 +1,28 @@ +use serde::{Deserialize, Serialize}; use std::collections::HashMap; +use std::sync::Arc; use pony::config::settings::ApiSettings; +use pony::memory::connection::Connections; use pony::memory::node::Node; +use pony::memory::subscription::Subscription; +use pony::memory::subscription::Subscriptions; +use pony::metrics::storage::MetricStorage; use pony::Connection; use pony::ConnectionApiOp; use pony::ConnectionBaseOp; -use pony::MemoryCache; use pony::NodeStorageOp; -use pony::Subscription; use pony::SubscriptionOp; -use crate::core::clickhouse::ChContext; use crate::core::sync::MemSync; -pub(crate) mod clickhouse; pub(crate) mod http; pub(crate) mod metrics; pub(crate) mod postgres; pub(crate) mod sync; pub(crate) mod tasks; -pub type ApiState = MemoryCache>, Connection, Subscription>; +pub type ApiState = Cache>, Connection, Subscription>; pub struct Api where @@ -29,8 +31,8 @@ where S: SubscriptionOp + Send + Sync + Clone + 'static, { pub sync: MemSync, - pub ch: ChContext, pub settings: ApiSettings, + pub metrics: Arc, } impl Api @@ -39,7 +41,49 @@ where C: ConnectionBaseOp + ConnectionApiOp + Send + Sync + Clone + 'static + std::cmp::PartialEq, S: SubscriptionOp + Send + Sync + Clone + 'static, { - pub fn new(ch: ChContext, sync: MemSync, settings: ApiSettings) -> Self { - Self { ch, sync, settings } + pub fn new(sync: MemSync, settings: ApiSettings, metrics: Arc) -> Self { + Self { + sync, + settings, + metrics, + } + } +} + +#[derive(Deserialize, Serialize, Debug, Clone)] +pub struct Cache +where + T: Send + Sync + Clone + 'static, + C: Send + Sync + Clone + 'static, + S: Send + Sync + Clone + 'static, +{ + pub connections: Connections, + pub subscriptions: Subscriptions, + pub nodes: T, +} + +impl Default for Cache +where + T: NodeStorageOp + Sync + Send + Clone + 'static, + C: ConnectionApiOp + ConnectionBaseOp + Clone + Send + Sync + 'static + PartialEq, + S: SubscriptionOp + Clone + Send + Sync + 'static, +{ + fn default() -> Self { + Self::new() + } +} + +impl Cache +where + T: NodeStorageOp + Sync + Send + Clone + 'static, + C: ConnectionApiOp + ConnectionBaseOp + Clone + Send + Sync + 'static + PartialEq, + S: SubscriptionOp + Clone + Send + Sync + 'static, +{ + pub fn new() -> Self { + Cache { + nodes: T::default(), + connections: Connections::default(), + subscriptions: Subscriptions::default(), + } } } diff --git a/src/bin/api/core/postgres/connection.rs b/src/bin/api/core/postgres/connection.rs index 0750d177..e9a68de9 100644 --- a/src/bin/api/core/postgres/connection.rs +++ b/src/bin/api/core/postgres/connection.rs @@ -8,7 +8,6 @@ use std::sync::Arc; use tokio::sync::Mutex; use pony::Connection as Conn; -use pony::ConnectionStat; use pony::Proto; use pony::Tag; use pony::WgKeys; @@ -27,7 +26,6 @@ pub struct ConnRow { pub modified_at: DateTime, pub expires_at: Option>, pub subscription_id: Option, - pub stat: ConnectionStat, pub wg: Option, pub node_id: Option, pub proto: Tag, @@ -37,12 +35,6 @@ pub struct ConnRow { impl From<(uuid::Uuid, Conn)> for ConnRow { fn from((conn_id, conn): (uuid::Uuid, Conn)) -> Self { - let conn_stat = ConnectionStat { - online: conn.stat.online, - uplink: conn.stat.uplink, - downlink: conn.stat.downlink, - }; - ConnRow { conn_id, password: conn.get_password(), @@ -51,7 +43,6 @@ impl From<(uuid::Uuid, Conn)> for ConnRow { modified_at: conn.modified_at, expires_at: conn.expires_at, subscription_id: conn.subscription_id, - stat: conn_stat, wg: conn.get_wireguard().cloned(), node_id: conn.get_wireguard_node_id(), proto: conn.get_proto().proto(), @@ -99,7 +90,6 @@ impl TryFrom for Conn { Ok(Self { env: row.env, proto, - stat: row.stat, subscription_id: row.subscription_id, created_at: row.created_at, modified_at: row.modified_at, @@ -132,9 +122,6 @@ impl PgConn { modified_at, expires_at, subscription_id, - online, - uplink, - downlink, proto, node_id, wg_privkey, @@ -160,9 +147,6 @@ impl PgConn { let expires_at: Option> = row.get("expires_at"); let subscription_id: Option = row.get("subscription_id"); let token: Option = row.get("token"); - let online: i64 = row.get("online"); - let uplink: i64 = row.get("uplink"); - let downlink: i64 = row.get("downlink"); let proto: Tag = row.get("proto"); let node_id: Option = row.get("node_id"); let wg_privkey: Option = row.get("wg_privkey"); @@ -189,11 +173,6 @@ impl PgConn { modified_at, expires_at, subscription_id, - stat: ConnectionStat { - online, - uplink, - downlink, - }, proto, wg, node_id, @@ -203,25 +182,6 @@ impl PgConn { .collect() } - pub async fn update_stat(&self, conn_id: &uuid::Uuid, stat: ConnectionStat) -> Result<()> { - let mut manager = self.manager.lock().await; - let client = manager.get_client().await?; - - let query = " - UPDATE connections - SET downlink = $1, uplink = $2, online = $3 - WHERE id = $4"; - - client - .execute( - query, - &[&stat.downlink, &stat.uplink, &stat.online, &conn_id], - ) - .await?; - - Ok(()) - } - pub async fn delete(&self, conn_id: &uuid::Uuid) -> Result<()> { let mut manager = self.manager.lock().await; let client = manager.get_client().await?; @@ -257,9 +217,6 @@ impl PgConn { modified_at, expires_at, subscription_id, - online, - uplink, - downlink, proto, is_deleted, wg_privkey, @@ -270,7 +227,7 @@ impl PgConn { ) VALUES ( $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, - $11, $12, $13, $14, $15, $16, $17 + $11, $12, $13, $14 ) "; @@ -285,9 +242,6 @@ impl PgConn { &conn.modified_at, &conn.expires_at, &conn.subscription_id, - &conn.stat.online, - &conn.stat.uplink, - &conn.stat.downlink, &conn.proto, &conn.is_deleted, &conn.wg.as_ref().map(|w| &w.keys.privkey), diff --git a/src/bin/api/core/postgres/mod.rs b/src/bin/api/core/postgres/mod.rs index 73df44f5..4573d76f 100644 --- a/src/bin/api/core/postgres/mod.rs +++ b/src/bin/api/core/postgres/mod.rs @@ -8,7 +8,6 @@ use pony::config::settings::PostgresConfig; use pony::memory::node::Node; use pony::Connection; use pony::ConnectionStorageApiOp; -use pony::MemoryCache; use pony::NodeStorageOp; use pony::OperationStatus; use pony::Result; @@ -20,6 +19,7 @@ use crate::core::postgres::connection::PgConn; use crate::core::postgres::keys::PgKey; use crate::core::postgres::node::PgNode; use crate::core::postgres::subscription::PgSubscription; +use crate::Cache; pub mod connection; pub mod keys; @@ -115,7 +115,7 @@ pub trait Tasks { } #[async_trait::async_trait] -impl Tasks for MemoryCache>, Connection, Subscription> { +impl Tasks for Cache>, Connection, Subscription> { async fn add_conn(&mut self, db_conn: ConnRow) -> Result { let conn_id = db_conn.conn_id; let conn: Connection = db_conn.try_into()?; diff --git a/src/bin/api/core/postgres/node.rs b/src/bin/api/core/postgres/node.rs index 22d8718f..8d7b875e 100644 --- a/src/bin/api/core/postgres/node.rs +++ b/src/bin/api/core/postgres/node.rs @@ -163,7 +163,7 @@ impl PgNode { .query( "SELECT n.id AS node_id, n.uuid, n.env, n.hostname, n.address, n.status, - n.created_at, n.modified_at, n.label, n.interface, n.cores, n.max_bandwidth_bps, + n.created_at, n.modified_at, n.label, n.interface, n.cores, n.max_bandwidth_bps, n.country, i.id AS inbound_id, i.tag, i.port, i.stream_settings, i.uplink, i.downlink, i.conn_count, i.wg_pubkey, i.wg_privkey, i.wg_interface, i.wg_network, i.wg_address, i.dns, i.h2, i.mtproto_secret FROM nodes n @@ -186,6 +186,7 @@ impl PgNode { let label: String = row.get("label"); let interface: String = row.get("interface"); let cores: i32 = row.get("cores"); + let country: String = row.get("country"); let max_bandwidth_bps: i64 = row.get("max_bandwidth_bps"); let wg_network: Option = row @@ -224,6 +225,7 @@ impl PgNode { inbounds: HashMap::new(), cores: cores as usize, max_bandwidth_bps, + country, }); if let Some(_inbound_id) = inbound_id { diff --git a/src/bin/api/core/sync/mod.rs b/src/bin/api/core/sync/mod.rs index c475a59c..f99f7cf2 100644 --- a/src/bin/api/core/sync/mod.rs +++ b/src/bin/api/core/sync/mod.rs @@ -5,11 +5,11 @@ use tokio::sync::RwLock; use pony::Connection as Conn; use pony::ConnectionApiOp; use pony::ConnectionBaseOp; -use pony::MemoryCache; use pony::NodeStorageOp; use pony::SubscriptionOp; use super::postgres::PgContext; +use super::Cache; pub(crate) mod tasks; @@ -20,7 +20,7 @@ where C: Send + Sync + Clone + 'static, S: Send + Sync + Clone + 'static, { - pub memory: Arc>>, + pub memory: Arc>>, pub db: PgContext, pub publisher: Publisher, } @@ -31,11 +31,7 @@ where C: ConnectionBaseOp + ConnectionApiOp + Send + Sync + Clone + 'static + From + PartialEq, S: SubscriptionOp + Send + Sync + Clone + 'static, { - pub fn new( - memory: Arc>>, - db: PgContext, - publisher: Publisher, - ) -> Self { + pub fn new(memory: Arc>>, db: PgContext, publisher: Publisher) -> Self { Self { memory, db, diff --git a/src/bin/api/core/sync/tasks.rs b/src/bin/api/core/sync/tasks.rs index 9a10e881..aee99762 100644 --- a/src/bin/api/core/sync/tasks.rs +++ b/src/bin/api/core/sync/tasks.rs @@ -6,7 +6,6 @@ use pony::Connection as Conn; use pony::Connection; use pony::ConnectionApiOp; use pony::ConnectionBaseOp; -use pony::ConnectionStat; use pony::ConnectionStorageApiOp as ApiOp; use pony::ConnectionStorageBaseOp; use pony::NodeStorageOp; @@ -83,7 +82,6 @@ where env: &str, status: NodeStatus, ) -> SyncResult<()>; - async fn update_conn_stat(&self, conn_id: &uuid::Uuid, stat: ConnectionStat) -> SyncResult<()>; async fn update_sub(&self, sub_id: &uuid::Uuid, sub_req: SubReq) -> SyncResult; async fn add_days(&self, sub_id: &uuid::Uuid, days: i64) -> SyncResult; @@ -520,38 +518,6 @@ where Ok(()) } - async fn update_conn_stat( - &self, - conn_id: &uuid::Uuid, - conn_stat: ConnectionStat, - ) -> SyncResult<()> { - log::debug!("Updating connection stats: {}", conn_id); - - // Update database first - if let Err(e) = self.db.conn().update_stat(conn_id, conn_stat.clone()).await { - log::error!( - "Failed to update connection {} stats in database: {}", - conn_id, - e - ); - return Err(SyncError::Database(e)); - } - - { - let mut memory = self.memory.write().await; - if let Err(e) = memory.connections.update_stats(conn_id, conn_stat) { - log::warn!( - "Failed to update connection {} stats in memory: {}", - conn_id, - e - ); - } - } - - log::debug!("Successfully updated connection {} stats", conn_id); - Ok(()) - } - async fn update_sub( &self, sub_id: &uuid::Uuid, diff --git a/src/bin/api/core/tasks.rs b/src/bin/api/core/tasks.rs index 5378667e..6ca9bcd5 100644 --- a/src/bin/api/core/tasks.rs +++ b/src/bin/api/core/tasks.rs @@ -1,38 +1,28 @@ use async_trait::async_trait; -use chrono::NaiveTime; -use chrono::TimeZone; use chrono::Utc; -use futures::future::join_all; use rand::Rng; use std::collections::HashMap; use std::time::Duration; use pony::memory::node::Node; -use pony::memory::node::Status as NodeStatus; use pony::utils::measure_time; use pony::Connection; use pony::ConnectionBaseOp; -use pony::ConnectionStat; use pony::ConnectionStorageApiOp; -use pony::MemoryCache; -use pony::NodeStorageOp; use pony::OperationStatus as StorageOperationStatus; -use pony::PonyError; use pony::Result; use pony::Subscription; use pony::SubscriptionOp; -use crate::core::clickhouse::query::Queries; use crate::core::postgres::Tasks as MemoryCacheTasks; use crate::core::sync::tasks::SyncOp; use crate::Api; +use crate::Cache; #[async_trait] pub trait Tasks { async fn get_state_from_db(&self) -> Result<()>; async fn periodic_db_sync(&self, interval_sec: u64); - async fn node_healthcheck(&self) -> Result<()>; - async fn collect_conn_stat(&self) -> Result<()>; async fn cleanup_expired_connections(&self, interval_sec: u64); async fn cleanup_expired_subscriptions(&self, interval_sec: u64); async fn restore_subscriptions(&self, interval_sec: u64); @@ -177,12 +167,7 @@ impl Tasks for Api>, Connection, Subscription> { let interval_sec_with_jitter = base + Duration::from_secs(jitter); tokio::time::sleep(interval_sec_with_jitter).await; - if let Err(e) = measure_time( - self.get_state_from_db(), - format!("Periodic DB Sync: interval + jitter {interval_sec}, {jitter}"), - ) - .await - { + if let Err(e) = measure_time(self.get_state_from_db(), "Periodic DB Sync").await { log::error!("Periodic DB sync failed: {:?}", e); } else { log::info!("Periodic DB sync completed successfully"); @@ -194,8 +179,7 @@ impl Tasks for Api>, Connection, Subscription> { let db = self.sync.db.clone(); let mut mem = self.sync.memory.write().await; - let mut tmp_mem: MemoryCache>, Connection, Subscription> = - MemoryCache::new(); + let mut tmp_mem: Cache>, Connection, Subscription> = Cache::new(); let node_repo = db.node(); let conn_repo = db.conn(); @@ -218,129 +202,4 @@ impl Tasks for Api>, Connection, Subscription> { Ok(()) } - - async fn node_healthcheck(&self) -> Result<()> { - let mem = self.sync.memory.read().await; - let nodes = match mem.nodes.all() { - Some(n) => n.clone(), - None => return Ok(()), - }; - drop(mem); - - let timeout = chrono::Duration::seconds(self.settings.api.node_health_check_timeout as i64); - let now = Utc::now(); - - let tasks = nodes.into_iter().map(|node| { - let ch = self.ch.clone(); - let sync = self.sync.clone(); - let uuid = node.uuid; - let env = node.env.clone(); - let hostname = node.hostname; - - async move { - let status = - match ch.fetch_node_heartbeat(&env, &uuid, &hostname).await { - Some(hb) => { - let ts = Utc.timestamp_opt(hb.timestamp, 0); - match ts { - chrono::LocalResult::Single(dt) => { - let diff = now - dt; - if diff <= timeout { - log::debug!( - "[ONLINE] Heartbeat OK: now: {}, dt: {}, diff: {}s <= {}s", - now, dt, diff.num_seconds(), timeout.num_seconds() - ); - NodeStatus::Online - } else { - log::debug!( - "[OFFLINE] Heartbeat too old: now: {}, dt: {}, diff: {}s > {}s", - now, dt, diff.num_seconds(), timeout.num_seconds() - ); - NodeStatus::Offline - } - } - _ => { - log::debug!( - "Could not parse timestamp from heartbeat: {:?}", - hb.timestamp - ); - return Err(PonyError::Custom(format!( - "Could not parse timestamp from heartbeat: {:?}", - hb.timestamp - ))); - } - } - } - None => { - log::debug!( - "No heartbeat data found for node {} ({}) in env {}", - uuid, - hostname, - env - ); - return Err(PonyError::Custom(format!( - "No heartbeat data found for node {} ({}) in env {}", - uuid, hostname, env - ))); - } - }; - - if let Err(e) = SyncOp::update_node_status(&sync, &uuid, &env, status).await { - log::error!("Failed to update_node_status {} database: {}", uuid, e); - return Err(PonyError::Custom(e.to_string())); - } - Ok::<_, PonyError>(()) - } - }); - - for result in join_all(tasks).await { - if let Err(e) = result { - log::error!("Healthcheck task error: {:?}", e); - } - } - - Ok(()) - } - - async fn collect_conn_stat(&self) -> Result<()> { - let conns_map = { - let mem = self.sync.memory.read().await; - mem.connections - .iter() - .filter(|(_, conn)| !conn.get_deleted()) - .map(|(k, v)| (*k, v.clone())) - .collect::>() - }; - - let today = Utc::now().date_naive(); - let start_of_day = - Utc.from_utc_datetime(&today.and_time(NaiveTime::from_hms_opt(0, 0, 0).unwrap())); - - let tasks = conns_map.into_iter().map(|(conn_id, _)| { - let ch = self.ch.clone(); - let sync = self.sync.clone(); - - async move { - if let Some(metrics) = ch.fetch_conn_stats(conn_id, start_of_day).await { - log::debug!("metrics {:?}", metrics); - let stat = ConnectionStat::from_metrics(metrics); - log::debug!("Stat to update - {}", stat); - SyncOp::update_conn_stat(&sync, &conn_id, stat).await - } else { - log::warn!("No metrics found for conn_id {}", conn_id); - Ok(()) - } - } - }); - - let results = join_all(tasks).await; - - for result in results { - if let Err(e) = result { - log::error!("Error during stat update: {:?}", e); - } - } - - Ok(()) - } } diff --git a/src/bin/api/main.rs b/src/bin/api/main.rs index 1f95842d..1cd75863 100644 --- a/src/bin/api/main.rs +++ b/src/bin/api/main.rs @@ -1,26 +1,24 @@ use fern::Dispatch; -use std::net::Ipv4Addr; use std::sync::Arc; use tokio::sync::RwLock; -use tokio::time::sleep; use tokio::time::Duration; use pony::config::settings::ApiSettings; use pony::config::settings::Settings; -use pony::http::debug; -use pony::metrics::Metrics; +use pony::metrics::storage::MetricStorage; use pony::utils::*; use pony::zmq::publisher::Publisher as ZmqPublisher; -use pony::MemoryCache; use pony::Result; +use pony::Subscriber; -use crate::core::clickhouse::ChContext; use crate::core::http::routes::Http; +use crate::core::metrics::MetricWorker; use crate::core::postgres::PgContext; use crate::core::sync::MemSync; use crate::core::tasks::Tasks; use crate::core::Api; use crate::core::ApiState; +use crate::core::Cache; mod core; @@ -55,8 +53,6 @@ async fn main() -> Result<()> { .apply() .unwrap(); - let debug = settings.debug.enabled; - let db = match PgContext::init(&settings.pg).await { Ok(db) => db, Err(err) => { @@ -65,15 +61,18 @@ async fn main() -> Result<()> { } }; - let ch = ChContext::new(&settings.clickhouse.address); let publisher = ZmqPublisher::new(&settings.zmq.endpoint).await; - let mem: Arc> = Arc::new(RwLock::new(MemoryCache::new())); + let mem: Arc> = Arc::new(RwLock::new(Cache::new())); let mem_sync = MemSync::new(mem.clone(), db.clone(), publisher.clone()); + let metric_storage = Arc::new(MetricStorage::new( + settings.api.max_points, + settings.api.retention_seconds, + )); - let api = Arc::new(Api::new(ch.clone(), mem_sync.clone(), settings.clone())); + let api = Arc::new(Api::new(mem_sync.clone(), settings.clone(), metric_storage)); - measure_time(api.get_state_from_db(), "Init PG DB".to_string()).await?; + measure_time(api.get_state_from_db(), "Init PG DB").await?; let api_clone = api.clone(); tokio::spawn(async move { @@ -82,77 +81,24 @@ async fn main() -> Result<()> { .await; }); - if settings.api.metrics_enabled { - log::info!("Running metrics send task"); - tokio::spawn({ - let settings = settings.clone(); - let api = api.clone(); - - async move { - loop { - sleep(Duration::from_secs(settings.api.metrics_interval)).await; - let _ = api.send_metrics(&settings.carbon.address).await; - } - } - }); - } - - if settings.api.metrics_enabled { - log::info!("Running HB metrics send task"); - tokio::spawn({ - let settings = settings.clone(); - let api = api.clone(); - - async move { - loop { - sleep(Duration::from_secs(settings.api.metrics_hb_interval)).await; - let _ = api.send_hb_metric(&settings.carbon.address).await; - } - } - }); - } - - let token = settings.api.token.clone(); - if debug { - let token = Arc::new(token); - tokio::spawn(debug::start_ws_server( - mem.clone(), - settings - .debug - .web_server - .unwrap_or(Ipv4Addr::new(127, 0, 0, 1)), - settings.debug.web_port, - token, - )); - } + log::debug!("Running metrics reciever task"); + let subscriber = Subscriber::new_bound(&settings.metrics.reciever, settings.metrics.topic); - tokio::spawn({ - log::info!("node_healthcheck task started"); - let job_interval = Duration::from_secs(settings.api.healthcheck_interval); - let api = api.clone(); + MetricWorker::start(api.metrics.clone(), subscriber).await; - async move { - loop { - if let Err(e) = api.node_healthcheck().await { - log::error!("node_healthcheck task failed: {:?}", e); - } - tokio::time::sleep(job_interval).await; - } - } - }); + log::info!("Metrics system initialized via MetricWorker"); - tokio::spawn({ - log::info!("collect_conn_stat task started"); - let api = api.clone(); - let job_interval = Duration::from_secs(settings.api.collect_conn_stat_interval); - - async move { - loop { - tokio::time::sleep(job_interval).await; - if let Err(e) = api.collect_conn_stat().await { - log::error!("collect_conn_stat task failed: {:?}", e); - } - } + let metrics_storage = api.metrics.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(60)); + loop { + interval.tick().await; + let s = metrics_storage.clone(); + tokio::task::spawn_blocking(move || { + log::info!("Starting MetricStorage GC..."); + s.perform_gc(); + log::info!("MetricStorage GC finished."); + }); } }); diff --git a/src/bin/auth/core/handlers.rs b/src/bin/auth/core/handlers.rs index a2c84974..57cb5419 100644 --- a/src/bin/auth/core/handlers.rs +++ b/src/bin/auth/core/handlers.rs @@ -1,9 +1,8 @@ +use pony::memory::connection::Connections; use std::sync::Arc; use tokio::sync::RwLock; -use crate::core::helpers::{activate_key, validate_key}; - -use super::helpers::{create_connection, create_subscription}; +use super::helpers::{activate_key, create_connection, create_subscription, validate_key}; use super::request; use super::response; use super::EmailStore; @@ -11,11 +10,9 @@ use super::HttpClient; use pony::config::settings::ApiAccessConfig; use pony::http::helpers as http; +use pony::http::response::Instance; use pony::ConnectionBaseOp; use pony::ConnectionStorageBaseOp; -use pony::MemoryCache as Cache; -use pony::NodeStorageOp; -use pony::SubscriptionOp; use super::Env; use super::DEFAULT_DAYS; @@ -140,7 +137,7 @@ pub async fn activate_key_handler( Ok(http::success_response( "Key code activated.".to_string(), Some(key.id), - http::Instance::Subscription(updated_sub), + Instance::Subscription(updated_sub), )) } @@ -252,22 +249,20 @@ pub async fn trial_handler( Ok(http::success_response( "Trial activated. Check your email".to_string(), Some(sub.id), - http::Instance::None, + Instance::None, )) } -pub async fn auth_handler( +pub async fn auth_handler( req: request::Auth, - memory: Arc>>, + memory: Arc>>, ) -> Result where - N: NodeStorageOp + Sync + Send + Clone + 'static, - S: SubscriptionOp + Sync + Send + Clone + 'static + std::cmp::PartialEq + serde::Serialize, - C: ConnectionBaseOp + Sync + Send + Clone + 'static + std::fmt::Display, + C: ConnectionBaseOp + Sync + Send + Clone + 'static, { log::debug!("Auth req {} {} {}", req.auth, req.addr, req.tx); let mem = memory.read().await; - if let Some(id) = mem.connections.validate_token(&req.auth) { + if let Some(id) = mem.validate_token(&req.auth) { Ok(warp::reply::json(&response::Auth { ok: true, id: Some(id.to_string()), diff --git a/src/bin/auth/core/helpers.rs b/src/bin/auth/core/helpers.rs index 5a560ea3..7f2b0a77 100644 --- a/src/bin/auth/core/helpers.rs +++ b/src/bin/auth/core/helpers.rs @@ -1,7 +1,7 @@ -use super::response; use super::Env; use super::HttpClient; -use pony::http::helpers::{Instance, InstanceWithId}; +use pony::http::response::ResponseMessage; +use pony::http::response::{Instance, InstanceWithId}; use pony::memory::key::Code; use pony::Subscription; @@ -28,7 +28,7 @@ pub async fn validate_key( let text = res.text().await?; if status.is_success() { - let parsed: response::Api> = serde_json::from_str(&text)?; + let parsed: ResponseMessage> = serde_json::from_str(&text)?; log::debug!("Response validate_key {} {}", parsed.status, parsed.message); @@ -69,7 +69,7 @@ pub async fn activate_key( let text = res.text().await?; if status.is_success() { - let parsed: response::Api> = serde_json::from_str(&text)?; + let parsed: ResponseMessage> = serde_json::from_str(&text)?; log::debug!("Response activate_key {} {}", parsed.status, parsed.message); @@ -111,7 +111,7 @@ pub async fn create_subscription( let text = res.text().await?; if status.is_success() { - let parsed: response::Api> = serde_json::from_str(&text)?; + let parsed: ResponseMessage> = serde_json::from_str(&text)?; match parsed.response.instance { Instance::Subscription(sub) => Ok(sub), @@ -172,7 +172,7 @@ pub async fn create_connection( anyhow::bail!("empty connection response, status = {}", status); } - let parsed: response::Api> = serde_json::from_str(&text)?; + let parsed: ResponseMessage> = serde_json::from_str(&text)?; log::debug!( "Response create_connection {} {}", diff --git a/src/bin/auth/core/http.rs b/src/bin/auth/core/http.rs index 2f9df24a..e0a3973b 100644 --- a/src/bin/auth/core/http.rs +++ b/src/bin/auth/core/http.rs @@ -1,9 +1,8 @@ use async_trait::async_trait; +use reqwest::StatusCode; use serde::{Deserialize, Serialize}; use pony::ConnectionBaseOp; -use pony::NodeStorageOp; -use pony::SubscriptionOp; use pony::Tag; use pony::{PonyError, Result as PonyResult}; @@ -12,6 +11,8 @@ use super::HttpClient; use reqwest::Url; +use pony::http::response::{Instance, InstanceWithId, ResponseMessage}; + #[derive(Clone, Debug, Deserialize, Serialize)] pub struct ConnTypeParam { pub proto: Tag, @@ -32,11 +33,9 @@ pub trait ApiRequests { } #[async_trait] -impl ApiRequests for AuthService +impl ApiRequests for AuthService where - T: NodeStorageOp + Send + Sync + Clone, C: ConnectionBaseOp + Send + Sync + Clone + 'static, - S: SubscriptionOp + Send + Sync + Clone + 'static, { async fn get_connections( &self, @@ -45,16 +44,7 @@ where proto: Tag, last_update: Option, ) -> PonyResult<()> { - let node = { - let mem = self.memory.read().await; - mem.nodes - .get_self() - .expect("No node available to register") - .clone() - }; - - let id = node.uuid; - let env = node.env; + let id = self.node.uuid; let conn_type_param = ConnTypeParam { proto, @@ -81,7 +71,22 @@ where let status = res.status(); let body = res.text().await?; if status.is_success() { - log::debug!("Connections Request Accepted: {:?}", status); + let result: ResponseMessage> = serde_json::from_str(&body)?; + let count = match result.response.instance { + Instance::Count(count) => count, + _ => { + return Err(PonyError::Custom("Unexpected instance type".into())); + } + }; + log::debug!( + "Connections Request Accepted for {}: {} Count: {} ", + proto, + status, + count, + ); + Ok(()) + } else if status == StatusCode::NOT_MODIFIED { + log::debug!("Connections Request Accepted for {}: {} ", proto, status,); Ok(()) } else { log::error!("Connections Request failed: {} - {}", status, body); diff --git a/src/bin/auth/core/metrics.rs b/src/bin/auth/core/metrics.rs new file mode 100644 index 00000000..ce43862c --- /dev/null +++ b/src/bin/auth/core/metrics.rs @@ -0,0 +1,19 @@ +use pony::memory::node::Node; +use pony::metrics::storage::HasMetrics; +use pony::metrics::storage::MetricBuffer; +use pony::ConnectionBaseOp; + +use super::AuthService; + +impl HasMetrics for AuthService +where + C: ConnectionBaseOp + Send + Sync + Clone + 'static, +{ + fn metrics(&self) -> &MetricBuffer { + &self.metrics + } + + fn node_settings(&self) -> &Node { + &self.node + } +} diff --git a/src/bin/auth/core/mod.rs b/src/bin/auth/core/mod.rs index 77c0dcde..122de773 100644 --- a/src/bin/auth/core/mod.rs +++ b/src/bin/auth/core/mod.rs @@ -7,14 +7,11 @@ use warp::Filter; use pony::config::settings::ApiAccessConfig; use pony::http::filters as pony_filters; +use pony::memory::connection::Connections; use pony::memory::node::Node; +use pony::metrics::storage::MetricBuffer; use pony::zmq::subscriber::Subscriber as ZmqSubscriber; -use pony::BaseConnection as Connection; use pony::ConnectionBaseOp; -use pony::MemoryCache; -use pony::NodeStorageOp; -use pony::Subscription; -use pony::SubscriptionOp; use crate::core::email::EmailStore; use crate::core::handlers::trial_handler; @@ -25,12 +22,12 @@ pub mod filters; pub mod handlers; pub mod helpers; pub mod http; +pub mod metrics; pub mod request; pub mod response; pub mod service; pub mod tasks; -pub type AuthServiceState = MemoryCache; pub type HttpClient = Client; const PROTOS: [&str; 4] = [ @@ -58,44 +55,45 @@ impl Display for Env { const DEFAULT_DAYS: i64 = 1; -pub struct AuthService +pub struct AuthService where - N: NodeStorageOp + Send + Sync + Clone + 'static, C: ConnectionBaseOp + Send + Sync + Clone + 'static, - S: SubscriptionOp + Send + Sync + Clone + 'static, { - pub memory: Arc>>, + pub memory: Arc>>, + pub metrics: Arc, + pub node: Node, pub subscriber: ZmqSubscriber, pub email_store: EmailStore, pub http_client: HttpClient, - pub ipaddr: Ipv4Addr, - pub port: u16, pub api: ApiAccessConfig, + pub listen: Ipv4Addr, + pub port: u16, } -impl AuthService +impl AuthService where - N: NodeStorageOp + Send + Sync + Clone + 'static, - C: ConnectionBaseOp + Send + Sync + Clone + 'static + Display + PartialEq, - S: SubscriptionOp + Send + Sync + Clone + 'static + PartialEq + serde::Serialize, + C: ConnectionBaseOp + Send + Sync + Clone + 'static, { pub fn new( - memory: Arc>>, + metrics: Arc, + node: Node, subscriber: ZmqSubscriber, email_store: EmailStore, http_client: HttpClient, - addr: Ipv4Addr, - port: u16, api: ApiAccessConfig, + listen: (Ipv4Addr, u16), ) -> Self { + let memory = Arc::new(RwLock::new(Connections::default())); Self { memory, + metrics, + node, subscriber, email_store, http_client, - ipaddr: addr, - port, api, + listen: listen.0, + port: listen.1, } } @@ -146,7 +144,7 @@ where .or(activate_route); warp::serve(routes) - .run(SocketAddr::new(IpAddr::V4(self.ipaddr), self.port)) + .run(SocketAddr::new(IpAddr::V4(self.listen), self.port)) .await; } } diff --git a/src/bin/auth/core/response.rs b/src/bin/auth/core/response.rs index 048c6e2d..9e9c73b9 100644 --- a/src/bin/auth/core/response.rs +++ b/src/bin/auth/core/response.rs @@ -1,4 +1,4 @@ -use serde::{Deserialize, Serialize}; +use serde::Serialize; #[derive(Serialize)] pub struct Auth { @@ -6,10 +6,3 @@ pub struct Auth { #[serde(skip_serializing_if = "Option::is_none")] pub id: Option, } - -#[derive(Deserialize)] -pub struct Api { - pub status: u16, - pub message: String, - pub response: T, -} diff --git a/src/bin/auth/core/service.rs b/src/bin/auth/core/service.rs index 79980077..2037f543 100644 --- a/src/bin/auth/core/service.rs +++ b/src/bin/auth/core/service.rs @@ -3,16 +3,16 @@ use std::path::Path; use std::sync::Arc; use tokio::signal; use tokio::sync::broadcast; -use tokio::sync::RwLock; use tokio::task::JoinHandle; use tokio::time::sleep; use tokio::time::Duration; use pony::config::settings::AuthServiceSettings; use pony::config::settings::NodeConfig; -use pony::http::debug; use pony::memory::node::Node; -use pony::MemoryCache; +use pony::metrics::storage::MetricBuffer; +use pony::BaseConnection as Connection; +use pony::Publisher; use pony::Result; use pony::SnapshotManager; use pony::Subscriber as ZmqSubscriber; @@ -20,7 +20,6 @@ use pony::Subscriber as ZmqSubscriber; use super::AuthService; use crate::core::http::ApiRequests; use crate::core::tasks::Tasks; -use crate::core::AuthServiceState; use crate::core::EmailStore; use crate::core::HttpClient; @@ -31,9 +30,6 @@ pub async fn run(settings: AuthServiceSettings) -> Result<()> { let node_config = NodeConfig::from_raw(settings.node.clone()); let node = Node::new(node_config?, None, None, None, None); - let memory: Arc> = - Arc::new(RwLock::new(MemoryCache::with_node(node.clone()))); - let subscriber = ZmqSubscriber::new( &settings.zmq.endpoint, &settings.node.uuid, @@ -50,21 +46,30 @@ pub async fn run(settings: AuthServiceSettings) -> Result<()> { email_store.load_trials().await?; let http_client = HttpClient::new(); - let auth = Arc::new(AuthService::new( - memory.clone(), + let listen_addr = settings + .auth + .web_server + .unwrap_or(Ipv4Addr::from_octets([127, 0, 0, 1])); + + let metric_publisher = Publisher::connect(&settings.metrics.publisher).await; + + let metrics = MetricBuffer { + batch: parking_lot::Mutex::new(Vec::new()), + publisher: metric_publisher, + }; + + let auth = Arc::new(AuthService::::new( + Arc::new(metrics), + node, subscriber, email_store, http_client, - settings - .auth - .web_server - .unwrap_or(Ipv4Addr::from_octets([127, 0, 0, 1])), - settings.auth.web_port, settings.api.clone(), + (listen_addr, settings.auth.web_port), )); let snapshot_manager = - SnapshotManager::new(settings.clone().auth.snapshot_path, memory.clone()); + SnapshotManager::new(settings.clone().auth.snapshot_path, auth.memory.clone()); let snapshot_timestamp = if Path::new(&snapshot_manager.snapshot_path).exists() { match snapshot_manager.load_snapshot().await { @@ -94,6 +99,11 @@ pub async fn run(settings: AuthServiceSettings) -> Result<()> { let snapshot_manager = snapshot_manager.clone(); let snapshot_handle = tokio::spawn(async move { + log::info!( + "Running snapshot task, interval {}", + settings.auth.snapshot_interval + ); + let mut interval = tokio::time::interval(std::time::Duration::from_secs( settings.auth.snapshot_interval, )); @@ -102,7 +112,8 @@ pub async fn run(settings: AuthServiceSettings) -> Result<()> { if let Err(e) = snapshot_manager.create_snapshot().await { log::error!("Failed to create snapshot: {}", e); } else { - log::debug!("Auth snapshot saved successfully"); + let len = snapshot_manager.len().await; + log::debug!("Auth snapshot saved successfully; {} connections", len); } } }); @@ -129,6 +140,7 @@ pub async fn run(settings: AuthServiceSettings) -> Result<()> { { let settings = settings.clone(); + let auth = auth.clone(); loop { let api_token = settings.api.token.clone(); @@ -154,6 +166,7 @@ pub async fn run(settings: AuthServiceSettings) -> Result<()> { { let mut shutdown = shutdown_tx.subscribe(); + let auth = auth.clone(); let auth_handle = tokio::spawn(async move { tokio::select! { @@ -164,31 +177,41 @@ pub async fn run(settings: AuthServiceSettings) -> Result<()> { tasks.push(auth_handle); }; - let api_token = settings.api.token.clone(); + log::info!("Running metrics task"); - let token = Arc::new(api_token.clone()); - if settings.debug.enabled { - log::debug!( - "Running debug server: localhost:{}", - settings.debug.web_port - ); + let auth_for_collect = auth.clone(); + let metrics_handle = tokio::spawn({ let mut shutdown = shutdown_tx.subscribe(); - let memory = memory.clone(); - let addr = settings - .debug - .web_server - .unwrap_or(Ipv4Addr::new(127, 0, 0, 1)); - let port = settings.debug.web_port; - let token = token.clone(); - - let debug_handle = tokio::spawn(async move { - tokio::select! { - _ = debug::start_ws_server(memory, addr, port, token) => {}, - _ = shutdown.recv() => {}, + async move { + loop { + tokio::select! { + _ = sleep(Duration::from_secs(settings.metrics.interval)) => { + auth_for_collect.collect_metrics().await; + }, + _ = shutdown.recv() => break, + } } - }); - tasks.push(debug_handle); - } + } + }); + + let auth_for_flush = auth.clone(); + log::info!("Running flush metrics task"); + let metrics_flush_handle = tokio::spawn({ + let mut shutdown = shutdown_tx.subscribe(); + async move { + loop { + tokio::select! { + _ = sleep(Duration::from_secs(settings.metrics.interval + 2)) => { + auth_for_flush.metrics.flush_to_zmq().await; + }, + _ = shutdown.recv() => break, + } + } + } + }); + + tasks.push(metrics_handle); + tasks.push(metrics_flush_handle); wait_all_tasks_or_ctrlc(tasks, shutdown_tx).await; Ok(()) diff --git a/src/bin/auth/core/tasks.rs b/src/bin/auth/core/tasks.rs index 66b6bd5e..9fce518f 100644 --- a/src/bin/auth/core/tasks.rs +++ b/src/bin/auth/core/tasks.rs @@ -5,12 +5,11 @@ use rkyv::AlignedVec; use rkyv::Infallible; use tokio::time::Duration; +use pony::metrics::Metrics; use pony::Action; use pony::BaseConnection as Connection; use pony::ConnectionBaseOp; use pony::Message; -use pony::NodeStorageOp; -use pony::SubscriptionOp; use pony::Topic; use pony::{PonyError, Result}; @@ -22,14 +21,13 @@ use super::AuthService; pub trait Tasks { async fn run_subscriber(&self) -> Result<()>; async fn handle_messages_batch(&self, msg: Vec) -> Result<()>; + async fn collect_metrics(&self); } #[async_trait] -impl Tasks for AuthService +impl Tasks for AuthService where - T: NodeStorageOp + Send + Sync + Clone, C: ConnectionBaseOp + Send + Sync + Clone + 'static + From, - S: SubscriptionOp + Send + Sync + Clone + 'static + std::cmp::PartialEq, { async fn run_subscriber(&self) -> Result<()> { let sub = self.subscriber.clone(); @@ -116,15 +114,9 @@ where msg.expires_at.map(Into::into), msg.subscription_id, ); - mem.connections - .add(&conn_id, conn.into()) - .map(|_| ()) - .map_err(|err| { - PonyError::Custom(format!( - "Failed to add conn {}: {}", - conn_id, err - )) - }) + mem.add(&conn_id, conn.into()).map(|_| ()).map_err(|err| { + PonyError::Custom(format!("Failed to add conn {}: {}", conn_id, err)) + }) } else { log::debug!("Skipped message {:?}", msg); Ok(()) @@ -132,7 +124,7 @@ where } Action::Delete => { - let _ = mem.connections.remove(&conn_id); + let _ = mem.remove(&conn_id); Ok(()) } @@ -144,4 +136,13 @@ where Ok(()) } + + async fn collect_metrics(&self) { + self.heartbeat().await; + self.bandwidth().await; + self.cpu_usage().await; + self.loadavg().await; + self.memory().await; + self.disk_usage().await; + } } diff --git a/src/config/settings.rs b/src/config/settings.rs index fdda2c44..e369dc85 100644 --- a/src/config/settings.rs +++ b/src/config/settings.rs @@ -9,267 +9,53 @@ use std::fs; use std::net::Ipv4Addr; use uuid::Uuid; -fn default_enabled() -> bool { - true -} fn default_disabled() -> bool { false } fn default_env() -> String { "dev".to_string() } -fn default_carbon_server() -> String { - "localhost:2003".to_string() -} -fn default_clickhouse_server() -> String { - "http://localhost:8123".to_string() -} -fn default_loglevel() -> String { - "debug".to_string() -} -fn default_zmq_sub_endpoint() -> String { - "tcp://localhost:3000".to_string() -} -fn default_zmq_pub_endpoint() -> String { - "tcp://*:3000".to_string() -} -fn default_xray_config_path() -> String { - "dev/xray-config.json".to_string() -} -fn default_wg_port() -> u16 { - 51820 -} -fn default_wg_interface() -> String { - "wg0".to_string() -} -fn default_pg_address() -> String { - "localhost".to_string() -} -fn default_pg_port() -> u16 { - 5432 -} -fn default_pg_db() -> String { - "postgres".to_string() -} -fn default_pg_username() -> String { - "postgres".to_string() -} -fn default_pg_password() -> String { - "password".to_string() -} -fn default_api_endpoint_address() -> String { - "http://localhost:3005".to_string() -} -fn default_uuid() -> Uuid { - Uuid::parse_str("9658b391-01cb-4031-a3f5-6cbdd749bcff").unwrap() -} -fn default_debug_web_server() -> Option { - Some(Ipv4Addr::new(127, 0, 0, 1)) -} -fn default_debug_web_port() -> u16 { - 3001 -} -fn default_auth_web_server() -> Option { - Some(Ipv4Addr::new(127, 0, 0, 1)) -} -fn default_auth_web_port() -> u16 { - 3005 -} -fn default_api_web_listen() -> Option { - Some(Ipv4Addr::new(127, 0, 0, 1)) -} -fn default_api_web_port() -> u16 { - 3005 -} -fn default_api_token() -> String { - "token".to_string() -} - -fn default_key_sign_token() -> Vec { - b"sign-token".to_vec() -} - -fn default_label() -> String { - "πŸ΄β€β˜ οΈπŸ΄β€β˜ οΈπŸ΄β€β˜ οΈ dev".to_string() -} - -fn default_stat_job_interval() -> u64 { - 60 -} - -fn default_snapshot_interval() -> u64 { - 30 -} - -fn default_snapshot_path() -> String { - "snapshots/snapshot.bin".to_string() -} - -fn default_metrics_interval() -> u64 { - 60 -} - -fn default_metrics_hb_interval() -> u64 { - 1 -} - -fn default_healthcheck_interval() -> u64 { - 60 -} -fn default_collect_conn_stat_interval() -> u64 { - 60 -} -fn default_db_sync_interval_sec() -> u64 { - 300 -} - -fn default_subscription_restore_interval_sec() -> u64 { - 60 -} - -fn default_subscription_expire_interval_sec() -> u64 { - 60 -} - -fn default_node_healthcheck_timeout() -> i16 { - 60 -} - -fn default_max_bandwidth_bps() -> i64 { - 100_000_000 -} - -fn default_web_host() -> String { - "http://localhost:8000".to_string() -} - -fn default_api_web_host() -> String { - "http://localhost:5005".to_string() -} - -fn default_h2_config_path() -> String { - "dev/h2.yaml".to_string() -} - -fn default_mtproto_port() -> u16 { - 8443 -} - -fn default_smtp_username() -> String { - "user".to_string() -} - -fn default_smtp_password() -> String { - "password".to_string() -} - -fn default_smtp_server() -> String { - "smtp.example.com".to_string() -} - -fn default_smtp_port() -> u16 { - 587 -} - -fn default_smtp_from() -> String { - "noreply@example.com".to_string() -} - -fn default_email_file() -> String { - "trials.csv".to_string() -} - -fn default_email_sign_token() -> Vec { - b"email-sign-token".to_vec() -} - -fn default_bonus_days() -> i64 { - 7 -} - -fn default_promo_codes() -> Vec { - vec![ - "FRKN.ORG".to_string(), - "mobile".to_string(), - "mobile-dev".to_string(), - ] -} #[derive(Clone, Debug, Deserialize, Default)] pub struct ApiServiceConfig { - #[serde(default = "default_api_web_listen")] pub address: Option, - #[serde(default = "default_api_web_port")] pub port: u16, - #[serde(default = "default_collect_conn_stat_interval")] - pub collect_conn_stat_interval: u64, - #[serde(default = "default_healthcheck_interval")] - pub healthcheck_interval: u64, - #[serde(default = "default_api_token")] pub token: String, - #[serde(default = "default_enabled")] - pub metrics_enabled: bool, - #[serde(default = "default_metrics_interval")] - pub metrics_interval: u64, - #[serde(default = "default_metrics_hb_interval")] - pub metrics_hb_interval: u64, - #[serde(default = "default_db_sync_interval_sec")] pub db_sync_interval_sec: u64, - #[serde(default = "default_web_host")] pub web_host: String, - #[serde(default = "default_api_web_host")] pub api_web_host: String, - #[serde(default = "default_subscription_restore_interval_sec")] pub subscription_restore_interval: u64, - #[serde(default = "default_subscription_expire_interval_sec")] pub subscription_expire_interval: u64, - #[serde(default = "default_node_healthcheck_timeout")] - pub node_health_check_timeout: i16, - #[serde(default = "default_key_sign_token")] pub key_sign_token: Vec, - #[serde(default = "default_bonus_days")] pub bonus_days: i64, - #[serde(default = "default_promo_codes")] pub promo_codes: Vec, + pub max_points: usize, + pub retention_seconds: i64, } #[derive(Clone, Debug, Deserialize, Default)] pub struct ApiAccessConfig { - #[serde(default = "default_api_endpoint_address")] pub endpoint: String, - #[serde(default = "default_api_token")] pub token: String, } #[derive(Clone, Debug, Deserialize, Default)] pub struct AuthServiceConfig { - #[serde(default = "default_snapshot_interval")] pub snapshot_interval: u64, - #[serde(default = "default_snapshot_path")] pub snapshot_path: String, - #[serde(default = "default_web_host")] pub web_host: String, - #[serde(default = "default_auth_web_server")] pub web_server: Option, - #[serde(default = "default_auth_web_port")] pub web_port: u16, - #[serde(default = "default_email_file")] pub email_file: String, - #[serde(default = "default_email_sign_token")] pub email_sign_token: Vec, } #[derive(Clone, Debug, Deserialize, Default)] pub struct SmtpConfig { - #[serde(default = "default_smtp_server")] pub server: String, - #[serde(default = "default_smtp_port")] pub port: u16, - #[serde(default = "default_smtp_username")] pub username: String, - #[serde(default = "default_smtp_password")] pub password: String, - #[serde(default = "default_smtp_from")] pub from: String, } @@ -277,47 +63,24 @@ pub struct SmtpConfig { pub struct AgentConfig { #[serde(default = "default_disabled")] pub local: bool, - #[serde(default = "default_enabled")] - pub metrics_enabled: bool, - #[serde(default = "default_metrics_interval")] - pub metrics_interval: u64, - #[serde(default = "default_metrics_hb_interval")] - pub metrics_hb_interval: u64, - #[serde(default = "default_enabled")] - pub stat_enabled: bool, - #[serde(default = "default_stat_job_interval")] - pub stat_job_interval: u64, - #[serde(default = "default_snapshot_interval")] pub snapshot_interval: u64, - #[serde(default = "default_snapshot_path")] pub snapshot_path: String, } -#[derive(Clone, Debug, Deserialize, Default)] -pub struct CarbonConfig { - #[serde(default = "default_carbon_server")] - pub address: String, -} - -#[derive(Clone, Debug, Deserialize, Default)] -pub struct ClickhouseConfig { - #[serde(default = "default_clickhouse_server")] - pub address: String, +#[derive(Clone, Default, Debug, Deserialize)] +pub struct MetricsTxConfig { + pub publisher: String, + pub interval: u64, } -#[derive(Clone, Debug, Deserialize, Default)] -pub struct DebugConfig { - #[serde(default = "default_enabled")] - pub enabled: bool, - #[serde(default = "default_debug_web_server")] - pub web_server: Option, - #[serde(default = "default_debug_web_port")] - pub web_port: u16, +#[derive(Clone, Default, Debug, Deserialize)] +pub struct MetricsRxConfig { + pub reciever: String, + pub topic: Vec, } #[derive(Clone, Debug, Deserialize, Default)] pub struct LoggingConfig { - #[serde(default = "default_loglevel")] pub level: String, } @@ -331,6 +94,7 @@ pub struct NodeConfig { pub label: String, pub max_bandwidth_bps: i64, pub cores: usize, + pub country: String, } #[derive(Clone, Debug, Deserialize, Default)] @@ -340,12 +104,10 @@ pub struct NodeConfigRaw { pub hostname: Option, pub default_interface: Option, pub address: Option, - #[serde(default = "default_uuid")] pub uuid: Uuid, - #[serde(default = "default_label")] pub label: String, - #[serde(default = "default_max_bandwidth_bps")] pub max_bandwidth_bps: i64, + pub country: String, } impl NodeConfig { @@ -406,7 +168,6 @@ impl NodeConfig { ))); } } else { - // Ни адрСс, Π½ΠΈ интСрфСйс Π½Π΅ ΡƒΠΊΠ°Π·Π°Π½Ρ‹ - ΠΈΡΠΏΠΎΠ»ΡŒΠ·ΡƒΠ΅ΠΌ Π΄Π΅Ρ„ΠΎΠ»Ρ‚Π½Ρ‹ΠΉ интСрфСйс match get_default_interface() { Ok(interface) => { if interface.ipv4.is_empty() { @@ -434,29 +195,24 @@ impl NodeConfig { label: raw.label, max_bandwidth_bps: raw.max_bandwidth_bps, cores: num_cpus, + country: raw.country, }) } } #[derive(Clone, Debug, Deserialize, Default)] pub struct PostgresConfig { - #[serde(default = "default_pg_address")] pub host: String, - #[serde(default = "default_pg_port")] pub port: u16, - #[serde(default = "default_pg_db")] pub db: String, - #[serde(default = "default_pg_username")] pub username: String, - #[serde(default = "default_pg_password")] pub password: String, } #[derive(Clone, Debug, Deserialize, Default)] pub struct XrayConfig { - #[serde(default = "default_enabled")] + #[serde(default = "default_disabled")] pub enabled: bool, - #[serde(default = "default_xray_config_path")] pub xray_config_path: String, } @@ -464,9 +220,7 @@ pub struct XrayConfig { pub struct WgConfig { #[serde(default = "default_disabled")] pub enabled: bool, - #[serde(default = "default_wg_port")] pub port: u16, - #[serde(default = "default_wg_interface")] pub interface: String, pub network: Option, pub privkey: Option, @@ -479,7 +233,6 @@ pub struct WgConfig { pub struct H2Config { #[serde(default = "default_disabled")] pub enabled: bool, - #[serde(default = "default_h2_config_path")] pub path: String, } @@ -487,14 +240,12 @@ pub struct H2Config { pub struct MtprotoConfig { #[serde(default = "default_disabled")] pub enabled: bool, - #[serde(default = "default_mtproto_port")] pub port: u16, pub secret: String, } #[derive(Clone, Debug, Deserialize, Default)] pub struct ZmqSubscriberConfig { - #[serde(default = "default_zmq_sub_endpoint")] pub endpoint: String, } @@ -511,7 +262,6 @@ impl ZmqSubscriberConfig { #[derive(Clone, Debug, Deserialize, Default)] pub struct ZmqPublisherConfig { - #[serde(default = "default_zmq_pub_endpoint")] pub endpoint: String, } @@ -551,25 +301,19 @@ pub struct ApiSettings { #[serde(default)] pub api: ApiServiceConfig, #[serde(default)] - pub debug: DebugConfig, - #[serde(default)] pub node: NodeConfigRaw, #[serde(default)] pub logging: LoggingConfig, #[serde(default)] pub zmq: ZmqPublisherConfig, #[serde(default)] - pub clickhouse: ClickhouseConfig, - #[serde(default)] pub pg: PostgresConfig, #[serde(default)] - pub carbon: CarbonConfig, + pub metrics: MetricsRxConfig, } #[derive(Clone, Debug, Deserialize)] pub struct AuthServiceSettings { - #[serde(default)] - pub debug: DebugConfig, #[serde(default)] pub logging: LoggingConfig, #[serde(default)] @@ -582,14 +326,12 @@ pub struct AuthServiceSettings { pub api: ApiAccessConfig, #[serde(default)] pub smtp: SmtpConfig, + #[serde(default)] + pub metrics: MetricsTxConfig, } #[derive(Clone, Debug, Deserialize)] pub struct AgentSettings { - #[serde(default)] - pub debug: DebugConfig, - #[serde(default)] - pub carbon: CarbonConfig, #[serde(default)] pub logging: LoggingConfig, #[serde(default)] @@ -608,6 +350,8 @@ pub struct AgentSettings { pub node: NodeConfigRaw, #[serde(default)] pub api: ApiAccessConfig, + #[serde(default)] + pub metrics: MetricsTxConfig, } impl Settings for AgentSettings { diff --git a/src/http/debug.rs b/src/http/debug.rs deleted file mode 100644 index f68e9c33..00000000 --- a/src/http/debug.rs +++ /dev/null @@ -1,198 +0,0 @@ -use core::fmt; -use futures::{SinkExt, StreamExt}; -use serde::{Deserialize, Serialize}; -use std::net::{IpAddr, Ipv4Addr, SocketAddr}; -use std::sync::Arc; -use tokio::sync::RwLock; -use warp::http::header::SEC_WEBSOCKET_PROTOCOL; -use warp::ws::Message; -use warp::ws::Ws; -use warp::Filter; -use warp::Rejection; - -use crate::http::Unauthorized; - -use super::super::memory::cache::Cache; -use super::super::memory::connection::op::base::Operations as ConnectionBaseOp; -use super::super::memory::storage::connection::BaseOp as ConnectionStorageBaseOp; -use super::super::memory::storage::node::Operations as NodeStorageOp; -use crate::SubscriptionOp; - -enum Kind { - Conn, - Conns, - Nodes, - Subs, -} - -impl fmt::Display for Kind { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Kind::Conn => write!(f, "conn"), - Kind::Conns => write!(f, "conns"), - Kind::Nodes => write!(f, "nodes"), - Kind::Subs => write!(f, "subs"), - } - } -} - -#[derive(Deserialize, Serialize, Debug)] -pub struct Request { - pub kind: String, - pub message: String, - pub conn_id: Option, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct Response { - pub kind: String, - pub data: serde_json::Value, - pub len: usize, -} - -pub async fn start_ws_server( - memory: Arc>>, - ipaddr: Ipv4Addr, - port: u16, - expected_token: Arc, -) where - N: NodeStorageOp + Sync + Send + Clone + 'static, - S: SubscriptionOp + Sync + Send + Clone + 'static + std::cmp::PartialEq + serde::Serialize, - C: ConnectionBaseOp + Sync + Send + Clone + 'static + std::fmt::Display, -{ - let health_check = warp::path("health-check").map(|| "Server OK"); - - let ws_route = { - let expected_token = expected_token.clone(); - warp::path("ws") - .and(warp::ws()) - .and(warp::header::optional::( - SEC_WEBSOCKET_PROTOCOL.as_str(), - )) - .and_then(move |ws: Ws, token: Option| { - let memory = memory.clone(); - let expected_token = expected_token.clone(); - async move { - match token { - Some(t) if t == *expected_token => { - Ok::<_, Rejection>(ws.on_upgrade(move |socket| { - handle_debug_connection::(socket, memory) - })) - } - _ => { - log::warn!( - "Unauthorized WebSocket connection attempt with token: {:?}", - token - ); - Err(warp::reject::custom(Unauthorized)) - } - } - } - }) - }; - - let routes = health_check - .or(ws_route) - .with(warp::cors().allow_any_origin()); - - warp::serve(routes) - .run(SocketAddr::new(IpAddr::V4(ipaddr), port)) - .await; -} - -pub async fn handle_debug_connection( - socket: warp::ws::WebSocket, - memory: Arc>>, -) where - N: NodeStorageOp + Sync + Send + Clone + 'static, - C: ConnectionBaseOp + Sync + Send + Clone + 'static + fmt::Display, - S: SubscriptionOp + Sync + Send + Clone + 'static + std::cmp::PartialEq + serde::Serialize, -{ - let (mut sender, mut receiver) = socket.split(); - - while let Some(Ok(msg)) = receiver.next().await { - let message = match msg.to_str() { - Ok(s) => s, - Err(_) => continue, - }; - - let req: Request = match serde_json::from_str(message) { - Ok(r) => r, - Err(err) => { - log::error!("Invalid request format: {}", err); - continue; - } - }; - - // COMMENT(@qezz): A `match` would probably work better here. - if req.kind == "get_connections" { - let memory = memory.read().await; - let conns: Vec<_> = memory - .connections - .iter() - .map(|(k, v)| (k, v.get_proto().proto())) - .collect(); - let data = serde_json::to_string(&conns).unwrap(); - let response = Response { - kind: Kind::Conns.to_string(), - data: serde_json::Value::String(data), - len: memory.connections.len(), - }; - let response_str = serde_json::to_string(&response).unwrap(); - sender.send(Message::text(response_str)).await.unwrap(); - } else if req.kind == "get_nodes" { - let memory = memory.read().await; - let nodes = memory.nodes.all_json(); - - let data = match serde_json::to_string(&nodes) { - Ok(json) => json, - Err(err) => { - log::error!("Failed to serialize nodes: {}", err); - continue; - } - }; - - let response = Response { - kind: Kind::Nodes.to_string(), - data: serde_json::Value::String(data), - len: memory.connections.len(), - }; - let response_str = serde_json::to_string(&response).unwrap(); - sender.send(Message::text(response_str)).await.unwrap(); - } else if req.kind == "get_conn_info" { - if let Some(conn_id) = req.conn_id { - let memory = memory.read().await; - if let Some(conn) = memory.connections.get(&conn_id) { - let response = Response { - kind: Kind::Conn.to_string(), - data: serde_json::Value::String(conn.to_string()), - len: 1, - }; - let response_str = serde_json::to_string(&response).unwrap(); - sender.send(Message::text(response_str)).await.unwrap(); - } - } - } else if req.kind == "get_subscriptions" { - let memory = memory.read().await; - - let subs: Vec<_> = memory.subscriptions.values().cloned().collect(); - - let data = match serde_json::to_value(&subs) { - Ok(v) => v, - Err(err) => { - log::error!("Failed to serialize subscriptions: {}", err); - continue; - } - }; - - let response = Response { - kind: Kind::Subs.to_string(), - len: subs.len(), - data, - }; - - let response_str = serde_json::to_string(&response).unwrap(); - sender.send(Message::text(response_str)).await.unwrap(); - } - } -} diff --git a/src/http/helpers.rs b/src/http/helpers.rs index 1868389c..5cf2b256 100644 --- a/src/http/helpers.rs +++ b/src/http/helpers.rs @@ -1,12 +1,7 @@ -use serde::{Deserialize, Serialize}; use warp::http::StatusCode; use warp::reply::Json; -use crate::memory::connection::conn::Conn as Connection; -use crate::memory::connection::stat::Stat as ConnectionStat; -use crate::memory::key::Key; -use crate::memory::subscription::Subscription; -use crate::memory::tag::ProtoTag as Tag; +use super::response::{Instance, InstanceWithId, ResponseMessage}; pub fn bad_request(msg: &str) -> warp::reply::WithStatus { let resp = ResponseMessage::> { @@ -66,26 +61,3 @@ pub fn success_response( }; warp::reply::with_status(warp::reply::json(&resp), StatusCode::OK) } - -#[derive(Serialize)] -struct ResponseMessage { - status: u16, - message: String, - response: T, -} - -#[derive(Serialize, Deserialize, Debug)] -pub struct InstanceWithId { - pub id: uuid::Uuid, - pub instance: T, -} - -#[derive(Debug, Serialize, Deserialize)] -pub enum Instance { - Connection(Connection), - Subscription(Subscription), - Stat(Vec<(uuid::Uuid, ConnectionStat, Tag)>), - Connections(Vec<(uuid::Uuid, Connection)>), - Key(Key), - None, -} diff --git a/src/http/mod.rs b/src/http/mod.rs index 31c6510d..b4252188 100644 --- a/src/http/mod.rs +++ b/src/http/mod.rs @@ -2,9 +2,9 @@ use serde::{Deserialize, Serialize}; use warp::reject; use warp::reject::Reject; -pub mod debug; pub mod filters; pub mod helpers; +pub mod response; #[derive(Debug)] pub struct AuthError(pub String); @@ -26,11 +26,6 @@ pub struct IdResponse { pub id: uuid::Uuid, } -#[derive(Debug)] -struct Unauthorized; - -impl Reject for Unauthorized {} - #[derive(Debug)] pub struct IpParseError(pub std::net::AddrParseError); diff --git a/src/http/response.rs b/src/http/response.rs new file mode 100644 index 00000000..ecc1e34c --- /dev/null +++ b/src/http/response.rs @@ -0,0 +1,30 @@ +use crate::memory::connection::conn::Conn as Connection; +use crate::memory::connection::stat::Stat as ConnectionStat; +use crate::memory::key::Key; +use crate::memory::subscription::Subscription; +use crate::memory::tag::ProtoTag as Tag; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, Serialize)] +pub struct ResponseMessage { + pub status: u16, + pub message: String, + pub response: T, +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct InstanceWithId { + pub id: uuid::Uuid, + pub instance: T, +} + +#[derive(Debug, Serialize, Deserialize)] +pub enum Instance { + Connection(Connection), + Subscription(Subscription), + Stat(Vec<(uuid::Uuid, ConnectionStat, Tag)>), + Connections(Vec<(uuid::Uuid, Connection)>), + Key(Key), + Count(usize), + None, +} diff --git a/src/lib.rs b/src/lib.rs index 61f70432..9f101e9c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,7 +13,6 @@ pub mod zmq; pub use crate::error::{PonyError, Result, SyncError}; -pub use memory::cache::Cache as MemoryCache; pub use memory::connection::base::Base as BaseConnection; pub use memory::connection::conn::Conn as Connection; pub use memory::subscription::Subscription; diff --git a/src/memory/cache.rs b/src/memory/cache.rs deleted file mode 100644 index 72892496..00000000 --- a/src/memory/cache.rs +++ /dev/null @@ -1,96 +0,0 @@ -use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize}; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::fmt; -use std::ops::{Deref, DerefMut}; - -use super::connection::op::api::Operations as ConnectionApiOp; -use super::connection::op::base::Operations as ConnectionBaseOp; -use super::node::Node; -use super::storage::node::Operations as NodeStorageOp; -use super::subscription::Operations as SubscriptionOp; -use super::subscription::Subscriptions; - -#[derive(Archive, Deserialize, Serialize, RkyvDeserialize, RkyvSerialize, Debug, Clone)] -#[archive(check_bytes)] -pub struct Connections(pub HashMap); - -impl Default for Connections { - fn default() -> Self { - Connections(HashMap::new()) - } -} - -impl fmt::Display for Connections { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - for (uuid, conn) in &self.0 { - writeln!(f, "{} => {}", uuid, conn)?; - } - Ok(()) - } -} - -impl Deref for Connections { - type Target = HashMap; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} - -impl DerefMut for Connections { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.0 - } -} - -#[derive(Deserialize, Serialize, Debug, Clone)] -pub struct Cache -where - T: Send + Sync + Clone + 'static, - C: Send + Sync + Clone + 'static, - S: Send + Sync + Clone + 'static, -{ - pub connections: Connections, - pub subscriptions: Subscriptions, - pub nodes: T, -} - -impl Default for Cache -where - T: NodeStorageOp + Sync + Send + Clone + 'static, - C: ConnectionApiOp + ConnectionBaseOp + Clone + Send + Sync + 'static + PartialEq, - S: SubscriptionOp + Clone + Send + Sync + 'static, -{ - fn default() -> Self { - Self::new() - } -} - -impl Cache -where - T: NodeStorageOp + Sync + Send + Clone + 'static, - C: ConnectionApiOp + ConnectionBaseOp + Clone + Send + Sync + 'static + PartialEq, - S: SubscriptionOp + Clone + Send + Sync + 'static, -{ - pub fn new() -> Self { - Cache { - nodes: T::default(), - connections: Connections::default(), - subscriptions: Subscriptions::default(), - } - } -} -impl Cache -where - C: ConnectionBaseOp + Send + Sync + Clone + 'static + PartialEq, - S: SubscriptionOp + Clone + Send + Sync + 'static, -{ - pub fn with_node(node: Node) -> Self { - Self { - nodes: node, - connections: Connections::default(), - subscriptions: Subscriptions::default(), - } - } -} diff --git a/src/memory/connection/base.rs b/src/memory/connection/base.rs index 3d23b517..f15a2430 100644 --- a/src/memory/connection/base.rs +++ b/src/memory/connection/base.rs @@ -6,7 +6,6 @@ use std::fmt; use super::super::connection::conn::Conn; use super::super::connection::proto::Proto; -use super::super::connection::stat::Stat as ConnectionStat; use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize}; @@ -14,7 +13,6 @@ use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize}; Archive, Deserialize, Serialize, RkyvDeserialize, RkyvSerialize, Debug, Clone, PartialEq, )] pub struct Base { - pub stat: ConnectionStat, pub created_at: DateTime, pub modified_at: DateTime, pub expires_at: Option>, @@ -32,7 +30,6 @@ impl Base { let now = Utc::now(); Self { - stat: ConnectionStat::default(), created_at: now, modified_at: now, expires_at, @@ -45,13 +42,7 @@ impl Base { impl From for Base { fn from(conn: Conn) -> Self { - let conn_stat = ConnectionStat { - online: conn.stat.online, - uplink: conn.stat.uplink, - downlink: conn.stat.downlink, - }; Base { - stat: conn_stat, created_at: conn.created_at, modified_at: conn.modified_at, expires_at: conn.expires_at, @@ -64,13 +55,7 @@ impl From for Base { impl From<&Conn> for Base { fn from(conn: &Conn) -> Self { - let conn_stat = ConnectionStat { - online: conn.stat.online, - uplink: conn.stat.uplink, - downlink: conn.stat.downlink, - }; Base { - stat: conn_stat, created_at: conn.created_at, modified_at: conn.modified_at, expires_at: conn.expires_at, @@ -85,7 +70,6 @@ impl fmt::Display for Base { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { writeln!(f, "Connection Base {{")?; - writeln!(f, " conn stat: {}", self.stat)?; writeln!(f, " created_at: {},", self.created_at)?; writeln!(f, " modified_at: {},", self.modified_at)?; writeln!(f, " expires_at: {:?},", self.expires_at)?; diff --git a/src/memory/connection/conn.rs b/src/memory/connection/conn.rs index 02c3e266..70ef8c35 100644 --- a/src/memory/connection/conn.rs +++ b/src/memory/connection/conn.rs @@ -9,13 +9,11 @@ use serde::Serialize; use super::op::api::Operations as ApiOps; use super::op::base::Operations as BasOps; use super::proto::Proto; -use super::stat::Stat; #[derive(Serialize, Deserialize, Debug, Clone)] pub struct Conn { pub env: String, pub proto: Proto, - pub stat: Stat, pub subscription_id: Option, pub created_at: DateTime, pub modified_at: DateTime, @@ -33,7 +31,6 @@ impl fmt::Display for Conn { writeln!(f, " subscription_id: None,")?; } writeln!(f, " env: {},", self.env)?; - writeln!(f, " conn stat: {},", self.stat)?; writeln!(f, " created_at: {},", self.created_at)?; writeln!(f, " modified_at: {},", self.modified_at)?; writeln!(f, " expires_at: {:?},", self.expires_at)?; @@ -57,7 +54,6 @@ impl Conn { pub fn new( env: &str, subscription_id: Option, - stat: Stat, proto: Proto, expires_at: Option>, ) -> Self { @@ -65,7 +61,6 @@ impl Conn { Self { env: env.to_string(), - stat, created_at: now, modified_at: now, expires_at, @@ -87,7 +82,6 @@ pub struct ConnWithId { pub struct ConnPatch { pub env: Option, pub proto: Option, - pub stat: Option, pub subscription_id: Option, pub created_at: Option, pub modified_at: Option, diff --git a/src/memory/connection/mod.rs b/src/memory/connection/mod.rs index 18dee4ea..c66e2195 100644 --- a/src/memory/connection/mod.rs +++ b/src/memory/connection/mod.rs @@ -1,6 +1,45 @@ +use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::fmt; +use std::ops::{Deref, DerefMut}; + pub(crate) mod base; pub(crate) mod conn; pub(crate) mod op; pub(crate) mod proto; pub mod stat; pub mod wireguard; + +#[derive(Archive, Deserialize, Serialize, RkyvDeserialize, RkyvSerialize, Debug, Clone)] +#[archive(check_bytes)] +pub struct Connections(pub HashMap); + +impl Default for Connections { + fn default() -> Self { + Connections(HashMap::new()) + } +} + +impl fmt::Display for Connections { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for (uuid, conn) in &self.0 { + writeln!(f, "{} => {}", uuid, conn)?; + } + Ok(()) + } +} + +impl Deref for Connections { + type Target = HashMap; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for Connections { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} diff --git a/src/memory/connection/op/base.rs b/src/memory/connection/op/base.rs index c21778ad..61e114db 100644 --- a/src/memory/connection/op/base.rs +++ b/src/memory/connection/op/base.rs @@ -4,22 +4,10 @@ use chrono::Utc; use super::super::base::Base; use super::super::conn::Conn; use super::super::proto::Proto; -use super::super::stat::Stat; use super::super::wireguard::Param as WgParam; use crate::error::{PonyError, Result}; pub trait Operations { - fn get_uplink(&self) -> i64; - fn set_uplink(&mut self, v: i64); - fn reset_uplink(&mut self); - - fn get_downlink(&self) -> i64; - fn set_downlink(&mut self, v: i64); - fn reset_downlink(&mut self); - - fn get_online(&self) -> i64; - fn set_online(&mut self, v: i64); - fn get_modified_at(&self) -> DateTime; fn set_modified_at(&mut self); @@ -32,8 +20,6 @@ pub trait Operations { fn set_deleted(&mut self, v: bool); fn get_deleted(&self) -> bool; - fn as_conn_stat(&self) -> Stat; - fn get_wireguard(&self) -> Option<&WgParam>; fn get_wireguard_node_id(&self) -> Option; fn get_password(&self) -> Option; @@ -42,33 +28,6 @@ pub trait Operations { } impl Operations for Base { - fn get_uplink(&self) -> i64 { - self.stat.uplink - } - fn set_uplink(&mut self, v: i64) { - self.stat.uplink = v; - } - fn reset_uplink(&mut self) { - self.stat.uplink = 0; - } - - fn get_downlink(&self) -> i64 { - self.stat.downlink - } - fn set_downlink(&mut self, v: i64) { - self.stat.downlink = v; - } - fn reset_downlink(&mut self) { - self.stat.downlink = 0; - } - - fn get_online(&self) -> i64 { - self.stat.online - } - fn set_online(&mut self, v: i64) { - self.stat.online = v; - } - fn get_modified_at(&self) -> DateTime { self.modified_at } @@ -91,14 +50,6 @@ impl Operations for Base { self.proto.clone() } - fn as_conn_stat(&self) -> Stat { - Stat { - uplink: self.stat.uplink, - downlink: self.stat.downlink, - online: self.stat.online, - } - } - fn set_deleted(&mut self, v: bool) { self.is_deleted = v; } @@ -148,33 +99,6 @@ impl Operations for Base { } impl Operations for Conn { - fn get_uplink(&self) -> i64 { - self.stat.uplink - } - fn set_uplink(&mut self, v: i64) { - self.stat.uplink = v; - } - fn reset_uplink(&mut self) { - self.stat.uplink = 0; - } - - fn get_downlink(&self) -> i64 { - self.stat.downlink - } - fn set_downlink(&mut self, v: i64) { - self.stat.downlink = v; - } - fn reset_downlink(&mut self) { - self.stat.downlink = 0; - } - - fn get_online(&self) -> i64 { - self.stat.online - } - fn set_online(&mut self, v: i64) { - self.stat.online = v; - } - fn get_modified_at(&self) -> DateTime { self.modified_at } @@ -204,14 +128,6 @@ impl Operations for Conn { self.is_deleted } - fn as_conn_stat(&self) -> Stat { - Stat { - uplink: self.stat.uplink, - downlink: self.stat.downlink, - online: self.stat.online, - } - } - fn get_wireguard_node_id(&self) -> Option { match &self.proto { Proto::Wireguard { node_id, .. } => Some(*node_id), diff --git a/src/memory/connection/stat.rs b/src/memory/connection/stat.rs index 9b59ff7e..06a9dc58 100644 --- a/src/memory/connection/stat.rs +++ b/src/memory/connection/stat.rs @@ -2,9 +2,6 @@ use serde::Deserialize; use serde::Serialize; use std::fmt; -use super::super::stat::Kind; -use crate::metrics::Metric; - use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize}; #[derive( @@ -24,22 +21,6 @@ pub struct Stat { pub online: i64, } -impl Stat { - pub fn from_metrics + Clone + Copy>(metrics: Vec>) -> Self { - let mut stat = Stat::default(); - for metric in metrics { - let value = metric.value; - match metric.stat_type() { - Kind::Uplink => stat.uplink = value.into(), - Kind::Downlink => stat.downlink = value.into(), - Kind::Online => stat.online = value.into(), - Kind::Unknown => {} - } - } - stat - } -} - impl fmt::Display for Stat { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( diff --git a/src/memory/key.rs b/src/memory/key.rs index 2d41385a..f22f1efc 100644 --- a/src/memory/key.rs +++ b/src/memory/key.rs @@ -303,9 +303,9 @@ mod tests { let mut chars: Vec = code_str.chars().collect(); - for i in 0..chars.len() { - if chars[i] != '-' { - chars[i] = if chars[i] == 'A' { 'B' } else { 'A' }; + for c in chars.iter_mut() { + if *c != '-' { + *c = if *c == 'A' { 'B' } else { 'A' }; break; } } diff --git a/src/memory/mod.rs b/src/memory/mod.rs index 618fa465..21d145c4 100644 --- a/src/memory/mod.rs +++ b/src/memory/mod.rs @@ -1,4 +1,3 @@ -pub mod cache; pub mod connection; pub mod key; pub mod node; diff --git a/src/memory/node.rs b/src/memory/node.rs index ef030c7f..54ef68c4 100644 --- a/src/memory/node.rs +++ b/src/memory/node.rs @@ -1,6 +1,9 @@ use std::fmt; use std::str::FromStr; -use std::{collections::HashMap, net::Ipv4Addr}; +use std::{ + collections::{BTreeMap, HashMap}, + net::Ipv4Addr, +}; use chrono::DateTime; use chrono::Utc; @@ -11,7 +14,6 @@ use super::tag::ProtoTag as Tag; use crate::config::h2::H2Settings; use crate::config::settings::{MtprotoConfig, NodeConfig}; use crate::config::wireguard::WireguardSettings; -use crate::config::xray::InboundResponse; use crate::config::xray::{Config as XrayConfig, Inbound}; #[derive(Clone, Debug, Deserialize, Serialize, Copy, ToSql, FromSql)] @@ -57,11 +59,20 @@ pub struct NodeResponse { pub env: String, pub hostname: String, pub address: Ipv4Addr, - pub inbounds: HashMap, + pub inbounds: Vec, pub status: Status, pub label: String, pub cores: usize, pub max_bandwidth_bps: i64, + pub metrics: Vec, + pub country: String, +} + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub struct NodeMetricInfo { + pub key: String, + pub name: String, + pub tags: BTreeMap, } #[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] @@ -78,6 +89,7 @@ pub struct Node { pub inbounds: HashMap, pub cores: usize, pub max_bandwidth_bps: i64, + pub country: String, } impl Node { @@ -167,27 +179,42 @@ impl Node { inbounds, cores: settings.cores, max_bandwidth_bps: settings.max_bandwidth_bps, + country: settings.country, } } + pub fn get_base_tags(&self) -> BTreeMap { + let mut tags = BTreeMap::new(); + tags.insert("env".to_string(), self.env.clone()); + tags.insert("hostname".to_string(), self.hostname.clone()); + tags.insert("label".to_string(), self.label.clone()); + tags.insert("address".to_string(), self.address.to_string()); + tags.insert("label".to_string(), self.label.clone()); + tags.insert("interface".to_string(), self.interface.clone()); + tags.insert("cores".to_string(), self.cores.to_string()); + tags.insert( + "max_bandwidth_bps".to_string(), + self.max_bandwidth_bps.to_string(), + ); + tags.insert("country".to_string(), self.country.clone()); + tags + } + pub fn as_node_response(&self) -> NodeResponse { - let inbound_response = self - .inbounds - .clone() - .into_iter() - .map(|inbound| (inbound.0, inbound.1.as_inbound_response())) - .collect(); + let tags: Vec = self.inbounds.keys().cloned().collect(); NodeResponse { env: self.env.clone(), hostname: self.hostname.clone(), address: self.address, uuid: self.uuid, - inbounds: inbound_response, + inbounds: tags, status: self.status, label: self.label.clone(), cores: self.cores, max_bandwidth_bps: self.max_bandwidth_bps, + metrics: [].to_vec(), + country: self.country.clone(), } } diff --git a/src/memory/snapshot.rs b/src/memory/snapshot.rs index 818bb0ac..add744ac 100644 --- a/src/memory/snapshot.rs +++ b/src/memory/snapshot.rs @@ -11,9 +11,8 @@ use tokio::sync::RwLock; use crate::error::Result; -use super::cache::Cache; -use super::cache::Connections; use super::connection::conn::Conn; +use super::connection::Connections; #[derive(Archive, Deserialize, Serialize, SerdeDeserialize, SerdeSerialize, Debug, Clone)] #[archive(check_bytes)] @@ -26,24 +25,21 @@ where pub version: u32, } -pub struct SnapshotManager +pub struct SnapshotManager where C: Archive + Send + Sync + Clone + 'static, - N: Send + Sync + Clone + 'static, - S: Send + Sync + Clone + 'static, { pub snapshot_path: String, - pub memory: Arc>>, + pub memory: Arc>, } -impl Clone for SnapshotManager +impl Clone for SnapshotManager where C: Archive + Send + Sync + Clone + 'static - + std::convert::From + rkyv::Serialize< rkyv::ser::serializers::CompositeSerializer< rkyv::ser::serializers::AlignedSerializer, @@ -54,8 +50,6 @@ where rkyv::ser::serializers::SharedSerializeMap, >, >, - N: Send + Sync + Clone, - S: Send + Sync + Clone + 'static, { fn clone(&self) -> Self { SnapshotManager { @@ -65,7 +59,7 @@ where } } -impl SnapshotManager +impl SnapshotManager> where C: Archive + Send @@ -83,10 +77,8 @@ where rkyv::ser::serializers::SharedSerializeMap, >, >, - N: Send + Sync + Clone + 'static, - S: Send + Sync + Clone + 'static, { - pub fn new(snapshot_path: String, memory: Arc>>) -> Self { + pub fn new(snapshot_path: String, memory: Arc>>) -> Self { Self { snapshot_path, memory, @@ -95,12 +87,13 @@ where pub async fn create_snapshot(&self) -> Result<()> { let memory_guard = self.memory.read().await; - let connections = memory_guard.connections.clone(); + let memory = memory_guard.clone(); + let timestamp = Utc::now().timestamp() as u64; drop(memory_guard); let snapshot = SnapshotData { - timestamp: Utc::now().timestamp() as u64, - memory: connections, + timestamp, + memory, version: 1, }; @@ -127,7 +120,7 @@ where let snapshot: SnapshotData = with.into_inner(); let mut memory_guard = self.memory.write().await; - memory_guard.connections = snapshot.memory.clone(); + memory_guard.0 = snapshot.memory.0.clone(); Ok(Some(snapshot.timestamp)) } @@ -145,12 +138,12 @@ where pub async fn len(&self) -> usize { let mem = self.memory.read().await; - mem.connections.0.len() + mem.0.len() } pub async fn is_empty(&self) -> bool { let mem = self.memory.read().await; - mem.connections.0.len() == 0 + mem.0.len() == 0 } } diff --git a/src/memory/storage/connection.rs b/src/memory/storage/connection.rs index 2f7cf1e4..7e7d7a99 100644 --- a/src/memory/storage/connection.rs +++ b/src/memory/storage/connection.rs @@ -1,12 +1,10 @@ use std::collections::hash_map::Entry; -use super::super::cache::Connections; use super::super::connection::conn::Conn as Connection; use super::super::connection::conn::ConnPatch as ConnectionPatch; use super::super::connection::op::api::Operations as ConnectionApiOp; use super::super::connection::op::base::Operations as ConnectionBaseOp; -use super::super::connection::stat::Stat as ConnectionStat; -use super::super::stat::Kind as StatKind; +use super::super::connection::Connections; use super::super::storage::Status as OperationStatus; use super::super::tag::ProtoTag as Tag; use crate::error::{PonyError, Result}; @@ -17,7 +15,6 @@ where C: Clone + Send + Sync + 'static, { fn add(&mut self, conn_id: &uuid::Uuid, new_conn: C) -> Result; - fn reset_stat(&mut self, conn_id: &uuid::Uuid, stat: StatKind); fn get_by_subscription_id(&self, subscription_id: &uuid::Uuid) -> Option>; fn apply_update(conn: &mut Connection, patch: ConnectionPatch) -> Option; } @@ -31,12 +28,6 @@ where fn add(&mut self, conn_id: &uuid::Uuid, new_conn: C) -> Result; fn remove(&mut self, conn_id: &uuid::Uuid) -> Result<()>; fn get(&self, conn_id: &uuid::Uuid) -> Option; - fn update_stat(&mut self, conn_id: &uuid::Uuid, stat: ConnectionStat) -> Result<()>; - fn reset_stat(&mut self, conn_id: &uuid::Uuid, stat: StatKind); - fn update_uplink(&mut self, conn_id: &uuid::Uuid, new_uplink: i64) -> Result<()>; - fn update_downlink(&mut self, conn_id: &uuid::Uuid, new_downlink: i64) -> Result<()>; - fn update_online(&mut self, conn_id: &uuid::Uuid, new_online: i64) -> Result<()>; - fn update_stats(&mut self, conn_id: &uuid::Uuid, stats: ConnectionStat) -> Result<()>; fn validate_token(&self, token: &uuid::Uuid) -> Option; } @@ -83,69 +74,6 @@ where fn get(&self, conn_id: &uuid::Uuid) -> Option { self.0.get(conn_id).cloned() } - - fn update_stats(&mut self, conn_id: &uuid::Uuid, stats: ConnectionStat) -> Result<()> { - let conn = self - .get_mut(conn_id) - .ok_or(PonyError::Custom("Conn not found".into()))?; - - conn.set_online(stats.online); - conn.set_uplink(stats.uplink); - conn.set_downlink(stats.downlink); - - Ok(()) - } - - fn update_stat(&mut self, conn_id: &uuid::Uuid, stat: ConnectionStat) -> Result<()> { - let conn = self - .get_mut(conn_id) - .ok_or(PonyError::Custom("Conn not found".into()))?; - conn.set_uplink(stat.uplink); - conn.set_downlink(stat.downlink); - conn.set_online(stat.online); - conn.set_modified_at(); - Ok(()) - } - - fn reset_stat(&mut self, conn_id: &uuid::Uuid, stat: StatKind) { - if let Some(conn) = self.get_mut(conn_id) { - match stat { - StatKind::Uplink => conn.reset_uplink(), - StatKind::Downlink => conn.reset_downlink(), - StatKind::Online | StatKind::Unknown => {} - } - } - } - - fn update_uplink(&mut self, conn_id: &uuid::Uuid, new_uplink: i64) -> Result<()> { - if let Some(conn) = self.get_mut(conn_id) { - conn.set_uplink(new_uplink); - conn.set_modified_at(); - Ok(()) - } else { - Err(PonyError::Custom(format!("Conn not found: {}", conn_id))) - } - } - - fn update_downlink(&mut self, conn_id: &uuid::Uuid, new_downlink: i64) -> Result<()> { - if let Some(conn) = self.get_mut(conn_id) { - conn.set_downlink(new_downlink); - conn.set_modified_at(); - Ok(()) - } else { - Err(PonyError::Custom(format!("Conn not found: {}", conn_id))) - } - } - - fn update_online(&mut self, conn_id: &uuid::Uuid, new_online: i64) -> Result<()> { - if let Some(conn) = self.get_mut(conn_id) { - conn.set_online(new_online); - conn.set_modified_at(); - Ok(()) - } else { - Err(PonyError::Custom(format!("Conn not found: {}", conn_id))) - } - } } impl ApiOp for Connections @@ -234,14 +162,4 @@ where Some(conns) } } - - fn reset_stat(&mut self, conn_id: &uuid::Uuid, stat: StatKind) { - if let Some(conn) = self.get_mut(conn_id) { - match stat { - StatKind::Uplink => conn.reset_uplink(), - StatKind::Downlink => conn.reset_downlink(), - StatKind::Online | StatKind::Unknown => {} - } - } - } } diff --git a/src/memory/storage/node.rs b/src/memory/storage/node.rs index 50f7a7cc..07c287c9 100644 --- a/src/memory/storage/node.rs +++ b/src/memory/storage/node.rs @@ -3,8 +3,8 @@ use rand::thread_rng; use serde_json::json; use std::collections::HashMap; -use super::super::cache::Connections; use super::super::connection::op::base::Operations as ConnectionBaseOp; +use super::super::connection::Connections; use super::super::node::Node; use super::super::storage::Status as OperationStatus; use super::super::tag::ProtoTag as Tag; @@ -19,7 +19,6 @@ pub trait Operations { fn get_by_env(&self, env: &str) -> Option>; fn get_by_id(&self, id: &uuid::Uuid) -> Option; fn get(&self, env: &str, uuid: &uuid::Uuid) -> Option<&Node>; - fn get_self(&self) -> Option; fn get_mut(&mut self, env: &str, uuid: &uuid::Uuid) -> Option<&mut Node>; fn update_node_uplink( &mut self, @@ -52,96 +51,6 @@ pub trait Operations { C: ConnectionBaseOp; } -impl Operations for Node { - fn iter_nodes(&self) -> Box + '_> { - Box::new(std::iter::empty()) - } - fn add(&mut self, _new_node: Node) -> Result { - Err(PonyError::Custom( - "Cannot add node to single Node instance".into(), - )) - } - fn get_by_env(&self, _env: &str) -> Option> { - Some(vec![self.clone()]) - } - fn all(&self) -> Option> { - Some(vec![self.clone()]) - } - fn all_json(&self) -> serde_json::Value { - serde_json::to_value(vec![self]).unwrap_or_else(|_| serde_json::json!([])) - } - fn get_self(&self) -> Option { - Some(self.clone()) - } - fn get(&self, _env: &str, _uuid: &uuid::Uuid) -> Option<&Node> { - None - } - - fn get_mut(&mut self, _env: &str, _uuid: &uuid::Uuid) -> Option<&mut Node> { - None - } - - fn update_node_uplink( - &mut self, - tag: &Tag, - new_uplink: i64, - _env: &str, - node_id: &uuid::Uuid, - ) -> Result<()> { - if &self.uuid == node_id { - self.update_uplink(tag, new_uplink)?; - Ok(()) - } else { - Err(PonyError::Custom("Node ID does not match".into())) - } - } - - fn update_node_downlink( - &mut self, - tag: &Tag, - new_downlink: i64, - _env: &str, - node_id: &uuid::Uuid, - ) -> Result<()> { - if self.uuid == *node_id { - self.update_downlink(tag, new_downlink)?; - Ok(()) - } else { - Err(PonyError::Custom("Node ID does not match".into())) - } - } - - fn update_node_conn_count( - &mut self, - tag: &Tag, - conn_count: i64, - _env: &str, - node_id: &uuid::Uuid, - ) -> Result<()> { - if self.uuid == *node_id { - self.update_conn_count(tag, conn_count)?; - Ok(()) - } else { - Err(PonyError::Custom("Node ID does not match".into())) - } - } - fn get_by_id(&self, _: &uuid::Uuid) -> std::option::Option { - None - } - - fn select_least_loaded_node( - &self, - _env: &str, - _proto: &Tag, - _connections: &Connections, - ) -> Option { - Some(self.uuid) - } - fn clear(&mut self) -> Result<()> { - Err(PonyError::Custom("Cannot clear Node".to_string())) - } -} - impl Operations for HashMap> { fn clear(&mut self) -> Result<()> { self.clear(); @@ -181,9 +90,6 @@ impl Operations for HashMap> { Ok(OperationStatus::Ok(uuid)) } - fn get_self(&self) -> Option { - None - } fn get(&self, env: &str, uuid: &uuid::Uuid) -> Option<&Node> { self.get(env)?.iter().find(|n| &n.uuid == uuid) } diff --git a/src/memory/subscription.rs b/src/memory/subscription.rs index b8272e4f..8d83871d 100644 --- a/src/memory/subscription.rs +++ b/src/memory/subscription.rs @@ -125,24 +125,6 @@ pub struct SubscriptionStats { pub is_active: bool, } -impl Subscription { - pub fn stats(&self) -> SubscriptionStats { - let now = Utc::now(); - let days_remaining = if let Some(expires_at) = self.expires_at { - (expires_at - now).num_days() - } else { - 99999 - }; - - SubscriptionStats { - id: self.id, - expires_at: self.expires_at, - days_remaining, - is_active: days_remaining > 0 && !self.is_deleted, - } - } -} - pub trait Operations { fn extend(&mut self, days: i64); fn id(&self) -> uuid::Uuid; @@ -155,9 +137,25 @@ pub trait Operations { fn days_remaining(&self) -> Option; fn set_expires_at(&mut self, expires_at: DateTime) -> Result<(), String>; fn mark_deleted(&mut self); + fn stats(&self) -> SubscriptionStats; } impl Operations for Subscription { + fn stats(&self) -> SubscriptionStats { + let now = Utc::now(); + let days_remaining = if let Some(expires_at) = self.expires_at { + (expires_at - now).num_days() + } else { + 99999 + }; + + SubscriptionStats { + id: self.id, + expires_at: self.expires_at, + days_remaining, + is_active: days_remaining > 0 && !self.is_deleted, + } + } fn extend(&mut self, days: i64) { if let Some(expires_at) = self.expires_at { let new_exp = expires_at + chrono::Duration::days(days); diff --git a/src/metrics/bandwidth.rs b/src/metrics/bandwidth.rs deleted file mode 100644 index 2fb59ff7..00000000 --- a/src/metrics/bandwidth.rs +++ /dev/null @@ -1,93 +0,0 @@ -use std::fmt; -use std::thread::sleep; -use std::time::Duration; - -use log::error; -use sysinfo::Networks; - -use super::{AsMetric, Metric, MetricType}; -use crate::utils::current_timestamp; - -#[derive(Debug, Default)] -struct Bandwidth { - rx_bps: u64, - tx_bps: u64, - rx_err: u64, - tx_err: u64, -} - -impl fmt::Display for Bandwidth { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "Bandwidth {{ rx_bps: {}, tx_bps: {}, rx_err: {}, tx_err: {} }}", - self.rx_bps, self.tx_bps, self.rx_err, self.tx_err - ) - } -} - -impl AsMetric for Bandwidth { - type Output = u64; - - fn as_metric(&self, interface: &str, env: &str, hostname: &str) -> Vec> { - let timestamp = current_timestamp(); - - vec![ - Metric { - //dev.localhost.network.eth0.rx_bps (bytes per second) - metric: format!("{env}.{hostname}.network.{interface}.rx_bps"), - value: self.rx_bps, - timestamp, - }, - Metric { - //dev.localhost.network.eth0.tx_bps (bytes per second) - metric: format!("{env}.{hostname}.network.{interface}.tx_bps"), - value: self.tx_bps, - timestamp, - }, - Metric { - //dev.localhost.network.eth0.rx_err - metric: format!("{env}.{hostname}.network.{interface}.rx_err"), - value: self.rx_err, - timestamp, - }, - Metric { - //dev.localhost.network.eth0.tx_err - metric: format!("{env}.{hostname}.network.{interface}.tx_err"), - value: self.tx_err, - timestamp, - }, - ] - } -} - -pub fn bandwidth_metrics(env: &str, hostname: &str, target_interface: &str) -> Vec { - let mut networks = Networks::new_with_refreshed_list(); - sleep(Duration::from_secs(1)); - - networks.refresh(true); - let res = networks - .iter() - .find(|&(interface, _)| interface == target_interface); - - match res { - Some((interface, data)) => { - let bandwidth = Bandwidth { - rx_bps: data.received(), - tx_bps: data.transmitted(), - rx_err: data.errors_on_received(), - tx_err: data.errors_on_transmitted(), - }; - let bandwidth_metrics = bandwidth.as_metric(interface, env, hostname); - - bandwidth_metrics - .iter() - .map(|metric| MetricType::U64(metric.clone())) - .collect() - } - None => { - error!("Cannot find interface: {}", target_interface); - vec![] - } - } -} diff --git a/src/metrics/cpuusage.rs b/src/metrics/cpuusage.rs deleted file mode 100644 index 70b0e527..00000000 --- a/src/metrics/cpuusage.rs +++ /dev/null @@ -1,70 +0,0 @@ -use std::fmt; -use std::thread::sleep; -use std::time::Duration; -use sysinfo::{CpuRefreshKind, RefreshKind, System}; - -use super::{AsMetric, Metric, MetricType}; -use crate::utils::{current_timestamp, round_to_two_decimal_places}; - -struct CpuUsage<'a> { - name: &'a str, - usage: f32, -} - -impl Default for CpuUsage<'_> { - fn default() -> Self { - CpuUsage { - name: "0", - usage: 0.0, - } - } -} - -impl fmt::Display for CpuUsage<'_> { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "CpuUsage {{ name: {}, usage: {} }}", - self.name, self.usage - ) - } -} - -impl AsMetric for CpuUsage<'_> { - type Output = f32; - - fn as_metric(&self, name: &str, env: &str, hostname: &str) -> Vec> { - let timestamp = current_timestamp(); - - vec![Metric { - //dev.localhost.cpu.processor1.percentage - metric: format!("{env}.{hostname}.cpu.{name}.percentage"), - value: self.usage, - timestamp, - }] - } -} - -pub fn cpu_metrics(env: &str, hostname: &str) -> Vec { - let mut s = - System::new_with_specifics(RefreshKind::nothing().with_cpu(CpuRefreshKind::everything())); - - s.refresh_cpu_all(); - sleep(Duration::from_millis(1000)); - s.refresh_cpu_all(); - - let cpu_metrics: Vec<_> = s - .cpus() - .iter() - .map(|cpu| CpuUsage { - name: cpu.name(), - usage: round_to_two_decimal_places(cpu.cpu_usage()), - }) - .flat_map(|cpu| cpu.as_metric(cpu.name, env, hostname)) - .collect(); - - cpu_metrics - .iter() - .map(|metric| MetricType::F32(metric.clone())) - .collect() -} diff --git a/src/metrics/heartbeat.rs b/src/metrics/heartbeat.rs deleted file mode 100644 index 51e6cdb4..00000000 --- a/src/metrics/heartbeat.rs +++ /dev/null @@ -1,27 +0,0 @@ -use super::{Metric, MetricType}; -use crate::utils::current_timestamp; - -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::OnceLock; - -static BEAT_INDEX: OnceLock = OnceLock::new(); - -pub fn heartbeat_metrics(env: &str, uuid: &uuid::Uuid, hostname: &str) -> MetricType { - let timestamp = current_timestamp(); - let metric_name = format!("{env}.{hostname}.{uuid}.heartbeat"); - - let pattern: &[f64] = &[ - 1.0, 1.0, 2.5, 0.4, 0.6, 1.3, 1.0, 1.0, 1.1, 1.0, 1.0, 0.8, 1.0, 1.9, 0.3, 0.5, 1.0, - ]; - - let value = { - let index = BEAT_INDEX - .get_or_init(|| AtomicUsize::new(0)) - .fetch_add(1, Ordering::Relaxed) - % pattern.len(); - pattern[index] - }; - - let metric = Metric::new(metric_name, value, timestamp); - MetricType::F64(metric) -} diff --git a/src/metrics/impls.rs b/src/metrics/impls.rs new file mode 100644 index 00000000..f7278bc1 --- /dev/null +++ b/src/metrics/impls.rs @@ -0,0 +1,184 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::OnceLock; +use sysinfo::Disks; +use sysinfo::Networks; +use sysinfo::{CpuRefreshKind, MemoryRefreshKind, RefreshKind, System}; + +use super::storage::HasMetrics; +use super::storage::MetricSink; +use super::Metrics; + +static BEAT_INDEX: OnceLock = OnceLock::new(); + +impl Metrics for T +where + T: HasMetrics + Sync + Send, +{ + async fn heartbeat(&self) { + let pattern: &[f64] = &[ + 1.0, 1.0, 2.5, 0.4, 0.6, 1.3, 1.0, 1.0, 1.1, 1.0, 1.0, 0.8, 1.0, 1.9, 0.3, 0.5, 1.0, + ]; + let value = { + let index = BEAT_INDEX + .get_or_init(|| AtomicUsize::new(0)) + .fetch_add(1, Ordering::Relaxed) + % pattern.len(); + pattern[index] + }; + + let node = self.node_settings(); + let tags = node.get_base_tags(); + let node_uuid = node.uuid; + + self.metrics() + .write(&node_uuid, "sys.heartbeat", value, tags); + } + + async fn memory(&self) { + let mut system = System::new_with_specifics( + RefreshKind::nothing().with_memory(MemoryRefreshKind::everything()), + ); + system.refresh_memory(); + + let free = system.free_memory() as f64; + let total = system.total_memory() as f64; + let used = system.used_memory() as f64; + let available = system.available_memory() as f64; + + let node = self.node_settings(); + let tags = node.get_base_tags(); + let node_uuid = node.uuid; + + self.metrics() + .write(&node_uuid, "sys.mem_free", free, tags.clone()); + self.metrics() + .write(&node_uuid, "sys.mem_total", total, tags.clone()); + self.metrics() + .write(&node_uuid, "sys.mem_used", used, tags.clone()); + self.metrics() + .write(&node_uuid, "sys.mem_available", available, tags); + } + + async fn loadavg(&self) { + let load = System::load_average(); + + let node = self.node_settings(); + let tags = node.get_base_tags(); + let node_uuid = node.uuid; + + self.metrics() + .write(&node_uuid, "sys.loadavg_1", load.one, tags.clone()); + self.metrics() + .write(&node_uuid, "sys.loadavg_5", load.five, tags.clone()); + self.metrics() + .write(&node_uuid, "sys.loadavg_15", load.fifteen, tags); + } + + async fn bandwidth(&self) { + let mut networks = Networks::new_with_refreshed_list(); + std::thread::sleep(std::time::Duration::from_secs(1)); + networks.refresh(true); + + let node = self.node_settings(); + let tags = node.get_base_tags(); + let node_uuid = node.uuid; + + for (interface, data) in networks.iter() { + self.metrics().write( + &node_uuid, + &format!("net.{interface}.total_rx_bps"), + data.total_received() as f64, + tags.clone(), + ); + + self.metrics().write( + &node_uuid, + &format!("net.{interface}.total_tx_bps"), + data.total_transmitted() as f64, + tags.clone(), + ); + + self.metrics().write( + &node_uuid, + &format!("net.{interface}.rx_bps"), + data.received() as f64, + tags.clone(), + ); + self.metrics().write( + &node_uuid, + &format!("net.{interface}.tx_bps"), + data.transmitted() as f64, + tags.clone(), + ); + } + } + async fn cpu_usage(&self) { + let mut sys = System::new_with_specifics( + RefreshKind::nothing().with_cpu(CpuRefreshKind::everything()), + ); + + sys.refresh_cpu_all(); + std::thread::sleep(std::time::Duration::from_millis(1000)); + sys.refresh_cpu_all(); + + let node = self.node_settings(); + let tags = node.get_base_tags(); + let node_uuid = node.uuid; + + for cpu in sys.cpus() { + let usage = (cpu.cpu_usage() * 100.0).round() / 100.0; + let metric_name = format!("sys.cpu.{}", cpu.name()); + + self.metrics() + .write(&node_uuid, &metric_name, usage as f64, tags.clone()); + } + } + + async fn disk_usage(&self) { + let disks = Disks::new_with_refreshed_list(); + + let node = self.node_settings(); + let tags = node.get_base_tags(); + let node_uuid = node.uuid; + + for disk in &disks { + let mount_point = disk.mount_point().to_str().unwrap_or("unknown"); + let safe_mount = if mount_point == "/" { + "root" + } else { + mount_point.trim_start_matches('/') + }; + + let total = disk.total_space(); + let available = disk.available_space(); + let used = total - available; + + let usage_percent = if total > 0 { + (used as f64 / total as f64 * 100.0).round() / 100.0 + } else { + 0.0 + }; + + self.metrics().write( + &node_uuid, + &format!("sys.disk.{}.used_bytes", safe_mount), + used as f64, + tags.clone(), + ); + + self.metrics().write( + &node_uuid, + &format!("sys.disk.{}.total_bytes", safe_mount), + total as f64, + tags.clone(), + ); + + self.metrics().write( + &node_uuid, + &format!("sys.disk.{}.usage_percent", safe_mount), + usage_percent, + tags.clone(), + ); + } + } +} diff --git a/src/metrics/loadavg.rs b/src/metrics/loadavg.rs deleted file mode 100644 index 5bea0a3f..00000000 --- a/src/metrics/loadavg.rs +++ /dev/null @@ -1,48 +0,0 @@ -use sysinfo::{LoadAvg, System}; - -use super::{AsMetric, Metric, MetricType}; -use crate::utils::current_timestamp; - -struct LoadAvgWrapper { - load_avg: LoadAvg, -} - -impl AsMetric for LoadAvgWrapper { - type Output = f64; - - fn as_metric(&self, name: &str, env: &str, hostname: &str) -> Vec> { - let timestamp = current_timestamp(); - vec![ - Metric { - //dev.localhost.loadavg.1m - metric: format!("{env}.{hostname}.{name}.1m"), - value: self.load_avg.one, - timestamp, - }, - Metric { - //dev.localhost.loadavg.5m - metric: format!("{env}.{hostname}.{name}.5m"), - value: self.load_avg.five, - timestamp, - }, - Metric { - //dev.localhost.loadavg.15m - metric: format!("{env}.{hostname}.{name}.15m"), - value: self.load_avg.fifteen, - timestamp, - }, - ] - } -} - -pub fn loadavg_metrics(env: &str, hostname: &str) -> Vec { - let load_avg = System::load_average(); - let wrapper = LoadAvgWrapper { load_avg }; - - let load_metrics = wrapper.as_metric("loadavg", env, hostname); - - load_metrics - .iter() - .map(|metric| MetricType::F64(metric.clone())) - .collect() -} diff --git a/src/metrics/memory.rs b/src/metrics/memory.rs deleted file mode 100644 index 31adc0e2..00000000 --- a/src/metrics/memory.rs +++ /dev/null @@ -1,64 +0,0 @@ -use sysinfo::{MemoryRefreshKind, System}; - -use super::{AsMetric, Metric, MetricType}; -use crate::utils::current_timestamp; - -struct MemUsage { - free: u64, - total: u64, - used: u64, - available: u64, -} - -impl AsMetric for MemUsage { - type Output = u64; - fn as_metric(&self, name: &str, env: &str, hostname: &str) -> Vec> { - let timestamp = current_timestamp(); - - vec![ - Metric { - //dev.localhost.mem.total - metric: format!("{env}.{hostname}.{name}.total"), - value: self.total, - timestamp, - }, - Metric { - //dev.localhost.mem.free - metric: format!("{env}.{hostname}.{name}.free"), - value: self.free, - timestamp, - }, - Metric { - //dev.localhost.mem.used - metric: format!("{env}.{hostname}.{name}.used"), - value: self.used, - timestamp, - }, - Metric { - //dev.localhost.mem.available - metric: format!("{env}.{hostname}.{name}.available"), - value: self.available, - timestamp, - }, - ] - } -} - -pub fn mem_metrics(env: &str, hostname: &str) -> Vec { - let mut system = System::new(); - - system.refresh_memory_specifics(MemoryRefreshKind::nothing().with_ram()); - let mem = MemUsage { - free: system.free_memory(), - total: system.total_memory(), - used: system.used_memory(), - available: system.available_memory(), - }; - - let mem_metrics = mem.as_metric("mem", env, hostname); - - mem_metrics - .iter() - .map(|metric| MetricType::U64(metric.clone())) - .collect() -} diff --git a/src/metrics/mod.rs b/src/metrics/mod.rs index 69d10eca..0c232015 100644 --- a/src/metrics/mod.rs +++ b/src/metrics/mod.rs @@ -1,131 +1,33 @@ -use crate::{PonyError, Result as PonyResult}; - -pub mod bandwidth; -pub mod cpuusage; -pub mod heartbeat; -pub mod loadavg; -pub mod memory; -pub mod xray; - -use clickhouse::Row; +use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize}; use serde::{Deserialize, Serialize}; -use std::fmt; -use std::io; -use tokio::{io::AsyncWriteExt, net::TcpStream}; - -use crate::memory::stat::Kind as StatKind; +use std::collections::BTreeMap; -pub trait AsMetric { - type Output; - fn as_metric(&self, name: &str, env: &str, hostname: &str) -> Vec>; -} +pub mod impls; +pub mod storage; -#[derive(Serialize, Deserialize, Row, Clone, Debug)] -pub struct Metric { - pub metric: String, - pub value: T, +#[derive(Archive, RkyvSerialize, RkyvDeserialize, Serialize, Deserialize, Clone, Debug)] +pub struct MetricPoint { + #[serde(rename = "x")] pub timestamp: i64, + #[serde(rename = "y")] + pub value: f64, } -impl Metric { - pub fn new(metric: String, value: T, timestamp: i64) -> Self { - Metric { - metric, - value, - timestamp, - } - } - pub fn stat_type(&self) -> StatKind { - StatKind::from_path(&self.metric) - } -} - -impl Default for Metric { - fn default() -> Self { - Metric { - value: T::default(), - metric: String::from(""), - timestamp: 0, - } - } -} - -impl fmt::Display for Metric { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "{} {} {}", self.metric, self.value, self.timestamp) - } -} - -impl Metric { - pub async fn send(&self, server: &str) -> Result<(), io::Error> { - let metric_string = format!("{self}\n"); - - log::debug!("Sending carbon metric string: {:?}", metric_string); - - match TcpStream::connect(server).await { - Ok(mut stream) => { - if let Err(e) = stream.write_all(metric_string.as_bytes()).await { - log::warn!("Failed to send metric: {}", e); - return Err(e); - } - - if let Err(e) = stream.flush().await { - log::warn!("Failed to flush stream: {}", e); - return Err(e); - } - - if let Err(e) = stream.shutdown().await { - log::warn!("Failed to shutdown stream: {}", e); - return Err(e); - } - - Ok(()) - } - Err(e) => { - log::error!("Failed to connect to Carbon server: {}", e); - Err(e) - } - } - } -} - -#[derive(Debug, Clone)] -pub enum MetricType { - F32(Metric), - F64(Metric), - I64(Metric), - U64(Metric), - U8(Metric), +#[derive(Archive, RkyvSerialize, RkyvDeserialize, Serialize, Deserialize, Clone, Debug)] +#[archive(check_bytes)] +pub struct MetricEnvelope { + pub node_id: uuid::Uuid, + pub name: String, + pub value: f64, + pub timestamp: i64, + pub tags: BTreeMap, } -#[async_trait::async_trait] -pub trait Metrics { - async fn collect_metrics(&self) -> Vec; - async fn collect_hb_metrics(&self) -> MetricType; - - async fn send_metrics(&self, carbon_address: &str) -> PonyResult<()> { - let metrics = self.collect_metrics::().await; - - for metric in metrics { - match metric { - MetricType::F32(m) => m.send(carbon_address).await?, - MetricType::F64(m) => m.send(carbon_address).await?, - MetricType::I64(m) => m.send(carbon_address).await?, - MetricType::U64(m) => m.send(carbon_address).await?, - MetricType::U8(m) => m.send(carbon_address).await?, - } - } - Ok(()) - } - async fn send_hb_metric(&self, carbon_address: &str) -> PonyResult<()> { - let metric = self.collect_hb_metrics::().await; - - match metric { - MetricType::F64(m) => { - m.send(carbon_address).await?; - Ok(()) - } - _ => Err(PonyError::Custom("Doesn't support metric type".to_string())), - } - } +pub trait Metrics { + fn heartbeat(&self) -> impl std::future::Future + Send; + fn memory(&self) -> impl std::future::Future + Send; + fn bandwidth(&self) -> impl std::future::Future + Send; + fn cpu_usage(&self) -> impl std::future::Future + Send; + fn loadavg(&self) -> impl std::future::Future + Send; + fn disk_usage(&self) -> impl std::future::Future + Send; } diff --git a/src/metrics/storage.rs b/src/metrics/storage.rs new file mode 100644 index 00000000..00ac0f73 --- /dev/null +++ b/src/metrics/storage.rs @@ -0,0 +1,190 @@ +use dashmap::DashMap; +use std::collections::{BTreeMap, VecDeque}; + +use super::MetricEnvelope; +use super::MetricPoint; +use crate::memory::node::Node; +use crate::zmq::publisher::Publisher; + +pub trait HasMetrics { + fn metrics(&self) -> &MetricBuffer; + fn node_settings(&self) -> &Node; +} + +pub trait MetricSink { + fn write(&self, node_id: &uuid::Uuid, metric: &str, value: f64, tags: BTreeMap); +} + +impl MetricSink for MetricBuffer { + fn write( + &self, + node_id: &uuid::Uuid, + metric: &str, + value: f64, + tags: BTreeMap, + ) { + let mut b = self.batch.lock(); + b.push(MetricEnvelope { + node_id: *node_id, + name: metric.to_string(), + value, + timestamp: chrono::Utc::now().timestamp_millis(), + tags, + }); + } +} + +pub struct MetricBuffer { + pub batch: parking_lot::Mutex>, + pub publisher: Publisher, +} + +impl MetricBuffer { + pub fn push( + &self, + node_id: uuid::Uuid, + name: &str, + value: f64, + tags: BTreeMap, + ) { + let mut b = self.batch.lock(); + b.push(MetricEnvelope { + node_id, + name: name.to_string(), + value, + timestamp: chrono::Utc::now().timestamp_millis(), + tags, + }); + } + + pub async fn flush_to_zmq(&self) { + let metrics = { + let mut batch = self.batch.lock(); + if batch.is_empty() { + return; + } + std::mem::take(&mut *batch) + }; + + let bytes = rkyv::to_bytes::<_, 65536>(&metrics).expect("Failed to serialize batch"); + + if let Err(e) = &self + .publisher + .send_binary("metrics", bytes.as_slice()) + .await + { + log::error!("Batch publish failed: {}", e); + } + } +} + +pub struct MetricStorage { + pub inner: DashMap>>, + pub metadata: DashMap>, + + pub max_points: usize, + pub retention_seconds: i64, +} + +impl MetricStorage { + pub fn new(max_points: usize, retention_seconds: i64) -> Self { + Self { + inner: DashMap::new(), + metadata: DashMap::new(), + max_points, + retention_seconds, + } + } + fn make_series_key(name: &str, tags: &BTreeMap) -> String { + if tags.is_empty() { + return name.to_string(); + } + let tags_part: Vec = tags.iter().map(|(k, v)| format!("{}={}", k, v)).collect(); + format!("{}:{{{}}}", name, tags_part.join(",")) + } + + pub fn get_range( + &self, + node_id: &uuid::Uuid, + series_key: &str, + from: i64, + to: i64, + ) -> Vec { + self.inner + .get(node_id) + .and_then(|node_data| { + node_data.get(series_key).map(|deque| { + deque + .iter() + .filter(|p| p.timestamp >= from && p.timestamp <= to) + .cloned() + .collect() + }) + }) + .unwrap_or_default() + } + + pub fn insert_envelope(&self, env: MetricEnvelope) { + let series_key = Self::make_series_key(&env.name, &env.tags); + + if !self.metadata.contains_key(&series_key) { + self.metadata.insert(series_key.clone(), env.tags); + } + + let node_map = self.inner.entry(env.node_id).or_default(); + let mut entry = node_map.entry(series_key).or_default(); + + // env.timestamp ΡƒΠΆΠ΅ Π² миллисСкундах (ΠΈΠ· MetricBuffer) + entry.push_back(MetricPoint { + timestamp: env.timestamp, + value: env.value, + }); + + // 1. ΠžΡ‡ΠΈΡΡ‚ΠΊΠ° ΠΏΠΎ количСству Ρ‚ΠΎΡ‡Π΅ΠΊ + while entry.len() > self.max_points { + entry.pop_front(); + } + + // 2. ΠžΡ‡ΠΈΡΡ‚ΠΊΠ° ΠΏΠΎ Π²Ρ€Π΅ΠΌΠ΅Π½ΠΈ (Retention) + // ΠŸΠ΅Ρ€Π΅Π²ΠΎΠ΄ΠΈΠΌ сСкунды ΠΊΠΎΠ½Ρ„ΠΈΠ³Π° Π² мс для сравнСния с env.timestamp + let retention_ms = self.retention_seconds * 1000; + let min_ts = env.timestamp - retention_ms; + + while let Some(front) = entry.front() { + if front.timestamp < min_ts { + entry.pop_front(); + } else { + break; + } + } + } + + pub fn perform_gc(&self) { + // Π˜ΡΠΏΠΎΠ»ΡŒΠ·ΡƒΠ΅ΠΌ встроСнный ΠΌΠ΅Ρ‚ΠΎΠ΄ для мс, Ρ‡Ρ‚ΠΎΠ±Ρ‹ Π½Π΅ ΠΌΠ½ΠΎΠΆΠΈΡ‚ΡŒ Π½Π° 1000 Π²Ρ€ΡƒΡ‡Π½ΡƒΡŽ + let now_ms = chrono::Utc::now().timestamp_millis(); + let retention_ms = self.retention_seconds * 1000; + let min_ts = now_ms - retention_ms; + + log::debug!("Starting MetricStorage GC. Min timestamp: {}", min_ts); + + self.inner.retain(|_node_id, node_map| { + node_map.retain(|_series_key, deque| { + while let Some(front) = deque.front() { + if front.timestamp < min_ts { + deque.pop_front(); + } else { + break; + } + } + !deque.is_empty() + }); + !node_map.is_empty() + }); + + self.metadata.retain(|series_key, _| { + self.inner + .iter() + .any(|node_ref| node_ref.value().contains_key(series_key)) + }); + } +} diff --git a/src/metrics/xray.rs b/src/metrics/xray.rs deleted file mode 100644 index a43a4ab4..00000000 --- a/src/metrics/xray.rs +++ /dev/null @@ -1,100 +0,0 @@ -use super::{AsMetric, Metric, MetricType}; -use crate::memory::cache::Connections; -use crate::memory::connection::op::base::Operations as ConnectionBaseOp; -use crate::memory::node::Node; -use crate::memory::{connection::stat::Stat as ConnectionStat, node::Stat as InboundStat}; -use crate::utils::current_timestamp; - -impl AsMetric for InboundStat { - type Output = i64; - fn as_metric(&self, name: &str, env: &str, hostname: &str) -> Vec> { - let timestamp = current_timestamp(); - - vec![ - Metric { - //dev.localhost.vmess.inbound_stat.uplink - metric: format!("{env}.{hostname}.{name}.inbound_stat.uplink"), - value: self.uplink, - timestamp, - }, - Metric { - //dev.localhost.vmess.inbound_stat.downlink - metric: format!("{env}.{hostname}.{name}.inbound_stat.downlink"), - value: self.downlink, - timestamp, - }, - Metric { - // dev.localhost.vmess.inbound_stat.user_count - metric: format!("{env}.{hostname}.{name}.inbound_stat.user_count"), - value: self.conn_count, - timestamp, - }, - ] - } -} - -impl AsMetric for ConnectionStat { - type Output = i64; - fn as_metric(&self, name: &str, env: &str, hostname: &str) -> Vec> { - let timestamp = current_timestamp(); - - vec![ - Metric { - //dev.localhost.user_id.uplink - metric: format!("{env}.{hostname}.{name}.conn_stat.uplink"), - value: self.uplink, - timestamp, - }, - Metric { - //dev.localhost.user_id.downlink - metric: format!("{env}.{hostname}.{name}.conn_stat.downlink"), - value: self.downlink, - timestamp, - }, - Metric { - // dev.localhost.user_id.online - metric: format!("{env}.{hostname}.{name}.conn_stat.online"), - value: self.online, - timestamp, - }, - ] - } -} - -pub fn xray_stat_metrics(node: Node) -> Vec { - let xray_stat_metrics: Vec<_> = node - .inbounds - .clone() - .into_iter() - .flat_map(|(tag, inbound)| { - inbound - .as_inbound_stat() - .as_metric(&tag.to_string(), &node.env, &node.hostname) - }) - .collect(); - - xray_stat_metrics.into_iter().map(MetricType::I64).collect() -} - -pub fn xray_conn_metrics( - connections: Connections, - env: &str, - hostname: &str, -) -> Vec -where - C: ConnectionBaseOp + Send + Sync + Clone + 'static, -{ - let conn_stat_metrics: Vec<_> = connections - .clone() - .iter() - .flat_map(|(conn_id, conn)| { - let metric_name = format!("{}.{}", conn.get_proto().proto(), &conn_id); - conn.as_conn_stat().as_metric(&metric_name, env, hostname) - }) - .collect(); - - conn_stat_metrics - .iter() - .map(|metric| MetricType::I64(metric.clone())) - .collect() -} diff --git a/src/utils.rs b/src/utils.rs index 1cf389f7..307d3c48 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -104,14 +104,14 @@ pub fn to_ipv4(ip: IpAddr) -> Option { } } -pub async fn measure_time(task: F, name: String) -> T +pub async fn measure_time(task: F, name: &str) -> T where F: std::future::Future, { let start_time = Instant::now(); let result = task.await; let duration = start_time.elapsed(); - log::info!("Task {} completed in {:?}", name, duration); + log::debug!("Task {} completed in {:?}", name, duration); result } diff --git a/src/zmq/publisher.rs b/src/zmq/publisher.rs index 2e619f5a..be3aa463 100644 --- a/src/zmq/publisher.rs +++ b/src/zmq/publisher.rs @@ -10,32 +10,46 @@ pub struct Publisher { } impl Publisher { - pub async fn new(endpoint: &str) -> Self { + pub async fn bind(endpoint: &str) -> Self { + Self::init(endpoint, true).await + } + + pub async fn connect(endpoint: &str) -> Self { + Self::init(endpoint, false).await + } + + async fn init(endpoint: &str, is_bind: bool) -> Self { let context = zmq::Context::new(); let publisher = context.socket(zmq::PUB).expect("Failed to create socket"); let mut i = 0; loop { - match publisher.bind(endpoint) { + let result = if is_bind { + publisher.bind(endpoint) + } else { + publisher.connect(endpoint) + }; + + match result { Ok(_) => { - log::debug!("Connected to Pub socket at {}", endpoint); + let mode = if is_bind { "Bound to" } else { "Connected to" }; + log::debug!("PUB: {} {}", mode, endpoint); break; } Err(err) => { if i >= 5 { panic!( - "Failed to connect to Pub socket at {}: {}. Giving up.", + "Failed to setup PUB socket at {}: {}. Giving up.", endpoint, err ); } - log::debug!("Trying to connect Pub socket, attempt {}: {}", i + 1, err); + log::warn!("PUB: Setup attempt {} failed: {}", i + 1, err); i += 1; sleep(Duration::from_secs(5)).await; } } } - // Ugly hack for zmq pub/sub sync sleep(Duration::from_millis(1000)).await; Self { @@ -43,6 +57,10 @@ impl Publisher { } } + pub async fn new(endpoint: &str) -> Self { + Self::bind(endpoint).await + } + pub async fn send_binary(&self, topic: &str, payload: &[u8]) -> zmq::Result<()> { let socket = self.socket.lock().await; @@ -52,20 +70,4 @@ impl Publisher { log::debug!("PUB: Message sent: {} | {} bytes", topic, payload.len()); Ok(()) } - - pub async fn send(&self, topic: &str, message: impl ToString) -> zmq::Result<()> { - let full_message = format!("{} {}", topic, message.to_string()); - let socket = self.socket.lock().await; - - match socket.send(full_message.as_bytes(), 0) { - Ok(_) => { - log::debug!("PUB: Message sent: {}", full_message); - Ok(()) - } - Err(e) => { - log::error!("PUB: Failed to send message: {}", e); - Err(e) - } - } - } } diff --git a/src/zmq/subscriber.rs b/src/zmq/subscriber.rs index d808e74b..520161dc 100644 --- a/src/zmq/subscriber.rs +++ b/src/zmq/subscriber.rs @@ -77,6 +77,34 @@ impl Subscriber { } } } + + pub fn new_bound(endpoint: &str, topics: Vec) -> Self { + let context = zmq::Context::new(); + let socket = context + .socket(zmq::SUB) + .expect("Failed to create SUB socket"); + + socket.bind(endpoint).expect("Failed to bind SUB socket"); + + if topics.is_empty() { + socket + .set_subscribe(b"") + .expect("Failed to subscribe to all"); + log::info!("Subscribed to all topics (wildcard)"); + } else { + for topic in &topics { + socket + .set_subscribe(topic.as_bytes()) + .expect("Failed to subscribe to topic"); + } + log::info!("Subscribed to topics: {:?}", topics); + } + + Self { + socket: Arc::new(Mutex::new(socket)), + topics, + } + } } impl Clone for Subscriber {