transbeam/src/file.rs

use std::{fs::File, io::Write, path::PathBuf, task::Waker};

pub trait LiveWriter: Write {
    fn add_waker(&mut self, waker: Waker);
}

/// A simple wrapper for a file that can be read while we're still appending data
pub struct LiveFileWriter {
    file: File,
    /// Wake handles for contexts that are waiting for us to write more
    wakers: Vec<Waker>,
}

impl LiveFileWriter {
    pub fn new(path: &PathBuf) -> std::io::Result<Self> {
        Ok(Self {
            file: File::options().write(true).create_new(true).open(path)?,
            wakers: Vec::new(),
        })
    }
}

impl LiveWriter for LiveFileWriter {
    fn add_waker(&mut self, waker: Waker) {
        self.wakers.push(waker);
    }
}

impl Write for LiveFileWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let result = self.file.write(buf);
        if let Ok(n) = result {
            if n > 0 {
                for waker in self.wakers.drain(..) {
                    waker.wake();
                }
            }
        }
        result
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.file.flush()
    }
}