diff --git a/crates/udp/src/config.rs b/crates/udp/src/config.rs index 0d3e6bc..3851b31 100644 --- a/crates/udp/src/config.rs +++ b/crates/udp/src/config.rs @@ -183,6 +183,12 @@ pub struct StatisticsConfig { pub write_html_to_file: bool, /// Path to save HTML file to pub html_file_path: PathBuf, + /// Save statistics as JSON to a file + pub write_json_to_file: bool, + /// Path to dump JSON info-hash IPv4 to + pub json_info_hash_ipv4_file_path: PathBuf, + /// Path to dump JSON info-hash IPv6 to + pub json_info_hash_ipv6_file_path: PathBuf, /// Run a prometheus endpoint #[cfg(feature = "prometheus")] pub run_prometheus_endpoint: bool, @@ -223,6 +229,9 @@ impl Default for StatisticsConfig { print_to_stdout: false, write_html_to_file: false, html_file_path: "tmp/statistics.html".into(), + write_json_to_file: false, + json_info_hash_ipv4_file_path: "tmp/info_hash_v4.json".into(), + json_info_hash_ipv6_file_path: "tmp/info_hash_v6.json".into(), #[cfg(feature = "prometheus")] run_prometheus_endpoint: false, #[cfg(feature = "prometheus")] diff --git a/crates/udp/src/swarm.rs b/crates/udp/src/swarm.rs index 2d422f9..1496540 100644 --- a/crates/udp/src/swarm.rs +++ b/crates/udp/src/swarm.rs @@ -1,3 +1,4 @@ +use std::io::Write; use std::iter::repeat_with; use std::net::IpAddr; use std::ops::DerefMut; @@ -133,6 +134,43 @@ impl TorrentMaps { ::log::error!("couldn't send statistics message: {:#}", err); } } + + if config.statistics.write_json_to_file { + fn save_to_file( + path: &std::path::PathBuf, + info_hashes: &Vec, + ) -> anyhow::Result<()> { + let mut file = + anyhow::Context::with_context(std::fs::File::create(path), || { + format!("File path: {}", path.to_string_lossy()) + })?; + write!(file, "[")?; + if !info_hashes.is_empty() { + write!(file, "\"{}\"", info_hashes[0])?; + if let Some(i) = info_hashes.get(1..) { + for info_hash in i { + write!(file, ",\"{info_hash}\"")?; + } + } + } + write!(file, "]")?; // @TODO serialize with serde_json? + Ok(()) + } + if config.network.ipv4_active() { + if let Err(err) = + save_to_file(&config.statistics.json_info_hash_ipv4_file_path, &ipv4.3) + { + ::log::error!("Couldn't dump IPv4 info-hash table to file: {:#}", err) + } + } + if config.network.ipv6_active() { + if let Err(err) = + save_to_file(&config.statistics.json_info_hash_ipv6_file_path, &ipv6.3) + { + ::log::error!("Couldn't dump IPv6 info-hash table to file: {:#}", err) + } + } + } } } } @@ -219,9 +257,10 @@ impl TorrentMapShards { access_list_cache: &mut AccessListCache, access_list_mode: AccessListMode, now: SecondsSinceServerStart, - ) -> (usize, usize, Option>) { + ) -> (usize, usize, Option>, Vec) { let mut total_num_torrents = 0; let mut total_num_peers = 0; + let mut info_hashes: Vec = Vec::new(); let mut opt_histogram: Option> = config .statistics @@ -297,9 +336,21 @@ impl TorrentMapShards { torrent_map_shard.shrink_to_fit(); total_num_torrents += torrent_map_shard.len(); + + if config.statistics.write_json_to_file { + info_hashes.reserve(total_num_torrents); + for (k, _) in torrent_map_shard.iter() { + info_hashes.push(*k) + } + } } - (total_num_torrents, total_num_peers, opt_histogram) + ( + total_num_torrents, + total_num_peers, + opt_histogram, + info_hashes, + ) } fn get_shard(&self, info_hash: &InfoHash) -> &RwLock> {