diff --git a/native/jni-bridge/src/errors.rs b/native/jni-bridge/src/errors.rs index 179a214ed92..2b78195300f 100644 --- a/native/jni-bridge/src/errors.rs +++ b/native/jni-bridge/src/errors.rs @@ -24,13 +24,7 @@ use jni::errors::{Exception, ToException}; use regex::Regex; use std::{ - any::Any, - convert, - fmt::Write, - panic::UnwindSafe, - result, str, - str::Utf8Error, - sync::{Arc, Mutex}, + any::Any, cell::Cell, convert, fmt::Write, panic::UnwindSafe, result, str, str::Utf8Error, }; // This is just a pointer. We'll be returning it from our function. We @@ -44,8 +38,38 @@ use lazy_static::lazy_static; use parquet::errors::ParquetError; use thiserror::Error; -lazy_static! { - static ref PANIC_BACKTRACE: Arc>> = Arc::new(Mutex::new(None)); +/// A panic's message, when it is a string, and its backtrace, as the panic hook captured them. +struct CapturedPanic { + message: Option, + backtrace: String, +} + +thread_local! { + /// The last panic on this thread, captured by the panic hook and taken at the boundary that + /// catches it. Kept per thread because Spark runs many tasks at once. + static LAST_PANIC: Cell> = const { Cell::new(None) }; +} + +/// Forgets the panic captured on this thread. Every boundary starts with this, so a panic that +/// was caught and handled before the call is never taken for one raised during it. +fn clear_captured_panic() { + let _ = LAST_PANIC.try_with(Cell::take); +} + +/// The backtrace of the panic with `message` that reached the boundary, if the panic hook +/// captured it on this thread. A panic resumed from another thread has none: DataFusion runs +/// some operators' inputs on spawned Tokio tasks and resumes their panics on the consuming +/// thread with `resume_unwind`, which does not run the hook. Matching the message keeps such a +/// panic from taking the backtrace of another panic that was caught and handled on this thread +/// during the same call, for instance by a `catch_unwind` outside Comet's control. +fn take_panic_backtrace(message: &str) -> Option { + LAST_PANIC + .try_with(Cell::take) + .ok() + .flatten() + // `CometError::Panic` reports a panic that is not a string as "unknown panic". + .filter(|captured| captured.message.as_deref().unwrap_or("unknown panic") == message) + .map(|captured| captured.backtrace) } /// Error returned during executing operators. @@ -187,9 +211,13 @@ pub fn init() { // Log the panic message and location to stderr so it is visible in CI logs // even if the exception message is lost crossing the FFI boundary eprintln!("Comet native panic: {panic_info}"); - // Capture the backtrace for a panic - *PANIC_BACKTRACE.lock().unwrap() = - Some(std::backtrace::Backtrace::force_capture().to_string()); + // Capture the backtrace for a panic. `try_with` because a panic in the hook aborts the + // process, and the thread-local is gone while the thread is being torn down. + let captured = CapturedPanic { + message: panic_info.payload_as_str().map(str::to_string), + backtrace: std::backtrace::Backtrace::force_capture().to_string(), + }; + let _ = LAST_PANIC.try_with(|slot| slot.set(Some(captured))); })); } @@ -643,8 +671,8 @@ pub struct NativeFailure { impl From for NativeFailure { fn from(error: CometError) -> Self { - let backtrace = match error { - CometError::Panic { .. } => PANIC_BACKTRACE.lock().unwrap().take(), + let backtrace = match &error { + CometError::Panic { msg } => take_panic_backtrace(msg), _ => None, }; NativeFailure { error, backtrace } @@ -664,6 +692,7 @@ pub fn catch_native(f: F) -> Result where F: FnOnce() -> CometResult + UnwindSafe, { + clear_captured_panic(); match std::panic::catch_unwind(f) { Ok(Ok(value)) => Ok(value), Ok(Err(error)) => Err(NativeFailure::from(error)), @@ -1080,6 +1109,7 @@ where T: JNIDefault, F: FnOnce(&mut Env) -> Result + UnwindSafe, { + clear_captured_panic(); let raw = env.as_raw(); let mut env1 = unsafe { EnvUnowned::from_raw(raw) }; match env1.with_env(f).into_outcome() { @@ -1615,6 +1645,127 @@ mod tests { std::error::Error::source(error).unwrap().to_string() } + /// Waits twice on a barrier when dropped. Dropped while its thread unwinds, it runs after that + /// thread's panic hook and before `catch_native` takes the backtrace. + struct WaitOnDrop(Arc); + + impl Drop for WaitOnDrop { + fn drop(&mut self) { + // Both panic hooks have run once both threads pass the first wait. + self.0.wait(); + self.0.wait(); + } + } + + #[inline(never)] + fn panic_in_first_thread(barrier: Arc) -> CometResult<()> { + let _wait = WaitOnDrop(barrier); + panic!("first thread") + } + + #[inline(never)] + fn panic_in_second_thread(barrier: Arc) -> CometResult<()> { + let _wait = WaitOnDrop(barrier); + panic!("second thread") + } + + /// Two threads panic at once, and each takes its backtrace only after both panic hooks have + /// run. Each must still get the backtrace of its own panic. + #[test] + fn panic_backtrace_stays_with_its_thread() { + init(); + let barrier = Arc::new(std::sync::Barrier::new(2)); + let first = { + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || catch_native(move || panic_in_first_thread(barrier))) + }; + let second = { + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || catch_native(move || panic_in_second_thread(barrier))) + }; + + for (handle, own, other) in [ + (first, "panic_in_first_thread", "panic_in_second_thread"), + (second, "panic_in_second_thread", "panic_in_first_thread"), + ] { + let failure = handle.join().unwrap().unwrap_err(); + let backtrace = failure + .backtrace + .unwrap_or_else(|| panic!("no backtrace for {own}")); + assert!(backtrace.contains(own), "{own} got another backtrace"); + assert!(!backtrace.contains(other), "{own} got {other}'s backtrace"); + } + } + + #[inline(never)] + fn panic_in_spawned_thread() { + panic!("spawned thread") + } + + /// Resumes on this thread a panic raised on another one, as DataFusion does with the panics + /// of its spawned tasks. `resume_unwind` does not run the panic hook. + fn resume_panic_from_another_thread() -> CometResult<()> { + let payload = std::thread::spawn(panic_in_spawned_thread) + .join() + .expect_err("the spawned thread panics"); + std::panic::resume_unwind(payload) + } + + #[test] + fn panic_resumed_from_another_thread_has_no_backtrace() { + init(); + let failure = catch_native(resume_panic_from_another_thread).unwrap_err(); + assert!( + matches!(&failure.error, CometError::Panic { msg } if msg == "spawned thread"), + "{:?}", + failure.error + ); + assert!(failure.backtrace.is_none(), "{:?}", failure.backtrace); + } + + /// A panic caught and handled during the call leaves its backtrace on this thread. A panic + /// resumed from another thread later in the same call must not take it. + #[test] + fn resumed_panic_does_not_take_a_handled_panic_backtrace() { + init(); + let failure = catch_native(|| { + let handled = std::panic::catch_unwind(|| panic!("handled during the call")); + assert!(handled.is_err()); + resume_panic_from_another_thread() + }) + .unwrap_err(); + assert!(failure.backtrace.is_none(), "{:?}", failure.backtrace); + } + + /// A panic handled before the call, here with the same message as the one the call resumes, + /// must not lend its backtrace either. + #[test] + fn resumed_panic_does_not_take_a_backtrace_from_before_the_call() { + init(); + let handled = std::panic::catch_unwind(|| panic!("spawned thread")); + assert!(handled.is_err()); + let failure = catch_native(resume_panic_from_another_thread).unwrap_err(); + assert!(failure.backtrace.is_none(), "{:?}", failure.backtrace); + } + + #[inline(never)] + fn panic_with_formatted_message(value: i32) -> CometResult<()> { + panic!("formatted message {value}") + } + + #[test] + fn panic_with_formatted_message_keeps_its_backtrace() { + init(); + let failure = catch_native(|| panic_with_formatted_message(7)).unwrap_err(); + assert!( + matches!(&failure.error, CometError::Panic { msg } if msg == "formatted message 7"), + "{:?}", + failure.error + ); + let backtrace = failure.backtrace.expect("the panic's own backtrace"); + assert!(backtrace.contains("panic_with_formatted_message")); + } + #[test] fn catch_native_reports_errors_and_panics() { assert_eq!(catch_native(|| Ok(7)).unwrap(), 7);