starbase_console/
buffer.rs

1use crate::stream::ConsoleStreamType;
2use parking_lot::Mutex;
3use std::io::{self, Write};
4use std::mem;
5use std::sync::{Arc, mpsc};
6use std::thread::sleep;
7use std::time::Duration;
8
9pub struct ConsoleBuffer {
10    buffer: Arc<Mutex<Vec<u8>>>,
11    stream: ConsoleStreamType,
12}
13
14impl ConsoleBuffer {
15    pub fn new(buffer: Arc<Mutex<Vec<u8>>>, stream: ConsoleStreamType) -> Self {
16        Self { buffer, stream }
17    }
18}
19
20impl Write for ConsoleBuffer {
21    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
22        self.buffer.lock().extend_from_slice(data);
23
24        Ok(data.len())
25    }
26
27    fn flush(&mut self) -> io::Result<()> {
28        flush(&mut self.buffer.lock(), self.stream)
29    }
30}
31
32pub fn flush(buffer: &mut Vec<u8>, stream: ConsoleStreamType) -> io::Result<()> {
33    if buffer.is_empty() {
34        return Ok(());
35    }
36
37    let data = mem::take(buffer);
38
39    match stream {
40        ConsoleStreamType::Stderr => io::stderr().lock().write_all(&data),
41        ConsoleStreamType::Stdout => io::stdout().lock().write_all(&data),
42    }
43}
44
45pub fn flush_on_loop(
46    buffer: Arc<Mutex<Vec<u8>>>,
47    stream: ConsoleStreamType,
48    receiver: mpsc::Receiver<bool>,
49) {
50    loop {
51        sleep(Duration::from_millis(100));
52
53        let _ = flush(&mut buffer.lock(), stream);
54
55        // Has the thread been closed?
56        match receiver.try_recv() {
57            Ok(true) | Err(mpsc::TryRecvError::Disconnected) => {
58                break;
59            }
60            _ => {}
61        }
62    }
63}