From f89d9d1334a554e0c36d4b064af378ad2dfec826 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" Date: Thu, 23 Jul 2026 16:20:29 +0800 Subject: [PATCH] fix(mito2): notify pick finished even when compaction planning panics The worker only leaves the picking phase after receiving the CompactionPickFinished notification. Previously the planning task was spawned fire-and-forget: if it panicked before sending the notification, the region would be stuck in the picking phase forever, blocking all future compactions and pending DDLs (e.g. entering staging) of the region. Wrap the planning future with catch_unwind so a panic is converted into a CompactionPlanningResult::Error and the notification is always sent, letting the worker run the normal error cleanup path. Signed-off-by: Lei, HUANG --- src/mito2/src/compaction.rs | 105 +++++++++++++++++++++++++++++------- 1 file changed, 87 insertions(+), 18 deletions(-) diff --git a/src/mito2/src/compaction.rs b/src/mito2/src/compaction.rs index f055fa15d7..265b432b56 100644 --- a/src/mito2/src/compaction.rs +++ b/src/mito2/src/compaction.rs @@ -23,8 +23,10 @@ mod test_util; mod twcs; mod window; +use std::any::Any; use std::collections::{HashMap, HashSet}; use std::fmt; +use std::future::Future; use std::sync::{Arc, Mutex}; use std::time::Instant; @@ -42,6 +44,7 @@ use datafusion_common::ScalarValue; use datafusion_expr::Expr; use datatypes::extension::json::is_structured_json_field; use datatypes::types::json_type::JsonNativeType; +use futures::FutureExt; use parquet::arrow::parquet_to_arrow_schema; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; use serde::{Deserialize, Serialize}; @@ -62,7 +65,7 @@ use crate::error::{ CompactRegionSnafu, CompactionCancelledSnafu, DataTypeMismatchSnafu, Error, GetSchemaMetadataSnafu, JoinSnafu, ManualCompactionOverrideSnafu, ParquetToArrowSchemaSnafu, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, RemoteCompactionSnafu, Result, - TimeRangePredicateOverflowSnafu, TimeoutSnafu, + TimeRangePredicateOverflowSnafu, TimeoutSnafu, UnexpectedSnafu, }; use crate::metrics::{ COMPACTION_MEMORY_REJECTED, COMPACTION_STAGE_ELAPSED, INFLIGHT_COMPACTION_COUNT, @@ -606,27 +609,69 @@ impl CompactionScheduler { common_runtime::spawn_compact(async move { let region_id = request.region_id(); let request_sender = request.request_sender.clone(); - let result = - Self::prepare_compaction(request, options, plugins, max_background_compactions) - .await; - if let CompactionPlanningResult::Error(err) = &result { - error!(err; "Compaction planning failed for region {}, plan_id: {}", region_id, plan_id); - } - let request = WorkerRequestWithTime::new(WorkerRequest::Background { + let planning = + Self::prepare_compaction(request, options, plugins, max_background_compactions); + Self::notify_planning_result( region_id, - notify: BackgroundNotify::CompactionPickFinished(CompactionPickFinished { - region_id, - plan_id, - version_control, - result, - }), - }); - if request_sender.send(request).await.is_err() { - warn!("Failed to send compaction planning result for region {region_id}"); - } + plan_id, + version_control, + request_sender, + planning, + ) + .await; }); } + /// Runs the planning future and always sends the planning result back to + /// the worker, even if the planning panics. + /// + /// The worker only leaves the picking phase after it receives the + /// `CompactionPickFinished` notification. If a panicked planning task + /// swallowed the notification, the region would be stuck in the picking + /// phase forever, blocking all future compactions and pending DDLs (e.g. + /// entering staging) of the region. + async fn notify_planning_result( + region_id: RegionId, + plan_id: u64, + version_control: VersionControlRef, + request_sender: Sender, + planning: impl Future + Send, + ) { + // The idiomatic way to handle a panic result. + let result = std::panic::AssertUnwindSafe(planning).catch_unwind().await.unwrap_or_else(|payload| { + let reason = if let Some(message) = payload.as_ref().downcast_ref::<&str>() { + message.to_string() + } else if let Some(message) = payload.as_ref().downcast_ref::() { + message.clone() + } else { + "unknown panic".to_string() + }; + CompactionPlanningResult::Error(Arc::new( + UnexpectedSnafu { + reason: format!( + "Compaction planning panicked for region {region_id}, plan_id {plan_id}: {reason}" + ), + } + .build(), + )) + }); + if let CompactionPlanningResult::Error(err) = &result { + error!(err; "Compaction planning failed for region {}, plan_id: {}", region_id, plan_id); + } + let request = WorkerRequestWithTime::new(WorkerRequest::Background { + region_id, + notify: BackgroundNotify::CompactionPickFinished(CompactionPickFinished { + region_id, + plan_id, + version_control, + result, + }), + }); + if request_sender.send(request).await.is_err() { + warn!("Failed to send compaction planning result for region {region_id}"); + } + } + async fn prepare_compaction( request: CompactionRequest, options: compact_request::Options, @@ -2750,6 +2795,30 @@ mod tests { assert!(!scheduler.region_status.contains_key(®ion_id)); } + #[tokio::test] + async fn test_planning_panic_still_sends_pick_finished() { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let (tx, mut rx) = mpsc::channel(4); + + CompactionScheduler::notify_planning_result(region_id, 7, version_control, tx, async { + panic!("planning boom") + }) + .await; + + let finished = recv_compaction_pick_finished(&mut rx).await; + assert_eq!(region_id, finished.region_id); + assert_eq!(7, finished.plan_id); + let CompactionPlanningResult::Error(err) = finished.result else { + panic!("expected planning error, got {:?}", finished.result); + }; + assert!( + err.to_string().contains("planning boom"), + "unexpected error: {err}" + ); + } + #[tokio::test] async fn test_picking_lifecycle_close_drop_truncate_fail_waiters_and_ignore_completion() { for lifecycle in ["close", "drop", "truncate"] {