Skip to main content

mtlog_core/
log_writer.rs

1use std::{
2    collections::HashMap,
3    fs::File,
4    io::{BufWriter, Seek, SeekFrom, Write},
5};
6
7use uuid::Uuid;
8
9pub trait LogWriter {
10    fn regular(&mut self, line: &str);
11    fn progress(&mut self, line: &str, id: Uuid);
12    fn finished(&mut self, id: Uuid);
13    fn flush(&mut self);
14}
15
16pub(crate) fn replace_line_in_file(file: &mut BufWriter<File>, line: &str, pos: u64) {
17    file.seek(SeekFrom::Start(pos)).unwrap();
18    write!(file, "{line}").unwrap();
19    file.seek(SeekFrom::End(0)).unwrap();
20}
21
22pub struct LogFile {
23    file: BufWriter<File>,
24    progress_positions: HashMap<Uuid, u64>,
25}
26
27impl LogFile {
28    pub fn new<P: AsRef<std::path::Path>>(path: P) -> Result<Self, std::io::Error> {
29        let mut file = File::options()
30            .create(true)
31            .truncate(false)
32            .write(true)
33            .open(&path)?;
34        file.seek(SeekFrom::End(0)).unwrap();
35        Ok(Self {
36            file: BufWriter::new(file),
37            progress_positions: HashMap::new(),
38        })
39    }
40}
41
42impl LogWriter for LogFile {
43    fn regular(&mut self, line: &str) {
44        writeln!(self.file, "{line}").unwrap()
45    }
46
47    fn progress(&mut self, line: &str, id: Uuid) {
48        self.flush();
49        if let Some(pos) = self.progress_positions.get(&id) {
50            replace_line_in_file(&mut self.file, line, *pos);
51        } else {
52            let pos = self.file.get_ref().metadata().unwrap().len();
53            self.progress_positions.insert(id, pos);
54            writeln!(self.file, "{line}").unwrap();
55        }
56    }
57
58    fn finished(&mut self, id: Uuid) {
59        self.progress_positions.remove(&id);
60        self.flush();
61    }
62
63    fn flush(&mut self) {
64        self.file.flush().unwrap();
65    }
66}
67
68#[derive(Default, Debug)]
69pub struct LogStdout {
70    progress_positions: HashMap<Uuid, usize>,
71    line_counter: usize,
72}
73
74impl LogWriter for LogStdout {
75    fn regular(&mut self, line: &str) {
76        if !self.progress_positions.is_empty() {
77            self.line_counter += 1;
78        }
79        println!("{line}");
80        std::io::stdout().flush().unwrap();
81    }
82
83    fn progress(&mut self, line: &str, id: Uuid) {
84        if let Some(pos) = self.progress_positions.get(&id) {
85            let offset = self.line_counter + 1 - pos;
86            // Move up, clear line, write content, move back down
87            print!("\x1B[{offset}A\x1B[2K\r{line}\x1B[{offset}B\r");
88            std::io::stdout().flush().unwrap();
89        } else {
90            println!("{line}");
91            std::io::stdout().flush().unwrap();
92            self.line_counter += 1;
93            self.progress_positions.insert(id, self.line_counter);
94        }
95    }
96
97    fn finished(&mut self, id: Uuid) {
98        if let Some(removed_pos) = self.progress_positions.remove(&id) {
99            let offset = self.line_counter + 1 - removed_pos;
100            // Move up to the line, delete it (shifts content below up), move back down
101            if offset > 1 {
102                print!("\x1B[{offset}A\x1B[M\x1B[{}B", offset - 1);
103            } else {
104                print!("\x1B[{offset}A\x1B[M");
105            }
106            std::io::stdout().flush().unwrap();
107
108            // Update positions of progress bars that were below the removed one
109            for pos in self.progress_positions.values_mut() {
110                if *pos > removed_pos {
111                    *pos -= 1;
112                }
113            }
114
115            // Decrement line counter
116            self.line_counter = self.line_counter.saturating_sub(1);
117        }
118        if self.progress_positions.is_empty() {
119            self.line_counter = 0;
120        }
121    }
122
123    fn flush(&mut self) {
124        std::io::stdout().flush().unwrap();
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    #[test]
133    fn test_log_file() {
134        std::fs::remove_file("/tmp/test_log_file.log").ok();
135        let mut log_file = LogFile::new("/tmp/test_log_file.log").unwrap();
136        let uuid = Uuid::default();
137        log_file.regular("Hello, world!");
138        log_file.progress("lorem ipsum", uuid);
139        log_file.regular("rust is awesome !");
140        log_file.progress("LOREM IPSUM", uuid);
141        log_file.finished(uuid);
142        log_file.regular("test");
143        log_file.flush();
144        assert_eq!(
145            std::fs::read_to_string("/tmp/test_log_file.log").unwrap(),
146            "Hello, world!\nLOREM IPSUM\nrust is awesome !\ntest\n"
147        );
148    }
149
150    #[test]
151    fn test_log_stdout() {
152        let mut log_stdout = LogStdout::default();
153        let uuid_1 = Uuid::new_v4();
154        let uuid_2 = Uuid::new_v4();
155        log_stdout.regular("Hello, world!");
156        log_stdout.progress("lorem ipsum", uuid_1);
157        log_stdout.progress("ipsum lorem", uuid_2);
158        log_stdout.regular("rust is awesome !");
159        log_stdout.progress("LOREM IPSUM", uuid_2);
160        log_stdout.finished(uuid_2);
161        log_stdout.regular("test");
162        log_stdout.progress("LOREM IPSUM", uuid_1);
163        log_stdout.finished(uuid_1);
164    }
165}