From 0a2b10e476635ea9ffd2ac7fd21e45d48a1ecb24 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 10 Aug 2026 13:30:13 +0800 Subject: [PATCH] fix(mito2): avoid region worker panic when building a WAL entry fails (#8810) When `add_wal_entry` fails for a region, the worker only sets the error on the write context, which stays in `region_ctxs`. The region's entry is not in the batch, so a successful `write_to_wal` returns no last entry id for it and the success branch panics on `response.last_entry_ids.get(region_id).unwrap()`, killing the region worker. When the failed region is the only one in the batch, the batch is empty and `append_batch` always returns an empty response, so the panic is guaranteed. No in-tree log store can fail to build an entry at runtime today (the provider/log store combination is validated when the region opens), so this is a latent panic rather than a reachable crash. Skip contexts already marked as failed when updating next entry ids; their waiters are already notified with the error. Extract the WAL phase of `handle_write_requests` into `write_wal` and cover the failure paths with unit tests. Signed-off-by: jeremyhi --- src/mito2/src/region_write_ctx.rs | 11 + src/mito2/src/worker/handle_write.rs | 287 ++++++++++++++++++++++++--- 2 files changed, 269 insertions(+), 29 deletions(-) diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index c6abdb0872..b71e3e5ab1 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -218,11 +218,22 @@ impl RegionWriteCtx { self.failed = true; } + /// Returns whether the write operation is already marked as failed. + pub(crate) fn is_failed(&self) -> bool { + self.failed + } + /// Updates next entry id. pub(crate) fn set_next_entry_id(&mut self, next_entry_id: EntryId) { self.next_entry_id = next_entry_id } + /// Returns the next entry id to write. + #[cfg(test)] + pub(crate) fn next_entry_id(&self) -> EntryId { + self.next_entry_id + } + /// Consumes mutations and writes them into mutable memtable. pub(crate) async fn write_memtable(&mut self) { debug_assert_eq!(self.notifiers.len(), self.wal_entry.mutations.len()); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index c52572df11..3be4055314 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -35,6 +35,7 @@ use crate::metrics::{ use crate::region::{RegionLeaderState, RegionRoleState}; use crate::region_write_ctx::RegionWriteCtx; use crate::request::{SenderBulkRequest, SenderWriteRequest, WriteRequest}; +use crate::wal::Wal; use crate::worker::RegionWorkerLoop; impl RegionWorkerLoop { @@ -110,35 +111,9 @@ impl RegionWorkerLoop { let _timer = WRITE_STAGE_ELAPSED .with_label_values(&["write_wal"]) .start_timer(); - let mut wal_writer = self.wal.writer(); - for region_ctx in region_ctxs.values_mut() { - if region_ctx.skip_wal() { - continue; - } - if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { - region_ctx.set_error(e); - } - } - match wal_writer.write_to_wal().await.map_err(Arc::new) { - Ok(response) => { - for (region_id, region_ctx) in region_ctxs.iter_mut() { - if region_ctx.skip_wal() { - continue; - } - - // Safety: the log store implementation ensures that either the `write_to_wal` fails and no - // response is returned or the last entry ids for each region do exist. - let last_entry_id = response.last_entry_ids.get(region_id).unwrap(); - region_ctx.set_next_entry_id(last_entry_id + 1); - } - } - Err(e) => { - // Failed to write wal. - for mut region_ctx in region_ctxs.into_values() { - region_ctx.set_error(e.clone()); - } - return; - } + if !write_wal(&self.wal, &mut region_ctxs).await { + // Failed to write to the WAL, all waiters are notified with the error. + return; } } @@ -625,6 +600,55 @@ impl RegionWorkerLoop { } } +/// Writes WAL entries of all region contexts to the WAL in one batch and updates +/// the next entry id of each region on success. +/// +/// Returns `false` if the batch fails to be written to the WAL. In this case all +/// contexts are consumed and their waiters are notified with the error, so the +/// caller should skip the memtable phase. +async fn write_wal( + wal: &Wal, + region_ctxs: &mut HashMap, +) -> bool { + let mut wal_writer = wal.writer(); + for region_ctx in region_ctxs.values_mut() { + if region_ctx.skip_wal() { + continue; + } + if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { + region_ctx.set_error(e); + } + } + match wal_writer.write_to_wal().await.map_err(Arc::new) { + Ok(response) => { + for (region_id, region_ctx) in region_ctxs.iter_mut() { + if region_ctx.skip_wal() { + continue; + } + // The entry of a failed region (e.g. failed to build its WAL entry) is + // not in the batch so the response has no last entry id for it. Its + // waiters are already notified with the error. + if region_ctx.is_failed() { + continue; + } + + // Safety: the log store implementation ensures that either the `write_to_wal` fails and no + // response is returned or the last entry ids for each region in the batch do exist. + let last_entry_id = response.last_entry_ids.get(region_id).unwrap(); + region_ctx.set_next_entry_id(last_entry_id + 1); + } + true + } + Err(e) => { + // Failed to write wal. + for (_, mut region_ctx) in region_ctxs.drain() { + region_ctx.set_error(e.clone()); + } + false + } + } +} + /// Send rejected error to all `write_requests`. fn reject_write_requests( write_requests: &mut Vec, @@ -707,3 +731,208 @@ fn check_partition_expr_version( } Ok(()) } + +#[cfg(test)] +mod tests { + use api::v1::{Row, Rows}; + use futures::stream; + use log_store::error::{ + Error as LogStoreError, IllegalStateSnafu, InvalidProviderSnafu, Result as LogStoreResult, + }; + use store_api::logstore::entry::{Entry, NaiveEntry}; + use store_api::logstore::provider::Provider; + use store_api::logstore::{AppendBatchResponse, EntryId, SendableEntryStream, WalIndex}; + use store_api::region_request::AffectedRows; + use tokio::sync::oneshot; + + use super::*; + use crate::request::OptionOutputTx; + use crate::test_util::version_util::VersionControlBuilder; + + /// A log store that fails to build entries for `failing_region` and fails the + /// whole batch when `fail_append` is true. + #[derive(Debug, Default)] + struct MockLogStore { + failing_region: Option, + fail_append: bool, + } + + #[async_trait::async_trait] + impl LogStore for MockLogStore { + type Error = LogStoreError; + + async fn stop(&self) -> LogStoreResult<()> { + Ok(()) + } + + async fn append_batch(&self, entries: Vec) -> LogStoreResult { + if self.fail_append { + return IllegalStateSnafu {}.fail(); + } + let mut last_entry_ids = HashMap::new(); + for entry in &entries { + let last_entry_id = last_entry_ids.entry(entry.region_id()).or_insert(0); + *last_entry_id = entry.entry_id().max(*last_entry_id); + } + Ok(AppendBatchResponse { last_entry_ids }) + } + + async fn read( + &self, + _provider: &Provider, + _entry_id: EntryId, + _index: Option, + ) -> LogStoreResult> { + Ok(Box::pin(stream::empty())) + } + + async fn create_namespace(&self, _ns: &Provider) -> LogStoreResult<()> { + Ok(()) + } + + async fn delete_namespace(&self, _ns: &Provider) -> LogStoreResult<()> { + Ok(()) + } + + async fn list_namespaces(&self) -> LogStoreResult> { + Ok(vec![]) + } + + async fn obsolete( + &self, + _provider: &Provider, + _region_id: RegionId, + _entry_id: EntryId, + ) -> LogStoreResult<()> { + Ok(()) + } + + async fn obsolete_all( + &self, + _provider: &Provider, + _region_id: RegionId, + ) -> LogStoreResult<()> { + Ok(()) + } + + fn entry( + &self, + data: Vec, + entry_id: EntryId, + region_id: RegionId, + provider: &Provider, + ) -> LogStoreResult { + if self.failing_region == Some(region_id) { + return InvalidProviderSnafu { + expected: "raft_engine", + actual: "mock", + } + .fail(); + } + Ok(Entry::Naive(NaiveEntry { + provider: provider.clone(), + region_id, + entry_id, + data, + })) + } + + fn latest_entry_id(&self, _provider: &Provider) -> LogStoreResult { + Ok(0) + } + } + + /// Creates a write context for `region_id` with one pending mutation of one row. + fn new_region_ctx( + region_id: RegionId, + ) -> (RegionWriteCtx, oneshot::Receiver>) { + let version_control = Arc::new(VersionControlBuilder::new().build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + let (tx, rx) = oneshot::channel(); + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![], + rows: vec![Row { values: vec![] }], + }), + None, + OptionOutputTx::from(tx), + None, + ); + (ctx, rx) + } + + #[tokio::test] + async fn test_write_wal_skips_region_failed_to_build_entry() { + let failing_region = RegionId::new(1, 1); + let ok_region = RegionId::new(1, 2); + let wal = Wal::new(Arc::new(MockLogStore { + failing_region: Some(failing_region), + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, failing_rx) = new_region_ctx(failing_region); + region_ctxs.insert(failing_region, ctx); + let (ctx, ok_rx) = new_region_ctx(ok_region); + region_ctxs.insert(ok_region, ctx); + let entry_id = region_ctxs[&ok_region].next_entry_id(); + + // The failed region must not fail the batch or panic the worker. + assert!(write_wal(&wal, &mut region_ctxs).await); + + assert!(region_ctxs[&failing_region].is_failed()); + assert!(!region_ctxs[&ok_region].is_failed()); + assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); + + // Waiters of the failed region get the error while others get the result. + drop(region_ctxs); + assert!(failing_rx.await.unwrap().is_err()); + assert_eq!(1, ok_rx.await.unwrap().unwrap()); + } + + #[tokio::test] + async fn test_write_wal_all_regions_failed_to_build_entries() { + let failing_region = RegionId::new(1, 1); + let wal = Wal::new(Arc::new(MockLogStore { + failing_region: Some(failing_region), + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, rx) = new_region_ctx(failing_region); + region_ctxs.insert(failing_region, ctx); + + // Writing an empty batch to the WAL succeeds, the failed region must not panic + // the worker. + assert!(write_wal(&wal, &mut region_ctxs).await); + + assert!(region_ctxs[&failing_region].is_failed()); + drop(region_ctxs); + assert!(rx.await.unwrap().is_err()); + } + + #[tokio::test] + async fn test_write_wal_append_batch_failure() { + let region_id = RegionId::new(1, 1); + let wal = Wal::new(Arc::new(MockLogStore { + fail_append: true, + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, rx) = new_region_ctx(region_id); + region_ctxs.insert(region_id, ctx); + + assert!(!write_wal(&wal, &mut region_ctxs).await); + + // All contexts are consumed and waiters are notified with the error. + assert!(region_ctxs.is_empty()); + assert!(rx.await.unwrap().is_err()); + } +}