xy/crates/xy-supervisor/src/logs.rs

use std::collections::VecDeque;
use std::fs::{File, OpenOptions};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use tokio::sync::broadcast;
use xy_protocol::rpc::{LogLine, LogStream};

pub struct RotatingLogWriter {
    base: PathBuf,
    max_bytes: u64,
    keep: usize,
    file: File,
    written: u64,
}

impl RotatingLogWriter {
    pub fn open(base: &Path, max_bytes: u64, keep: usize) -> std::io::Result<Self> {
        if let Some(parent) = base.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let file = OpenOptions::new().create(true).append(true).open(base)?;
        let written = file.metadata()?.len();
        Ok(Self {
            base: base.to_path_buf(),
            max_bytes,
            keep,
            file,
            written,
        })
    }

    pub fn write_line(&mut self, tag: &str, line: &str) -> std::io::Result<()> {
        let bytes = format!("{tag} {line}\n");
        self.file.write_all(bytes.as_bytes())?;
        self.written += bytes.len() as u64;
        if self.written >= self.max_bytes {
            self.rotate()?;
        }
        Ok(())
    }

    fn rotate(&mut self) -> std::io::Result<()> {
        // Drop the current handle by replacing with /dev/null briefly.
        self.file = OpenOptions::new().read(true).open("/dev/null")?;
        for i in (1..self.keep).rev() {
            let src = self.gen_path(i);
            let dst = self.gen_path(i + 1);
            if src.exists() {
                let _ = std::fs::rename(&src, &dst);
            }
        }
        if self.base.exists() {
            let _ = std::fs::rename(&self.base, self.gen_path(1));
        }
        self.file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.base)?;
        self.written = 0;
        Ok(())
    }

    fn gen_path(&self, n: usize) -> PathBuf {
        let mut s = self.base.as_os_str().to_os_string();
        s.push(format!(".{n}"));
        PathBuf::from(s)
    }
}

impl std::io::Write for RotatingLogWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.file.write_all(buf)?;

        self.written += buf.len() as u64;

        if self.written >= self.max_bytes {
            self.rotate()?;
        }

        Ok(buf.len())
    }

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

#[derive(Clone)]
pub struct RingBuffer {
    inner: Arc<Mutex<RingBufferInner>>,
    capacity_bytes: usize,
}

struct RingBufferInner {
    lines: VecDeque<RecordedLine>,
    bytes: usize,
}

#[derive(Debug, Clone)]
pub struct RecordedLine {
    pub stream: xy_protocol::rpc::LogStream,
    pub line: String,
    pub ts_unix_ms: u64,
}

impl RingBuffer {
    pub fn new(capacity_bytes: usize) -> Self {
        Self {
            inner: Arc::new(Mutex::new(RingBufferInner {
                lines: VecDeque::new(),
                bytes: 0,
            })),
            capacity_bytes,
        }
    }

    pub fn push(&self, line: RecordedLine) {
        let mut g = self.inner.lock().unwrap();
        g.bytes += line.line.len();
        g.lines.push_back(line);
        while g.bytes > self.capacity_bytes {
            if let Some(removed) = g.lines.pop_front() {
                g.bytes -= removed.line.len();
            } else {
                break;
            }
        }
    }

    pub fn snapshot_tail(&self, n: Option<u32>) -> Vec<RecordedLine> {
        let g = self.inner.lock().unwrap();
        match n {
            None => g.lines.iter().cloned().collect(),
            Some(n) => {
                let take = (n as usize).min(g.lines.len());
                let start = g.lines.len() - take;
                g.lines.iter().skip(start).cloned().collect()
            }
        }
    }
}

const LOG_BROADCAST_CAP: usize = 256;

#[derive(Clone)]
pub struct LogSink {
    pub server_name: String,
    writer: Arc<Mutex<RotatingLogWriter>>,
    pub ring: RingBuffer,
    pub broadcast: broadcast::Sender<LogLine>,
}

impl LogSink {
    pub fn new(server_name: String, writer: RotatingLogWriter, ring_capacity_bytes: usize) -> Self {
        let (tx, _) = broadcast::channel(LOG_BROADCAST_CAP);
        Self {
            server_name,
            writer: Arc::new(Mutex::new(writer)),
            ring: RingBuffer::new(ring_capacity_bytes),
            broadcast: tx,
        }
    }

