diff --git a/AGENTS.md b/AGENTS.md
index ce25569aff1..fd2ec44aba1 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -77,8 +77,10 @@ change-coupling points, and gotchas:
2. `make clippy`
3. `make test` (or `cargo nextest run`)
4. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies).
-5. If you changed sample config under `config/`, run `make config-docs` (needs
- Docker) and commit the regenerated `config/config.md`.
+5. If you added or changed a public configuration option, update the applicable
+ example TOMLs, configuration-loading and serialized-config snapshot tests,
+ and related user-facing documentation. Run `make config-docs` (needs Docker)
+ and commit the regenerated `config/config.md`.
6. If you changed a persisted or wire format, add a compatibility test case (see
`.agents/architecture-invariants.md`).
7. Use a conventional-commit title, sign off commits (`git commit -s`), and sign
diff --git a/config/config.md b/config/config.md
index b0da5722f98..be287826fbc 100644
--- a/config/config.md
+++ b/config/config.md
@@ -157,6 +157,7 @@
| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. |
| `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. |
| `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size`. |
+| `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. |
| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. |
| `region_engine.mito.vector_cache_size` | String | Auto | Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. |
| `region_engine.mito.page_cache_size` | String | Auto | Cache size for pages of SST row groups. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory. |
@@ -560,6 +561,7 @@
| `region_engine.mito.auto_flush_interval` | String | `1h` | Interval to auto flush a region if it has not flushed yet. |
| `region_engine.mito.global_write_buffer_size` | String | Auto | Global write buffer size for all regions. If not set, it's default to 1/8 of OS memory with a max limitation of 1GB. |
| `region_engine.mito.global_write_buffer_reject_size` | String | Auto | Global write buffer size threshold to reject write requests. If not set, it's default to 2 times of `global_write_buffer_size` |
+| `region_engine.mito.default_region_write_buffer_size` | String | `0` | Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`. |
| `region_engine.mito.sst_meta_cache_size` | String | Auto | Cache size for SST metadata. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory with a max limitation of 512MB. |
| `region_engine.mito.vector_cache_size` | String | Auto | Cache size for vectors and arrow arrays. Setting it to 0 to disable the cache.
If not set, it's default to 1/16 of OS memory with a max limitation of 512MB. |
| `region_engine.mito.page_cache_size` | String | Auto | Cache size for pages of SST row groups. Setting it to 0 to disable the cache.
If not set, it's default to 1/8 of OS memory. |
diff --git a/config/datanode.example.toml b/config/datanode.example.toml
index f823d3114be..cfa9f695326 100644
--- a/config/datanode.example.toml
+++ b/config/datanode.example.toml
@@ -468,6 +468,9 @@ auto_flush_interval = "1h"
## @toml2docs:none-default="Auto"
#+ global_write_buffer_reject_size = "2GB"
+## Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`.
+#+ default_region_write_buffer_size = "0"
+
## Cache size for SST metadata. Setting it to 0 to disable the cache.
## If not set, it's default to 1/8 of OS memory with a max limitation of 512MB.
## @toml2docs:none-default="Auto"
diff --git a/config/standalone.example.toml b/config/standalone.example.toml
index 822f9ba568d..439c947c15c 100644
--- a/config/standalone.example.toml
+++ b/config/standalone.example.toml
@@ -597,6 +597,9 @@ auto_flush_interval = "1h"
## @toml2docs:none-default="Auto"
#+ global_write_buffer_reject_size = "2GB"
+## Default write buffer size for each region. Regions stall at this size and reject writes at twice this size. Setting it to 0 disables both limits unless the table specifies `write_buffer_size`.
+#+ default_region_write_buffer_size = "0"
+
## Cache size for SST metadata. Setting it to 0 to disable the cache.
## If not set, it's default to 1/8 of OS memory with a max limitation of 512MB.
## @toml2docs:none-default="Auto"
diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs
index 3570e88cc1e..40188dd0489 100644
--- a/src/cmd/tests/load_config_test.rs
+++ b/src/cmd/tests/load_config_test.rs
@@ -16,6 +16,7 @@ use std::time::Duration;
use cmd::options::GreptimeOptions;
use common_base::memory_limit::MemoryLimit;
+use common_base::readable_size::ReadableSize;
use common_config::{Configurable, DEFAULT_DATA_HOME, ENV_VAR_SEP};
use common_options::datanode::{ClientOptions, DatanodeClientOptions};
use common_telemetry::logging::{DEFAULT_LOGGING_DIR, DEFAULT_OTLP_HTTP_ENDPOINT, LoggingOptions};
@@ -90,6 +91,7 @@ fn test_load_datanode_example_config() {
region_engine: vec![
RegionEngineConfig::Mito(MitoConfig {
auto_flush_interval: Duration::from_secs(3600),
+ default_region_write_buffer_size: ReadableSize::mb(0),
write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)),
scan_memory_limit: MemoryLimit::Unlimited,
..Default::default()
@@ -299,6 +301,7 @@ fn test_load_standalone_example_config() {
region_engine: vec![
RegionEngineConfig::Mito(MitoConfig {
auto_flush_interval: Duration::from_secs(3600),
+ default_region_write_buffer_size: ReadableSize::mb(0),
write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)),
scan_memory_limit: MemoryLimit::Unlimited,
..Default::default()
diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs
index 49e43f41fd4..f802601df30 100644
--- a/src/mito2/src/compaction/window.rs
+++ b/src/mito2/src/compaction/window.rs
@@ -254,6 +254,7 @@ mod tests {
merge_mode: None,
sst_format: None,
primary_key_encoding: None,
+ write_buffer_size: None,
},
compaction_time_window: None,
}
diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs
index 407fedc955d..f7489709ee3 100644
--- a/src/mito2/src/config.rs
+++ b/src/mito2/src/config.rs
@@ -109,6 +109,10 @@ pub struct MitoConfig {
pub global_write_buffer_size: ReadableSize,
/// Global write buffer size threshold to reject write requests.
pub global_write_buffer_reject_size: ReadableSize,
+ /// Default write buffer size for each region. Regions stall at this size and
+ /// reject writes at twice this size. Setting it to 0 disables both limits
+ /// unless the table specifies `write_buffer_size`.
+ pub default_region_write_buffer_size: ReadableSize,
// Cache configs:
/// Cache size for SST metadata. Setting it to 0 to disable the cache.
@@ -206,6 +210,7 @@ impl Default for MitoConfig {
auto_flush_interval: Duration::from_secs(30 * 60),
global_write_buffer_size: ReadableSize::gb(1),
global_write_buffer_reject_size: ReadableSize::gb(2),
+ default_region_write_buffer_size: ReadableSize::mb(0),
sst_meta_cache_size: ReadableSize::mb(128),
vector_cache_size: ReadableSize::mb(512),
page_cache_size: ReadableSize::mb(512),
diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs
index 6d68a33abda..ccac8982970 100644
--- a/src/mito2/src/engine/flush_test.rs
+++ b/src/mito2/src/engine/flush_test.rs
@@ -14,6 +14,7 @@
//! Flush tests for mito engine.
+use std::assert_matches;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};
@@ -22,6 +23,7 @@ use std::time::Duration;
use api::v1::Rows;
use async_trait::async_trait;
use common_base::Plugins;
+use common_base::readable_size::ReadableSize;
use common_recordbatch::RecordBatches;
use common_time::util::current_time_millis;
use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions};
@@ -29,17 +31,20 @@ use rstest::rstest;
use rstest_reuse::{self, apply};
use store_api::ManifestVersion;
use store_api::metadata::RegionMetadataRef;
-use store_api::region_engine::RegionEngine;
+use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY;
+use store_api::region_engine::{RegionEngine, RegionRole};
use store_api::region_request::{
- PathType, RegionCloseRequest, RegionFlushRequest, RegionOpenRequest, RegionRequest,
- ReplayCheckpoint,
+ AlterKind, PathType, RegionAlterRequest, RegionCloseRequest, RegionFlushRequest,
+ RegionOpenRequest, RegionPutRequest, RegionRequest, ReplayCheckpoint, SetRegionOption,
};
use store_api::storage::{RegionId, ScanRequest};
use tokio::sync::Notify;
use crate::config::MitoConfig;
-use crate::engine::listener::{FlushListener, StallListener};
+use crate::engine::MitoEngine;
+use crate::engine::listener::{EventListener, FlushListener, StallListener};
use crate::engine::region_hook::{RegionHook, RegionHookRef, SstFileInfo};
+use crate::error::Error;
use crate::manifest::action::RegionMetaActionList;
use crate::test_util::{
CreateRequestBuilder, LogStoreFactory, MockWriteBufferManager, TestEnv, build_rows,
@@ -50,6 +55,73 @@ use crate::test_util::{
use crate::time_provider::TimeProvider;
use crate::worker::MAX_INITIAL_CHECK_DELAY_SECS;
+async fn set_write_buffer_size_to_current_usage(engine: &MitoEngine, region_id: RegionId) {
+ let version = engine.get_region(region_id).unwrap().version();
+ let memory_usage = version.memtables.mutable_usage() + version.memtables.immutables_usage();
+ assert!(memory_usage > 0);
+ engine
+ .handle_request(
+ region_id,
+ RegionRequest::Alter(RegionAlterRequest {
+ kind: AlterKind::SetRegionOptions {
+ options: vec![SetRegionOption::WriteBufferSize(Some(ReadableSize(
+ memory_usage as u64,
+ )))],
+ },
+ }),
+ )
+ .await
+ .unwrap();
+}
+
+struct RegionFlushGate {
+ blocked_region_id: RegionId,
+ observed_region_id: RegionId,
+ blocked_flush_started: Notify,
+ resume_blocked_flush: Notify,
+ observed_flush_finished: Notify,
+}
+
+impl RegionFlushGate {
+ fn new(blocked_region_id: RegionId, observed_region_id: RegionId) -> Self {
+ Self {
+ blocked_region_id,
+ observed_region_id,
+ blocked_flush_started: Notify::new(),
+ resume_blocked_flush: Notify::new(),
+ observed_flush_finished: Notify::new(),
+ }
+ }
+
+ async fn wait_blocked_flush_started(&self) {
+ self.blocked_flush_started.notified().await;
+ }
+
+ fn resume_blocked_flush(&self) {
+ self.resume_blocked_flush.notify_one();
+ }
+
+ async fn wait_observed_flush_finished(&self) {
+ self.observed_flush_finished.notified().await;
+ }
+}
+
+#[async_trait]
+impl EventListener for RegionFlushGate {
+ fn on_flush_success(&self, region_id: RegionId) {
+ if region_id == self.observed_region_id {
+ self.observed_flush_finished.notify_one();
+ }
+ }
+
+ async fn on_flush_begin(&self, region_id: RegionId) {
+ if region_id == self.blocked_region_id {
+ self.blocked_flush_started.notify_one();
+ self.resume_blocked_flush.notified().await;
+ }
+ }
+}
+
#[tokio::test]
async fn test_manual_flush() {
test_manual_flush_with_format(false).await;
@@ -273,6 +345,503 @@ async fn test_write_stall_with_format(flat_format: bool) {
assert_eq!(expected, batches.pretty_print().unwrap());
}
+#[tokio::test]
+async fn test_region_write_buffer_limit_flushes_hot_region() {
+ let mut env = TestEnv::new().await;
+ let listener = Arc::new(FlushListener::default());
+ let engine = env
+ .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
+ .await;
+
+ let hot_region_id = RegionId::new(1, 1);
+ let normal_region_id = RegionId::new(2, 1);
+ let idle_region_id = RegionId::new(3, 1);
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ hot_region_id.table_id(),
+ "hot_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ idle_region_id.table_id(),
+ "idle_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ normal_region_id.table_id(),
+ "normal_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+
+ // The hot region sets its limit to its memory usage after the initial write.
+ let hot_request = CreateRequestBuilder::new().build();
+ let hot_schema = rows_schema(&hot_request);
+ engine
+ .handle_request(hot_region_id, RegionRequest::Create(hot_request))
+ .await
+ .unwrap();
+
+ // An explicit table option takes precedence over the engine-level default.
+ let normal_request = CreateRequestBuilder::new()
+ .table_dir("normal")
+ .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB")
+ .build();
+ let normal_schema = rows_schema(&normal_request);
+ engine
+ .handle_request(normal_region_id, RegionRequest::Create(normal_request))
+ .await
+ .unwrap();
+
+ let idle_request = CreateRequestBuilder::new()
+ .table_dir("idle")
+ .insert_option(WRITE_BUFFER_SIZE_KEY, "1B")
+ .build();
+ let idle_schema = rows_schema(&idle_request);
+ engine
+ .handle_request(idle_region_id, RegionRequest::Create(idle_request))
+ .await
+ .unwrap();
+
+ put_rows(
+ &engine,
+ hot_region_id,
+ Rows {
+ schema: hot_schema.clone(),
+ rows: build_rows_for_key("hot", 0, 2, 0),
+ },
+ )
+ .await;
+ set_write_buffer_size_to_current_usage(&engine, hot_region_id).await;
+ put_rows(
+ &engine,
+ normal_region_id,
+ Rows {
+ schema: normal_schema,
+ rows: build_rows_for_key("normal", 0, 2, 0),
+ },
+ )
+ .await;
+ put_rows(
+ &engine,
+ idle_region_id,
+ Rows {
+ schema: idle_schema,
+ rows: build_rows_for_key("idle", 0, 2, 0),
+ },
+ )
+ .await;
+
+ // The next write triggers a flush after the first write crosses the threshold.
+ put_rows(
+ &engine,
+ hot_region_id,
+ Rows {
+ schema: hot_schema,
+ rows: build_rows_for_key("hot", 2, 2, 0),
+ },
+ )
+ .await;
+ listener.wait().await;
+
+ let scanner = engine
+ .scanner(hot_region_id, ScanRequest::default())
+ .await
+ .unwrap();
+ assert_eq!(1, scanner.num_files());
+
+ let scanner = engine
+ .scanner(normal_region_id, ScanRequest::default())
+ .await
+ .unwrap();
+ assert_eq!(0, scanner.num_files());
+ assert_eq!(1, scanner.num_memtables());
+
+ let scanner = engine
+ .scanner(idle_region_id, ScanRequest::default())
+ .await
+ .unwrap();
+ assert_eq!(0, scanner.num_files());
+ assert_eq!(1, scanner.num_memtables());
+}
+
+#[tokio::test]
+async fn test_region_write_buffer_stall_does_not_block_other_region() {
+ let mut env = TestEnv::new().await;
+ let listener = Arc::new(StallListener::default());
+ let engine = env
+ .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
+ .await;
+
+ let stalled_region_id = RegionId::new(1, 1);
+ let normal_region_id = RegionId::new(2, 1);
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ stalled_region_id.table_id(),
+ "stalled_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ normal_region_id.table_id(),
+ "normal_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+
+ let stalled_request = CreateRequestBuilder::new().build();
+ let stalled_schema = rows_schema(&stalled_request);
+ engine
+ .handle_request(stalled_region_id, RegionRequest::Create(stalled_request))
+ .await
+ .unwrap();
+
+ let normal_request = CreateRequestBuilder::new()
+ .table_dir("normal")
+ .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB")
+ .build();
+ let normal_schema = rows_schema(&normal_request);
+ engine
+ .handle_request(normal_region_id, RegionRequest::Create(normal_request))
+ .await
+ .unwrap();
+
+ put_rows(
+ &engine,
+ stalled_region_id,
+ Rows {
+ schema: stalled_schema.clone(),
+ rows: build_rows_for_key("stalled", 0, 2, 0),
+ },
+ )
+ .await;
+ set_write_buffer_size_to_current_usage(&engine, stalled_region_id).await;
+
+ let engine_cloned = engine.clone();
+ let stalled_write = tokio::spawn(async move {
+ put_rows(
+ &engine_cloned,
+ stalled_region_id,
+ Rows {
+ schema: stalled_schema,
+ rows: build_rows_for_key("stalled", 2, 2, 0),
+ },
+ )
+ .await;
+ });
+
+ listener.wait().await;
+
+ put_rows(
+ &engine,
+ normal_region_id,
+ Rows {
+ schema: normal_schema,
+ rows: build_rows_for_key("normal", 0, 2, 0),
+ },
+ )
+ .await;
+
+ flush_region(&engine, stalled_region_id, None).await;
+ stalled_write.await.unwrap();
+
+ let scanner = engine
+ .scanner(normal_region_id, ScanRequest::default())
+ .await
+ .unwrap();
+ assert_eq!(0, scanner.num_files());
+ assert_eq!(1, scanner.num_memtables());
+}
+
+#[tokio::test]
+async fn test_region_write_buffer_does_not_stall_follower_write() {
+ let mut env = TestEnv::new().await;
+ let engine = env
+ .create_engine(MitoConfig {
+ default_region_write_buffer_size: ReadableSize(1),
+ ..Default::default()
+ })
+ .await;
+
+ let region_id = RegionId::new(1, 1);
+ let request = CreateRequestBuilder::new().build();
+ let schema = rows_schema(&request);
+ engine
+ .handle_request(region_id, RegionRequest::Create(request))
+ .await
+ .unwrap();
+
+ put_rows(
+ &engine,
+ region_id,
+ Rows {
+ schema: schema.clone(),
+ rows: build_rows_for_key("follower", 0, 2, 0),
+ },
+ )
+ .await;
+
+ engine
+ .set_region_role(region_id, RegionRole::Follower)
+ .unwrap();
+
+ let err = tokio::time::timeout(
+ Duration::from_secs(1),
+ engine.handle_request(
+ region_id,
+ RegionRequest::Put(RegionPutRequest {
+ rows: Rows {
+ schema,
+ rows: build_rows_for_key("follower", 2, 4, 0),
+ },
+ hint: None,
+ partition_expr_version: None,
+ }),
+ ),
+ )
+ .await
+ .expect("write to follower should not remain stalled")
+ .unwrap_err();
+ assert_matches!(
+ err.into_inner().as_any().downcast_ref::().unwrap(),
+ Error::RegionState { .. }
+ );
+}
+
+#[tokio::test]
+async fn test_unrelated_flush_does_not_resume_region_stalled_write() {
+ let mut env = TestEnv::new().await;
+ let stalled_region_id = RegionId::new(1, 1);
+ let normal_region_id = RegionId::new(2, 1);
+ let listener = Arc::new(RegionFlushGate::new(stalled_region_id, normal_region_id));
+ let engine = env
+ .create_engine_with(
+ MitoConfig {
+ max_background_flushes: 2,
+ ..Default::default()
+ },
+ None,
+ Some(listener.clone()),
+ None,
+ )
+ .await;
+
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ stalled_region_id.table_id(),
+ "stalled_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ normal_region_id.table_id(),
+ "normal_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+
+ let stalled_request = CreateRequestBuilder::new().build();
+ let stalled_schema = rows_schema(&stalled_request);
+ engine
+ .handle_request(stalled_region_id, RegionRequest::Create(stalled_request))
+ .await
+ .unwrap();
+
+ let normal_request = CreateRequestBuilder::new()
+ .table_dir("normal")
+ .insert_option(WRITE_BUFFER_SIZE_KEY, "1GiB")
+ .build();
+ let normal_schema = rows_schema(&normal_request);
+ engine
+ .handle_request(normal_region_id, RegionRequest::Create(normal_request))
+ .await
+ .unwrap();
+
+ put_rows(
+ &engine,
+ stalled_region_id,
+ Rows {
+ schema: stalled_schema.clone(),
+ rows: build_rows_for_key("stalled", 0, 2, 0),
+ },
+ )
+ .await;
+ set_write_buffer_size_to_current_usage(&engine, stalled_region_id).await;
+ put_rows(
+ &engine,
+ normal_region_id,
+ Rows {
+ schema: normal_schema,
+ rows: build_rows_for_key("normal", 0, 2, 0),
+ },
+ )
+ .await;
+
+ let engine_cloned = engine.clone();
+ let mut stalled_write = tokio::spawn(async move {
+ put_rows(
+ &engine_cloned,
+ stalled_region_id,
+ Rows {
+ schema: stalled_schema,
+ rows: build_rows_for_key("stalled", 2, 2, 0),
+ },
+ )
+ .await;
+ });
+
+ listener.wait_blocked_flush_started().await;
+ flush_region(&engine, normal_region_id, None).await;
+ listener.wait_observed_flush_finished().await;
+
+ assert!(
+ tokio::time::timeout(Duration::from_secs(1), &mut stalled_write)
+ .await
+ .is_err(),
+ "an unrelated region flush resumed the region-stalled write"
+ );
+
+ listener.resume_blocked_flush();
+ tokio::time::timeout(Duration::from_secs(10), &mut stalled_write)
+ .await
+ .expect("the stalled write should resume after its region flushes")
+ .unwrap();
+}
+
+#[tokio::test]
+async fn test_region_write_buffer_rejects_only_full_region_queue() {
+ let mut env = TestEnv::new().await;
+ let hot_region_id = RegionId::new(1, 1);
+ let normal_region_id = RegionId::new(2, 1);
+ let listener = Arc::new(RegionFlushGate::new(hot_region_id, hot_region_id));
+ let engine = env
+ .create_engine_with(
+ MitoConfig {
+ max_background_flushes: 2,
+ ..Default::default()
+ },
+ None,
+ Some(listener.clone()),
+ None,
+ )
+ .await;
+
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ hot_region_id.table_id(),
+ "hot_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+ env.get_schema_metadata_manager()
+ .register_region_table_info(
+ normal_region_id.table_id(),
+ "normal_table",
+ "test_catalog",
+ "test_schema",
+ None,
+ env.get_kv_backend(),
+ )
+ .await;
+
+ let hot_request = CreateRequestBuilder::new().build();
+ let hot_schema = rows_schema(&hot_request);
+ engine
+ .handle_request(hot_region_id, RegionRequest::Create(hot_request))
+ .await
+ .unwrap();
+ let normal_request = CreateRequestBuilder::new().table_dir("normal").build();
+ let normal_schema = rows_schema(&normal_request);
+ engine
+ .handle_request(normal_region_id, RegionRequest::Create(normal_request))
+ .await
+ .unwrap();
+
+ put_rows(
+ &engine,
+ hot_region_id,
+ Rows {
+ schema: hot_schema.clone(),
+ rows: build_rows_for_key("hot", 0, 2, 0),
+ },
+ )
+ .await;
+ set_write_buffer_size_to_current_usage(&engine, hot_region_id).await;
+
+ let engine_cloned = engine.clone();
+ let stalled_schema = hot_schema.clone();
+ let stalled_write = tokio::spawn(async move {
+ engine_cloned
+ .handle_request(
+ hot_region_id,
+ RegionRequest::Put(RegionPutRequest {
+ rows: Rows {
+ schema: stalled_schema,
+ rows: build_rows_for_key("hot", 2, 1026, 2),
+ },
+ hint: None,
+ partition_expr_version: None,
+ }),
+ )
+ .await
+ });
+
+ listener.wait_blocked_flush_started().await;
+
+ put_rows(
+ &engine,
+ normal_region_id,
+ Rows {
+ schema: normal_schema,
+ rows: build_rows_for_key("normal", 0, 2, 0),
+ },
+ )
+ .await;
+
+ listener.resume_blocked_flush();
+ listener.wait_observed_flush_finished().await;
+
+ let err = tokio::time::timeout(Duration::from_secs(1), stalled_write)
+ .await
+ .expect("the stalled queue should be rejected at the region hard limit")
+ .unwrap()
+ .unwrap_err();
+ assert_matches!(
+ err.into_inner().as_any().downcast_ref::().unwrap(),
+ Error::RejectWrite { .. }
+ );
+}
+
#[tokio::test]
async fn test_flush_empty() {
test_flush_empty_with_format(false).await;
diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs
index e6aa15f9d48..71abc0b0a13 100644
--- a/src/mito2/src/flush.rs
+++ b/src/mito2/src/flush.rs
@@ -211,6 +211,8 @@ impl WriteBufferManager for WriteBufferManagerImpl {
pub enum FlushReason {
/// Engine reaches flush threshold.
EngineFull,
+ /// Region reaches its write buffer threshold.
+ RegionFull,
/// Manual flush.
Manual,
/// Flush to alter table.
diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs
index 92dfd460fe6..2ac881c4072 100644
--- a/src/mito2/src/region/options.rs
+++ b/src/mito2/src/region/options.rs
@@ -107,6 +107,11 @@ pub struct RegionOptions {
/// Internal primary key encoding override used by metric-engine.
#[serde(skip_serializing_if = "Option::is_none")]
pub primary_key_encoding: Option,
+ /// Per-region write buffer size. A positive size flushes/stalls this region
+ /// independently of the global write buffer limit and rejects writes at twice
+ /// the configured size; zero disables both limits.
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pub write_buffer_size: Option,
}
impl RegionOptions {
@@ -256,6 +261,7 @@ impl RegionOptions {
merge_mode: options.merge_mode,
sst_format,
primary_key_encoding,
+ write_buffer_size: options.write_buffer_size,
};
opts.validate()?;
@@ -354,6 +360,7 @@ impl Default for TwcsOptions {
#[derive(Debug, Deserialize)]
#[serde(default)]
struct RegionOptionsWithoutEnum {
+ write_buffer_size: Option,
/// Region SST files TTL.
ttl: Option,
#[serde(with = "humantime_serde")]
@@ -371,6 +378,7 @@ impl Default for RegionOptionsWithoutEnum {
fn default() -> Self {
let options = RegionOptions::default();
RegionOptionsWithoutEnum {
+ write_buffer_size: options.write_buffer_size,
ttl: options.ttl,
auto_flush_interval: options.auto_flush_interval,
storage: options.storage,
@@ -520,6 +528,7 @@ mod tests {
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
use common_wal::options::KafkaWalOptions;
+ use store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY;
use super::*;
@@ -569,6 +578,28 @@ mod tests {
);
}
+ #[test]
+ fn test_with_write_buffer_size() {
+ let map = make_map(&[(WRITE_BUFFER_SIZE_KEY, "128MiB")]);
+ let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap();
+ let expect = RegionOptions {
+ write_buffer_size: Some(ReadableSize::mb(128)),
+ ..Default::default()
+ };
+ assert_eq!(expect, options);
+ }
+
+ #[test]
+ fn test_with_zero_write_buffer_size() {
+ let map = make_map(&[(WRITE_BUFFER_SIZE_KEY, "0")]);
+ let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap();
+ let expect = RegionOptions {
+ write_buffer_size: Some(ReadableSize::mb(0)),
+ ..Default::default()
+ };
+ assert_eq!(expect, options);
+ }
+
#[test]
fn test_with_storage() {
let map = make_map(&[("storage", "S3")]);
@@ -898,6 +929,7 @@ mod tests {
merge_mode: Some(MergeMode::LastNonNull),
sst_format: Some(FormatType::Flat),
primary_key_encoding: None,
+ write_buffer_size: None,
};
assert_eq!(expect, options);
}
@@ -928,10 +960,18 @@ mod tests {
merge_mode: Some(MergeMode::LastNonNull),
sst_format: None,
primary_key_encoding: None,
+ write_buffer_size: Some(ReadableSize::mb(128)),
};
let region_options_json_str = serde_json::to_string(&options).unwrap();
let got: RegionOptions = serde_json::from_str(®ion_options_json_str).unwrap();
assert_eq!(options, got);
+
+ let old_region_options_json_str = r#"{"ttl":null}"#;
+ let got: RegionOptions = serde_json::from_str(old_region_options_json_str).unwrap();
+ assert_eq!(None, got.write_buffer_size);
+
+ let default_json = serde_json::to_value(RegionOptions::default()).unwrap();
+ assert!(default_json.get(WRITE_BUFFER_SIZE_KEY).is_none());
}
#[test]
@@ -985,6 +1025,7 @@ mod tests {
merge_mode: Some(MergeMode::LastNonNull),
sst_format: None,
primary_key_encoding: None,
+ write_buffer_size: None,
};
assert_eq!(options, got);
}
diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs
index 23b6dd4d7b8..a3602d29ff6 100644
--- a/src/mito2/src/worker.rs
+++ b/src/mito2/src/worker.rs
@@ -771,6 +771,14 @@ pub(crate) struct StalledRequests {
}
impl StalledRequests {
+ /// Returns the estimated size of stalled requests for a region.
+ pub(crate) fn estimated_size(&self, region_id: &RegionId) -> usize {
+ self.requests
+ .get(region_id)
+ .map(|(size, _, _)| *size)
+ .unwrap_or_default()
+ }
+
/// Appends stalled requests.
pub(crate) fn append(
&mut self,
diff --git a/src/mito2/src/worker/handle_alter.rs b/src/mito2/src/worker/handle_alter.rs
index e022d166315..87044d768de 100644
--- a/src/mito2/src/worker/handle_alter.rs
+++ b/src/mito2/src/worker/handle_alter.rs
@@ -188,6 +188,20 @@ impl RegionWorkerLoop {
let mut current_options = version.options.clone();
for option in options {
match option {
+ SetRegionOption::WriteBufferSize(new_write_buffer_size) => {
+ info!(
+ "Update region write_buffer_size: {}, previous: {:?} new: {:?}",
+ region.region_id, current_options.write_buffer_size, new_write_buffer_size
+ );
+ current_options.write_buffer_size = new_write_buffer_size;
+ current_options.validate().map_err(|e| {
+ store_api::metadata::InvalidRegionRequestSnafu {
+ region_id: region.region_id,
+ err: e.to_string(),
+ }
+ .build()
+ })?;
+ }
SetRegionOption::Ttl(new_ttl) => {
info!(
"Update region ttl: {}, previous: {:?} new: {:?}",
@@ -273,7 +287,8 @@ fn new_region_options_on_empty_memtable(
let mut current_options = current_options.clone();
for option in options {
match option {
- SetRegionOption::Ttl(_)
+ SetRegionOption::WriteBufferSize(_)
+ | SetRegionOption::Ttl(_)
| SetRegionOption::Twsc(_, _)
| SetRegionOption::AutoFlushInterval(_) => (),
SetRegionOption::Format(format_str) => {
diff --git a/src/mito2/src/worker/handle_flush.rs b/src/mito2/src/worker/handle_flush.rs
index c8e86699ea1..ccbe31c3a32 100644
--- a/src/mito2/src/worker/handle_flush.rs
+++ b/src/mito2/src/worker/handle_flush.rs
@@ -14,9 +14,11 @@
//! Handling flush related requests.
+use std::collections::HashSet;
use std::sync::Arc;
use std::sync::atomic::Ordering;
+use common_base::readable_size::ReadableSize;
use common_telemetry::{debug, error, info};
use store_api::logstore::LogStore;
use store_api::region_request::{RegionFlushReason, RegionFlushRequest};
@@ -26,10 +28,24 @@ use crate::config::{IndexBuildMode, MitoConfig};
use crate::error::{RegionNotFoundSnafu, Result};
use crate::flush::{FlushReason, RegionFlushTask};
use crate::region::MitoRegionRef;
+use crate::region::version::VersionRef;
use crate::request::{BuildIndexRequest, FlushFailed, FlushFinished, OnFailure, OptionOutputTx};
use crate::sst::index::IndexBuildType;
use crate::worker::RegionWorkerLoop;
+#[derive(Debug, Default)]
+pub(crate) struct RegionWriteBufferPressure {
+ pub(crate) stalled_region_ids: HashSet,
+ pub(crate) rejected_region_ids: HashSet,
+}
+
+#[derive(Debug, Default, PartialEq, Eq)]
+struct RegionWriteBufferStatus {
+ should_flush: bool,
+ should_stall: bool,
+ should_reject: bool,
+}
+
fn resolve_flush_reason(
request_reason: Option,
is_downgrading: bool,
@@ -85,8 +101,7 @@ impl RegionWorkerLoop {
}
let version = region.version();
- let region_memtable_size =
- version.memtables.mutable_usage() + version.memtables.immutables_usage();
+ let region_memtable_size = region_memtable_usage(&version);
let auto_flush_interval = version
.options
@@ -152,6 +167,53 @@ impl RegionWorkerLoop {
Ok(())
}
+ /// Flushes write regions that exceed their flush threshold and returns region pressure.
+ pub(crate) fn maybe_flush_write_regions(
+ &mut self,
+ region_ids: HashSet,
+ ) -> RegionWriteBufferPressure {
+ let mut pressure = RegionWriteBufferPressure::default();
+
+ for region_id in region_ids {
+ let Some(region) = self.regions.get_region(region_id) else {
+ continue;
+ };
+ if !region.is_writable() {
+ continue;
+ }
+
+ let status = region_write_buffer_status(
+ ®ion.version(),
+ self.config.default_region_write_buffer_size,
+ self.stalled_requests.estimated_size(®ion_id),
+ );
+
+ if status.should_flush && !self.flush_scheduler.is_flush_requested(region.region_id) {
+ let task = self.new_flush_task(
+ ®ion,
+ FlushReason::RegionFull,
+ None,
+ self.config.clone(),
+ );
+ if let Err(e) = self.flush_scheduler.schedule_flush(
+ region.region_id,
+ ®ion.version_control,
+ task,
+ ) {
+ error!(e; "Failed to schedule flush task for region {}", region.region_id);
+ }
+ }
+
+ if status.should_reject {
+ pressure.rejected_region_ids.insert(region_id);
+ } else if status.should_stall {
+ pressure.stalled_region_ids.insert(region_id);
+ }
+ }
+
+ pressure
+ }
+
/// Creates a flush task with specific `reason` for the `region`.
pub(crate) fn new_flush_task(
&self,
@@ -179,6 +241,61 @@ impl RegionWorkerLoop {
}
}
+pub(crate) fn region_memtable_usage(version: &VersionRef) -> usize {
+ version.memtables.mutable_usage() + version.memtables.immutables_usage()
+}
+
+pub(crate) fn region_write_buffer_size(
+ version: &VersionRef,
+ default_region_write_buffer_size: ReadableSize,
+) -> Option {
+ let size = version
+ .options
+ .write_buffer_size
+ .unwrap_or(default_region_write_buffer_size);
+ (size.as_bytes() > 0).then_some(size.as_bytes() as usize)
+}
+
+/// Returns the pressure state of a region's write buffer.
+fn region_write_buffer_status(
+ version: &VersionRef,
+ default_region_write_buffer_size: ReadableSize,
+ stalled_request_size: usize,
+) -> RegionWriteBufferStatus {
+ let Some(write_buffer_size) =
+ region_write_buffer_size(version, default_region_write_buffer_size)
+ else {
+ return RegionWriteBufferStatus::default();
+ };
+
+ let mutable_usage = version.memtables.mutable_usage();
+ let memory_usage = region_memtable_usage(version);
+ region_write_buffer_status_from_usage(
+ write_buffer_size,
+ mutable_usage,
+ memory_usage,
+ stalled_request_size,
+ )
+}
+
+fn region_write_buffer_status_from_usage(
+ write_buffer_size: usize,
+ mutable_usage: usize,
+ memory_usage: usize,
+ stalled_request_size: usize,
+) -> RegionWriteBufferStatus {
+ let mutable_limit = std::cmp::max(1, write_buffer_size / 2);
+ let should_stall = memory_usage >= write_buffer_size;
+ let reject_limit = write_buffer_size.saturating_mul(2);
+ let should_reject = memory_usage.saturating_add(stalled_request_size) >= reject_limit;
+
+ RegionWriteBufferStatus {
+ should_flush: mutable_usage >= mutable_limit || should_stall,
+ should_stall,
+ should_reject,
+ }
+}
+
impl RegionWorkerLoop {
/// Handles manual flush request.
pub(crate) fn handle_flush_request(
@@ -438,4 +555,48 @@ mod tests {
assert_eq!(resolve_flush_reason(None, true), FlushReason::Downgrading);
assert_eq!(resolve_flush_reason(None, false), FlushReason::Manual);
}
+
+ #[test]
+ fn test_region_write_buffer_status_boundaries() {
+ assert_eq!(
+ RegionWriteBufferStatus::default(),
+ region_write_buffer_status_from_usage(100, 49, 99, 0)
+ );
+ assert_eq!(
+ RegionWriteBufferStatus {
+ should_flush: true,
+ should_stall: false,
+ should_reject: false,
+ },
+ region_write_buffer_status_from_usage(100, 50, 99, 0)
+ );
+ assert_eq!(
+ RegionWriteBufferStatus {
+ should_flush: true,
+ should_stall: true,
+ should_reject: false,
+ },
+ region_write_buffer_status_from_usage(100, 50, 100, 99)
+ );
+ assert_eq!(
+ RegionWriteBufferStatus {
+ should_flush: true,
+ should_stall: true,
+ should_reject: true,
+ },
+ region_write_buffer_status_from_usage(100, 50, 100, 100)
+ );
+ }
+
+ #[test]
+ fn test_region_write_buffer_status_saturates_reject_limit() {
+ assert_eq!(
+ RegionWriteBufferStatus {
+ should_flush: true,
+ should_stall: true,
+ should_reject: true,
+ },
+ region_write_buffer_status_from_usage(usize::MAX, usize::MAX, usize::MAX, usize::MAX,)
+ );
+ }
}
diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs
index dd0acd63000..0eb4c7bc39e 100644
--- a/src/mito2/src/worker/handle_write.rs
+++ b/src/mito2/src/worker/handle_write.rs
@@ -14,7 +14,7 @@
//! Handling write requests.
-use std::collections::{HashMap, hash_map};
+use std::collections::{HashMap, HashSet, hash_map};
use std::sync::Arc;
use api::v1::OpType;
@@ -50,8 +50,11 @@ impl RegionWorkerLoop {
return;
}
- // Flush this worker if the engine needs to flush.
+ let write_region_ids = write_region_ids(write_requests, bulk_requests);
+
+ // Check region pressure before writes to match the global write buffer behavior.
self.maybe_flush_worker();
+ let pressure = self.maybe_flush_write_regions(write_region_ids);
if self.should_reject_write() {
// The memory pressure is still too high, reject write requests.
@@ -61,6 +64,20 @@ impl RegionWorkerLoop {
return;
}
+ if !pressure.rejected_region_ids.is_empty() {
+ reject_region_write_requests(
+ &pressure.rejected_region_ids,
+ write_requests,
+ bulk_requests,
+ );
+ for region_id in &pressure.rejected_region_ids {
+ self.reject_region_stalled_requests(region_id);
+ }
+ if write_requests.is_empty() && bulk_requests.is_empty() {
+ return;
+ }
+ }
+
if self.write_buffer_manager.should_stall() && allow_stall {
let stalled_count = (write_requests.len() + bulk_requests.len()) as i64;
self.stalling_count.add(stalled_count);
@@ -70,6 +87,17 @@ impl RegionWorkerLoop {
return;
}
+ if allow_stall {
+ self.stall_region_write_requests(
+ &pressure.stalled_region_ids,
+ write_requests,
+ bulk_requests,
+ );
+ if write_requests.is_empty() && bulk_requests.is_empty() {
+ return;
+ }
+ }
+
// Prepare write context.
let mut region_ctxs = {
let _timer = WRITE_STAGE_ELAPSED
@@ -163,15 +191,31 @@ impl RegionWorkerLoop {
.inc_by(delete_rows as u64);
}
- /// Handles all stalled write requests.
+ /// Handles stalled write requests whose regions no longer need to stall.
pub(crate) async fn handle_stalled_requests(&mut self) {
- // Handle stalled requests.
- let stalled = std::mem::take(&mut self.stalled_requests);
- self.stalling_count.sub(stalled.stalled_count() as i64);
- // We already stalled these requests, don't stall them again.
- for (_, (_, mut requests, mut bulk)) in stalled.requests {
- self.handle_write_requests(&mut requests, &mut bulk, false)
- .await;
+ let region_ids = self
+ .stalled_requests
+ .requests
+ .keys()
+ .copied()
+ .collect::>();
+ let pressure = self.maybe_flush_write_regions(region_ids);
+ for region_id in &pressure.rejected_region_ids {
+ self.reject_region_stalled_requests(region_id);
+ }
+ let ready_region_ids = self
+ .stalled_requests
+ .requests
+ .keys()
+ .filter(|region_id| !pressure.stalled_region_ids.contains(region_id))
+ .copied()
+ .collect::>();
+
+ // These requests have already been stalled. Retry ready regions without stalling the
+ // same requests again. Regions that still exceed their limit remain in the queue until
+ // their own flush releases the pressure.
+ for region_id in ready_region_ids {
+ self.handle_region_stalled_requests(®ion_id, false).await;
}
}
@@ -553,6 +597,34 @@ impl RegionWorkerLoop {
self.write_buffer_manager.memory_usage() + self.stalled_requests.estimated_size
>= self.config.global_write_buffer_reject_size.as_bytes() as usize
}
+
+ fn stall_region_write_requests(
+ &mut self,
+ stalled_region_ids: &HashSet,
+ write_requests: &mut Vec,
+ bulk_requests: &mut Vec,
+ ) {
+ let mut stalled_count = 0;
+ let mut stalled_write_requests = write_requests
+ .extract_if(.., |req| {
+ stalled_region_ids.contains(&req.request.region_id)
+ })
+ .collect::>();
+ let mut stalled_bulk_requests = bulk_requests
+ .extract_if(.., |req| stalled_region_ids.contains(&req.region_id))
+ .collect::>();
+
+ stalled_count += stalled_write_requests.len() + stalled_bulk_requests.len();
+ self.stalled_requests
+ .append(&mut stalled_write_requests, &mut stalled_bulk_requests);
+
+ if stalled_count > 0 {
+ let stalled_count = stalled_count as i64;
+ self.stalling_count.add(stalled_count);
+ WRITE_STALL_TOTAL.inc_by(stalled_count as u64);
+ self.listener.on_write_stall();
+ }
+ }
}
/// Send rejected error to all `write_requests`.
@@ -576,6 +648,33 @@ fn reject_write_requests(
}
}
+fn reject_region_write_requests(
+ rejected_region_ids: &HashSet,
+ write_requests: &mut Vec,
+ bulk_requests: &mut Vec,
+) {
+ let mut rejected_write_requests = write_requests
+ .extract_if(.., |req| {
+ rejected_region_ids.contains(&req.request.region_id)
+ })
+ .collect::>();
+ let mut rejected_bulk_requests = bulk_requests
+ .extract_if(.., |req| rejected_region_ids.contains(&req.region_id))
+ .collect::>();
+ reject_write_requests(&mut rejected_write_requests, &mut rejected_bulk_requests);
+}
+
+fn write_region_ids(
+ write_requests: &[SenderWriteRequest],
+ bulk_requests: &[SenderBulkRequest],
+) -> HashSet {
+ write_requests
+ .iter()
+ .map(|req| req.request.region_id)
+ .chain(bulk_requests.iter().map(|req| req.region_id))
+ .collect()
+}
+
/// Rejects delete request under append mode.
fn check_op_type(append_mode: bool, request: &WriteRequest) -> Result<()> {
if append_mode {
diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs
index 8762ed69779..4320a40d8de 100644
--- a/src/store-api/src/mito_engine_options.rs
+++ b/src/store-api/src/mito_engine_options.rs
@@ -19,6 +19,8 @@
pub use common_wal::options::WAL_OPTIONS_KEY;
/// Option key for append mode.
pub const APPEND_MODE_KEY: &str = "append_mode";
+/// Option key for the per-region write buffer size.
+pub const WRITE_BUFFER_SIZE_KEY: &str = "write_buffer_size";
/// Option key for merge mode.
pub const MERGE_MODE_KEY: &str = "merge_mode";
/// Option key for TTL(time-to-live)
@@ -72,6 +74,7 @@ pub const SST_FORMAT_KEY: &str = "sst_format";
pub fn is_mito_engine_option_key(key: &str) -> bool {
[
"ttl",
+ WRITE_BUFFER_SIZE_KEY,
AUTO_FLUSH_INTERVAL_KEY,
COMPACTION_TYPE,
COMPACTION_OVERRIDE,
@@ -107,6 +110,7 @@ mod tests {
#[test]
fn test_is_mito_engine_option_key() {
assert!(is_mito_engine_option_key("ttl"));
+ assert!(is_mito_engine_option_key("write_buffer_size"));
assert!(is_mito_engine_option_key("auto_flush_interval"));
assert!(is_mito_engine_option_key("compaction.type"));
assert!(is_mito_engine_option_key("compaction.override"));
diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs
index 6aa3e3f49d0..d63c9058a33 100644
--- a/src/store-api/src/region_request.rs
+++ b/src/store-api/src/region_request.rs
@@ -33,6 +33,7 @@ use api::v1::{
SemanticType, SkippingIndexType as PbSkippingIndexType, WriteHint,
};
pub use common_base::AffectedRows;
+use common_base::readable_size::ReadableSize;
use common_grpc::flight::FlightDecoder;
use common_recordbatch::DfRecordBatch;
use common_time::{TimeToLive, Timestamp};
@@ -54,7 +55,7 @@ use crate::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY;
use crate::metrics;
use crate::mito_engine_options::{
APPEND_MODE_KEY, AUTO_FLUSH_INTERVAL_KEY, SST_FORMAT_KEY, TTL_KEY, TWCS_MAX_OUTPUT_FILE_SIZE,
- TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM,
+ TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, WRITE_BUFFER_SIZE_KEY,
};
use crate::path_utils::table_dir;
use crate::storage::{ColumnId, RegionId, ScanRequest};
@@ -1405,6 +1406,7 @@ impl From for ModifyColumnType {
#[derive(Debug, Eq, PartialEq, Clone, Serialize, Deserialize)]
pub enum SetRegionOption {
+ WriteBufferSize(Option),
Ttl(Option),
// Modifying TwscOptions with values as (option name, new value).
Twsc(String, String),
@@ -1422,6 +1424,12 @@ impl TryFrom<&PbOption> for SetRegionOption {
fn try_from(value: &PbOption) -> std::result::Result {
let PbOption { key, value } = value;
match key.as_str() {
+ WRITE_BUFFER_SIZE_KEY => {
+ let size = value
+ .parse::()
+ .map_err(|_| InvalidSetRegionOptionRequestSnafu { key, value }.build())?;
+ Ok(Self::WriteBufferSize(Some(size)))
+ }
TTL_KEY => {
let ttl = TimeToLive::from_humantime_or_str(value)
.map_err(|_| InvalidSetRegionOptionRequestSnafu { key, value }.build())?;
@@ -1470,6 +1478,7 @@ impl From<&UnsetRegionOption> for SetRegionOption {
SetRegionOption::Twsc(unset_option.to_string(), String::new())
}
UnsetRegionOption::Ttl => SetRegionOption::Ttl(Default::default()),
+ UnsetRegionOption::WriteBufferSize => SetRegionOption::WriteBufferSize(None),
}
}
}
@@ -1480,6 +1489,7 @@ impl TryFrom<&str> for UnsetRegionOption {
fn try_from(key: &str) -> Result {
match key.to_ascii_lowercase().as_str() {
TTL_KEY => Ok(Self::Ttl),
+ WRITE_BUFFER_SIZE_KEY => Ok(Self::WriteBufferSize),
TWCS_TRIGGER_FILE_NUM => Ok(Self::TwcsTriggerFileNum),
TWCS_MAX_OUTPUT_FILE_SIZE => Ok(Self::TwcsMaxOutputFileSize),
TWCS_TIME_WINDOW => Ok(Self::TwcsTimeWindow),
@@ -1494,12 +1504,14 @@ pub enum UnsetRegionOption {
TwcsMaxOutputFileSize,
TwcsTimeWindow,
Ttl,
+ WriteBufferSize,
}
impl UnsetRegionOption {
pub fn as_str(&self) -> &str {
match self {
Self::Ttl => TTL_KEY,
+ Self::WriteBufferSize => WRITE_BUFFER_SIZE_KEY,
Self::TwcsTriggerFileNum => TWCS_TRIGGER_FILE_NUM,
Self::TwcsMaxOutputFileSize => TWCS_MAX_OUTPUT_FILE_SIZE,
Self::TwcsTimeWindow => TWCS_TIME_WINDOW,
@@ -1831,6 +1843,41 @@ mod tests {
);
}
+ #[test]
+ fn test_write_buffer_size_region_options() {
+ let option = PbOption {
+ key: WRITE_BUFFER_SIZE_KEY.to_string(),
+ value: "128MiB".to_string(),
+ };
+ assert_eq!(
+ SetRegionOption::WriteBufferSize(Some(ReadableSize::mb(128))),
+ SetRegionOption::try_from(&option).unwrap()
+ );
+
+ let option = PbOption {
+ key: WRITE_BUFFER_SIZE_KEY.to_string(),
+ value: "invalid".to_string(),
+ };
+ SetRegionOption::try_from(&option).unwrap_err();
+
+ // Clearing the option uses UNSET. Empty SET values, including SQL NULL,
+ // remain invalid because the SQL layer does not distinguish NULL from ''.
+ let option = PbOption {
+ key: WRITE_BUFFER_SIZE_KEY.to_string(),
+ value: String::new(),
+ };
+ SetRegionOption::try_from(&option).unwrap_err();
+
+ assert_eq!(
+ UnsetRegionOption::WriteBufferSize,
+ UnsetRegionOption::try_from(WRITE_BUFFER_SIZE_KEY).unwrap()
+ );
+ assert_eq!(
+ SetRegionOption::WriteBufferSize(None),
+ SetRegionOption::from(&UnsetRegionOption::WriteBufferSize)
+ );
+ }
+
/// Returns a new region metadata for testing. Metadata:
/// `[(ts, ms, 1), (tag_0, string, 2), (field_0, string, 3), (field_1, bool, 4)]`
fn new_metadata() -> RegionMetadata {
diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs
index 006cfb34962..b2801dd8654 100644
--- a/src/table/src/metadata.rs
+++ b/src/table/src/metadata.rs
@@ -350,6 +350,9 @@ impl TableMeta {
for request in requests {
match request {
+ SetRegionOption::WriteBufferSize(new_write_buffer_size) => {
+ new_options.write_buffer_size = *new_write_buffer_size;
+ }
SetRegionOption::Ttl(new_ttl) => {
new_options.ttl = *new_ttl;
}
diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs
index 12214a92f9d..12331bf8bb9 100644
--- a/src/table/src/requests.rs
+++ b/src/table/src/requests.rs
@@ -145,7 +145,7 @@ pub fn validate_table_option(key: &str) -> bool {
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
pub struct TableOptions {
- /// Memtable size of memtable.
+ /// Per-region write buffer stall threshold. Writes are rejected at twice this size.
pub write_buffer_size: Option,
/// Time-to-live of table. Expired data will be automatically purged.
pub ttl: Option,
@@ -155,7 +155,7 @@ pub struct TableOptions {
pub extra_options: HashMap,
}
-pub const WRITE_BUFFER_SIZE_KEY: &str = "write_buffer_size";
+pub const WRITE_BUFFER_SIZE_KEY: &str = store_api::mito_engine_options::WRITE_BUFFER_SIZE_KEY;
pub const TTL_KEY: &str = store_api::mito_engine_options::TTL_KEY;
pub const STORAGE_KEY: &str = "storage";
pub const COMMENT_KEY: &str = "comment";
diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs
index 133afd7822d..ba6710da34b 100644
--- a/tests-integration/tests/http.rs
+++ b/tests-integration/tests/http.rs
@@ -1939,6 +1939,7 @@ compress_manifest = false
experimental_compaction_memory_limit = "unlimited"
experimental_compaction_on_exhausted = "wait"
auto_flush_interval = "30m"
+default_region_write_buffer_size = "0KiB"
enable_write_cache = false
write_cache_path = ""
write_cache_size = "5GiB"
diff --git a/tests/cases/standalone/common/alter/alter_table_options.result b/tests/cases/standalone/common/alter/alter_table_options.result
index f91fd9da88e..b169562c894 100644
--- a/tests/cases/standalone/common/alter/alter_table_options.result
+++ b/tests/cases/standalone/common/alter/alter_table_options.result
@@ -259,6 +259,79 @@ DROP TABLE ato;
Affected Rows: 0
+CREATE TABLE write_buffer_size_options(
+ i INTEGER,
+ ts TIMESTAMP TIME INDEX,
+ PRIMARY KEY(i)
+);
+
+Affected Rows: 0
+
+ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB';
+
+Affected Rows: 0
+
+SHOW CREATE TABLE write_buffer_size_options;
+
++---------------------------+----------------------------------------------------------+
+| Table | Create Table |
++---------------------------+----------------------------------------------------------+
+| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( |
+| | "i" INT NULL, |
+| | "ts" TIMESTAMP(3) NOT NULL, |
+| | TIME INDEX ("ts"), |
+| | PRIMARY KEY ("i") |
+| | ) |
+| | |
+| | ENGINE=mito |
+| | WITH( |
+| | write_buffer_size = '128.0MiB' |
+| | ) |
++---------------------------+----------------------------------------------------------+
+
+-- SQLNESS ARG restart=true
+SHOW CREATE TABLE write_buffer_size_options;
+
++---------------------------+----------------------------------------------------------+
+| Table | Create Table |
++---------------------------+----------------------------------------------------------+
+| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( |
+| | "i" INT NULL, |
+| | "ts" TIMESTAMP(3) NOT NULL, |
+| | TIME INDEX ("ts"), |
+| | PRIMARY KEY ("i") |
+| | ) |
+| | |
+| | ENGINE=mito |
+| | WITH( |
+| | write_buffer_size = '128.0MiB' |
+| | ) |
++---------------------------+----------------------------------------------------------+
+
+ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size';
+
+Affected Rows: 0
+
+SHOW CREATE TABLE write_buffer_size_options;
+
++---------------------------+----------------------------------------------------------+
+| Table | Create Table |
++---------------------------+----------------------------------------------------------+
+| write_buffer_size_options | CREATE TABLE IF NOT EXISTS "write_buffer_size_options" ( |
+| | "i" INT NULL, |
+| | "ts" TIMESTAMP(3) NOT NULL, |
+| | TIME INDEX ("ts"), |
+| | PRIMARY KEY ("i") |
+| | ) |
+| | |
+| | ENGINE=mito |
+| | |
++---------------------------+----------------------------------------------------------+
+
+DROP TABLE write_buffer_size_options;
+
+Affected Rows: 0
+
CREATE TABLE phy (ts timestamp time index, val double) engine=metric with ("physical_metric_table" = "");
Affected Rows: 0
diff --git a/tests/cases/standalone/common/alter/alter_table_options.sql b/tests/cases/standalone/common/alter/alter_table_options.sql
index badf1275cdb..c3f67faeee5 100644
--- a/tests/cases/standalone/common/alter/alter_table_options.sql
+++ b/tests/cases/standalone/common/alter/alter_table_options.sql
@@ -55,6 +55,25 @@ SHOW CREATE TABLE ato;
DROP TABLE ato;
+CREATE TABLE write_buffer_size_options(
+ i INTEGER,
+ ts TIMESTAMP TIME INDEX,
+ PRIMARY KEY(i)
+);
+
+ALTER TABLE write_buffer_size_options SET 'write_buffer_size' = '128MiB';
+
+SHOW CREATE TABLE write_buffer_size_options;
+
+-- SQLNESS ARG restart=true
+SHOW CREATE TABLE write_buffer_size_options;
+
+ALTER TABLE write_buffer_size_options UNSET 'write_buffer_size';
+
+SHOW CREATE TABLE write_buffer_size_options;
+
+DROP TABLE write_buffer_size_options;
+
CREATE TABLE phy (ts timestamp time index, val double) engine=metric with ("physical_metric_table" = "");
ALTER TABLE phy set ttl='2years';
diff --git a/tests/cases/standalone/common/information_schema/region_info.result b/tests/cases/standalone/common/information_schema/region_info.result
index ed9a07add60..91cab62336d 100644
--- a/tests/cases/standalone/common/information_schema/region_info.result
+++ b/tests/cases/standalone/common/information_schema/region_info.result
@@ -41,7 +41,7 @@ ADMIN FLUSH_TABLE('region_info_case');
+---------------------------------------+
-- SQLNESS REPLACE (\s+\d+\s+)
--- SQLNESS REPLACE (\{".*"\})
+-- SQLNESS REPLACE (\{.*\})
-- SQLNESS REPLACE (-{40,}) ----------------
-- SQLNESS REPLACE (region_options\s+\|) region_options |
SELECT region_id, state, role, writable, committed_sequence, flushed_sequence, manifest_version, compaction_time_window, region_options, sst_format
diff --git a/tests/cases/standalone/common/information_schema/region_info.sql b/tests/cases/standalone/common/information_schema/region_info.sql
index 1aa682393f2..6c901387ce9 100644
--- a/tests/cases/standalone/common/information_schema/region_info.sql
+++ b/tests/cases/standalone/common/information_schema/region_info.sql
@@ -11,7 +11,7 @@ INSERT INTO region_info_case VALUES (1, 1), (2, 2);
ADMIN FLUSH_TABLE('region_info_case');
-- SQLNESS REPLACE (\s+\d+\s+)
--- SQLNESS REPLACE (\{".*"\})
+-- SQLNESS REPLACE (\{.*\})
-- SQLNESS REPLACE (-{40,}) ----------------
-- SQLNESS REPLACE (region_options\s+\|) region_options |
SELECT region_id, state, role, writable, committed_sequence, flushed_sequence, manifest_version, compaction_time_window, region_options, sst_format
diff --git a/tests/cases/standalone/common/show/show_create.result b/tests/cases/standalone/common/show/show_create.result
index 72392a0997d..18879ed95e8 100644
--- a/tests/cases/standalone/common/show/show_create.result
+++ b/tests/cases/standalone/common/show/show_create.result
@@ -15,7 +15,7 @@ PARTITION ON COLUMNS (`id`) (
ENGINE=mito
WITH(
ttl = '7d',
- write_buffer_size = 1024
+ write_buffer_size = '1KiB'
);
Affected Rows: 0
diff --git a/tests/cases/standalone/common/show/show_create.sql b/tests/cases/standalone/common/show/show_create.sql
index eb0f90c6746..fd8b1c16ea2 100644
--- a/tests/cases/standalone/common/show/show_create.sql
+++ b/tests/cases/standalone/common/show/show_create.sql
@@ -15,7 +15,7 @@ PARTITION ON COLUMNS (`id`) (
ENGINE=mito
WITH(
ttl = '7d',
- write_buffer_size = 1024
+ write_buffer_size = '1KiB'
);
SHOW CREATE TABLE system_metrics;