Skip to main content

vtcode_core/utils/
trace_writer.rs

1//! Shared buffered trace log writer with flush-on-exit support.
2//!
3//! Provides a `BufWriter`-backed file writer wrapped in `Arc<Mutex<..>>` so the
4//! tracing `fmt::layer` can write efficiently (batched syscalls) while still
5//! allowing an explicit `flush_trace_log()` call on process exit or signal.
6
7use std::fs::{File, OpenOptions};
8use std::io::{BufWriter, Write};
9use std::path::Path;
10use std::sync::{Arc, Mutex, MutexGuard, OnceLock};
11
12use anyhow::{Context, Result};
13
14/// Capacity of the internal `BufWriter` (64 KiB — large enough to batch many
15/// log lines before issuing a single `write` syscall).
16const BUF_CAPACITY: usize = 64 * 1024;
17
18/// Global handle to the active trace log writer so `flush_trace_log` can be
19/// called from signal handlers / shutdown hooks without threading the writer
20/// through the entire call stack.
21static GLOBAL_WRITER: OnceLock<FlushableWriter> = OnceLock::new();
22
23/// A clonable, thread-safe buffered writer that implements `std::io::Write`
24/// so it can be passed directly to `tracing_subscriber::fmt::layer().with_writer(..)`.
25#[derive(Clone)]
26pub struct FlushableWriter {
27    inner: Arc<Mutex<BufWriter<File>>>,
28}
29
30impl FlushableWriter {
31    /// Open (or create) a log file and wrap it in a buffered writer.
32    pub fn open(path: &Path) -> Result<Self> {
33        let file = OpenOptions::new()
34            .create(true)
35            .append(true)
36            .open(path)
37            .with_context(|| format!("Failed to open trace log file: {}", path.display()))?;
38        let writer = BufWriter::with_capacity(BUF_CAPACITY, file);
39        let flushable = Self { inner: Arc::new(Mutex::new(writer)) };
40        // Store globally so `flush_trace_log` works from anywhere.
41        let _ = GLOBAL_WRITER.set(flushable.clone());
42        // Register the flush hook in vtcode-commons so crates that don't
43        // depend on vtcode-core (e.g. vtcode-ui) can still trigger a flush.
44        vtcode_commons::trace_flush::register_trace_flush_hook(flush_trace_log);
45        Ok(flushable)
46    }
47
48    /// Flush the internal buffer to disk.
49    pub fn flush(&self) {
50        let _ = self.flush_locked();
51    }
52
53    fn lock_writer(&self) -> std::io::Result<MutexGuard<'_, BufWriter<File>>> {
54        self.inner
55            .lock()
56            .map_err(|e| std::io::Error::other(format!("trace writer lock poisoned: {e}")))
57    }
58
59    fn flush_locked(&self) -> std::io::Result<()> {
60        let mut guard = self.lock_writer()?;
61        guard.flush()
62    }
63}
64
65impl Write for FlushableWriter {
66    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
67        let mut guard = self.lock_writer()?;
68        guard.write(buf)
69    }
70
71    fn flush(&mut self) -> std::io::Result<()> {
72        self.flush_locked()
73    }
74}
75
76/// Flush the global trace log writer to disk.
77///
78/// Safe to call from signal handlers, shutdown hooks, or `Drop` implementations.
79/// No-op if no trace writer has been initialized.
80pub fn flush_trace_log() {
81    if let Some(writer) = GLOBAL_WRITER.get() {
82        writer.flush();
83    }
84}