    pub fn record(&self, stream: LogStream, line: String) {
        let ts = now_unix_ms();
        let tag = match stream {
            LogStream::Stdout => "[out]",
            LogStream::Stderr => "[err]",
        };
        if let Err(e) = self.writer.lock().unwrap().write_line(tag, &line) {
            tracing::warn!(server = %self.server_name, error = %e, "log file write failed");
        }
        self.ring.push(RecordedLine {
            stream,
            line: line.clone(),
            ts_unix_ms: ts,
        });
        let _ = self.broadcast.send(LogLine {
            subscription_id: 0,
            name: self.server_name.clone(),
            stream,
            line,
            ts_unix_ms: ts,
        });
    }
}

fn now_unix_ms() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Read;
    use tempfile::tempdir;

    #[test]
    fn writes_lines_with_tags() {
        let dir = tempdir().unwrap();
        let base = dir.path().join("x.log");
        let mut w = RotatingLogWriter::open(&base, 1024, 3).unwrap();
        w.write_line("[out]", "hello").unwrap();
        w.write_line("[err]", "boom").unwrap();
        let mut s = String::new();
        File::open(&base).unwrap().read_to_string(&mut s).unwrap();
        assert_eq!(s, "[out] hello\n[err] boom\n");
    }

    #[test]
    fn rotates_at_threshold() {
        let dir = tempdir().unwrap();
        let base = dir.path().join("x.log");
        let mut w = RotatingLogWriter::open(&base, 20, 3).unwrap();
        for _ in 0..5 {
            w.write_line("[out]", "0123456789").unwrap();
        }
        assert!(base.exists());
        let rotated = dir.path().join("x.log.1");
        assert!(
            rotated.exists(),
            "expected rotated file at {}",
            rotated.display()
        );
    }

    use xy_protocol::rpc::LogStream;

    fn recorded(s: &str) -> RecordedLine {
        RecordedLine {
            stream: LogStream::Stdout,
            line: s.to_string(),
            ts_unix_ms: 0,
        }
    }

    #[test]
    fn ring_buffer_drops_oldest_when_full() {
        let rb = RingBuffer::new(10);
        rb.push(recorded("aaaaa"));
        rb.push(recorded("bbbbb"));
        rb.push(recorded("ccc"));
        let snap = rb.snapshot_tail(None);
        assert_eq!(snap.len(), 2);
        assert_eq!(snap[0].line, "bbbbb");
        assert_eq!(snap[1].line, "ccc");
    }

    #[test]
    fn ring_buffer_tail_n() {
        let rb = RingBuffer::new(1024);
        for i in 0..5 {
            rb.push(recorded(&format!("line{i}")));
        }
        let snap = rb.snapshot_tail(Some(2));
        assert_eq!(snap.len(), 2);
        assert_eq!(snap[0].line, "line3");
        assert_eq!(snap[1].line, "line4");
    }

    #[tokio::test]
    async fn log_sink_records_and_broadcasts() {
        let dir = tempdir().unwrap();
        let writer = RotatingLogWriter::open(&dir.path().join("s.log"), 1024, 3).unwrap();
        let sink = LogSink::new("s".to_string(), writer, 1024);
        let mut rx = sink.broadcast.subscribe();
        sink.record(LogStream::Stdout, "hello".to_string());
        let got = tokio::time::timeout(std::time::Duration::from_millis(100), rx.recv())
            .await
            .unwrap()
            .unwrap();
        assert_eq!(got.line, "hello");
        assert_eq!(got.stream, LogStream::Stdout);
        assert_eq!(sink.ring.snapshot_tail(None).len(), 1);
    }

    #[test]
    fn write_trait_appends_bytes() {
        use std::io::Write;

        let tmp = tempfile::tempdir().unwrap();
        let base = tmp.path().join("daemon.log");

        let mut writer = RotatingLogWriter::open(&base, 1024, 3).unwrap();

        writer.write_all(b"hello\n").unwrap();
        writer.flush().unwrap();

        let contents = std::fs::read_to_string(&base).unwrap();
        assert_eq!(contents, "hello\n");
    }

    #[test]
    fn write_trait_rotates_at_threshold() {
        use std::io::Write;

        let tmp = tempfile::tempdir().unwrap();
        let base = tmp.path().join("daemon.log");

        let mut writer = RotatingLogWriter::open(&base, 8, 3).unwrap();

        writer.write_all(b"0123456789").unwrap();
        writer.write_all(b"after\n").unwrap();
        writer.flush().unwrap();

        let rotated = tmp.path().join("daemon.log.1");
        assert!(rotated.exists());
        assert_eq!(std::fs::read_to_string(&rotated).unwrap(), "0123456789");
        assert_eq!(std::fs::read_to_string(&base).unwrap(), "after\n");
    }
}