irox_progress/
console.rs

1// SPDX-License-Identifier: MIT
2// Copyright ${YEAR} IROX Contributors
3//
4
5use std::io::{stdout, Error, Write};
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::mpsc::{Sender, TryRecvError};
8use std::sync::Arc;
9
10use log::error;
11
12use irox_time::epoch::UnixTimestamp;
13use irox_time::format::iso8601::ISO8601Duration;
14use irox_time::format::Format;
15use irox_time::Duration;
16
17use crate::{get_human, ProgressPrinter, Task};
18
19pub struct ConsoleProgressBar {
20    width: usize,
21}
22
23impl ConsoleProgressBar {
24    pub fn new(width: usize) -> Self {
25        ConsoleProgressBar { width }
26    }
27
28    pub fn print_infinite_progress(&self, task: &Task) -> Result<(), Error> {
29        let current = task.current_progress_count();
30        if let Some(started) = task.get_started() {
31            let current = current as f64;
32            let elapsed = started.elapsed().as_seconds_f64();
33            let avg_per_sec = current / elapsed;
34            let (avg_per_sec, avg_unit) = get_human!(avg_per_sec);
35            let (current, unit) = get_human!(current);
36
37            let time = (elapsed * 2.) as u64 & 0x7;
38            let state = match time {
39                0 => "\u{25b6}\u{25b9}\u{25b9}\u{25b9}\u{25b9}",
40                1 => "\u{25b9}\u{25b6}\u{25b9}\u{25b9}\u{25b9}",
41                2 => "\u{25b9}\u{25b9}\u{25b6}\u{25b9}\u{25b9}",
42                3 => "\u{25b9}\u{25b9}\u{25b9}\u{25b6}\u{25b9}",
43                4 => "\u{25b9}\u{25b9}\u{25b9}\u{25b9}\u{25b6}",
44                _ => "\u{25b9}\u{25b9}\u{25b9}\u{25b9}\u{25b9}",
45            };
46            let state2 = match time & 0x3 {
47                0 => "\u{25dc}",
48                1 => "\u{25dd}",
49                2 => "\u{25de}",
50                _ => "\u{25df}",
51            };
52
53            let out = format!("| ({current:.02}{unit})  {avg_per_sec:.02}{avg_unit}/s\r");
54            let spaces =
55                " ".repeat(((self.width as i32 - out.len() as i32 - 9).max(1) / 4) as usize);
56            let status = task.current_status().unwrap_or_default();
57            let out = format!("| ({current:.02}{unit}) {spaces}{state2}{spaces} {state} {spaces}{state2}{spaces}{avg_per_sec:.02}{avg_unit}/s {status}\r");
58
59            let mut stdio = stdout();
60            stdio.write_all(out.as_bytes())?;
61            return stdio.flush();
62        }
63        Ok(())
64    }
65
66    pub fn print_progress(&self, task: &Task) -> Result<(), Error> {
67        let pct = task.current_progress_frac();
68        let current = task.current_progress_count();
69        let max = task.max_elements();
70        if max == u64::MAX {
71            return self.print_infinite_progress(task);
72        }
73
74        let rem_str = ISO8601Duration.format(&task.get_remaining_time());
75
76        let w_pct = self.width as f64 * pct;
77        let whole = w_pct.floor() as usize;
78        let part = (8.0 * w_pct.fract()).round() as u8;
79        let mut char = match part {
80            0 => "\u{00A0}",
81            1 => "\u{258F}",
82            2 => "\u{258E}",
83            3 => "\u{258D}",
84            4 => "\u{258C}",
85            5 => "\u{258B}",
86            6 => "\u{258A}",
87            7 => "\u{2589}",
88            _ => "\u{2588}",
89        };
90
91        let mut rem = self.width - whole;
92        // println!("{current} {max} {pct} {whole} {part} {rem}");
93        if rem == 1 {
94            rem = 0;
95        } else if rem == 0 {
96            char = "";
97        } else {
98            rem -= 1;
99        }
100        let whole = "\u{2588}".repeat(whole);
101        let rem = " ".repeat(rem);
102        let status = task.current_status().unwrap_or_default();
103        let out = format!(
104            "{:>3.0}%|{whole}{char}{rem}| ({current}/{max}) {rem_str} {status}\r",
105            pct * 100.
106        );
107        let mut stdio = stdout();
108        stdio.write_all(out.as_bytes())?;
109        stdio.flush()
110    }
111}
112
113pub struct ConsoleProgressPrinter {
114    thread_handle: Option<std::thread::JoinHandle<()>>,
115    running_flag: Arc<AtomicBool>,
116    sender: Sender<Task>,
117}
118
119impl ConsoleProgressPrinter {
120    pub fn new_update_rate(update_rate: Duration) -> ConsoleProgressPrinter {
121        let running_flag = Arc::new(AtomicBool::new(true));
122        let running = running_flag.clone();
123
124        let (sender, receiver) = std::sync::mpsc::channel();
125
126        #[allow(unused_assignments)]
127        let thread_handle = std::thread::spawn(move || {
128            let mut tasks: Vec<Task> = Vec::new();
129            let mut last_run = UnixTimestamp::now();
130            while running.load(Ordering::Relaxed) {
131                last_run = UnixTimestamp::now();
132                let next_run = last_run + update_rate;
133                let _r = stdout().write_all(&[0x1B, b'[', b'2', b'K']);
134                for task in &tasks {
135                    let _res = ConsoleProgressBar::new(60).print_progress(task);
136                }
137
138                match receiver.try_recv() {
139                    Ok(task) => tasks.push(task),
140                    Err(e) => {
141                        if e == TryRecvError::Disconnected {
142                            return;
143                        }
144                    }
145                };
146
147                let delay = next_run - UnixTimestamp::now();
148                std::thread::sleep(delay.into());
149            }
150            let _r = stdout().write_all(&[0x1B, b'[', b'2', b'K']);
151            for task in &tasks {
152                let _res = ConsoleProgressBar::new(40).print_progress(task);
153            }
154        });
155
156        ConsoleProgressPrinter {
157            thread_handle: Some(thread_handle),
158            running_flag,
159            sender,
160        }
161    }
162}
163
164impl Drop for ConsoleProgressPrinter {
165    fn drop(&mut self) {
166        self.running_flag.store(false, Ordering::Relaxed);
167        if let Some(handle) = self.thread_handle.take() {
168            let _res = handle.join();
169        }
170    }
171}
172
173impl ProgressPrinter for ConsoleProgressPrinter {
174    fn track_task_progress(&self, task: &Task) {
175        if let Err(e) = self.sender.send(task.clone()) {
176            error!("Error sending task to printer: {e:?}");
177        }
178    }
179}