use std::io::{ErrorKind, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
use crate::agent::append_bootstrap_log;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Emitted {
Written,
Logged,
Dropped,
}
static FALLBACK_WORKDIR: Mutex<Option<PathBuf>> = Mutex::new(None);
pub fn set_diagnostic_workdir(workdir: &Path) {
let mut slot = FALLBACK_WORKDIR
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
*slot = Some(workdir.to_path_buf());
}
pub fn diagnostic_workdir() -> Option<PathBuf> {
FALLBACK_WORKDIR
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clone()
}
pub fn to_stdout(line: &str) -> Emitted {
let fallback = diagnostic_workdir();
emit_to(&mut std::io::stdout().lock(), line, fallback.as_deref())
}
pub fn credential_to_stdout(line: &str) -> Emitted {
emit_to(&mut std::io::stdout().lock(), line, None)
}
pub fn to_stderr(line: &str) -> Emitted {
let fallback = diagnostic_workdir();
emit_to(&mut std::io::stderr().lock(), line, fallback.as_deref())
}
pub fn raw_to_stdout(chunk: &str) -> Emitted {
let fallback = diagnostic_workdir();
emit_raw_to(&mut std::io::stdout().lock(), chunk, fallback.as_deref())
}
pub(crate) fn emit_to<W: Write>(out: &mut W, line: &str, fallback: Option<&Path>) -> Emitted {
let written = out
.write_all(line.as_bytes())
.and_then(|()| out.write_all(b"\n"))
.and_then(|()| out.flush());
match written {
Ok(()) => Emitted::Written,
Err(_) => fall_back(line, fallback),
}
}
pub(crate) fn emit_raw_to<W: Write>(out: &mut W, chunk: &str, fallback: Option<&Path>) -> Emitted {
let written = out.write_all(chunk.as_bytes()).and_then(|()| out.flush());
match written {
Ok(()) => Emitted::Written,
Err(_) => fall_back(chunk.trim_end_matches('\n'), fallback),
}
}
fn fall_back(line: &str, fallback: Option<&Path>) -> Emitted {
match fallback {
Some(workdir) => {
append_bootstrap_log(workdir, line);
Emitted::Logged
}
None => Emitted::Dropped,
}
}
#[macro_export]
macro_rules! runtime_out {
($($arg:tt)*) => {{
$crate::diagnostic::to_stdout(&::std::format!($($arg)*));
}};
}
#[macro_export]
macro_rules! runtime_err {
($($arg:tt)*) => {{
$crate::diagnostic::to_stderr(&::std::format!($($arg)*));
}};
}
static STDOUT_CLOSED: AtomicBool = AtomicBool::new(false);
static STDERR_CLOSED: AtomicBool = AtomicBool::new(false);
static STDOUT_FAILURE: Mutex<Option<String>> = Mutex::new(None);
pub fn report_to_stdout(bytes: &[u8]) -> Emitted {
if STDOUT_CLOSED.load(Ordering::Acquire) {
return Emitted::Dropped;
}
report_to(
&mut std::io::stdout().lock(),
bytes,
&STDOUT_CLOSED,
Some(&STDOUT_FAILURE),
)
}
pub fn report_to_stderr(bytes: &[u8]) -> Emitted {
if STDERR_CLOSED.load(Ordering::Acquire) {
return Emitted::Dropped;
}
report_to(&mut std::io::stderr().lock(), bytes, &STDERR_CLOSED, None)
}
pub fn stdout_closed() -> bool {
STDOUT_CLOSED.load(Ordering::Acquire)
}
pub fn stdout_failure() -> Option<String> {
STDOUT_FAILURE
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clone()
}
pub(crate) fn report_to<W: Write>(
out: &mut W,
bytes: &[u8],
closed: &AtomicBool,
failure: Option<&Mutex<Option<String>>>,
) -> Emitted {
if closed.load(Ordering::Acquire) {
return Emitted::Dropped;
}
match out.write_all(bytes).and_then(|()| out.flush()) {
Ok(()) => Emitted::Written,
Err(error) => {
closed.store(true, Ordering::Release);
if error.kind() != ErrorKind::BrokenPipe {
if let Some(slot) = failure {
slot.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.get_or_insert_with(|| error.to_string());
}
}
Emitted::Dropped
}
}
}
#[doc(hidden)]
#[macro_export]
macro_rules! __report_line {
() => {
::std::string::String::from("\n")
};
($($arg:tt)*) => {{
let mut line = ::std::format!($($arg)*);
line.push('\n');
line
}};
}
#[macro_export]
macro_rules! report_println {
($($arg:tt)*) => {{
$crate::diagnostic::report_to_stdout($crate::__report_line!($($arg)*).as_bytes());
}};
}
#[macro_export]
macro_rules! report_print {
($($arg:tt)*) => {{
$crate::diagnostic::report_to_stdout(::std::format!($($arg)*).as_bytes());
}};
}
#[macro_export]
macro_rules! report_eprintln {
($($arg:tt)*) => {{
$crate::diagnostic::report_to_stderr($crate::__report_line!($($arg)*).as_bytes());
}};
}
#[macro_export]
macro_rules! report_eprint {
($($arg:tt)*) => {{
$crate::diagnostic::report_to_stderr(::std::format!($($arg)*).as_bytes());
}};
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{Error, ErrorKind};
struct Refusing {
kind: ErrorKind,
calls: usize,
}
impl Refusing {
fn new(kind: ErrorKind) -> Self {
Self { kind, calls: 0 }
}
}
impl Write for Refusing {
fn write(&mut self, _buf: &[u8]) -> std::io::Result<usize> {
self.calls += 1;
Err(Error::new(self.kind, "refused"))
}
fn flush(&mut self) -> std::io::Result<()> {
self.calls += 1;
Err(Error::new(self.kind, "refused"))
}
}
#[test]
fn refused_line_lands_in_the_bootstrap_log() {
let workdir = tempfile::tempdir().expect("temp workdir");
let outcome = emit_to(
&mut Refusing::new(ErrorKind::BrokenPipe),
"{\"url\":\"http://127.0.0.1:1/\"}",
Some(workdir.path()),
);
assert_eq!(outcome, Emitted::Logged);
let log = workdir.path().join("logs").join("bootstrap.log");
let contents = std::fs::read_to_string(&log).expect("bootstrap.log was written");
assert_eq!(contents, "{\"url\":\"http://127.0.0.1:1/\"}\n");
}
#[test]
fn refused_line_with_no_workdir_is_dropped() {
let outcome = emit_to(
&mut Refusing::new(ErrorKind::BrokenPipe),
"a warning raised during staging",
None,
);
assert_eq!(outcome, Emitted::Dropped);
}
#[test]
fn accepted_line_gets_exactly_one_newline() {
let mut sink: Vec<u8> = Vec::new();
let outcome = emit_to(&mut sink, "session: ses_01", None);
assert_eq!(outcome, Emitted::Written);
assert_eq!(String::from_utf8(sink).unwrap(), "session: ses_01\n");
}
#[test]
fn raw_chunk_is_written_verbatim() {
let mut sink: Vec<u8> = Vec::new();
let outcome = emit_raw_to(&mut sink, "child line\n", None);
assert_eq!(outcome, Emitted::Written);
assert_eq!(String::from_utf8(sink).unwrap(), "child line\n");
}
#[test]
fn refused_raw_chunk_is_logged_without_a_doubled_newline() {
let workdir = tempfile::tempdir().expect("temp workdir");
let outcome = emit_raw_to(
&mut Refusing::new(ErrorKind::BrokenPipe),
"child line\n",
Some(workdir.path()),
);
assert_eq!(outcome, Emitted::Logged);
let log = workdir.path().join("logs").join("bootstrap.log");
let contents = std::fs::read_to_string(&log).expect("bootstrap.log was written");
assert_eq!(contents, "child line\n");
}
#[test]
fn report_over_a_broken_pipe_latches_and_stops_writing() {
let closed = AtomicBool::new(false);
let failure = Mutex::new(None);
let mut out = Refusing::new(ErrorKind::BrokenPipe);
let first = report_to(&mut out, b"row one\n", &closed, Some(&failure));
assert_eq!(first, Emitted::Dropped);
assert!(closed.load(Ordering::Acquire));
assert_eq!(*failure.lock().unwrap(), None);
let calls_after_first = out.calls;
assert!(calls_after_first > 0);
let second = report_to(&mut out, b"row two\n", &closed, Some(&failure));
assert_eq!(second, Emitted::Dropped);
assert_eq!(
out.calls, calls_after_first,
"a latched stream is not written"
);
}
#[test]
fn report_over_a_full_disk_latches_and_records_the_error() {
let closed = AtomicBool::new(false);
let failure = Mutex::new(None);
let mut out = Refusing::new(ErrorKind::StorageFull);
let outcome = report_to(&mut out, b"row\n", &closed, Some(&failure));
assert_eq!(outcome, Emitted::Dropped);
assert!(closed.load(Ordering::Acquire));
assert_eq!(failure.lock().unwrap().as_deref(), Some("refused"));
}
#[test]
fn report_over_an_open_stream_writes_the_formatted_bytes() {
let cases: [(String, &str); 4] = [
(crate::__report_line!("{} rows", 3), "3 rows\n"),
(crate::__report_line!(), "\n"),
(::std::format!("value for {}: ", "KEY"), "value for KEY: "),
(crate::__report_line!("{:>4}|", 7), " 7|\n"),
];
for (text, expected) in cases {
let closed = AtomicBool::new(false);
let mut sink: Vec<u8> = Vec::new();
let outcome = report_to(&mut sink, text.as_bytes(), &closed, None);
assert_eq!(outcome, Emitted::Written);
assert_eq!(String::from_utf8(sink).unwrap(), expected);
assert!(!closed.load(Ordering::Acquire));
}
}
}