feat: open region in background

src/mito2/src/region.rs

@@ -18,7 +18,7 @@ pub(crate) mod opener;
 pub mod options;
 pub(crate) mod version;
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::atomic::{AtomicI64, Ordering};
 use std::sync::{Arc, RwLock};
 
@@ -471,6 +471,34 @@ impl RegionMap {
 
 pub(crate) type RegionMapRef = Arc<RegionMap>;
 
+/// Opening regions
+#[derive(Debug, Default)]
+pub(crate) struct OpeningRegions {
+    regions: RwLock<HashSet<RegionId>>,
+}
+
+impl OpeningRegions {
+    /// 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(&region_id)
+    }
+
+    /// Inserts a new region into the map.
+    pub(crate) fn insert_region(&self, region: RegionId) {
+        let mut regions = self.regions.write().unwrap();
+        regions.insert(region);
+    }
+
+    /// Remove region by id.
+    pub(crate) fn remove_region(&self, region_id: RegionId) {
+        let mut regions = self.regions.write().unwrap();
+        regions.remove(&region_id);
+    }
+}
+
+pub(crate) type OpeningRegionsRef = Arc<OpeningRegions>;
+
 #[cfg(test)]
 mod tests {
     use crossbeam_utils::atomic::AtomicCell;

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,
 };
@@ -373,6 +373,7 @@ impl<S: LogStore> WorkerStarter<S> {
             config: self.config.clone(),
             regions: regions.clone(),
             dropping_regions: Arc::new(RegionMap::default()),
+            opening_regions: Arc::new(OpeningRegions::default()),
             sender: sender.clone(),
             receiver,
             wal: Wal::new(self.log_store),
@@ -531,6 +532,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 +701,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)

src/mito2/src/worker/handle_open.rs

@@ -20,24 +20,22 @@ use common_telemetry::info;
 use object_store::util::join_path;
 use snafu::{OptionExt, ResultExt};
 use store_api::logstore::LogStore;
-use store_api::region_request::{AffectedRows, RegionOpenRequest};
+use store_api::region_request::RegionOpenRequest;
 use store_api::storage::RegionId;
 
 use crate::error::{ObjectStoreNotFoundSnafu, OpenDalSnafu, RegionNotFoundSnafu, Result};
 use crate::metrics::REGION_COUNT;
 use crate::region::opener::RegionOpener;
+use crate::request::OptionOutputTx;
 use crate::worker::handle_drop::remove_region_dir_once;
 use crate::worker::{RegionWorkerLoop, DROPPING_MARKER_FILE};
 
 impl<S: LogStore> RegionWorkerLoop<S> {
-    pub(crate) async fn handle_open_request(
-        &mut self,
+    async fn check_and_cleanup_region(
+        &self,
         region_id: RegionId,
-        request: RegionOpenRequest,
-    ) -> Result<AffectedRows> {
-        if self.regions.is_region_exists(region_id) {
-            return Ok(0);
-        }
+        request: &RegionOpenRequest,
+    ) -> Result<()> {
         let object_store = if let Some(storage_name) = request.options.get("storage") {
             self.object_store_manager
                 .find(storage_name)
@@ -59,10 +57,31 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             return RegionNotFoundSnafu { region_id }.fail();
         }
 
+        Ok(())
+    }
+
+    pub(crate) async fn handle_open_request(
+        &mut self,
+        region_id: RegionId,
+        request: RegionOpenRequest,
+        sender: OptionOutputTx,
+    ) {
+        if self.opening_regions.is_region_exists(region_id) {
+            sender.send(Ok(0));
+            return;
+        };
+        if self.regions.is_region_exists(region_id) {
+            sender.send(Ok(0));
+            return;
+        }
+        if let Err(err) = self.check_and_cleanup_region(region_id, &request).await {
+            sender.send(Err(err));
+            return;
+        }
         info!("Try to open region {}", region_id);
 
         // Open region from specific region dir.
-        let region = RegionOpener::new(
+        let opener = match RegionOpener::new(
             region_id,
             &request.region_dir,
             self.memtable_builder_provider.clone(),
@@ -71,18 +90,37 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             self.intermediate_manager.clone(),
         )
         .skip_wal_replay(request.skip_wal_replay)
-        .parse_options(request.options)?
         .cache(Some(self.cache_manager.clone()))
-        .open(&self.config, &self.wal)
-        .await?;
-
-        info!("Region {} is opened", region_id);
+        .parse_options(request.options)
+        {
+            Ok(opener) => opener,
+            Err(err) => {
+                sender.send(Err(err));
+                return;
+            }
+        };
 
-        REGION_COUNT.inc();
+        let regions = self.regions.clone();
+        let wal = self.wal.clone();
+        let config = self.config.clone();
+        let opening_regions = self.opening_regions.clone();
+        opening_regions.insert_region(region_id);
+        common_runtime::spawn_bg(async move {
+            match opener.open(&config, &wal).await {
+                Ok(region) => {
+                    info!("Region {} is opened", region_id);
+                    REGION_COUNT.inc();
 
-        // Insert the MitoRegion into the RegionMap.
-        self.regions.insert_region(Arc::new(region));
+                    // Insert the MitoRegion into the RegionMap.
+                    regions.insert_region(Arc::new(region));
 
-        Ok(0)
+                    sender.send(Ok(0));
+                }
+                Err(err) => {
+                    sender.send(Err(err));
+                }
+            }
+            opening_regions.remove_region(region_id);
+        });
     }
 }