From b3b78bb6450c3c06da1ebda85990547a4d73545d Mon Sep 17 00:00:00 2001 From: Mingwei Zhang Date: Thu, 23 Jul 2026 10:00:57 -0700 Subject: [PATCH 1/3] feat: add collector-aware indexers for as2rel and pfx2as Add as2rel-collector-index and pfx2as-collector-index binaries that preserve per-collector provenance in latest-snapshot aggregates. - as2rel-collector-index: walks daily as2rel_*.bz2 files and produces as2rel-collector-latest.json.bz2 with per-edge collector breakdown - pfx2as-collector-index: same pattern for pfx2as data - Each output includes generated_at timestamp, input file count, total/cross-collector sums, and collector-level detail - Fix pre-existing clippy warnings in bootstrap.rs (useless borrows in format! macros) --- Cargo.toml | 8 ++ src/bin/bootstrap.rs | 10 +- src/bin/index-as2rel-collector.rs | 204 ++++++++++++++++++++++++++++++ src/bin/index-pfx2as-collector.rs | 190 ++++++++++++++++++++++++++++ 4 files changed, 407 insertions(+), 5 deletions(-) create mode 100644 src/bin/index-as2rel-collector.rs create mode 100644 src/bin/index-pfx2as-collector.rs diff --git a/Cargo.toml b/Cargo.toml index ade074a..ab40e29 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,3 +61,11 @@ path = "src/bin/index-as2rel.rs" [[bin]] name = "pfx2as-index" path = "src/bin/index-pfx2as.rs" + +[[bin]] +name = "as2rel-collector-index" +path = "src/bin/index-as2rel-collector.rs" + +[[bin]] +name = "pfx2as-collector-index" +path = "src/bin/index-pfx2as-collector.rs" diff --git a/src/bin/bootstrap.rs b/src/bin/bootstrap.rs index 63c26f0..2391437 100644 --- a/src/bin/bootstrap.rs +++ b/src/bin/bootstrap.rs @@ -156,7 +156,7 @@ fn main() { "{}/{}/{}/{:02}/{:02}", output_dir, "as2rel", - &item.collector_id, + item.collector_id, ts.year(), ts.month() ), @@ -164,7 +164,7 @@ fn main() { "{}/{}/{}/{:02}/{:02}", output_dir, data_type, - &item.collector_id, + item.collector_id, ts.year(), ts.month() ), @@ -172,13 +172,13 @@ fn main() { fs::create_dir_all(file_dir.as_str()).unwrap(); let output_path = format!( "{}/{}_{}_{}-{:02}-{:02}_{}.bz2", - &file_dir, + file_dir, data_type, - &item.collector_id, + item.collector_id, ts.year(), ts.month(), ts.day(), - ×tamp + timestamp ); if !opts.force && std::path::Path::new(output_path.as_str()).exists() { info!( diff --git a/src/bin/index-as2rel-collector.rs b/src/bin/index-as2rel-collector.rs new file mode 100644 index 0000000..05da0f3 --- /dev/null +++ b/src/bin/index-as2rel-collector.rs @@ -0,0 +1,204 @@ +use chrono::{NaiveDate, Utc}; +use clap::Parser; +use peer_stats::As2Rel; +use serde::Serialize; +use std::collections::HashMap; +use std::io::Read; +use std::path::PathBuf; +use tracing::info; +use walkdir::WalkDir; + +/// Index AS relationship data with per-collector provenance. +/// Produces as2rel-collector-latest.json.bz2 with collector-level breakdown. +#[derive(Parser, Debug)] +struct Opts { + /// Path to output directory + output_dir: PathBuf, + + /// Path to the data file directory + data_dir: PathBuf, + + /// Whether to print debug logs + #[clap(long)] + debug: bool, + + /// Allow processing files from the previous day + #[clap(long)] + allow_previous_day: bool, +} + +#[derive(Debug, Clone, Serialize)] +struct CollectorDetail { + project: String, + collector: String, + paths_count: usize, + peers_count: usize, +} + +#[derive(Debug, Clone, Serialize)] +struct As2RelCollectorEntry { + asn1: u32, + asn2: u32, + rel: u8, + /// Total paths across all collectors + total_paths_count: usize, + /// Total unique peers across all collectors (deduplicated by IP) + total_peers_count: usize, + /// Number of distinct collectors observing this relationship + collector_count: usize, + /// Per-collector breakdown + collectors: HashMap, +} + +#[derive(Debug, Clone, Serialize)] +struct As2RelCollectorOutput { + generated_at: String, + input_files: usize, + entries: Vec, +} + +fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { + let date_part = file_path.split('_').collect::>(); + let parts = date_part[date_part.len() - 2] + .split('-') + .collect::>(); + ( + parts[0].parse::().unwrap(), + parts[1].parse::().unwrap(), + parts[2].parse::().unwrap(), + ) +} + +fn main() { + let opts = Opts::parse(); + + if opts.debug { + tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .init(); + } + + for file_prefix in ["as2rel_", "as2rel-v4_", "as2rel-v6_"] { + let file_paths = WalkDir::new(opts.data_dir.to_str().unwrap()) + .follow_links(true) + .into_iter() + .filter_map(|e| match e.ok() { + Some(entry) => { + let path: String = entry.path().to_str().unwrap().to_string(); + let path_str = path.as_str(); + if path_str.contains(file_prefix) && path_str.ends_with(".bz2") { + let (year, month, day) = get_ymd_from_file(path.as_str()); + let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); + let ts = Utc::now().date_naive(); + if file_date == ts { + return Some(path); + } + if opts.allow_previous_day && file_date == ts.pred_opt().unwrap() { + return Some(path); + } + } + None + } + None => None, + }) + .collect::>(); + + if file_paths.is_empty() { + info!( + "no matching current date {} files found, skipping", + file_prefix + ); + continue; + } + + let input_file_count = file_paths.len(); + + // Key: (asn1, asn2, rel) — Value: per-collector details + let mut collector_map: HashMap<(u32, u32, u8), HashMap> = + HashMap::new(); + + for file in &file_paths { + info!("processing {}", file.as_str()); + let mut data = String::new(); + oneio::get_reader(file.as_str()) + .unwrap() + .read_to_string(&mut data) + .unwrap(); + let as2rel_info: As2Rel = serde_json::from_str(&data).unwrap(); + + let project = as2rel_info.project; + let collector = as2rel_info.collector; + + for as2rel in as2rel_info.as2rel { + let key = (as2rel.asn1, as2rel.asn2, as2rel.rel); + let per_collector = collector_map.entry(key).or_default(); + let detail = + per_collector + .entry(collector.clone()) + .or_insert_with(|| CollectorDetail { + project: project.clone(), + collector: collector.clone(), + paths_count: 0, + peers_count: 0, + }); + detail.paths_count += as2rel.paths_count; + detail.peers_count += as2rel.peers_count; + } + } + + let entries: Vec = collector_map + .into_iter() + .map(|((asn1, asn2, rel), per_collector)| { + let total_paths_count: usize = per_collector.values().map(|d| d.paths_count).sum(); + let total_peers_count: usize = per_collector.values().map(|d| d.peers_count).sum(); + let collector_count = per_collector.len(); + + As2RelCollectorEntry { + asn1, + asn2, + rel, + total_paths_count, + total_peers_count, + collector_count, + collectors: per_collector, + } + }) + .collect(); + + let output = As2RelCollectorOutput { + generated_at: Utc::now().to_rfc3339(), + input_files: input_file_count, + entries, + }; + + let output_file = format!( + "{}/{}-collector-latest.json.bz2", + opts.output_dir.to_str().unwrap(), + file_prefix.strip_suffix('_').unwrap() + ); + let mut writer = oneio::get_writer(output_file.as_str()).unwrap(); + let _ = writer.write_all( + serde_json::to_string_pretty(&serde_json::to_value(&output).unwrap()) + .unwrap() + .as_ref(), + ); + info!("wrote {} entries to {}", file_paths.len(), output_file); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_get_file_date() { + assert_eq!( + get_ymd_from_file("as2rel_rrc16_2022-02-01_1643673600.bz2"), + (2022, 2, 1) + ); + assert_eq!( + get_ymd_from_file("/aaa_bbb-ccc/as2rel_rrc16_2022-02-01_1643673600.bz2"), + (2022, 2, 1) + ); + } +} diff --git a/src/bin/index-pfx2as-collector.rs b/src/bin/index-pfx2as-collector.rs new file mode 100644 index 0000000..960b049 --- /dev/null +++ b/src/bin/index-pfx2as-collector.rs @@ -0,0 +1,190 @@ +use chrono::{NaiveDate, Utc}; +use clap::Parser; +use peer_stats::Prefix2As; +use serde::Serialize; +use std::collections::HashMap; +use std::io::Read; +use std::path::PathBuf; +use tracing::info; +use walkdir::WalkDir; + +/// Index prefix-to-AS mapping data with per-collector provenance. +/// Produces pfx2as-collector-latest.json.bz2 with collector-level breakdown. +#[derive(Parser, Debug)] +struct Opts { + /// Path to output directory (file named pfx2as-collector-latest.json.bz2) + output_dir: PathBuf, + + /// Path to the data file directory + data_dir: PathBuf, + + /// Whether to print debug logs + #[clap(long)] + debug: bool, + + /// Allow processing files from the previous day + #[clap(long)] + allow_previous_day: bool, +} + +#[derive(Debug, Clone, Serialize)] +struct Pfx2AsCollectorDetail { + project: String, + collector: String, + count: usize, +} + +#[derive(Debug, Clone, Serialize)] +struct Pfx2AsCollectorEntry { + prefix: String, + asn: u32, + /// Total count across all collectors + total_count: usize, + /// Number of distinct collectors seeing this mapping + collector_count: usize, + /// Per-collector breakdown + collectors: HashMap, +} + +#[derive(Debug, Clone, Serialize)] +struct Pfx2AsCollectorOutput { + generated_at: String, + input_files: usize, + entries: Vec, +} + +fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { + let date_part = file_path.split('_').collect::>(); + let parts = date_part[date_part.len() - 2] + .split('-') + .collect::>(); + ( + parts[0].parse::().unwrap(), + parts[1].parse::().unwrap(), + parts[2].parse::().unwrap(), + ) +} + +fn main() { + let opts = Opts::parse(); + + if opts.debug { + tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .init(); + } + + let file_paths = WalkDir::new(opts.data_dir.to_str().unwrap()) + .follow_links(true) + .into_iter() + .filter_map(|e| match e.ok() { + Some(entry) => { + let path: String = entry.path().to_str().unwrap().to_string(); + let path_str = path.as_str(); + if path_str.contains("pfx2as_") && path_str.ends_with(".bz2") { + let (year, month, day) = get_ymd_from_file(path.as_str()); + let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); + let ts = Utc::now().date_naive(); + if file_date == ts { + return Some(path); + } + if opts.allow_previous_day && file_date == ts.pred_opt().unwrap() { + return Some(path); + } + } + None + } + None => None, + }) + .collect::>(); + + if file_paths.is_empty() { + info!("no data files found, skipping"); + return; + } + + let input_file_count = file_paths.len(); + + // Key: (prefix, asn) — Value: per-collector details + let mut collector_map: HashMap<(String, u32), HashMap> = + HashMap::new(); + + for file in &file_paths { + info!("processing {}", file.as_str()); + let mut data = String::new(); + oneio::get_reader(file.as_str()) + .unwrap() + .read_to_string(&mut data) + .unwrap(); + let pfx2as_info: Prefix2As = serde_json::from_str(&data).unwrap(); + + let project = pfx2as_info.project; + let collector = pfx2as_info.collector; + + for pfx2as in pfx2as_info.pfx2as { + let key = (pfx2as.prefix.clone(), pfx2as.asn); + let per_collector = collector_map.entry(key).or_default(); + let detail = + per_collector + .entry(collector.clone()) + .or_insert_with(|| Pfx2AsCollectorDetail { + project: project.clone(), + collector: collector.clone(), + count: 0, + }); + detail.count += pfx2as.count; + } + } + + let entry_count = collector_map.len(); + let entries: Vec = collector_map + .into_iter() + .map(|((prefix, asn), per_collector)| { + let total_count: usize = per_collector.values().map(|d| d.count).sum(); + let collector_count = per_collector.len(); + + Pfx2AsCollectorEntry { + prefix, + asn, + total_count, + collector_count, + collectors: per_collector, + } + }) + .collect(); + + let output = Pfx2AsCollectorOutput { + generated_at: Utc::now().to_rfc3339(), + input_files: input_file_count, + entries, + }; + + let output_file = format!( + "{}/pfx2as-collector-latest.json.bz2", + opts.output_dir.to_str().unwrap() + ); + let mut writer = oneio::get_writer(output_file.as_str()).unwrap(); + let _ = writer.write_all( + serde_json::to_string_pretty(&serde_json::to_value(&output).unwrap()) + .unwrap() + .as_ref(), + ); + info!("wrote {} entries to {}", entry_count, output_file); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_get_file_date() { + assert_eq!( + get_ymd_from_file("pfx2as_rrc16_2022-02-01_1643673600.bz2"), + (2022, 2, 1) + ); + assert_eq!( + get_ymd_from_file("/aaa_bbb-ccc/pfx2as_rrc16_2022-02-01_1643673600.bz2"), + (2022, 2, 1) + ); + } +} From 4414571cc0e2e50f409e2c86791b74eabd129eb0 Mon Sep 17 00:00:00 2001 From: Mingwei Zhang Date: Thu, 23 Jul 2026 10:20:59 -0700 Subject: [PATCH 2/3] fix: replace nested HashMaps with flat vector + sort approach The previous implementation used HashMap<(u32,u32,u8), HashMap> which created per-edge sub-HashMaps causing 30GB+ memory usage with ~500K edges. Now uses: - Flat Vec<(u32,u32,u8,u16,usize,usize)> for as2rel - Flat Vec<(u32,u32,u16,usize)> for pfx2as (with prefix string interning) - Sort by key then group with O(1) extra memory - Manual streaming JSON writer avoids serde_json::to_value() DOM tree --- src/bin/index-as2rel-collector.rs | 322 +++++++++++++++++------------- src/bin/index-pfx2as-collector.rs | 215 ++++++++++++-------- 2 files changed, 320 insertions(+), 217 deletions(-) diff --git a/src/bin/index-as2rel-collector.rs b/src/bin/index-as2rel-collector.rs index 05da0f3..bf87377 100644 --- a/src/bin/index-as2rel-collector.rs +++ b/src/bin/index-as2rel-collector.rs @@ -1,15 +1,14 @@ use chrono::{NaiveDate, Utc}; use clap::Parser; use peer_stats::As2Rel; -use serde::Serialize; use std::collections::HashMap; -use std::io::Read; +use std::io::{BufWriter, Read, Write}; use std::path::PathBuf; use tracing::info; use walkdir::WalkDir; /// Index AS relationship data with per-collector provenance. -/// Produces as2rel-collector-latest.json.bz2 with collector-level breakdown. +/// Uses flat vector + sort + group to stay memory-efficient. #[derive(Parser, Debug)] struct Opts { /// Path to output directory @@ -27,36 +26,6 @@ struct Opts { allow_previous_day: bool, } -#[derive(Debug, Clone, Serialize)] -struct CollectorDetail { - project: String, - collector: String, - paths_count: usize, - peers_count: usize, -} - -#[derive(Debug, Clone, Serialize)] -struct As2RelCollectorEntry { - asn1: u32, - asn2: u32, - rel: u8, - /// Total paths across all collectors - total_paths_count: usize, - /// Total unique peers across all collectors (deduplicated by IP) - total_peers_count: usize, - /// Number of distinct collectors observing this relationship - collector_count: usize, - /// Per-collector breakdown - collectors: HashMap, -} - -#[derive(Debug, Clone, Serialize)] -struct As2RelCollectorOutput { - generated_at: String, - input_files: usize, - entries: Vec, -} - fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { let date_part = file_path.split('_').collect::>(); let parts = date_part[date_part.len() - 2] @@ -69,120 +38,201 @@ fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { ) } -fn main() { - let opts = Opts::parse(); +/// Flat record: (asn1, asn2, rel, collector_idx, paths_count, peers_count) +type FlatRecord = (u32, u32, u8, u16, usize, usize); - if opts.debug { - tracing_subscriber::fmt() - .with_max_level(tracing::Level::INFO) - .init(); - } +struct CollectorInfo { + name: String, + project: String, +} - for file_prefix in ["as2rel_", "as2rel-v4_", "as2rel-v6_"] { - let file_paths = WalkDir::new(opts.data_dir.to_str().unwrap()) - .follow_links(true) - .into_iter() - .filter_map(|e| match e.ok() { - Some(entry) => { - let path: String = entry.path().to_str().unwrap().to_string(); - let path_str = path.as_str(); - if path_str.contains(file_prefix) && path_str.ends_with(".bz2") { - let (year, month, day) = get_ymd_from_file(path.as_str()); - let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); - let ts = Utc::now().date_naive(); - if file_date == ts { - return Some(path); - } - if opts.allow_previous_day && file_date == ts.pred_opt().unwrap() { - return Some(path); - } +fn process_prefix(file_prefix: &str, opts: &Opts) { + let file_paths: Vec = WalkDir::new(opts.data_dir.to_str().unwrap()) + .follow_links(true) + .into_iter() + .filter_map(|e| match e.ok() { + Some(entry) => { + let path: String = entry.path().to_str().unwrap().to_string(); + if path.contains(file_prefix) && path.ends_with(".bz2") { + let (year, month, day) = get_ymd_from_file(path.as_str()); + let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); + let ts = Utc::now().date_naive(); + if file_date == ts { + return Some(path); + } + if opts.allow_previous_day && file_date == ts.pred_opt().unwrap() { + return Some(path); } - None } - None => None, - }) - .collect::>(); - - if file_paths.is_empty() { - info!( - "no matching current date {} files found, skipping", - file_prefix - ); - continue; + None + } + None => None, + }) + .collect(); + + if file_paths.is_empty() { + info!( + "no matching current date {} files found, skipping", + file_prefix + ); + return; + } + + let input_file_count = file_paths.len(); + + // Collector name → compact index + let mut collector_index: HashMap = HashMap::new(); + let mut collector_info: Vec = Vec::new(); + + // Flat vector of records — much more memory-efficient than nested HashMaps + let mut records: Vec = Vec::new(); + + for file in &file_paths { + info!("processing {}", file.as_str()); + let mut data = String::new(); + oneio::get_reader(file.as_str()) + .unwrap() + .read_to_string(&mut data) + .unwrap(); + // Drop the raw string memory as soon as parsing is done + let as2rel_info: As2Rel = { + let result = serde_json::from_str(&data); + drop(data); + result.unwrap() + }; + + let project = &as2rel_info.project; + let collector = &as2rel_info.collector; + + // Get or assign collector index + let cidx = if let Some(&idx) = collector_index.get(collector) { + idx + } else { + let idx = collector_info.len() as u16; + collector_index.insert(collector.clone(), idx); + collector_info.push(CollectorInfo { + name: collector.clone(), + project: project.clone(), + }); + idx + }; + + for as2rel in &as2rel_info.as2rel { + records.push(( + as2rel.asn1, + as2rel.asn2, + as2rel.rel, + cidx, + as2rel.paths_count, + as2rel.peers_count, + )); + } + } + + info!( + "collected {} flat records from {} files, {} unique collectors", + records.len(), + input_file_count, + collector_info.len() + ); + + // Sort by (asn1, asn2, rel, collector_idx) + records.sort_unstable_by_key(|r| (r.0, r.1, r.2, r.3)); + + // Group and write JSON manually (streaming — no intermediate Value tree) + let output_file = format!( + "{}/{}-collector-latest.json.bz2", + opts.output_dir.to_str().unwrap(), + file_prefix.strip_suffix('_').unwrap() + ); + let file = std::fs::File::create(&output_file).unwrap(); + let compressor = bzip2::write::BzEncoder::new(file, bzip2::Compression::best()); + let mut writer = BufWriter::with_capacity(256 * 1024, compressor); + + let generated_at = Utc::now().to_rfc3339(); + + // Write JSON header + write!( + writer, + "{{\"generated_at\":{},\"input_files\":{},\"entries\":[", + serde_json::to_string(&generated_at).unwrap(), + input_file_count + ) + .unwrap(); + + let mut first_group = true; + let mut i = 0; + while i < records.len() { + let key = (records[i].0, records[i].1, records[i].2); + + // Find all records for this key + let group_start = i; + let mut total_paths = 0usize; + let mut total_peers = 0usize; + while i < records.len() + && records[i].0 == key.0 + && records[i].1 == key.1 + && records[i].2 == key.2 + { + total_paths += records[i].4; + total_peers += records[i].5; + i += 1; } + let group_slice = &records[group_start..i]; + let collector_count = group_slice.len(); - let input_file_count = file_paths.len(); - - // Key: (asn1, asn2, rel) — Value: per-collector details - let mut collector_map: HashMap<(u32, u32, u8), HashMap> = - HashMap::new(); - - for file in &file_paths { - info!("processing {}", file.as_str()); - let mut data = String::new(); - oneio::get_reader(file.as_str()) - .unwrap() - .read_to_string(&mut data) - .unwrap(); - let as2rel_info: As2Rel = serde_json::from_str(&data).unwrap(); - - let project = as2rel_info.project; - let collector = as2rel_info.collector; - - for as2rel in as2rel_info.as2rel { - let key = (as2rel.asn1, as2rel.asn2, as2rel.rel); - let per_collector = collector_map.entry(key).or_default(); - let detail = - per_collector - .entry(collector.clone()) - .or_insert_with(|| CollectorDetail { - project: project.clone(), - collector: collector.clone(), - paths_count: 0, - peers_count: 0, - }); - detail.paths_count += as2rel.paths_count; - detail.peers_count += as2rel.peers_count; + // Write comma separator + if !first_group { + write!(writer, ",").unwrap(); + } + first_group = false; + + // Write entry header + write!( + writer, + "{{\"asn1\":{},\"asn2\":{},\"rel\":{},\"total_paths_count\":{},\"total_peers_count\":{},\"collector_count\":{},\"collectors\":{{", + key.0, key.1, key.2, total_paths, total_peers, collector_count + ) + .unwrap(); + + // Write per-collector breakdown + for (j, rec) in group_slice.iter().enumerate() { + let info = &collector_info[rec.3 as usize]; + if j > 0 { + write!(writer, ",").unwrap(); } + write!( + writer, + "{}:{{\"project\":{},\"collector\":{},\"paths_count\":{},\"peers_count\":{}}}", + serde_json::to_string(&info.name).unwrap(), + serde_json::to_string(&info.project).unwrap(), + serde_json::to_string(&info.name).unwrap(), + rec.4, + rec.5 + ) + .unwrap(); } - let entries: Vec = collector_map - .into_iter() - .map(|((asn1, asn2, rel), per_collector)| { - let total_paths_count: usize = per_collector.values().map(|d| d.paths_count).sum(); - let total_peers_count: usize = per_collector.values().map(|d| d.peers_count).sum(); - let collector_count = per_collector.len(); - - As2RelCollectorEntry { - asn1, - asn2, - rel, - total_paths_count, - total_peers_count, - collector_count, - collectors: per_collector, - } - }) - .collect(); + write!(writer, "}}}}").unwrap(); + } - let output = As2RelCollectorOutput { - generated_at: Utc::now().to_rfc3339(), - input_files: input_file_count, - entries, - }; + write!(writer, "]}}").unwrap(); + writer.flush().unwrap(); - let output_file = format!( - "{}/{}-collector-latest.json.bz2", - opts.output_dir.to_str().unwrap(), - file_prefix.strip_suffix('_').unwrap() - ); - let mut writer = oneio::get_writer(output_file.as_str()).unwrap(); - let _ = writer.write_all( - serde_json::to_string_pretty(&serde_json::to_value(&output).unwrap()) - .unwrap() - .as_ref(), - ); - info!("wrote {} entries to {}", file_paths.len(), output_file); + info!("wrote output to {}", output_file); +} + +fn main() { + let opts = Opts::parse(); + + if opts.debug { + tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .init(); + } + + for file_prefix in ["as2rel_", "as2rel-v4_", "as2rel-v6_"] { + process_prefix(file_prefix, &opts); } } diff --git a/src/bin/index-pfx2as-collector.rs b/src/bin/index-pfx2as-collector.rs index 960b049..a66e5fc 100644 --- a/src/bin/index-pfx2as-collector.rs +++ b/src/bin/index-pfx2as-collector.rs @@ -1,15 +1,14 @@ use chrono::{NaiveDate, Utc}; use clap::Parser; use peer_stats::Prefix2As; -use serde::Serialize; use std::collections::HashMap; -use std::io::Read; +use std::io::{BufWriter, Read, Write}; use std::path::PathBuf; use tracing::info; use walkdir::WalkDir; /// Index prefix-to-AS mapping data with per-collector provenance. -/// Produces pfx2as-collector-latest.json.bz2 with collector-level breakdown. +/// Uses flat vector + sort + group to stay memory-efficient. #[derive(Parser, Debug)] struct Opts { /// Path to output directory (file named pfx2as-collector-latest.json.bz2) @@ -27,32 +26,6 @@ struct Opts { allow_previous_day: bool, } -#[derive(Debug, Clone, Serialize)] -struct Pfx2AsCollectorDetail { - project: String, - collector: String, - count: usize, -} - -#[derive(Debug, Clone, Serialize)] -struct Pfx2AsCollectorEntry { - prefix: String, - asn: u32, - /// Total count across all collectors - total_count: usize, - /// Number of distinct collectors seeing this mapping - collector_count: usize, - /// Per-collector breakdown - collectors: HashMap, -} - -#[derive(Debug, Clone, Serialize)] -struct Pfx2AsCollectorOutput { - generated_at: String, - input_files: usize, - entries: Vec, -} - fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { let date_part = file_path.split('_').collect::>(); let parts = date_part[date_part.len() - 2] @@ -65,6 +38,15 @@ fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { ) } +/// Flat record: (prefix, asn, collector_idx, count) +/// We use a string interning approach: store prefix strings in a vec and reference by index. +type FlatRecord = (u32, u32, u16, usize); // (prefix_idx, asn, collector_idx, count) + +struct CollectorInfo { + name: String, + project: String, +} + fn main() { let opts = Opts::parse(); @@ -74,14 +56,13 @@ fn main() { .init(); } - let file_paths = WalkDir::new(opts.data_dir.to_str().unwrap()) + let file_paths: Vec = WalkDir::new(opts.data_dir.to_str().unwrap()) .follow_links(true) .into_iter() .filter_map(|e| match e.ok() { Some(entry) => { let path: String = entry.path().to_str().unwrap().to_string(); - let path_str = path.as_str(); - if path_str.contains("pfx2as_") && path_str.ends_with(".bz2") { + if path.contains("pfx2as_") && path.ends_with(".bz2") { let (year, month, day) = get_ymd_from_file(path.as_str()); let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); let ts = Utc::now().date_naive(); @@ -96,7 +77,7 @@ fn main() { } None => None, }) - .collect::>(); + .collect(); if file_paths.is_empty() { info!("no data files found, skipping"); @@ -105,9 +86,16 @@ fn main() { let input_file_count = file_paths.len(); - // Key: (prefix, asn) — Value: per-collector details - let mut collector_map: HashMap<(String, u32), HashMap> = - HashMap::new(); + // Collector name → compact index + let mut collector_index: HashMap = HashMap::new(); + let mut collector_info: Vec = Vec::new(); + + // Prefix string → compact index (interning) + let mut prefix_index: HashMap = HashMap::new(); + let mut prefix_strings: Vec = Vec::new(); + + // Flat vector of records + let mut records: Vec = Vec::new(); for file in &file_paths { info!("processing {}", file.as_str()); @@ -116,60 +104,125 @@ fn main() { .unwrap() .read_to_string(&mut data) .unwrap(); - let pfx2as_info: Prefix2As = serde_json::from_str(&data).unwrap(); - - let project = pfx2as_info.project; - let collector = pfx2as_info.collector; - - for pfx2as in pfx2as_info.pfx2as { - let key = (pfx2as.prefix.clone(), pfx2as.asn); - let per_collector = collector_map.entry(key).or_default(); - let detail = - per_collector - .entry(collector.clone()) - .or_insert_with(|| Pfx2AsCollectorDetail { - project: project.clone(), - collector: collector.clone(), - count: 0, - }); - detail.count += pfx2as.count; + let pfx2as_info: Prefix2As = { + let result = serde_json::from_str(&data); + drop(data); + result.unwrap() + }; + + let project = &pfx2as_info.project; + let collector = &pfx2as_info.collector; + + let cidx = if let Some(&idx) = collector_index.get(collector) { + idx + } else { + let idx = collector_info.len() as u16; + collector_index.insert(collector.clone(), idx); + collector_info.push(CollectorInfo { + name: collector.clone(), + project: project.clone(), + }); + idx + }; + + for pfx2as in &pfx2as_info.pfx2as { + let pidx = if let Some(&idx) = prefix_index.get(&pfx2as.prefix) { + idx + } else { + let idx = prefix_strings.len() as u32; + prefix_index.insert(pfx2as.prefix.clone(), idx); + prefix_strings.push(pfx2as.prefix.clone()); + idx + }; + + records.push((pidx, pfx2as.asn, cidx, pfx2as.count)); } } - let entry_count = collector_map.len(); - let entries: Vec = collector_map - .into_iter() - .map(|((prefix, asn), per_collector)| { - let total_count: usize = per_collector.values().map(|d| d.count).sum(); - let collector_count = per_collector.len(); - - Pfx2AsCollectorEntry { - prefix, - asn, - total_count, - collector_count, - collectors: per_collector, - } - }) - .collect(); + info!( + "collected {} flat records from {} files, {} unique collectors, {} unique prefixes", + records.len(), + input_file_count, + collector_info.len(), + prefix_strings.len() + ); - let output = Pfx2AsCollectorOutput { - generated_at: Utc::now().to_rfc3339(), - input_files: input_file_count, - entries, - }; + // Sort by (prefix_idx, asn, collector_idx) + records.sort_unstable_by_key(|r| (r.0, r.1, r.2)); + // Group and write JSON manually let output_file = format!( "{}/pfx2as-collector-latest.json.bz2", opts.output_dir.to_str().unwrap() ); - let mut writer = oneio::get_writer(output_file.as_str()).unwrap(); - let _ = writer.write_all( - serde_json::to_string_pretty(&serde_json::to_value(&output).unwrap()) - .unwrap() - .as_ref(), - ); - info!("wrote {} entries to {}", entry_count, output_file); + let file = std::fs::File::create(&output_file).unwrap(); + let compressor = bzip2::write::BzEncoder::new(file, bzip2::Compression::best()); + let mut writer = BufWriter::with_capacity(256 * 1024, compressor); + + let generated_at = Utc::now().to_rfc3339(); + + write!( + writer, + "{{\"generated_at\":{},\"input_files\":{},\"entries\":[", + serde_json::to_string(&generated_at).unwrap(), + input_file_count + ) + .unwrap(); + + let mut first_group = true; + let mut i = 0; + while i < records.len() { + let key = (records[i].0, records[i].1); // (prefix_idx, asn) + + let group_start = i; + let mut total_count = 0usize; + while i < records.len() && records[i].0 == key.0 && records[i].1 == key.1 { + total_count += records[i].3; + i += 1; + } + let group_slice = &records[group_start..i]; + let collector_count = group_slice.len(); + + if !first_group { + write!(writer, ",").unwrap(); + } + first_group = false; + + let prefix_str = &prefix_strings[key.0 as usize]; + + write!( + writer, + "{{\"prefix\":{},\"asn\":{},\"total_count\":{},\"collector_count\":{},\"collectors\":{{", + serde_json::to_string(prefix_str).unwrap(), + key.1, + total_count, + collector_count + ) + .unwrap(); + + for (j, rec) in group_slice.iter().enumerate() { + let info = &collector_info[rec.2 as usize]; + if j > 0 { + write!(writer, ",").unwrap(); + } + write!( + writer, + "{}:{{\"project\":{},\"collector\":{},\"count\":{}}}", + serde_json::to_string(&info.name).unwrap(), + serde_json::to_string(&info.project).unwrap(), + serde_json::to_string(&info.name).unwrap(), + rec.3 + ) + .unwrap(); + } + + write!(writer, "}}}}").unwrap(); + } + + write!(writer, "]}}").unwrap(); + writer.flush().unwrap(); + + info!("wrote output to {}", output_file); } #[cfg(test)] From ae4b102c2d5ebd131bc9d4beb49c8a55a8f45aad Mon Sep 17 00:00:00 2001 From: Mingwei Zhang Date: Thu, 23 Jul 2026 10:39:19 -0700 Subject: [PATCH 3/3] feat: add collector-aware as2rel output alongside classic aggregate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewrite as2rel-index to produce BOTH outputs from a single file pass: 1. Classic: {prefix}-latest.json.bz2 (unchanged v1 format — no breaking change) 2. Collector: {prefix}-collector-latest.json.bz2 (new per-collector provenance) Uses flat vector + sort approach for memory efficiency (~1.7GB peak RSS for 21M records). All three prefixes (as2rel, as2rel-v4, as2rel-v6) get both output files. Remove standalone as2rel-collector-index and pfx2as-collector-index binaries — collector output is now integrated into as2rel-index. --- Cargo.toml | 8 - src/bin/index-as2rel-collector.rs | 254 ------------------------ src/bin/index-as2rel.rs | 319 ++++++++++++++++++++++-------- src/bin/index-pfx2as-collector.rs | 243 ----------------------- 4 files changed, 236 insertions(+), 588 deletions(-) delete mode 100644 src/bin/index-as2rel-collector.rs delete mode 100644 src/bin/index-pfx2as-collector.rs diff --git a/Cargo.toml b/Cargo.toml index ab40e29..ade074a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,11 +61,3 @@ path = "src/bin/index-as2rel.rs" [[bin]] name = "pfx2as-index" path = "src/bin/index-pfx2as.rs" - -[[bin]] -name = "as2rel-collector-index" -path = "src/bin/index-as2rel-collector.rs" - -[[bin]] -name = "pfx2as-collector-index" -path = "src/bin/index-pfx2as-collector.rs" diff --git a/src/bin/index-as2rel-collector.rs b/src/bin/index-as2rel-collector.rs deleted file mode 100644 index bf87377..0000000 --- a/src/bin/index-as2rel-collector.rs +++ /dev/null @@ -1,254 +0,0 @@ -use chrono::{NaiveDate, Utc}; -use clap::Parser; -use peer_stats::As2Rel; -use std::collections::HashMap; -use std::io::{BufWriter, Read, Write}; -use std::path::PathBuf; -use tracing::info; -use walkdir::WalkDir; - -/// Index AS relationship data with per-collector provenance. -/// Uses flat vector + sort + group to stay memory-efficient. -#[derive(Parser, Debug)] -struct Opts { - /// Path to output directory - output_dir: PathBuf, - - /// Path to the data file directory - data_dir: PathBuf, - - /// Whether to print debug logs - #[clap(long)] - debug: bool, - - /// Allow processing files from the previous day - #[clap(long)] - allow_previous_day: bool, -} - -fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { - let date_part = file_path.split('_').collect::>(); - let parts = date_part[date_part.len() - 2] - .split('-') - .collect::>(); - ( - parts[0].parse::().unwrap(), - parts[1].parse::().unwrap(), - parts[2].parse::().unwrap(), - ) -} - -/// Flat record: (asn1, asn2, rel, collector_idx, paths_count, peers_count) -type FlatRecord = (u32, u32, u8, u16, usize, usize); - -struct CollectorInfo { - name: String, - project: String, -} - -fn process_prefix(file_prefix: &str, opts: &Opts) { - let file_paths: Vec = WalkDir::new(opts.data_dir.to_str().unwrap()) - .follow_links(true) - .into_iter() - .filter_map(|e| match e.ok() { - Some(entry) => { - let path: String = entry.path().to_str().unwrap().to_string(); - if path.contains(file_prefix) && path.ends_with(".bz2") { - let (year, month, day) = get_ymd_from_file(path.as_str()); - let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); - let ts = Utc::now().date_naive(); - if file_date == ts { - return Some(path); - } - if opts.allow_previous_day && file_date == ts.pred_opt().unwrap() { - return Some(path); - } - } - None - } - None => None, - }) - .collect(); - - if file_paths.is_empty() { - info!( - "no matching current date {} files found, skipping", - file_prefix - ); - return; - } - - let input_file_count = file_paths.len(); - - // Collector name → compact index - let mut collector_index: HashMap = HashMap::new(); - let mut collector_info: Vec = Vec::new(); - - // Flat vector of records — much more memory-efficient than nested HashMaps - let mut records: Vec = Vec::new(); - - for file in &file_paths { - info!("processing {}", file.as_str()); - let mut data = String::new(); - oneio::get_reader(file.as_str()) - .unwrap() - .read_to_string(&mut data) - .unwrap(); - // Drop the raw string memory as soon as parsing is done - let as2rel_info: As2Rel = { - let result = serde_json::from_str(&data); - drop(data); - result.unwrap() - }; - - let project = &as2rel_info.project; - let collector = &as2rel_info.collector; - - // Get or assign collector index - let cidx = if let Some(&idx) = collector_index.get(collector) { - idx - } else { - let idx = collector_info.len() as u16; - collector_index.insert(collector.clone(), idx); - collector_info.push(CollectorInfo { - name: collector.clone(), - project: project.clone(), - }); - idx - }; - - for as2rel in &as2rel_info.as2rel { - records.push(( - as2rel.asn1, - as2rel.asn2, - as2rel.rel, - cidx, - as2rel.paths_count, - as2rel.peers_count, - )); - } - } - - info!( - "collected {} flat records from {} files, {} unique collectors", - records.len(), - input_file_count, - collector_info.len() - ); - - // Sort by (asn1, asn2, rel, collector_idx) - records.sort_unstable_by_key(|r| (r.0, r.1, r.2, r.3)); - - // Group and write JSON manually (streaming — no intermediate Value tree) - let output_file = format!( - "{}/{}-collector-latest.json.bz2", - opts.output_dir.to_str().unwrap(), - file_prefix.strip_suffix('_').unwrap() - ); - let file = std::fs::File::create(&output_file).unwrap(); - let compressor = bzip2::write::BzEncoder::new(file, bzip2::Compression::best()); - let mut writer = BufWriter::with_capacity(256 * 1024, compressor); - - let generated_at = Utc::now().to_rfc3339(); - - // Write JSON header - write!( - writer, - "{{\"generated_at\":{},\"input_files\":{},\"entries\":[", - serde_json::to_string(&generated_at).unwrap(), - input_file_count - ) - .unwrap(); - - let mut first_group = true; - let mut i = 0; - while i < records.len() { - let key = (records[i].0, records[i].1, records[i].2); - - // Find all records for this key - let group_start = i; - let mut total_paths = 0usize; - let mut total_peers = 0usize; - while i < records.len() - && records[i].0 == key.0 - && records[i].1 == key.1 - && records[i].2 == key.2 - { - total_paths += records[i].4; - total_peers += records[i].5; - i += 1; - } - let group_slice = &records[group_start..i]; - let collector_count = group_slice.len(); - - // Write comma separator - if !first_group { - write!(writer, ",").unwrap(); - } - first_group = false; - - // Write entry header - write!( - writer, - "{{\"asn1\":{},\"asn2\":{},\"rel\":{},\"total_paths_count\":{},\"total_peers_count\":{},\"collector_count\":{},\"collectors\":{{", - key.0, key.1, key.2, total_paths, total_peers, collector_count - ) - .unwrap(); - - // Write per-collector breakdown - for (j, rec) in group_slice.iter().enumerate() { - let info = &collector_info[rec.3 as usize]; - if j > 0 { - write!(writer, ",").unwrap(); - } - write!( - writer, - "{}:{{\"project\":{},\"collector\":{},\"paths_count\":{},\"peers_count\":{}}}", - serde_json::to_string(&info.name).unwrap(), - serde_json::to_string(&info.project).unwrap(), - serde_json::to_string(&info.name).unwrap(), - rec.4, - rec.5 - ) - .unwrap(); - } - - write!(writer, "}}}}").unwrap(); - } - - write!(writer, "]}}").unwrap(); - writer.flush().unwrap(); - - info!("wrote output to {}", output_file); -} - -fn main() { - let opts = Opts::parse(); - - if opts.debug { - tracing_subscriber::fmt() - .with_max_level(tracing::Level::INFO) - .init(); - } - - for file_prefix in ["as2rel_", "as2rel-v4_", "as2rel-v6_"] { - process_prefix(file_prefix, &opts); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_get_file_date() { - assert_eq!( - get_ymd_from_file("as2rel_rrc16_2022-02-01_1643673600.bz2"), - (2022, 2, 1) - ); - assert_eq!( - get_ymd_from_file("/aaa_bbb-ccc/as2rel_rrc16_2022-02-01_1643673600.bz2"), - (2022, 2, 1) - ); - } -} diff --git a/src/bin/index-as2rel.rs b/src/bin/index-as2rel.rs index ec89191..3bd4764 100644 --- a/src/bin/index-as2rel.rs +++ b/src/bin/index-as2rel.rs @@ -1,26 +1,33 @@ use chrono::{NaiveDate, Utc}; use clap::Parser; -use peer_stats::{As2Rel, As2RelCount}; -use serde_json::json; +use peer_stats::As2Rel; use std::collections::HashMap; -use std::io::Read; +use std::io::{BufWriter, Read, Write}; use std::path::PathBuf; use tracing::info; use walkdir::WalkDir; -/// peer-stats is a CLI tool that collects peer information from a given RIB dump file. +/// Index AS relationship data from per-collector daily files. +/// +/// Produces TWO output files per prefix (as2rel, as2rel-v4, as2rel-v6): +/// 1. Classic: {prefix}-latest.json.bz2 — unchanged v1 aggregate +/// 2. Collector: {prefix}-collector-latest.json.bz2 — per-collector provenance +/// +/// Both are generated from a single pass over the daily files. +/// The classic output format is identical to the previous indexer — no breaking changes. #[derive(Parser, Debug)] struct Opts { - /// Path to output file + /// Path to output directory output_dir: PathBuf, /// Path to the data file directory data_dir: PathBuf, - /// whether to print debug + /// Whether to print debug logs #[clap(long)] debug: bool, + /// Allow processing files from the previous day #[clap(long)] allow_previous_day: bool, } @@ -37,93 +44,239 @@ fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { ) } -fn main() { - let opts = Opts::parse(); +/// Flat record: (asn1, asn2, rel, collector_idx, paths_count, peers_count) +type FlatRecord = (u32, u32, u8, u16, usize, usize); - if opts.debug { - tracing_subscriber::fmt() - .with_max_level(tracing::Level::INFO) - .init(); - } +struct CollectorInfo { + name: String, + project: String, +} - for file_prefix in ["as2rel_", "as2rel-v4_", "as2rel-v6_"] { - let file_paths = WalkDir::new(opts.data_dir.to_str().unwrap()) - .follow_links(true) - .into_iter() - .filter_map(|e| match e.ok() { - Some(entry) => { - let path: String = entry.path().to_str().unwrap().to_string(); - let path_str = path.as_str(); - if path_str.contains(file_prefix) && path_str.ends_with(".bz2") { - let (year, month, day) = get_ymd_from_file(path.as_str()); - let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); - let ts = Utc::now().date_naive(); - if file_date == ts { - return Some(path); - } - if opts.allow_previous_day && file_date == ts.pred_opt().unwrap() { - return Some(path); - } +fn process_prefix(file_prefix: &str, opts: &Opts) { + let file_paths: Vec = WalkDir::new(opts.data_dir.to_str().unwrap()) + .follow_links(true) + .into_iter() + .filter_map(|e| match e.ok() { + Some(entry) => { + let path: String = entry.path().to_str().unwrap().to_string(); + if path.contains(file_prefix) && path.ends_with(".bz2") { + let (year, month, day) = get_ymd_from_file(path.as_str()); + let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); + let ts = Utc::now().date_naive(); + if file_date == ts { + return Some(path); + } + if opts.allow_previous_day && file_date == ts.pred_opt().unwrap() { + return Some(path); } - None } - None => None, - }) - .collect::>(); - - if file_paths.is_empty() { - info!( - "no matching current date {} file found, skipping", - file_prefix - ); - return; + None + } + None => None, + }) + .collect(); + + if file_paths.is_empty() { + info!( + "no matching current date {} files found, skipping", + file_prefix + ); + return; + } + + let input_file_count = file_paths.len(); + + // Collector name → compact index + let mut collector_index: HashMap = HashMap::new(); + let mut collector_info: Vec = Vec::new(); + + // --- Phase 1: collect flat records --- + let mut records: Vec = Vec::new(); + + for file in &file_paths { + info!("reading {}", file.as_str()); + let mut data = String::new(); + oneio::get_reader(file.as_str()) + .unwrap() + .read_to_string(&mut data) + .unwrap(); + let as2rel_info: As2Rel = { + let result = serde_json::from_str(&data); + drop(data); + result.unwrap() + }; + + let collector = &as2rel_info.collector; + let project = &as2rel_info.project; + + let cidx = if let Some(&idx) = collector_index.get(collector) { + idx + } else { + let idx = collector_info.len() as u16; + collector_index.insert(collector.clone(), idx); + collector_info.push(CollectorInfo { + name: collector.clone(), + project: project.clone(), + }); + idx + }; + + for as2rel in &as2rel_info.as2rel { + records.push(( + as2rel.asn1, + as2rel.asn2, + as2rel.rel, + cidx, + as2rel.paths_count, + as2rel.peers_count, + )); } + } - let mut data_map: HashMap<(u32, u32, u8), (usize, usize)> = HashMap::new(); - - for file in file_paths { - info!("processing {}", file.as_str()); - let mut data = "".to_string(); - oneio::get_reader(file.as_str()) - .unwrap() - .read_to_string(&mut data) - .unwrap(); - let as2rel_info: As2Rel = serde_json::from_str(&data).unwrap(); - - for as2rel in as2rel_info.as2rel { - let (asn1, asn2, rel, paths_count, peers_count) = ( - as2rel.asn1, - as2rel.asn2, - as2rel.rel, - as2rel.paths_count, - as2rel.peers_count, - ); - let (count_1, count_2) = data_map.entry((asn1, asn2, rel)).or_insert((0, 0)); - *count_1 += paths_count; - *count_2 += peers_count; - } + info!( + "{}: collected {} records from {} files, {} collectors", + file_prefix, + records.len(), + input_file_count, + collector_info.len() + ); + + // --- Phase 2: sort by key + collector --- + records.sort_unstable_by_key(|r| (r.0, r.1, r.2, r.3)); + + // --- Phase 3: group and write both outputs --- + let base_name = file_prefix.strip_suffix('_').unwrap(); + + // Classic output + let classic_file = format!( + "{}/{}-latest.json.bz2", + opts.output_dir.to_str().unwrap(), + base_name + ); + let classic_raw = std::fs::File::create(&classic_file).unwrap(); + let classic_comp = bzip2::write::BzEncoder::new(classic_raw, bzip2::Compression::best()); + let mut classic_w = BufWriter::with_capacity(256 * 1024, classic_comp); + + // Collector output + let collector_file = format!( + "{}/{}-collector-latest.json.bz2", + opts.output_dir.to_str().unwrap(), + base_name + ); + let collector_raw = std::fs::File::create(&collector_file).unwrap(); + let collector_comp = bzip2::write::BzEncoder::new(collector_raw, bzip2::Compression::best()); + let mut collector_w = BufWriter::with_capacity(256 * 1024, collector_comp); + + let generated_at = Utc::now().to_rfc3339(); + + // Classic: opening bracket + write!(classic_w, "[").unwrap(); + + // Collector: header + write!( + collector_w, + "{{\"generated_at\":{},\"input_files\":{},\"entries\":[", + serde_json::to_string(&generated_at).unwrap(), + input_file_count + ) + .unwrap(); + + let mut first_classic = true; + let mut first_collector = true; + let mut i = 0; + let mut entry_count: usize = 0; + + while i < records.len() { + let key = (records[i].0, records[i].1, records[i].2); + + // Find all records for this (asn1, asn2, rel) key + let group_start = i; + let mut total_paths = 0usize; + let mut total_peers = 0usize; + while i < records.len() + && records[i].0 == key.0 + && records[i].1 == key.1 + && records[i].2 == key.2 + { + total_paths += records[i].4; + total_peers += records[i].5; + i += 1; + } + let group_slice = &records[group_start..i]; + let collector_count = group_slice.len(); + entry_count += 1; + + // --- Classic output: single As2RelCount --- + if !first_classic { + write!(classic_w, ",").unwrap(); } + first_classic = false; + write!( + classic_w, + "{{\"asn1\":{},\"asn2\":{},\"rel\":{},\"paths_count\":{},\"peers_count\":{}}}", + key.0, key.1, key.2, total_paths, total_peers + ) + .unwrap(); - let res: Vec = data_map - .into_iter() - .map( - |((asn1, asn2, rel), (paths_count, peers_count))| As2RelCount { - asn1, - asn2, - rel, - paths_count, - peers_count, - }, + // --- Collector output: full entry with collectors map --- + if !first_collector { + write!(collector_w, ",").unwrap(); + } + first_collector = false; + write!( + collector_w, + "{{\"asn1\":{},\"asn2\":{},\"rel\":{},\"total_paths_count\":{},\"total_peers_count\":{},\"collector_count\":{},\"collectors\":{{", + key.0, key.1, key.2, total_paths, total_peers, collector_count + ) + .unwrap(); + + for (j, rec) in group_slice.iter().enumerate() { + let info = &collector_info[rec.3 as usize]; + if j > 0 { + write!(collector_w, ",").unwrap(); + } + write!( + collector_w, + "{}:{{\"project\":{},\"collector\":{},\"paths_count\":{},\"peers_count\":{}}}", + serde_json::to_string(&info.name).unwrap(), + serde_json::to_string(&info.project).unwrap(), + serde_json::to_string(&info.name).unwrap(), + rec.4, + rec.5 ) - .collect(); + .unwrap(); + } - let output_file = format!( - "{}/{}-latest.json.bz2", - opts.output_dir.to_str().unwrap(), - file_prefix.strip_suffix('_').unwrap() - ); - let mut writer = oneio::get_writer(output_file.as_str()).unwrap(); - let _ = writer.write_all(serde_json::to_string_pretty(&json!(res)).unwrap().as_ref()); + write!(collector_w, "}}}}").unwrap(); + } + + // Classic: closing bracket + write!(classic_w, "]").unwrap(); + classic_w.flush().unwrap(); + drop(classic_w); + + // Collector: closing brackets + write!(collector_w, "]}}").unwrap(); + collector_w.flush().unwrap(); + drop(collector_w); + + info!( + "{}: wrote {} entries, classic={} collector={}", + file_prefix, entry_count, classic_file, collector_file + ); +} + +fn main() { + let opts = Opts::parse(); + + if opts.debug { + tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .init(); + } + + for file_prefix in ["as2rel_", "as2rel-v4_", "as2rel-v6_"] { + process_prefix(file_prefix, &opts); } } diff --git a/src/bin/index-pfx2as-collector.rs b/src/bin/index-pfx2as-collector.rs deleted file mode 100644 index a66e5fc..0000000 --- a/src/bin/index-pfx2as-collector.rs +++ /dev/null @@ -1,243 +0,0 @@ -use chrono::{NaiveDate, Utc}; -use clap::Parser; -use peer_stats::Prefix2As; -use std::collections::HashMap; -use std::io::{BufWriter, Read, Write}; -use std::path::PathBuf; -use tracing::info; -use walkdir::WalkDir; - -/// Index prefix-to-AS mapping data with per-collector provenance. -/// Uses flat vector + sort + group to stay memory-efficient. -#[derive(Parser, Debug)] -struct Opts { - /// Path to output directory (file named pfx2as-collector-latest.json.bz2) - output_dir: PathBuf, - - /// Path to the data file directory - data_dir: PathBuf, - - /// Whether to print debug logs - #[clap(long)] - debug: bool, - - /// Allow processing files from the previous day - #[clap(long)] - allow_previous_day: bool, -} - -fn get_ymd_from_file(file_path: &str) -> (i32, u32, u32) { - let date_part = file_path.split('_').collect::>(); - let parts = date_part[date_part.len() - 2] - .split('-') - .collect::>(); - ( - parts[0].parse::().unwrap(), - parts[1].parse::().unwrap(), - parts[2].parse::().unwrap(), - ) -} - -/// Flat record: (prefix, asn, collector_idx, count) -/// We use a string interning approach: store prefix strings in a vec and reference by index. -type FlatRecord = (u32, u32, u16, usize); // (prefix_idx, asn, collector_idx, count) - -struct CollectorInfo { - name: String, - project: String, -} - -fn main() { - let opts = Opts::parse(); - - if opts.debug { - tracing_subscriber::fmt() - .with_max_level(tracing::Level::INFO) - .init(); - } - - let file_paths: Vec = WalkDir::new(opts.data_dir.to_str().unwrap()) - .follow_links(true) - .into_iter() - .filter_map(|e| match e.ok() { - Some(entry) => { - let path: String = entry.path().to_str().unwrap().to_string(); - if path.contains("pfx2as_") && path.ends_with(".bz2") { - let (year, month, day) = get_ymd_from_file(path.as_str()); - let file_date = NaiveDate::from_ymd_opt(year, month, day).unwrap(); - let ts = Utc::now().date_naive(); - if file_date == ts { - return Some(path); - } - if opts.allow_previous_day && file_date == ts.pred_opt().unwrap() { - return Some(path); - } - } - None - } - None => None, - }) - .collect(); - - if file_paths.is_empty() { - info!("no data files found, skipping"); - return; - } - - let input_file_count = file_paths.len(); - - // Collector name → compact index - let mut collector_index: HashMap = HashMap::new(); - let mut collector_info: Vec = Vec::new(); - - // Prefix string → compact index (interning) - let mut prefix_index: HashMap = HashMap::new(); - let mut prefix_strings: Vec = Vec::new(); - - // Flat vector of records - let mut records: Vec = Vec::new(); - - for file in &file_paths { - info!("processing {}", file.as_str()); - let mut data = String::new(); - oneio::get_reader(file.as_str()) - .unwrap() - .read_to_string(&mut data) - .unwrap(); - let pfx2as_info: Prefix2As = { - let result = serde_json::from_str(&data); - drop(data); - result.unwrap() - }; - - let project = &pfx2as_info.project; - let collector = &pfx2as_info.collector; - - let cidx = if let Some(&idx) = collector_index.get(collector) { - idx - } else { - let idx = collector_info.len() as u16; - collector_index.insert(collector.clone(), idx); - collector_info.push(CollectorInfo { - name: collector.clone(), - project: project.clone(), - }); - idx - }; - - for pfx2as in &pfx2as_info.pfx2as { - let pidx = if let Some(&idx) = prefix_index.get(&pfx2as.prefix) { - idx - } else { - let idx = prefix_strings.len() as u32; - prefix_index.insert(pfx2as.prefix.clone(), idx); - prefix_strings.push(pfx2as.prefix.clone()); - idx - }; - - records.push((pidx, pfx2as.asn, cidx, pfx2as.count)); - } - } - - info!( - "collected {} flat records from {} files, {} unique collectors, {} unique prefixes", - records.len(), - input_file_count, - collector_info.len(), - prefix_strings.len() - ); - - // Sort by (prefix_idx, asn, collector_idx) - records.sort_unstable_by_key(|r| (r.0, r.1, r.2)); - - // Group and write JSON manually - let output_file = format!( - "{}/pfx2as-collector-latest.json.bz2", - opts.output_dir.to_str().unwrap() - ); - let file = std::fs::File::create(&output_file).unwrap(); - let compressor = bzip2::write::BzEncoder::new(file, bzip2::Compression::best()); - let mut writer = BufWriter::with_capacity(256 * 1024, compressor); - - let generated_at = Utc::now().to_rfc3339(); - - write!( - writer, - "{{\"generated_at\":{},\"input_files\":{},\"entries\":[", - serde_json::to_string(&generated_at).unwrap(), - input_file_count - ) - .unwrap(); - - let mut first_group = true; - let mut i = 0; - while i < records.len() { - let key = (records[i].0, records[i].1); // (prefix_idx, asn) - - let group_start = i; - let mut total_count = 0usize; - while i < records.len() && records[i].0 == key.0 && records[i].1 == key.1 { - total_count += records[i].3; - i += 1; - } - let group_slice = &records[group_start..i]; - let collector_count = group_slice.len(); - - if !first_group { - write!(writer, ",").unwrap(); - } - first_group = false; - - let prefix_str = &prefix_strings[key.0 as usize]; - - write!( - writer, - "{{\"prefix\":{},\"asn\":{},\"total_count\":{},\"collector_count\":{},\"collectors\":{{", - serde_json::to_string(prefix_str).unwrap(), - key.1, - total_count, - collector_count - ) - .unwrap(); - - for (j, rec) in group_slice.iter().enumerate() { - let info = &collector_info[rec.2 as usize]; - if j > 0 { - write!(writer, ",").unwrap(); - } - write!( - writer, - "{}:{{\"project\":{},\"collector\":{},\"count\":{}}}", - serde_json::to_string(&info.name).unwrap(), - serde_json::to_string(&info.project).unwrap(), - serde_json::to_string(&info.name).unwrap(), - rec.3 - ) - .unwrap(); - } - - write!(writer, "}}}}").unwrap(); - } - - write!(writer, "]}}").unwrap(); - writer.flush().unwrap(); - - info!("wrote output to {}", output_file); -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_get_file_date() { - assert_eq!( - get_ymd_from_file("pfx2as_rrc16_2022-02-01_1643673600.bz2"), - (2022, 2, 1) - ); - assert_eq!( - get_ymd_from_file("/aaa_bbb-ccc/pfx2as_rrc16_2022-02-01_1643673600.bz2"), - (2022, 2, 1) - ); - } -}