diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index dd5b75d778b6..dbf27b1a5036 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -31,6 +31,7 @@ use crate::access_layer::AccessLayerRef; use crate::error::{RegionNotFoundSnafu, RegionReadonlySnafu, Result}; use crate::manifest::manager::RegionManifestManager; use crate::region::version::{VersionControlRef, VersionRef}; +use crate::request::OnFailure; use crate::sst::file_purger::FilePurgerRef; /// Metadata and runtime status of a region. @@ -128,13 +129,22 @@ impl RegionMap { regions.insert(region.region_id, region); } - /// Get region by region id. + /// Gets region by region id. pub(crate) fn get_region(&self, region_id: RegionId) -> Option { let regions = self.regions.read().unwrap(); regions.get(®ion_id).cloned() } - /// Get writable region by region id. + /// Gets region by region id or call the failure callback. + pub(crate) fn get_region_or_fail(&self, region_id: RegionId, cb: &mut F) -> Option { + let region_opt = self.get_region(region_id); + if region_opt.is_none() { + cb.on_failure(RegionNotFoundSnafu { region_id }.build()); + } + region_opt + } + + /// Gets writable region by region id. /// /// Returns error if the region does not exist or is readonly. pub(crate) fn get_writable_region(&self, region_id: RegionId) -> Result { @@ -145,6 +155,19 @@ impl RegionMap { Ok(region) } + /// Gets writable region by region id. + /// + /// Calls the callback if the region does not exist or is readonly. + pub(crate) fn get_writable_region_or_fail(&self, region_id: RegionId, cb: &mut F) -> Option { + match self.get_writable_region(region_id) { + Ok(region) => Some(region), + Err(e) => { + cb.on_failure(e); + None + }, + } + } + /// Remove region by id. pub(crate) fn remove_region(&self, region_id: RegionId) { let mut regions = self.regions.write().unwrap(); diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 222723e9d040..1f635844b1da 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -451,6 +451,12 @@ impl From>> for OptionOutputTx { } } +/// Callback on failure. +pub(crate) trait OnFailure { + /// Handles `err` on failure. + fn on_failure(&mut self, err: Error); +} + /// Sender and write request. #[derive(Debug)] pub(crate) struct SenderWriteRequest {