use std::cell::RefCell;
use std::io;
use std::sync::OnceLock;
use tracing_subscriber::fmt::MakeWriter;
use tracing_subscriber::layer::SubscriberExt as _;
use tracing_subscriber::util::SubscriberInitExt as _;
use tracing_subscriber::Layer as _;
thread_local! {
static BUF: RefCell<Vec<u8>> = const { RefCell::new(Vec::new()) };
}
struct ThreadLocalWriter;
impl io::Write for ThreadLocalWriter {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
BUF.with(|buf| buf.borrow_mut().extend_from_slice(bytes));
Ok(bytes.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[derive(Clone, Copy)]
struct ThreadLocalMakeWriter;
impl<'a> MakeWriter<'a> for ThreadLocalMakeWriter {
type Writer = ThreadLocalWriter;
fn make_writer(&'a self) -> Self::Writer {
ThreadLocalWriter
}
}
static GLOBAL_INSTALLED: OnceLock<()> = OnceLock::new();
fn install_once() {
GLOBAL_INSTALLED.get_or_init(|| {
let layer = tracing_subscriber::fmt::layer()
.with_writer(ThreadLocalMakeWriter)
.with_target(true)
.with_ansi(false)
.without_time()
.with_filter(tracing_subscriber::filter::LevelFilter::DEBUG);
let _ = tracing_subscriber::registry().with(layer).try_init();
});
}
pub fn capture<F: FnOnce() + Send + 'static>(f: F) -> String {
install_once();
tracing::callsite::rebuild_interest_cache();
std::thread::spawn(move || {
BUF.with(|b| b.borrow_mut().clear());
f();
BUF.with(|b| String::from_utf8(b.borrow().clone()).expect("tracing output should be UTF-8"))
})
.join()
.expect("capture thread panicked")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn capture_collects_events_emitted_inside_the_closure() {
let out = capture(|| {
tracing::info!(target: "studio_worker::test_support_demo", marker = "alpha", "hello");
});
assert!(out.contains("INFO"), "missing INFO level: {out:?}");
assert!(
out.contains("studio_worker::test_support_demo"),
"missing target: {out:?}"
);
assert!(out.contains("marker=\"alpha\""), "missing field: {out:?}");
assert!(out.contains("hello"), "missing message: {out:?}");
}
#[test]
fn capture_isolates_between_invocations() {
let first = capture(|| tracing::info!("first message"));
let second = capture(|| tracing::info!("second message"));
assert!(first.contains("first message") && !first.contains("second message"));
assert!(second.contains("second message") && !second.contains("first message"));
}
#[test]
fn capture_isolates_between_threads() {
let handle = std::thread::spawn(|| {
for _ in 0..50 {
tracing::info!("sibling thread noise");
}
});
let out = capture(|| {
tracing::info!("primary capture message");
});
handle.join().unwrap();
assert!(out.contains("primary capture message"));
assert!(
!out.contains("sibling thread noise"),
"cross-thread leak: {out:?}"
);
}
#[test]
fn capture_works_inside_a_multi_thread_tokio_runtime() {
let rt = tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.build()
.unwrap();
let out = rt.block_on(async {
capture(|| {
tracing::info!("emitted from spawned capture thread");
})
});
assert!(out.contains("emitted from spawned capture thread"));
}
}