use super::*;
use std::sync::{Arc, Mutex, OnceLock};
use std::thread::ThreadId;
use tracing::Level;
use tracing::level_filters::LevelFilter;
use tracing_subscriber::Registry;
use tracing_subscriber::layer::{Context, Layer, SubscriberExt};
type Sink = Arc<Mutex<Vec<String>>>;
static ACTIVE_SINK: Mutex<Option<(ThreadId, Sink)>> = Mutex::new(None);
#[derive(Clone, Default)]
struct WarnCapture;
impl<S: tracing::Subscriber> Layer<S> for WarnCapture {
fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
struct Visit(String);
impl tracing::field::Visit for Visit {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
if field.name() == "message" {
self.0 = format!("{value:?}");
}
}
}
if *event.metadata().level() != Level::WARN {
return;
}
let sink = {
let guard = ACTIVE_SINK
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match guard.as_ref() {
Some((owner, sink)) if *owner == std::thread::current().id() => Arc::clone(sink),
_ => return,
}
};
let mut v = Visit(String::new());
event.record(&mut v);
sink.lock().expect("warn log").push(v.0);
}
}
struct HangingExecutor;
impl AgentExecutor for HangingExecutor {
fn execute<'a>(
&'a self,
_ctx: &'a RequestContext,
_queue: &'a dyn EventQueueWriter,
) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
Box::pin(async { Ok(()) })
}
fn on_shutdown<'a>(&'a self) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
Box::pin(async { std::future::pending::<()>().await })
}
}
static CAPTURE_LOCK: Mutex<()> = Mutex::new(());
static INSTALLED: OnceLock<Result<(), String>> = OnceLock::new();
const CANARY: &str = "warn-capture canary";
fn warnings_during<F>(f: F) -> Vec<String>
where
F: FnOnce(),
{
let _serial = CAPTURE_LOCK
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Err(why) = INSTALLED.get_or_init(|| {
let subscriber = Registry::default().with(WarnCapture);
tracing::subscriber::set_global_default(subscriber).map_err(|e| e.to_string())
}) {
panic!(
"the global tracing subscriber in this binary is not the one this file \
installs ({why}), so WarnCapture is not in the dispatcher and every \
capture here can only return []. Whatever installs one first has to \
stop, or these tests have to observe it instead."
);
}
let sink: Sink = Arc::new(Mutex::new(Vec::new()));
*ACTIVE_SINK
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) =
Some((std::thread::current().id(), Arc::clone(&sink)));
tracing::warn!("{CANARY}");
let canary_seen = {
let captured = sink.lock().expect("warn log");
captured.iter().any(|w| w.contains(CANARY))
};
assert!(
canary_seen,
"the WARN capture is not working: this file's own canary never reached \
the sink (LevelFilter::current() == {}). Any `got []` from these tests \
is about the capture, not about shutdown.",
LevelFilter::current()
);
sink.lock().expect("warn log").clear();
f();
*ACTIVE_SINK
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = None;
sink.lock().expect("warn log").clone()
}
fn mentions_cleanup(warnings: &[String]) -> bool {
warnings.iter().any(|w| w.contains("executor cleanup"))
}
#[test]
fn a_warning_from_another_thread_is_not_captured() {
let warnings = warnings_during(|| {
std::thread::spawn(|| {
tracing::warn!("executor cleanup did not finish within the shutdown timeout");
})
.join()
.expect("the emitting thread should not panic");
});
assert!(
warnings.is_empty(),
"a WARN from another thread must not reach this capture; got {warnings:?}"
);
}
#[test]
fn clean_shutdown_warns_about_nothing() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_time()
.build()
.expect("runtime");
let warnings = warnings_during(|| {
rt.block_on(async {
let handler = make_handler();
let report = handler
.shutdown_with_timeout(Duration::from_millis(50))
.await;
assert!(
report.executor_cleanup_completed,
"the no-op executor's cleanup returns immediately"
);
});
});
assert!(
!mentions_cleanup(&warnings),
"a clean shutdown must not warn about executor cleanup; got {warnings:?}"
);
}
#[test]
fn fixed_budget_shutdown_warns_only_when_cleanup_hangs() {
let rt = || {
tokio::runtime::Builder::new_current_thread()
.enable_time()
.start_paused(true)
.build()
.expect("runtime")
};
let clean = warnings_during(|| {
rt().block_on(async {
let handler = make_handler();
let report = handler.shutdown().await;
assert!(report.executor_cleanup_completed);
});
});
assert!(
!mentions_cleanup(&clean),
"a clean shutdown() must not warn about executor cleanup; got {clean:?}"
);
let hung = warnings_during(|| {
rt().block_on(async {
let handler = RequestHandlerBuilder::new(HangingExecutor)
.build()
.expect("builder should succeed");
let report = handler.shutdown().await;
assert!(!report.executor_cleanup_completed);
});
});
assert!(
mentions_cleanup(&hung),
"a hung cleanup under shutdown() must be warned about; got {hung:?}"
);
}
#[test]
fn hung_cleanup_is_warned_about() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_time()
.build()
.expect("runtime");
let warnings = warnings_during(|| {
rt.block_on(async {
let handler = RequestHandlerBuilder::new(HangingExecutor)
.build()
.expect("builder should succeed");
let report = handler
.shutdown_with_timeout(Duration::from_millis(50))
.await;
assert!(
!report.executor_cleanup_completed,
"a hanging cleanup must be reported as incomplete"
);
});
});
assert!(
mentions_cleanup(&warnings),
"a hung executor cleanup must be warned about; got {warnings:?}"
);
}