perf(mito): scan SSTs and memtables in parallel (#2852)

* feat: seq scan support parallelism

* feat: scan region by parallelism in config

* feat: enlarge channel size

* chore: parallel builder logs

* feat: use parallel reader accroding to source num

* chore: 128 channel size

* feat: add fetch cost metrics

* feat: add channel size to config

* feat: builder cost

* feat: logs

* feat: compiler error

* feat: fetch cost

* feat: convert cost

* chore: Revert "feat: logs"

This reverts commit 01e0df2.

* chore: fix compiler errors

* feat: reduce channel size to 32

* chore: use workspace tokio-stream

* test: test scan in parallel

* chore: comment typos

* refactor: build all sources first

* test: test 0 parallelism

* feat: use parallel scan by default

* docs: update config example

* feat: log parallelism

* refactor: keep config in engine inner

* refactor: rename parallelism method

* docs: update docs

* test: fix config api test

* docs: update parallel scan comment

* feat: 0 for default parallelism

Cargo.toml

@@ -127,6 +127,7 @@ sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "6
 strum = { version = "0.25", features = ["derive"] }
 tempfile = "3"
 tokio = { version = "1.28", features = ["full"] }
+tokio-stream = { version = "0.1" }
 tokio-util = { version = "0.7", features = ["io-util", "compat"] }
 toml = "0.7"
 tonic = { version = "0.10", features = ["tls"] }
@@ -168,7 +169,6 @@ frontend = { path = "src/frontend" }
 log-store = { path = "src/log-store" }
 meta-client = { path = "src/meta-client" }
 meta-srv = { path = "src/meta-srv" }
-mito = { path = "src/mito" }
 mito2 = { path = "src/mito2" }
 object-store = { path = "src/object-store" }
 operator = { path = "src/operator" }

config/datanode.example.toml

@@ -89,6 +89,13 @@ vector_cache_size = "512MB"
 page_cache_size = "512MB"
 # Buffer size for SST writing.
 sst_write_buffer_size = "8MB"
+# Parallelism to scan a region (default: 1/4 of cpu cores).
+# - 0: using the default value (1/4 of cpu cores).
+# - 1: scan in current thread.
+# - n: scan in parallelism n.
+scan_parallelism = 0
+# Capacity of the channel to send data from parallel scan tasks to the main task (default 32).
+parallel_scan_channel_size = 32
 
 # Log options, see `standalone.example.toml`
 # [logging]

config/standalone.example.toml

@@ -157,6 +157,13 @@ vector_cache_size = "512MB"
 page_cache_size = "512MB"
 # Buffer size for SST writing.
 sst_write_buffer_size = "8MB"
+# Parallelism to scan a region (default: 1/4 of cpu cores).
+# - 0: using the default value (1/4 of cpu cores).
+# - 1: scan in current thread.
+# - n: scan in parallelism n.
+scan_parallelism = 0
+# Capacity of the channel to send data from parallel scan tasks to the main task (default 32).
+parallel_scan_channel_size = 32
 
 # Log options
 # [logging]

src/client/Cargo.toml

@@ -35,7 +35,7 @@ prost.workspace = true
 rand.workspace = true
 session.workspace = true
 snafu.workspace = true
-tokio-stream = { version = "0.1", features = ["net"] }
+tokio-stream = { workspace = true, features = ["net"] }
 tokio.workspace = true
 tonic.workspace = true
 

src/datanode/Cargo.toml

@@ -62,7 +62,7 @@ sql.workspace = true
 store-api.workspace = true
 substrait.workspace = true
 table.workspace = true
-tokio-stream = { version = "0.1", features = ["net"] }
+tokio-stream = { workspace = true, features = ["net"] }
 tokio.workspace = true
 toml.workspace = true
 tonic.workspace = true

src/meta-client/Cargo.toml

@@ -20,7 +20,7 @@ serde.workspace = true
 serde_json.workspace = true
 snafu.workspace = true
 table.workspace = true
-tokio-stream = { version = "0.1", features = ["net"] }
+tokio-stream = { workspace = true, features = ["net"] }
 tokio.workspace = true
 tonic.workspace = true
 

src/meta-srv/Cargo.toml

@@ -49,7 +49,7 @@ snafu.workspace = true
 store-api.workspace = true
 strum.workspace = true
 table.workspace = true
-tokio-stream = { version = "0.1", features = ["net"] }
+tokio-stream = { workspace = true, features = ["net"] }
 tokio.workspace = true
 toml.workspace = true
 tonic.workspace = true

src/mito2/Cargo.toml

@@ -58,6 +58,7 @@ snafu.workspace = true
 store-api.workspace = true
 strum.workspace = true
 table.workspace = true
+tokio-stream.workspace = true
 tokio-util.workspace = true
 tokio.workspace = true
 uuid.workspace = true

src/mito2/src/config.rs

@@ -24,13 +24,16 @@ use serde::{Deserialize, Serialize};
 const DEFAULT_MAX_BG_JOB: usize = 4;
 
 const MULTIPART_UPLOAD_MINIMUM_SIZE: ReadableSize = ReadableSize::mb(5);
+/// Default channel size for parallel scan task.
+const DEFAULT_SCAN_CHANNEL_SIZE: usize = 32;
 
 /// Configuration for [MitoEngine](crate::engine::MitoEngine).
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
 #[serde(default)]
 pub struct MitoConfig {
     // Worker configs:
-    /// Number of region workers (default 1).
+    /// Number of region workers (default: 1/2 of cpu cores).
+    /// Sets to 0 to use the default value.
     pub num_workers: usize,
     /// Request channel size of each worker (default 128).
     pub worker_channel_size: usize,
@@ -68,12 +71,19 @@ pub struct MitoConfig {
     // Other configs:
     /// Buffer size for SST writing.
     pub sst_write_buffer_size: ReadableSize,
+    /// Parallelism to scan a region (default: 1/4 of cpu cores).
+    /// - 0: using the default value (1/4 of cpu cores).
+    /// - 1: scan in current thread.
+    /// - n: scan in parallelism n.
+    pub scan_parallelism: usize,
+    /// Capacity of the channel to send data from parallel scan tasks to the main task (default 32).
+    pub parallel_scan_channel_size: usize,
 }
 
 impl Default for MitoConfig {
     fn default() -> Self {
         MitoConfig {
-            num_workers: num_cpus::get() / 2,
+            num_workers: divide_num_cpus(2),
             worker_channel_size: 128,
             worker_request_batch_size: 64,
             manifest_checkpoint_distance: 10,
@@ -86,23 +96,18 @@ impl Default for MitoConfig {
             vector_cache_size: ReadableSize::mb(512),
             page_cache_size: ReadableSize::mb(512),
             sst_write_buffer_size: ReadableSize::mb(8),
+            scan_parallelism: divide_num_cpus(4),
+            parallel_scan_channel_size: DEFAULT_SCAN_CHANNEL_SIZE,
         }
     }
 }
 
 impl MitoConfig {
     /// Sanitize incorrect configurations.
     pub(crate) fn sanitize(&mut self) {
-        // Sanitize worker num.
-        let num_workers_before = self.num_workers;
+        // Use default value if `num_workers` is 0.
         if self.num_workers == 0 {
-            self.num_workers = (num_cpus::get() / 2).max(1);
-        }
-        if num_workers_before != self.num_workers {
-            warn!(
-                "Sanitize worker num {} to {}",
-                num_workers_before, self.num_workers
-            );
+            self.num_workers = divide_num_cpus(2);
         }
 
         // Sanitize channel size.
@@ -131,5 +136,27 @@ impl MitoConfig {
                 self.sst_write_buffer_size
             );
         }
+
+        // Use default value if `scan_parallelism` is 0.
+        if self.scan_parallelism == 0 {
+            self.scan_parallelism = divide_num_cpus(4);
+        }
+
+        if self.parallel_scan_channel_size < 1 {
+            self.parallel_scan_channel_size = DEFAULT_SCAN_CHANNEL_SIZE;
+            warn!(
+                "Sanitize scan channel size to {}",
+                self.parallel_scan_channel_size
+            );
+        }
     }
 }
+
+/// Divide cpu num by a non-zero `divisor` and returns at least 1.
+fn divide_num_cpus(divisor: usize) -> usize {
+    debug_assert!(divisor > 0);
+    let cores = num_cpus::get();
+    debug_assert!(cores > 0);
+
+    (cores + divisor - 1) / divisor
+}

src/mito2/src/engine.rs

@@ -33,13 +33,16 @@ pub mod listener;
 #[cfg(test)]
 mod open_test;
 #[cfg(test)]
+mod parallel_test;
+#[cfg(test)]
 mod projection_test;
 #[cfg(test)]
 mod prune_test;
 #[cfg(test)]
 mod set_readonly_test;
 #[cfg(test)]
 mod truncate_test;
+
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -56,7 +59,7 @@ use store_api::storage::{RegionId, ScanRequest};
 use crate::config::MitoConfig;
 use crate::error::{RecvSnafu, RegionNotFoundSnafu, Result};
 use crate::metrics::HANDLE_REQUEST_ELAPSED;
-use crate::read::scan_region::{ScanRegion, Scanner};
+use crate::read::scan_region::{ScanParallism, ScanRegion, Scanner};
 use crate::region::RegionUsage;
 use crate::request::WorkerRequest;
 use crate::worker::WorkerGroup;
@@ -114,6 +117,8 @@ impl MitoEngine {
 struct EngineInner {
     /// Region workers group.
     workers: WorkerGroup,
+    /// Config of the engine.
+    config: Arc<MitoConfig>,
 }
 
 impl EngineInner {
@@ -123,8 +128,10 @@ impl EngineInner {
         log_store: Arc<S>,
         object_store_manager: ObjectStoreManagerRef,
     ) -> EngineInner {
+        let config = Arc::new(config);
         EngineInner {
-            workers: WorkerGroup::start(config, log_store, object_store_manager),
+            workers: WorkerGroup::start(config.clone(), log_store, object_store_manager),
+            config,
         }
     }
 
@@ -171,12 +178,18 @@ impl EngineInner {
         let version = region.version();
         // Get cache.
         let cache_manager = self.workers.cache_manager();
+        let scan_parallelism = ScanParallism {
+            parallelism: self.config.scan_parallelism,
+            channel_size: self.config.parallel_scan_channel_size,
+        };
+
         let scan_region = ScanRegion::new(
             version,
             region.access_layer.clone(),
             request,
             Some(cache_manager),
-        );
+        )
+        .with_parallelism(scan_parallelism);
 
         scan_region.scanner()
     }
@@ -303,15 +316,17 @@ impl MitoEngine {
     ) -> MitoEngine {
         config.sanitize();
 
+        let config = Arc::new(config);
         MitoEngine {
             inner: Arc::new(EngineInner {
                 workers: WorkerGroup::start_for_test(
-                    config,
+                    config.clone(),
                     log_store,
                     object_store_manager,
                     write_buffer_manager,
                     listener,
                 ),
+                config,
             }),
         }
     }