initrd/src/dklog.rs

use std::io::Write;
use std::sync::{LazyLock, Mutex};
use std::time::SystemTime;
use tokio::sync::watch;
use tokio::task::JoinSet;

pub static LOG: LazyLock<Log> = LazyLock::new(Log::new);

pub fn init() {
    log::set_logger(&*LOG).expect("set_logger should not fail");
    log::set_max_level(log::LevelFilter::Info);
}

pub struct Log {
    start: SystemTime,
    log: Mutex<Vec<u8>>,
    tx: Mutex<Option<watch::Sender<usize>>>,
    rx: watch::Receiver<usize>,
    tasks: Mutex<Option<JoinSet<()>>>,
}

impl Log {
    pub fn new() -> Self {
        let (tx, rx) = watch::channel(0);
        Self {
            start: SystemTime::now(),
            log: Mutex::new(Vec::new()),
            tx: Mutex::new(Some(tx)),
            rx,
            tasks: Mutex::new(Some(JoinSet::new())),
        }
    }

    pub fn spawn(&self, task: impl Future<Output = ()> + Send + 'static) {
        if let Some(tasks) = self.tasks.lock().unwrap().as_mut() {
            tasks.spawn(task);
        }
    }

    pub fn subscribe(&self) -> LogWatch {
        LogWatch {
            log: self,
            pos: 0,
            rx: self.rx.clone(),
        }
    }

    pub async fn copy_to<W: tokio::io::AsyncWrite + Unpin>(&self, mut out: W) {
        let mut log = self.subscribe();
        use tokio::io::AsyncWriteExt;
        while let Some(chunk) = log.next().await {
            let _ = out.write_all(&chunk).await;
            let _ = out.flush().await;
        }
    }

    pub async fn close(&self) {
        self.tx.lock().unwrap().take();
        if let Some(tasks) = self.tasks.lock().unwrap().take() {
            tasks.join_all().await;
        }
    }
}

impl log::Log for Log {
    fn enabled(&self, metadata: &log::Metadata) -> bool {
        log::max_level().to_level().unwrap_or(log::Level::Info) >= metadata.level()
    }

    fn log(&self, record: &log::Record) {
        if !self.enabled(record.metadata()) {
            return;
        }

        let ts = self.start.elapsed().unwrap_or_default();

        use log::Level::*;
        let level = match record.level() {
            Trace => "\x1b[2mTRC\x1b[0m",
            Debug => "\x1b[34mDBG\x1b[0m",
            Info => "\x1b[32mINF\x1b[0m",
            Warn => "\x1b[93mWRN\x1b[0m",
            Error => "\x1b[91mERR\x1b[0m",
        };

        let Ok(mut log) = self.log.lock() else {
            return;
        };

        writeln!(
            log,
            "{:3}.{:03} {level} {}",
            ts.as_secs(),
            ts.subsec_millis(),
            record.args()
        )
        .expect("write to buf should not fail");

        if let Some(tx) = self.tx.lock().unwrap().as_mut() {
            tx.send_replace(log.len());
        }
    }

    fn flush(&self) {
        // always flushed
    }
}

pub struct LogWatch<'t> {
    log: &'t Log,
    pos: usize,
    rx: watch::Receiver<usize>,
}

impl<'t> LogWatch<'t> {
    pub async fn next(&mut self) -> Option<Vec<u8>> {
        loop {
            let new_pos = self.rx.borrow_and_update().clone();
            if new_pos <= self.pos {
                if self.rx.changed().await.is_err() {
                    return None; // finished
                }
                continue;
            }

            let mut chunk = Vec::new();
            chunk.extend(&self.log.log.lock().unwrap()[self.pos..new_pos]);
            self.pos = new_pos;
            return Some(chunk);
        }
    }
}