manycastr/cli/writer/
mod.rs

1use std::fs::File;
2use std::io;
3use std::io::Write;
4
5use bimap::BiHashMap;
6use csv::Writer;
7use tokio::sync::mpsc::UnboundedReceiver;
8
9use crate::cli::writer::catchment_row::get_catchment_csv_row;
10use crate::cli::writer::csv_writer::get_csv_metadata;
11use crate::cli::writer::laces_row::get_laces_row;
12use crate::cli::writer::latency_row::get_latency_row;
13use crate::cli::writer::trace_row::get_trace_row;
14use crate::custom_module;
15use crate::custom_module::manycastr::MeasurementType;
16use crate::custom_module::manycastr::reply::ReplyData;
17use custom_module::manycastr::{Configuration, Reply, ReplyBatch};
18use flate2::Compression;
19use flate2::write::GzEncoder;
20use log::error;
21use std::io::BufWriter;
22
23mod catchment_row;
24pub mod csv_writer;
25mod laces_row;
26mod latency_row;
27pub mod parquet_writer;
28mod trace_row;
29
30/// Configuration for the results writing process.
31pub struct WriteConfig<'a> {
32    /// Determines whether the results should also be printed to the command-line interface.
33    pub print_to_cli: bool,
34    /// The file handle to which the measurement results should be written.
35    pub output_file: File,
36    /// Metadata for the measurement, to be written at the beginning of the output file.
37    pub metadata_args: MetadataArgs<'a>,
38    /// Measurement type
39    pub m_type: MeasurementType,
40    /// Indicates whether the measurement involves multiple origins
41    pub is_multi_origin: bool,
42    /// A bidirectional map used to convert worker IDs (u16) to their corresponding hostnames (String).
43    pub worker_map: BiHashMap<u32, String>,
44    /// Indicate whether any Origin is for CHAOS
45    pub is_chaos: bool,
46    /// Whether feed sessions are enabled (--sessions; adds a 'session' column to feed output)
47    pub is_sessions: bool,
48}
49
50/// Holds all the arguments required to metadata for the output file.
51pub struct MetadataArgs<'a> {
52    /// Path to the hitlist used.
53    pub hitlist: &'a str,
54    /// Number of targets in the hitlist.
55    pub hitlist_length: usize,
56    /// Whether the hitlist was shuffled.
57    pub is_shuffle: bool,
58    /// The probing rate used.
59    pub probing_rate: u32,
60    /// The interval between subsequent workers.
61    pub interval: u32,
62    /// A bidirectional map of all possible worker IDs to their hostnames.
63    pub all_workers: &'a BiHashMap<u32, String>,
64    /// Optional configuration file used.
65    pub configurations: &'a Vec<Configuration>,
66    /// Whether this is a responsiveness-based measurement.
67    pub is_responsive: bool,
68    /// Measurement type
69    pub m_type: MeasurementType,
70    /// Measurement start time (Unix epoch seconds).
71    pub start_time: u64,
72    /// Record to send CHAOS (TXT) or A/AAAA requests for.
73    pub record: Option<&'a str>,
74    /// URL encoded in probes (e.g., opt-out link).
75    pub url: Option<&'a str>,
76    /// Interval between probes from/to the same origin,dst pair (seconds).
77    pub probe_interval: u32,
78    /// Number of probes sent per origin,dst pair.
79    pub number_of_probes: u32,
80}
81
82struct DualWriter<W1: Write, W2: Write> {
83    file: Writer<W1>,
84    cli: Option<Writer<W2>>,
85}
86
87impl<W1: Write, W2: Write> DualWriter<W1, W2> {
88    fn write_record<I, T>(&mut self, record: I) -> csv::Result<()>
89    where
90        I: IntoIterator<Item = T> + Clone,
91        T: AsRef<[u8]>,
92    {
93        if let Some(ref mut cli) = self.cli {
94            cli.write_record(record.clone())?;
95        }
96        self.file.write_record(record)?;
97        Ok(())
98    }
99
100    fn flush(&mut self) -> io::Result<()> {
101        self.file.flush()?;
102        if let Some(ref mut cli) = self.cli {
103            cli.flush()?;
104        }
105        Ok(())
106    }
107}
108
109/// Writes the results to a file (and optionally to the command-line)
110///
111/// # Arguments
112/// * `rx` - The receiver channel that receives the results
113/// * `config` - The configuration for writing results, including file handle, metadata, and measurement type
114pub fn write_results_csv(mut rx: UnboundedReceiver<ReplyBatch>, config: WriteConfig) {
115    // Create writers (file writer and optional CLI writer)
116    let buffered_file_writer = BufWriter::new(config.output_file);
117    let mut gz_encoder = GzEncoder::new(buffered_file_writer, Compression::default());
118
119    // Write metadata to file
120    let md_lines = get_csv_metadata(config.metadata_args, &config.worker_map);
121    for line in md_lines {
122        if let Err(e) = writeln!(gz_encoder, "{line}") {
123            error!("Failed to write metadata line to Gzip stream: {e}");
124        }
125    }
126
127    let mut dual_wtr = DualWriter {
128        file: Writer::from_writer(gz_encoder),
129        cli: config
130            .print_to_cli
131            .then(|| Writer::from_writer(io::stdout())),
132    };
133
134    // Write header and flush it immediately
135    let header = get_header(
136        config.is_chaos,
137        config.is_multi_origin,
138        config.m_type,
139        config.is_sessions,
140    );
141    dual_wtr
142        .write_record(header)
143        .expect("Failed to write header to file");
144    dual_wtr.flush().expect("Failed to flush header");
145
146    tokio::spawn(async move {
147        // Receive task results from the outbound channel
148        while let Some(task_result) = rx.recv().await {
149            if task_result == ReplyBatch::default() {
150                break;
151            }
152            let results: Vec<Reply> = task_result.results;
153            let rx_id = task_result.rx_id;
154            let origin_id = task_result.origin_id;
155
156            for result in results {
157                let row = match result.reply_data {
158                    Some(data) => match data {
159                        ReplyData::Measurement(reply) => match config.m_type {
160                            MeasurementType::AnycastLatency => {
161                                get_latency_row(reply, &rx_id, &config.worker_map, origin_id)
162                            }
163                            MeasurementType::Catchment => {
164                                get_catchment_csv_row(reply, &rx_id, &config.worker_map, origin_id)
165                            }
166                            MeasurementType::Laces => {
167                                get_laces_row(reply, &rx_id, &config.worker_map, origin_id)
168                            }
169                            MeasurementType::Feed => {
170                                // Write session IDs for attribution if enabled (0 = no session)
171                                let session_id = reply.session_id;
172                                let mut row =
173                                    get_laces_row(reply, &rx_id, &config.worker_map, origin_id);
174                                if config.is_sessions {
175                                    row.push(session_id.to_string());
176                                }
177                                row
178                            }
179                            MeasurementType::AnycastTraceroute
180                            | MeasurementType::Tracemap
181                            | MeasurementType::FeedTrace => {
182                                panic!("Received regular reply during a traceroute measurement")
183                            }
184                        },
185                        ReplyData::Trace(reply) => {
186                            get_trace_row(reply, &rx_id, &config.worker_map, origin_id)
187                        }
188                        ReplyData::Discovery(_) => panic!("Discovery result forwarded to CLI"),
189                    },
190                    None => {
191                        panic!("Reply contained no result data!");
192                    }
193                };
194                // Write to command-line
195                dual_wtr
196                    .write_record(row)
197                    .expect("Failed to write record to file");
198            }
199            dual_wtr.flush().expect("Failed to flush file");
200        }
201        rx.close();
202        dual_wtr.flush().expect("Failed to flush file");
203    });
204}
205
206/// Creates the appropriate CSV header for the results file (based on the measurement type)
207///
208/// # Arguments
209/// * `is_chaos` - Whether CHAOS queries are sent
210/// * `is_multi_origin` - A boolean that determines whether multiple origins are used
211/// * `m_type` - Measurement type performed
212/// * `is_sessions` - Whether feed sessions are enabled (--sessions)
213pub fn get_header(
214    is_chaos: bool,
215    is_multi_origin: bool,
216    m_type: MeasurementType,
217    is_sessions: bool,
218) -> Vec<&'static str> {
219    // Determine headers based on measurement type
220    let mut header = match m_type {
221        MeasurementType::AnycastTraceroute | MeasurementType::Tracemap => {
222            vec!["rx", "addr", "ttl", "tx", "trace_dst", "hop_count", "rtt"]
223        }
224        MeasurementType::FeedTrace => {
225            vec!["rx", "addr", "ttl", "tx", "trace_dst", "probe_ttl", "rtt"]
226        }
227        MeasurementType::AnycastLatency => {
228            vec!["rx", "addr", "ttl", "rtt"]
229        }
230        MeasurementType::Catchment => {
231            vec!["rx", "addr", "ttl"]
232        }
233        MeasurementType::Laces | MeasurementType::Feed => {
234            // CHAOS replies carry no transmit timestamp, so there is no RTT to report
235            if is_chaos {
236                vec!["rx", "addr", "ttl", "tx"]
237            } else {
238                vec!["rx", "addr", "ttl", "tx", "rtt"]
239            }
240        }
241    };
242
243    // Optional fields (trace replies carry no CHAOS data; trace DNS probes are always A queries)
244    let is_trace = matches!(
245        m_type,
246        MeasurementType::AnycastTraceroute | MeasurementType::Tracemap | MeasurementType::FeedTrace
247    );
248    if is_chaos && !is_trace {
249        header.push("chaos_data");
250    }
251    if is_multi_origin {
252        header.push("origin_id");
253    }
254    // With --sessions, feed replies are attributed to the session of their target
255    if m_type == MeasurementType::Feed && is_sessions {
256        header.push("session");
257    }
258
259    header
260}
261
262/// Format RTT (milliseconds) as a three-decimal string (for .csv compression)
263pub fn format_rtt(rtt: f32) -> String {
264    format!("{rtt:.3}")
265}
manycastr/cli/writer/mod.rs

manycastr/cli/writer/
mod.rs
