use std::io::{self, Write};
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use tracing_subscriber::fmt::MakeWriter;
const MAX_BUFFERED_BYTES: usize = 64 * 1024;
const TRUNCATION_MARKER: &[u8] = b"\n[startup diagnostics truncated]\n";
enum Sink {
Buffered(BoundedBuffer),
Stderr,
Discard,
}
#[derive(Default)]
struct BoundedBuffer {
bytes: Vec<u8>,
truncated: bool,
}
impl BoundedBuffer {
fn append(&mut self, buf: &[u8]) {
if self.truncated {
return;
}
let remaining = MAX_BUFFERED_BYTES.saturating_sub(self.bytes.len());
if buf.len() <= remaining {
self.bytes.extend_from_slice(buf);
return;
}
let content_limit = MAX_BUFFERED_BYTES.saturating_sub(TRUNCATION_MARKER.len());
if self.bytes.len() < content_limit {
let keep = content_limit.saturating_sub(self.bytes.len());
if let Some(head) = buf.get(..keep.min(buf.len())) {
self.bytes.extend_from_slice(head);
}
}
self.bytes.truncate(content_limit);
self.bytes.extend_from_slice(TRUNCATION_MARKER);
self.truncated = true;
}
fn take(&mut self) -> Vec<u8> {
self.truncated = false;
std::mem::take(&mut self.bytes)
}
#[cfg(test)]
fn as_slice(&self) -> &[u8] {
&self.bytes
}
}
#[derive(Clone)]
pub struct StartupWriter {
sink: Arc<Mutex<Sink>>,
}
impl StartupWriter {
#[must_use]
pub fn buffering() -> Self {
Self {
sink: Arc::new(Mutex::new(Sink::Buffered(BoundedBuffer::default()))),
}
}
fn lock(&self) -> MutexGuard<'_, Sink> {
self.sink.lock().unwrap_or_else(PoisonError::into_inner)
}
pub fn release_to_stderr(&self) -> io::Result<()> {
let mut sink = self.lock();
let buffered = match &mut *sink {
Sink::Buffered(buffer) => buffer.take(),
Sink::Stderr | Sink::Discard => Vec::new(),
};
*sink = Sink::Stderr;
drop(sink);
if buffered.is_empty() {
return Ok(());
}
io::stderr().write_all(&buffered)
}
pub fn discard(&self) {
let mut sink = self.lock();
*sink = Sink::Discard;
}
#[cfg(test)]
#[must_use]
pub fn buffered(&self) -> Vec<u8> {
match &*self.lock() {
Sink::Buffered(buffer) => buffer.as_slice().to_vec(),
Sink::Stderr | Sink::Discard => Vec::new(),
}
}
}
pub struct StartupWriterHandle {
sink: Arc<Mutex<Sink>>,
}
impl Write for StartupWriterHandle {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let mut sink = self.sink.lock().unwrap_or_else(PoisonError::into_inner);
match &mut *sink {
Sink::Buffered(buffer) => {
buffer.append(buf);
Ok(buf.len())
}
Sink::Stderr => {
drop(sink);
io::stderr().write(buf)
}
Sink::Discard => Ok(buf.len()),
}
}
fn flush(&mut self) -> io::Result<()> {
let sink = self.sink.lock().unwrap_or_else(PoisonError::into_inner);
if matches!(&*sink, Sink::Stderr) {
drop(sink);
return io::stderr().flush();
}
Ok(())
}
}
impl<'writer> MakeWriter<'writer> for StartupWriter {
type Writer = StartupWriterHandle;
fn make_writer(&'writer self) -> Self::Writer {
StartupWriterHandle {
sink: Arc::clone(&self.sink),
}
}
}
#[cfg(test)]
#[path = "startup_tracing_tests.rs"]
mod tests;