From 4092bbe6c1706188047f600bf42cae7d9f5befe9 Mon Sep 17 00:00:00 2001 From: kid Date: Wed, 22 Jul 2026 20:22:38 +0800 Subject: [PATCH 1/3] fix(mem-wal): propagate task shutdown failures --- rust/lance/src/dataset/mem_wal/write.rs | 127 +++++++++++++++++++++++- 1 file changed, 123 insertions(+), 4 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 50a0932be49..8efef232804 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -541,15 +541,28 @@ impl TaskExecutor { self.cancellation_token.cancel(); let tasks = std::mem::take(&mut *self.tasks.write().unwrap()); + let mut first_error = None; for (name, handle) in tasks { match handle.await { Ok(Ok(())) => debug!("Task '{}' completed successfully", name), - Ok(Err(e)) => warn!("Task '{}' completed with error: {}", name, e), - Err(e) => error!("Task '{}' panicked: {}", name, e), + Ok(Err(e)) => { + warn!("Task '{}' completed with error: {}", name, e); + if first_error.is_none() { + first_error = Some(e); + } + } + Err(e) => { + error!("Task '{}' panicked: {}", name, e); + if first_error.is_none() { + first_error = Some(Error::internal(format!( + "Task '{name}' panicked during shutdown: {e}" + ))); + } + } } } - Ok(()) + first_error.map_or(Ok(()), Err) } } @@ -4703,7 +4716,113 @@ mod tests { call_count.load(Ordering::SeqCst) ); - executor.shutdown_all().await.ok(); + executor + .shutdown_all() + .await + .expect("dispatcher should shut down successfully"); + } + + #[tokio::test] + async fn test_task_executor_shutdown_propagates_cleanup_error_and_joins_all_tasks() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct CleanupHandler { + cleanup_count: Arc, + error_message: Option<&'static str>, + } + + #[async_trait] + impl MessageHandler for CleanupHandler { + async fn handle(&mut self, _message: u32) -> Result<()> { + Ok(()) + } + + async fn cleanup(&mut self, _shutdown_ok: bool) -> Result<()> { + self.cleanup_count.fetch_add(1, Ordering::SeqCst); + match self.error_message { + Some(message) => Err(Error::io(message)), + None => Ok(()), + } + } + } + + let executor = TaskExecutor::new(); + let cleanup_count = Arc::new(AtomicUsize::new(0)); + let (_failing_tx, failing_rx) = mpsc::unbounded_channel::(); + executor + .add_handler( + "failing-cleanup".to_string(), + Box::new(CleanupHandler { + cleanup_count: cleanup_count.clone(), + error_message: Some("intentional cleanup failure"), + }), + failing_rx, + ) + .unwrap(); + let (_successful_tx, successful_rx) = mpsc::unbounded_channel::(); + executor + .add_handler( + "successful-cleanup".to_string(), + Box::new(CleanupHandler { + cleanup_count: cleanup_count.clone(), + error_message: None, + }), + successful_rx, + ) + .unwrap(); + + let error = executor + .shutdown_all() + .await + .expect_err("shutdown must propagate the handler cleanup failure"); + assert!(matches!(&error, Error::IO { .. })); + assert!( + error.to_string().contains("intentional cleanup failure"), + "unexpected error: {error}" + ); + assert_eq!( + cleanup_count.load(Ordering::SeqCst), + 2, + "shutdown must join and clean up every task after the first failure" + ); + assert!(executor.tasks.read().unwrap().is_empty()); + } + + #[tokio::test] + async fn test_task_executor_shutdown_propagates_task_panic() { + struct PanickingCleanupHandler; + + #[async_trait] + impl MessageHandler for PanickingCleanupHandler { + async fn handle(&mut self, _message: u32) -> Result<()> { + Ok(()) + } + + async fn cleanup(&mut self, _shutdown_ok: bool) -> Result<()> { + panic!("intentional cleanup panic"); + } + } + + let executor = TaskExecutor::new(); + let (_tx, rx) = mpsc::unbounded_channel::(); + executor + .add_handler( + "panicking-cleanup".to_string(), + Box::new(PanickingCleanupHandler), + rx, + ) + .unwrap(); + + let error = executor + .shutdown_all() + .await + .expect_err("shutdown must propagate the task panic"); + assert!(matches!(&error, Error::Internal { .. })); + assert!( + error.to_string().contains("panicking-cleanup") + && error.to_string().contains("panicked during shutdown"), + "unexpected error: {error}" + ); } /// Same as the local-fs test but against memory:// — closer to S3 From eabd39006e5096f280dbdfa0d6c2f42f76686f09 Mon Sep 17 00:00:00 2001 From: kid Date: Wed, 22 Jul 2026 22:21:38 +0800 Subject: [PATCH 2/3] docs(mem-wal): document task shutdown contract --- rust/lance/src/dataset/mem_wal/write.rs | 50 +++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 8efef232804..fafa233fb4c 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -536,6 +536,56 @@ impl TaskExecutor { Ok(()) } + /// Cancel and join every handler registered by [`Self::add_handler`]. + /// + /// Cancellation causes each handler's dispatcher to stop accepting messages and call + /// [`MessageHandler::cleanup`]. This method waits for every dispatcher to finish, even if + /// cleanup fails or a dispatcher panics, and then returns the first such failure. It returns + /// `Ok(())` only after every registered handler has been cleaned up successfully. + /// + /// # Example + /// + /// ``` + /// use std::sync::{ + /// Arc, + /// atomic::{AtomicBool, Ordering}, + /// }; + /// + /// use async_trait::async_trait; + /// use lance::dataset::mem_wal::write::{MessageHandler, TaskExecutor}; + /// use lance::Result; + /// use tokio::sync::mpsc; + /// + /// struct Handler(Arc); + /// + /// #[async_trait] + /// impl MessageHandler<()> for Handler { + /// async fn handle(&mut self, _message: ()) -> Result<()> { + /// Ok(()) + /// } + /// + /// async fn cleanup(&mut self, _shutdown_ok: bool) -> Result<()> { + /// self.0.store(true, Ordering::SeqCst); + /// Ok(()) + /// } + /// } + /// + /// # #[tokio::main] + /// # async fn main() -> Result<()> { + /// let executor = TaskExecutor::new(); + /// let cleaned_up = Arc::new(AtomicBool::new(false)); + /// let (_sender, receiver) = mpsc::unbounded_channel(); + /// executor.add_handler( + /// "example".to_string(), + /// Box::new(Handler(cleaned_up.clone())), + /// receiver, + /// )?; + /// + /// executor.shutdown_all().await?; + /// assert!(cleaned_up.load(Ordering::SeqCst)); + /// # Ok(()) + /// # } + /// ``` pub async fn shutdown_all(&self) -> Result<()> { info!("Shutting down all tasks"); self.cancellation_token.cancel(); From dd2c506f959550e6df02e155de12b2f2e7a7d9d9 Mon Sep 17 00:00:00 2001 From: kid Date: Thu, 23 Jul 2026 09:42:58 +0800 Subject: [PATCH 3/3] docs(mem-wal): remove shutdown example --- rust/lance/src/dataset/mem_wal/write.rs | 44 ------------------------- 1 file changed, 44 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index fafa233fb4c..ab87ec68393 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -542,50 +542,6 @@ impl TaskExecutor { /// [`MessageHandler::cleanup`]. This method waits for every dispatcher to finish, even if /// cleanup fails or a dispatcher panics, and then returns the first such failure. It returns /// `Ok(())` only after every registered handler has been cleaned up successfully. - /// - /// # Example - /// - /// ``` - /// use std::sync::{ - /// Arc, - /// atomic::{AtomicBool, Ordering}, - /// }; - /// - /// use async_trait::async_trait; - /// use lance::dataset::mem_wal::write::{MessageHandler, TaskExecutor}; - /// use lance::Result; - /// use tokio::sync::mpsc; - /// - /// struct Handler(Arc); - /// - /// #[async_trait] - /// impl MessageHandler<()> for Handler { - /// async fn handle(&mut self, _message: ()) -> Result<()> { - /// Ok(()) - /// } - /// - /// async fn cleanup(&mut self, _shutdown_ok: bool) -> Result<()> { - /// self.0.store(true, Ordering::SeqCst); - /// Ok(()) - /// } - /// } - /// - /// # #[tokio::main] - /// # async fn main() -> Result<()> { - /// let executor = TaskExecutor::new(); - /// let cleaned_up = Arc::new(AtomicBool::new(false)); - /// let (_sender, receiver) = mpsc::unbounded_channel(); - /// executor.add_handler( - /// "example".to_string(), - /// Box::new(Handler(cleaned_up.clone())), - /// receiver, - /// )?; - /// - /// executor.shutdown_all().await?; - /// assert!(cleaned_up.load(Ordering::SeqCst)); - /// # Ok(()) - /// # } - /// ``` pub async fn shutdown_all(&self) -> Result<()> { info!("Shutting down all tasks"); self.cancellation_token.cancel();