test: better handling of `loom` double panics

This ensures the current iteration's trace is *always* printed, even if
another panic occurs while unwinding.

src/loom.rs

@@ -31,87 +31,93 @@ mod inner {
     ) {
         use std::{
             env, io,
-            sync::atomic::{AtomicUsize, Ordering},
+            sync::{
+                atomic::{AtomicUsize, Ordering},
+                Once,
+            },
         };
         use tracing_subscriber::{filter::Targets, fmt, prelude::*};
 
-        // set up tracing for loom.
-        const LOOM_LOG: &str = "LOOM_LOG";
-
-        struct TracebufWriter;
-        impl io::Write for TracebufWriter {
-            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-                let len = buf.len();
-                let s = std::str::from_utf8(buf)
-                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
-                TRACE_BUF.with(|buf| buf.borrow_mut().push_str(s));
-                Ok(len)
-            }
+        static SETUP_TRACE: Once = Once::new();
 
-            fn flush(&mut self) -> io::Result<()> {
-                Ok(())
-            }
-        }
+        SETUP_TRACE.call_once(|| {
+            // set up tracing for loom.
+            const LOOM_LOG: &str = "LOOM_LOG";
 
-        let filter = env::var(LOOM_LOG)
-            .ok()
-            .and_then(|var| match var.parse::<Targets>() {
-                Err(e) => {
-                    eprintln!("invalid {}={:?}: {}", LOOM_LOG, var, e);
-                    None
+            struct TracebufWriter;
+            impl io::Write for TracebufWriter {
+                fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+                    let len = buf.len();
+                    let s = std::str::from_utf8(buf)
+                        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
+                    TRACE_BUF.with(|buf| buf.borrow_mut().push_str(s));
+                    Ok(len)
                 }
-                Ok(targets) => Some(targets),
-            })
-            .unwrap_or_else(|| Targets::new().with_target("loom", tracing::Level::INFO));
-        let _ = fmt::Subscriber::builder()
-            .with_writer(|| TracebufWriter)
-            .without_time()
-            .finish()
-            .with(filter)
-            .try_init();
 
-        // wrap the loom model with `catch_unwind` to avoid potentially losing
-        // test output on double panics.
-        let current_iteration = std::sync::Arc::new(AtomicUsize::new(1));
-        let result = {
-            let current_iteration = current_iteration.clone();
-            std::panic::catch_unwind(move || {
-                builder.check(move || {
-                    traceln(format_args!(
-                        "\n---- {} iteration {} ----",
-                        std::thread::current().name().unwrap_or("<unknown test>"),
-                        current_iteration.fetch_add(1, Ordering::Relaxed)
-                    ));
-
-                    model();
-                    // if this iteration succeeded, clear the buffer for the
-                    // next iteration...
-                    TRACE_BUF.with(|buf| buf.borrow_mut().clear());
-                })
-            })
-        };
+                fn flush(&mut self) -> io::Result<()> {
+                    Ok(())
+                }
+            }
 
-        if let Err(panic) = result {
-            TRACE_BUF
-                .try_with(|buf| {
-                    if let Ok(buf) = buf.try_borrow() {
-                        eprint!("{}", buf);
-                    } else {
-                        eprint!("trace buf already mutably borrowed?");
+            let filter = env::var(LOOM_LOG)
+                .ok()
+                .and_then(|var| match var.parse::<Targets>() {
+                    Err(e) => {
+                        eprintln!("invalid {}={:?}: {}", LOOM_LOG, var, e);
+                        None
                     }
+                    Ok(targets) => Some(targets),
                 })
-                .unwrap_or_else(|e| eprintln!("trace buf already torn down: {}", e));
-            eprintln!(
-                "test '{}' panicked after {} iterations!",
+                .unwrap_or_else(|| Targets::new().with_target("loom", tracing::Level::INFO));
+            fmt::Subscriber::builder()
+                .with_writer(|| TracebufWriter)
+                .without_time()
+                .finish()
+                .with(filter)
+                .init();
+
+            let default_hook = std::panic::take_hook();
+            std::panic::set_hook(Box::new(move |panic| {
+                // try to print the trace buffer.
+                TRACE_BUF
+                    .try_with(|buf| {
+                        if let Ok(mut buf) = buf.try_borrow_mut() {
+                            eprint!("{}", buf);
+                            buf.clear();
+                        } else {
+                            eprint!("trace buf already mutably borrowed?");
+                        }
+                    })
+                    .unwrap_or_else(|e| eprintln!("trace buf already torn down: {}", e));
+
+                // let the default panic hook do the rest...
+                default_hook(panic);
+            }))
+        });
+
+        // wrap the loom model with `catch_unwind` to avoid potentially losing
+        // test output on double panics.
+        let current_iteration = std::sync::Arc::new(AtomicUsize::new(1));
+        builder.check(move || {
+            traceln(format_args!(
+                "\n---- {} iteration {} ----",
                 std::thread::current().name().unwrap_or("<unknown test>"),
-                current_iteration.load(Ordering::Relaxed),
-            );
-            std::panic::resume_unwind(panic);
-        }
+                current_iteration.fetch_add(1, Ordering::Relaxed)
+            ));
+
+            model();
+            // if this iteration succeeded, clear the buffer for the
+            // next iteration...
+            TRACE_BUF.with(|buf| buf.borrow_mut().clear());
+        });
     }
 
     pub(crate) fn model(model: impl Fn() + std::panic::UnwindSafe + Sync + Send + 'static) {
-        run_builder(loom::model::Builder::default(), model)
+        let mut builder = loom::model::Builder::default();
+        // // A couple of our tests will hit the max number of branches riiiiight
+        // // before they should complete. Double it so this stops happening.
+        builder.max_branches *= 2;
+        run_builder(builder, model)
     }
 
     pub(crate) mod alloc {