Skip to main content

aqt_sim/simulation/
recorder.rs

1//! This module contains the `Recorder` trait and its implementations, which takes "snapshots" of
2//! the simulation/network and either prints them to the console or writes them to a file.
3
4use std::fs::OpenOptions;
5use std::io::prelude::*;
6use crate::network::Network;
7
8
9/// Trait implemented by all recorders.
10pub trait Recorder {
11    fn record(&mut self, rd: usize, prime_timestep: bool, network: &Network);
12    fn close(&mut self);
13}
14
15
16/// Print the whole network to the console.
17pub struct DebugPrintRecorder;
18
19impl DebugPrintRecorder {
20    pub fn new() -> Self {
21        DebugPrintRecorder {}
22    }
23}
24
25impl Recorder for DebugPrintRecorder {
26    fn record(&mut self, rd: usize, prime_timestep: bool, network: &Network) {
27        if prime_timestep { println!("{}':", rd) } else { println!("{}:", rd) };
28        println!("{}", network);
29    }
30
31    fn close(&mut self) {
32        println!("Simulation finished.");
33    }
34}
35
36
37const DEFAULT_LOCAL_LINE_LIMIT: usize = 5000;
38const BUFFER_LOAD_CSV_HEADER: &str = "rd,prime,buffer_from,buffer_to,load\n";
39
40/// Records the load in each buffer.
41pub struct BufferLoadRecorder {
42    local_line_limit: usize, // How many lines to keep in memory before dumping record to file. 
43    lines: Vec<String>,
44    output_filepath: String,
45}
46
47impl Recorder for BufferLoadRecorder {
48    fn record(&mut self, rd: usize, prime_timestep: bool, network: &Network) {
49        // format: rd,prime (0/1), buffer_from, buffer_to, load (see BUFFER_LOAD_CSV_HEADER)
50        let prime_flag = if prime_timestep { 1 } else { 0 };
51        for (from_id, to_id) in network.get_edgebuffers() {
52            let load = network.get_edgebuffer(from_id, to_id).unwrap().buffer.len();
53            self.write(format!("{},{},{},{},{}\n", rd, prime_flag, from_id, to_id, load));
54        }
55    }
56
57    fn close(&mut self) {
58        self.save()
59    }
60}
61
62impl BufferLoadRecorder {
63    /// Get a new `BufferLoadRecorder` with the default line limit.
64    pub fn new(output_filepath: String) -> Self {
65        BufferLoadRecorder {
66            local_line_limit: DEFAULT_LOCAL_LINE_LIMIT,
67            lines: vec![BUFFER_LOAD_CSV_HEADER.to_string()],
68            output_filepath,
69        }
70    }
71
72    /// Get a new `BufferLoadRecorder` with the specified line limit.
73    pub fn new_with_line_limit(output_filepath: String, local_line_limit: usize) -> Self {
74        BufferLoadRecorder {
75            local_line_limit,
76            lines: vec![BUFFER_LOAD_CSV_HEADER.to_string()],
77            output_filepath,
78        }
79    }
80
81    fn write(&mut self, line: String) {
82        if self.lines.len() >= self.local_line_limit {
83            // Write to file
84            self.save();
85
86            self.lines = Vec::new();
87        }
88        self.lines.push(line);
89    }
90
91    fn save(&mut self) {
92        let data = self.lines.concat();
93        let mut file = OpenOptions::new()
94            .write(true)
95            .create(true)
96            .append(true)
97            .open(&self.output_filepath)
98            .unwrap();
99
100        if let Err(e) = writeln!(file, "{}", data) {
101            eprintln!("Couldn't write to file: {}", e);
102        }
103    }
104}