Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 42 additions & 40 deletions crates/spectacular-macros/src/attr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,6 @@ pub(crate) fn expand(
let has_before = before_name.is_some();
let has_after = after_name.is_some();
let has_after_each = after_each_name.is_some();
let test_count = test_fns.len();

// Extract params from before_each (ref params come from before context)
let before_each_params = before_each_fn.map(extract_params).unwrap_or_default();
Expand Down Expand Up @@ -277,10 +276,36 @@ pub(crate) fn expand(
None
};

let countdown_static = has_after.then(|| {
// The `#[after]` hook runs once at process exit (registered via
// `register_teardown`), guarded by this `Once` so it registers a single time
// no matter which — or how many — of the group's tests libtest selects.
// Counting tests at expansion time can't work: name filters, `--skip`, and
// `#[ignore]` all change the run set at runtime (issue #1).
let after_at_exit = has_after.then(|| {
let name = after_name.unwrap();
let invocation = if after_params.is_empty() {
quote! { #name(); }
} else {
// In practice `#[after]` only takes `&T` references to the
// `#[before]` context, which lives in `__SPEC_BEFORE_CTX` for the
// whole program and is initialized by the first test that runs.
let args = after_params.iter().map(|_| quote! { __spec_after_ctx });
quote! {
// If `#[before]` never produced context (e.g. no test ran, or
// `#[before]` itself panicked) there is nothing to tear down.
let ::std::option::Option::Some(__spec_after_ctx) =
__SPEC_BEFORE_CTX.get()
else {
return;
};
#name(#(#args),*);
}
};
quote! {
static __SPEC_AFTER_REMAINING: ::std::sync::atomic::AtomicUsize =
::std::sync::atomic::AtomicUsize::new(#test_count);
static __SPEC_AFTER_ONCE: ::std::sync::Once = ::std::sync::Once::new();
fn __spec_after_at_exit() {
#invocation
}
}
});

Expand All @@ -306,6 +331,15 @@ pub(crate) fn expand(
let mut pre = proc_macro2::TokenStream::new();
let mut post = proc_macro2::TokenStream::new();

// --- Register group `after` teardown (runs once, at process exit) ---
if has_after {
pre.extend(quote! {
__SPEC_AFTER_ONCE.call_once(|| {
::spectacular::__internal::register_teardown(__spec_after_at_exit);
});
});
}

// --- Suite before ---
if has_suite {
pre.extend(quote! { super::__spectacular_suite::before(); });
Expand Down Expand Up @@ -530,42 +564,10 @@ pub(crate) fn expand(
post.extend(quote! { super::__spectacular_suite::after_each(); });
}

// --- after (countdown) ---
if let Some(name) = after_name {
let call_args: Vec<proc_macro2::TokenStream> = after_params
.iter()
.map(|p| {
if p.is_ref {
quote! { __before_ctx }
} else {
let pat = &p.pat;
quote! { #pat }
}
})
.collect();

if call_args.is_empty() {
post.extend(quote! {
if __SPEC_AFTER_REMAINING
.fetch_sub(1, ::std::sync::atomic::Ordering::SeqCst)
== 1
{
#name();
}
});
} else {
post.extend(quote! {
if __SPEC_AFTER_REMAINING
.fetch_sub(1, ::std::sync::atomic::Ordering::SeqCst)
== 1
{
#name(#(#call_args),*);
}
});
}
}
// NOTE: the group `after` hook is NOT run here. It runs once at
// process exit — see `after_at_exit` and the registration in `pre`.

let needs_catch = has_after || has_after_each || has_suite;
let needs_catch = has_after_each || has_suite;

if test_needs_async {
let rt = runtime.unwrap();
Expand Down Expand Up @@ -597,7 +599,7 @@ pub(crate) fn expand(
#vis mod #mod_name {
#(#cleaned_items)*
#once_static
#countdown_static
#after_at_exit
#(#test_fn_defs)*
}
})
Expand Down
79 changes: 39 additions & 40 deletions crates/spectacular-macros/src/spec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -386,7 +386,6 @@ pub(crate) fn expand(input: proc_macro2::TokenStream) -> syn::Result<proc_macro2
|| after_each_params.iter().any(|p| is_type_infer(&p.ty));
let before_each_needs_inline = !has_before_each_ctx && has_infer_consumers;
let after_each_needs_inline = after_each_params.iter().any(|p| is_type_infer(&p.ty));
let test_count = tests.len();

// Generate before fn
let before_fn = if let Some(body) = &before_body {
Expand Down Expand Up @@ -483,10 +482,33 @@ pub(crate) fn expand(input: proc_macro2::TokenStream) -> syn::Result<proc_macro2
None
};

let countdown_static = has_after.then(|| {
// The `after` hook runs once at process exit (registered via
// `register_teardown`), guarded by this `Once`. Counting tests at expansion
// time can't work: name filters, `--skip`, and `#[ignore]` change the run
// set at runtime (issue #1).
let after_at_exit = has_after.then(|| {
let invocation = if after_params.is_empty() {
quote! { __spec_after(); }
} else {
// `after` only takes `&T` references to the `before` context, which
// lives in `__SPEC_BEFORE_CTX` for the whole program.
let args = after_params.iter().map(|_| quote! { __spec_after_ctx });
quote! {
// If `before` never produced context (e.g. no test ran, or
// `before` itself panicked) there is nothing to tear down.
let ::std::option::Option::Some(__spec_after_ctx) =
__SPEC_BEFORE_CTX.get()
else {
return;
};
__spec_after(#(#args),*);
}
};
quote! {
static __SPEC_AFTER_REMAINING: ::std::sync::atomic::AtomicUsize =
::std::sync::atomic::AtomicUsize::new(#test_count);
static __SPEC_AFTER_ONCE: ::std::sync::Once = ::std::sync::Once::new();
fn __spec_after_at_exit() {
#invocation
}
}
});

Expand All @@ -499,6 +521,15 @@ pub(crate) fn expand(input: proc_macro2::TokenStream) -> syn::Result<proc_macro2
let mut pre = proc_macro2::TokenStream::new();
let mut post = proc_macro2::TokenStream::new();

// --- Register group `after` teardown (runs once, at process exit) ---
if has_after {
pre.extend(quote! {
__SPEC_AFTER_ONCE.call_once(|| {
::spectacular::__internal::register_teardown(__spec_after_at_exit);
});
});
}

// --- Suite before ---
if has_suite {
pre.extend(quote! { super::__spectacular_suite::before(); });
Expand Down Expand Up @@ -713,42 +744,10 @@ pub(crate) fn expand(input: proc_macro2::TokenStream) -> syn::Result<proc_macro2
post.extend(quote! { super::__spectacular_suite::after_each(); });
}

// --- after (countdown) ---
if has_after {
let call_args: Vec<proc_macro2::TokenStream> = after_params
.iter()
.map(|p| {
if p.is_ref {
quote! { __before_ctx }
} else {
let pat = &p.pat;
quote! { #pat }
}
})
.collect();

if call_args.is_empty() {
post.extend(quote! {
if __SPEC_AFTER_REMAINING
.fetch_sub(1, ::std::sync::atomic::Ordering::SeqCst)
== 1
{
__spec_after();
}
});
} else {
post.extend(quote! {
if __SPEC_AFTER_REMAINING
.fetch_sub(1, ::std::sync::atomic::Ordering::SeqCst)
== 1
{
__spec_after(#(#call_args),*);
}
});
}
}
// NOTE: the group `after` hook is NOT run here. It runs once at
// process exit — see `after_at_exit` and the registration in `pre`.

let needs_catch = has_after || has_after_each || has_suite;
let needs_catch = has_after_each || has_suite;

if test_needs_async {
let rt = runtime.unwrap();
Expand Down Expand Up @@ -778,7 +777,7 @@ pub(crate) fn expand(input: proc_macro2::TokenStream) -> syn::Result<proc_macro2
#vis mod #mod_name {
#(#other_items)*
#once_static
#countdown_static
#after_at_exit
#before_fn
#after_fn
#before_each_fn
Expand Down
72 changes: 68 additions & 4 deletions crates/spectacular/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
//! TEST
//! group::after_each
//! suite::after_each
//! group::after (countdownlast test in group triggers it)
//! group::after (onceat process exit, if any test in the group ran)
//! ```
//!
//! Groups without `suite;` skip the suite layer entirely.
Expand Down Expand Up @@ -598,9 +598,14 @@ pub use spectacular_macros::before;
/// Marks a function as a once-per-group teardown hook inside a
/// [`#[test_suite]`](macro@test_suite) module.
///
/// The function runs exactly once after the last test in the group completes,
/// using an atomic countdown. Only one `#[after]` per module is allowed.
/// Must be sync.
/// The function runs exactly once at process exit, after every test the
/// harness selected has finished — provided at least one test in the group
/// ran. Running teardown at exit (rather than counting tests) keeps it correct
/// under name filters, `--skip`, and `#[ignore]`/`#[cfg]`, which change the run
/// set at runtime. Only one `#[after]` per module is allowed. Must be sync.
///
/// If the process is killed (e.g. `SIGKILL`) or aborts, `after` does not run —
/// no in-process teardown hook can guarantee that.
///
/// When `#[before]` returns context, `after` can receive it as `&T` via a
/// reference parameter: `fn cleanup(pool: &PgPool)`. Without parameters,
Expand Down Expand Up @@ -800,6 +805,7 @@ pub mod __internal {
use std::any::Any;
use std::future::Future;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::{Mutex, Once};
use std::task::Poll;

/// Like `std::panic::catch_unwind` but for async blocks.
Expand All @@ -817,6 +823,64 @@ pub mod __internal {
})
.await
}

// --- Group teardown (`after`) registry ---
//
// A group's `after` hook must run exactly once, after every *selected* test
// in the binary has finished. libtest decides which tests run at runtime
// (name filters, `--skip`, `--ignored`), so no compile-time count can know
// that set. Instead we register the group's teardown the first time any of
// its tests runs and execute it at process exit, when the run is provably
// over. This is robust to filtering, `--skip`, and `#[ignore]`/`#[cfg]`.

unsafe extern "C" {
/// C runtime `atexit`, linked by every Rust binary. The callback runs
/// after `main` returns or `std::process::exit` is called.
fn atexit(cb: extern "C" fn()) -> core::ffi::c_int;
}

static TEARDOWNS: Mutex<Vec<fn()>> = Mutex::new(Vec::new());
static ATEXIT_REGISTERED: Once = Once::new();

extern "C" fn run_teardowns() {
let pending: Vec<fn()> = {
let mut guard = TEARDOWNS.lock().unwrap_or_else(|e| e.into_inner());
std::mem::take(&mut *guard)
};
// Reverse registration order: the last group set up is torn down first.
// Each teardown is a normal Rust `fn`, so a panic (e.g. a failing
// assertion in `after`) unwinds safely into this `catch_unwind` rather
// than crossing the `extern "C"` boundary.
let mut any_panicked = false;
for teardown in pending.into_iter().rev() {
if catch_unwind(AssertUnwindSafe(teardown)).is_err() {
any_panicked = true;
}
}
if any_panicked {
eprintln!("spectacular: an `after` teardown hook panicked during shutdown; aborting");
std::process::abort();
}
}

/// Register a once-per-group teardown to run at process exit.
///
/// Generated code calls this — guarded by a per-group [`Once`] — the first
/// time one of the group's tests runs, so the group's `after` hook fires
/// exactly once regardless of how libtest filtered the run.
pub fn register_teardown(teardown: fn()) {
ATEXIT_REGISTERED.call_once(|| {
// SAFETY: `run_teardowns` is a valid `extern "C" fn()`; `atexit` is
// provided by the C runtime linked into every Rust binary.
unsafe {
atexit(run_teardowns);
}
});
TEARDOWNS
.lock()
.unwrap_or_else(|e| e.into_inner())
.push(teardown);
}
}

/// Convenience re-export of all spectacular macros.
Expand Down
Loading
Loading