feat: open region in background

src/mito2/src/engine.rs

@@ -104,6 +104,11 @@ impl MitoEngine {
         self.inner.workers.is_region_exists(region_id)
     }
 
+    /// Returns true if the specific region exists.
+    pub fn is_region_opening(&self, region_id: RegionId) -> bool {
+        self.inner.workers.is_region_opening(region_id)
+    }
+
     /// Returns the region disk/memory usage information.
     pub async fn get_region_usage(&self, region_id: RegionId) -> Result<RegionUsage> {
         let region = self

src/mito2/src/engine/open_test.rs

@@ -24,8 +24,10 @@ use store_api::region_request::{
     RegionCloseRequest, RegionOpenRequest, RegionPutRequest, RegionRequest,
 };
 use store_api::storage::{RegionId, ScanRequest};
+use tokio::sync::oneshot;
 
 use crate::config::MitoConfig;
+use crate::error;
 use crate::test_util::{
     build_rows, flush_region, put_rows, reopen_region, rows_schema, CreateRequestBuilder, TestEnv,
 };
@@ -319,3 +321,87 @@ async fn test_open_region_skip_wal_replay() {
 +-------+---------+---------------------+";
     assert_eq!(expected, batches.pretty_print().unwrap());
 }
+
+#[tokio::test]
+async fn test_open_region_wait_for_opening_region_ok() {
+    let mut env = TestEnv::with_prefix("wait-for-opening-region-ok");
+    let engine = env.create_engine(MitoConfig::default()).await;
+    let region_id = RegionId::new(1, 1);
+    let worker = engine.inner.workers.worker(region_id);
+    let (tx, rx) = oneshot::channel();
+    let opening_regions = worker.opening_regions().clone();
+    opening_regions.insert_sender(region_id, tx.into());
+    assert!(engine.is_region_opening(region_id));
+
+    let handle_open = tokio::spawn(async move {
+        engine
+            .handle_request(
+                region_id,
+                RegionRequest::Open(RegionOpenRequest {
+                    engine: String::new(),
+                    region_dir: "empty".to_string(),
+                    options: HashMap::default(),
+                    skip_wal_replay: false,
+                }),
+            )
+            .await
+    });
+
+    // Wait for conditions
+    while opening_regions.sender_len(region_id) != 2 {
+        tokio::time::sleep(Duration::from_millis(100)).await;
+    }
+
+    let senders = opening_regions.remove_sender(region_id);
+    for sender in senders {
+        sender.send(Ok(0));
+    }
+
+    assert_eq!(handle_open.await.unwrap().unwrap().affected_rows, 0);
+    assert_eq!(rx.await.unwrap().unwrap(), 0);
+}
+
+#[tokio::test]
+async fn test_open_region_wait_for_opening_region_err() {
+    let mut env = TestEnv::with_prefix("wait-for-opening-region-err");
+    let engine = env.create_engine(MitoConfig::default()).await;
+    let region_id = RegionId::new(1, 1);
+    let worker = engine.inner.workers.worker(region_id);
+    let (tx, rx) = oneshot::channel();
+    let opening_regions = worker.opening_regions().clone();
+    opening_regions.insert_sender(region_id, tx.into());
+    assert!(engine.is_region_opening(region_id));
+
+    let handle_open = tokio::spawn(async move {
+        engine
+            .handle_request(
+                region_id,
+                RegionRequest::Open(RegionOpenRequest {
+                    engine: String::new(),
+                    region_dir: "empty".to_string(),
+                    options: HashMap::default(),
+                    skip_wal_replay: false,
+                }),
+            )
+            .await
+    });
+
+    // Wait for conditions
+    while opening_regions.sender_len(region_id) != 2 {
+        tokio::time::sleep(Duration::from_millis(100)).await;
+    }
+
+    let senders = opening_regions.remove_sender(region_id);
+    for sender in senders {
+        sender.send(Err(error::RegionNotFoundSnafu { region_id }.build()));
+    }
+
+    assert_eq!(
+        handle_open.await.unwrap().unwrap_err().status_code(),
+        StatusCode::RegionNotFound
+    );
+    assert_eq!(
+        rx.await.unwrap().unwrap_err().status_code(),
+        StatusCode::RegionNotFound
+    );
+}

src/mito2/src/error.rs

@@ -722,6 +722,13 @@ pub enum Error {
         #[snafu(implicit)]
         location: Location,
     },
+
+    #[snafu(display("Failed to open region"))]
+    OpenRegion {
+        #[snafu(implicit)]
+        location: Location,
+        source: Arc<Error>,
+    },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -783,6 +790,7 @@ impl ErrorExt for Error {
             | Recv { .. }
             | EncodeWal { .. }
             | DecodeWal { .. } => StatusCode::Internal,
+            OpenRegion { source, .. } => source.status_code(),
             WriteBuffer { source, .. } => source.status_code(),
             WriteGroup { source, .. } => source.status_code(),
             FieldTypeMismatch { source, .. } => source.status_code(),

src/mito2/src/region.rs

@@ -18,6 +18,7 @@ pub(crate) mod opener;
 pub mod options;
 pub(crate) mod version;
 
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicI64, Ordering};
 use std::sync::{Arc, RwLock};
@@ -35,7 +36,7 @@ use crate::manifest::action::{RegionMetaAction, RegionMetaActionList};
 use crate::manifest::manager::RegionManifestManager;
 use crate::memtable::MemtableBuilderRef;
 use crate::region::version::{VersionControlRef, VersionRef};
-use crate::request::OnFailure;
+use crate::request::{OnFailure, OptionOutputTx};
 use crate::sst::file_purger::FilePurgerRef;
 use crate::time_provider::TimeProviderRef;
 
@@ -471,6 +472,60 @@ impl RegionMap {
 
 pub(crate) type RegionMapRef = Arc<RegionMap>;
 
+/// Opening regions
+#[derive(Debug, Default)]
+pub(crate) struct OpeningRegions {
+    regions: RwLock<HashMap<RegionId, Vec<OptionOutputTx>>>,
+}
+
+impl OpeningRegions {
+    /// Registers `sender` for an opening region; Otherwise, it returns `None`.
+    pub(crate) fn wait_for_opening_region(
+        &self,
+        region_id: RegionId,
+        sender: OptionOutputTx,
+    ) -> Option<OptionOutputTx> {
+        let mut regions = self.regions.write().unwrap();
+        match regions.entry(region_id) {
+            Entry::Occupied(mut senders) => {
+                senders.get_mut().push(sender);
+                None
+            }
+            Entry::Vacant(_) => Some(sender),
+        }
+    }
+
+    /// Returns true if the region exists.
+    pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool {
+        let regions = self.regions.read().unwrap();
+        regions.contains_key(&region_id)
+    }
+
+    /// Inserts a new region into the map.
+    pub(crate) fn insert_sender(&self, region: RegionId, sender: OptionOutputTx) {
+        let mut regions = self.regions.write().unwrap();
+        regions.insert(region, vec![sender]);
+    }
+
+    /// Remove region by id.
+    pub(crate) fn remove_sender(&self, region_id: RegionId) -> Vec<OptionOutputTx> {
+        let mut regions = self.regions.write().unwrap();
+        regions.remove(&region_id).unwrap_or_default()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn sender_len(&self, region_id: RegionId) -> usize {
+        let regions = self.regions.read().unwrap();
+        if let Some(senders) = regions.get(&region_id) {
+            senders.len()
+        } else {
+            0
+        }
+    }
+}
+
+pub(crate) type OpeningRegionsRef = Arc<OpeningRegions>;
+
 #[cfg(test)]
 mod tests {
     use crossbeam_utils::atomic::AtomicCell;

src/mito2/src/request.rs

@@ -384,11 +384,11 @@ pub(crate) fn validate_proto_value(
 
 /// Oneshot output result sender.
 #[derive(Debug)]
-pub(crate) struct OutputTx(Sender<Result<AffectedRows>>);
+pub(crate) struct OutputTx(Sender<Result<AffectedRows, Error>>);
 
 impl OutputTx {
     /// Creates a new output sender.
-    pub(crate) fn new(sender: Sender<Result<AffectedRows>>) -> OutputTx {
+    pub(crate) fn new(sender: Sender<Result<AffectedRows, Error>>) -> OutputTx {
         OutputTx(sender)
     }
 

src/mito2/src/wal.rs

@@ -49,7 +49,7 @@ pub type WalEntryStream<'a> = BoxStream<'a, Result<(EntryId, WalEntry)>>;
 /// Write ahead log.
 ///
 /// All regions in the engine shares the same WAL instance.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct Wal<S> {
     /// The underlying log store.
     store: Arc<S>,
@@ -62,6 +62,14 @@ impl<S> Wal<S> {
     }
 }
 
+impl<S> Clone for Wal<S> {
+    fn clone(&self) -> Self {
+        Self {
+            store: self.store.clone(),
+        }
+    }
+}
+
 impl<S: LogStore> Wal<S> {
     /// Returns a writer to write to the WAL.
     pub fn writer(&self) -> WalWriter<S> {

src/mito2/src/worker.rs

@@ -51,7 +51,7 @@ use crate::error::{JoinSnafu, Result, WorkerStoppedSnafu};
 use crate::flush::{FlushScheduler, WriteBufferManagerImpl, WriteBufferManagerRef};
 use crate::memtable::MemtableBuilderProvider;
 use crate::metrics::WRITE_STALL_TOTAL;
-use crate::region::{MitoRegionRef, RegionMap, RegionMapRef};
+use crate::region::{MitoRegionRef, OpeningRegions, OpeningRegionsRef, RegionMap, RegionMapRef};
 use crate::request::{
     BackgroundNotify, DdlRequest, SenderDdlRequest, SenderWriteRequest, WorkerRequest,
 };
@@ -212,6 +212,11 @@ impl WorkerGroup {
         self.worker(region_id).is_region_exists(region_id)
     }
 
+    /// Returns true if the specific region is opening.
+    pub(crate) fn is_region_opening(&self, region_id: RegionId) -> bool {
+        self.worker(region_id).is_region_opening(region_id)
+    }
+
     /// Returns region of specific `region_id`.
     ///
     /// This method should not be public.
@@ -225,7 +230,7 @@ impl WorkerGroup {
     }
 
     /// Get worker for specific `region_id`.
-    fn worker(&self, region_id: RegionId) -> &RegionWorker {
+    pub(crate) fn worker(&self, region_id: RegionId) -> &RegionWorker {
         let index = region_id_to_index(region_id, self.workers.len());
 
         &self.workers[index]
@@ -364,6 +369,7 @@ impl<S: LogStore> WorkerStarter<S> {
     /// Starts a region worker and its background thread.
     fn start(self) -> RegionWorker {
         let regions = Arc::new(RegionMap::default());
+        let opening_regions = Arc::new(OpeningRegions::default());
         let (sender, receiver) = mpsc::channel(self.config.worker_channel_size);
 
         let running = Arc::new(AtomicBool::new(true));
@@ -373,6 +379,7 @@ impl<S: LogStore> WorkerStarter<S> {
             config: self.config.clone(),
             regions: regions.clone(),
             dropping_regions: Arc::new(RegionMap::default()),
+            opening_regions: opening_regions.clone(),
             sender: sender.clone(),
             receiver,
             wal: Wal::new(self.log_store),
@@ -409,6 +416,7 @@ impl<S: LogStore> WorkerStarter<S> {
         RegionWorker {
             id: self.id,
             regions,
+            opening_regions,
             sender,
             handle: Mutex::new(Some(handle)),
             running,
@@ -422,6 +430,8 @@ pub(crate) struct RegionWorker {
     id: WorkerId,
     /// Regions bound to the worker.
     regions: RegionMapRef,
+    /// The opening regions.
+    opening_regions: OpeningRegionsRef,
     /// Request sender.
     sender: Sender<WorkerRequest>,
     /// Handle to the worker thread.
@@ -481,10 +491,21 @@ impl RegionWorker {
         self.regions.is_region_exists(region_id)
     }
 
+    /// Returns true if the region is opening.
+    fn is_region_opening(&self, region_id: RegionId) -> bool {
+        self.opening_regions.is_region_exists(region_id)
+    }
+
     /// Returns region of specific `region_id`.
     fn get_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
         self.regions.get_region(region_id)
     }
+
+    #[cfg(test)]
+    /// Returns the [OpeningRegionsRef].
+    pub(crate) fn opening_regions(&self) -> &OpeningRegionsRef {
+        &self.opening_regions
+    }
 }
 
 impl Drop for RegionWorker {
@@ -531,6 +552,8 @@ struct RegionWorkerLoop<S> {
     regions: RegionMapRef,
     /// Regions that are not yet fully dropped.
     dropping_regions: RegionMapRef,
+    /// Regions that are opening.
+    opening_regions: OpeningRegionsRef,
     /// Request sender.
     sender: Sender<WorkerRequest>,
     /// Request receiver.
@@ -698,7 +721,11 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             let res = match ddl.request {
                 DdlRequest::Create(req) => self.handle_create_request(ddl.region_id, req).await,
                 DdlRequest::Drop(_) => self.handle_drop_request(ddl.region_id).await,
-                DdlRequest::Open(req) => self.handle_open_request(ddl.region_id, req).await,
+                DdlRequest::Open(req) => {
+                    self.handle_open_request(ddl.region_id, req, ddl.sender)
+                        .await;
+                    continue;
+                }
                 DdlRequest::Close(_) => self.handle_close_request(ddl.region_id).await,
                 DdlRequest::Alter(req) => {
                     self.handle_alter_request(ddl.region_id, req, ddl.sender)