use std::fmt;
use std::future::Future;
use std::io::{self, Write};
use std::sync::{Arc, Mutex, RwLock};
pub type Sink = Arc<Mutex<dyn Write + Send>>;
static SINK: RwLock<Option<Sink>> = RwLock::new(None);
pub fn notice(args: fmt::Arguments<'_>) {
let _ = Stderr.write_fmt(format_args!("{}\n", args));
}
tokio::task_local! {
static SCOPED: Sink;
}
pub fn scoped<F: Future>(sink: Sink, fut: F) -> impl Future<Output = F::Output> {
SCOPED.scope(sink, fut)
}
pub fn redirect(sink: Sink) -> RedirectGuard {
let mut slot = SINK.write().unwrap_or_else(|p| p.into_inner());
let previous = slot.replace(sink);
RedirectGuard { previous }
}
#[must_use = "diagnostics return to stderr as soon as the guard is dropped"]
pub struct RedirectGuard {
previous: Option<Sink>,
}
impl Drop for RedirectGuard {
fn drop(&mut self) {
let mut slot = SINK.write().unwrap_or_else(|p| p.into_inner());
*slot = self.previous.take();
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct Stderr;
impl Write for Stderr {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if let Ok(written) =
SCOPED.try_with(|sink| sink.lock().unwrap_or_else(|p| p.into_inner()).write(buf))
{
return written;
}
let slot = SINK.read().unwrap_or_else(|p| p.into_inner());
match slot.as_ref() {
Some(sink) => sink.lock().unwrap_or_else(|p| p.into_inner()).write(buf),
None => io::stderr().write(buf),
}
}
fn flush(&mut self) -> io::Result<()> {
if let Ok(flushed) =
SCOPED.try_with(|sink| sink.lock().unwrap_or_else(|p| p.into_inner()).flush())
{
return flushed;
}
let slot = SINK.read().unwrap_or_else(|p| p.into_inner());
match slot.as_ref() {
Some(sink) => sink.lock().unwrap_or_else(|p| p.into_inner()).flush(),
None => io::stderr().flush(),
}
}
}
#[macro_export]
macro_rules! notice {
($($arg:tt)*) => {
$crate::diag::notice(format_args!($($arg)*))
};
}
#[cfg(test)]
mod tests {
use super::*;
static REDIRECT_LOCK: Mutex<()> = Mutex::new(());
#[test]
fn test_stderr_writer_flushes_with_and_without_redirect() {
let _lock = REDIRECT_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let mut w = Stderr;
assert!(w.flush().is_ok());
let buf: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
let _g = redirect(buf.clone());
assert!(w.flush().is_ok());
}
#[test]
fn test_a_scoped_future_keeps_its_notices_from_the_redirect() {
let _lock = REDIRECT_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let rt = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
let pane: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
let watch: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
let _g = redirect(pane.clone());
let watching = scoped(watch.clone(), async {
for i in 0..3 {
notice(format_args!("watch {i}"));
tokio::task::yield_now().await;
}
});
let command = async {
for i in 0..3 {
notice(format_args!("pane {i}"));
tokio::task::yield_now().await;
}
};
rt.block_on(async { tokio::join!(watching, command) });
let pane = String::from_utf8_lossy(&pane.lock().unwrap()).into_owned();
let watch = String::from_utf8_lossy(&watch.lock().unwrap()).into_owned();
assert!(
watch.contains(
"watch 0
watch 1
watch 2
"
),
"{watch}"
);
assert!(!watch.contains("pane"), "{watch}");
assert!(
pane.contains(
"pane 0
"
) && pane.contains(
"pane 2
"
),
"{pane}"
);
assert!(!pane.contains("watch"), "{pane}");
}
#[test]
fn test_redirect_captures_and_restores() {
let _lock = REDIRECT_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let outer: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
let inner: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
{
let _g1 = redirect(outer.clone());
notice(format_args!("one"));
{
let _g2 = redirect(inner.clone());
crate::notice!("two {}", 2);
}
notice(format_args!("three"));
}
let outer = String::from_utf8_lossy(&outer.lock().unwrap()).into_owned();
let inner = String::from_utf8_lossy(&inner.lock().unwrap()).into_owned();
let (one, three) = (outer.find("one\n").unwrap(), outer.find("three\n").unwrap());
assert!(one < three);
assert!(!outer.contains("two 2"));
assert!(inner.contains("two 2\n"));
assert!(!inner.contains("one") && !inner.contains("three"));
}
}