From 0a002f4d65f83aec561567d8e90497fd0a99c0b1 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Wed, 15 Jul 2026 15:31:25 +0800 Subject: [PATCH] feat: support per-region write buffer limits (#8473) * feat(mito): add per-region write buffer limit Signed-off-by: evenyag * feat(mito): add default region write buffer size Signed-off-by: evenyag * docs: expand configuration change checklist Signed-off-by: evenyag * test: cover table write buffer size option Signed-off-by: evenyag * fix(mito): refine region write buffer checks Signed-off-by: evenyag * fix(mito): keep region-stalled writes queued Signed-off-by: evenyag * fix(mito): allow zero region write buffer size Signed-off-by: evenyag * fix(mito): validate region state before write stall Signed-off-by: evenyag * fix(mito): reject writes beyond region hard limit Signed-off-by: evenyag * docs: update example Signed-off-by: evenyag --------- Signed-off-by: evenyag --- AGENTS.md | 6 +- config/config.md | 2 + config/datanode.example.toml | 3 + config/standalone.example.toml | 3 + src/cmd/tests/load_config_test.rs | 3 + src/mito2/src/compaction/window.rs | 1 + src/mito2/src/config.rs | 5 + src/mito2/src/engine/flush_test.rs | 577 +++++++++++++++++- src/mito2/src/flush.rs | 2 + src/mito2/src/region/options.rs | 41 ++ src/mito2/src/worker.rs | 8 + src/mito2/src/worker/handle_alter.rs | 17 +- src/mito2/src/worker/handle_flush.rs | 165 ++++- src/mito2/src/worker/handle_write.rs | 119 +++- src/store-api/src/mito_engine_options.rs | 4 + src/store-api/src/region_request.rs | 49 +- src/table/src/metadata.rs | 3 + src/table/src/requests.rs | 4 +- tests-integration/tests/http.rs | 1 + .../common/alter/alter_table_options.result | 73 +++ .../common/alter/alter_table_options.sql | 19 + .../information_schema/region_info.result | 2 +- .../common/information_schema/region_info.sql | 2 +- .../standalone/common/show/show_create.result | 2 +- .../standalone/common/show/show_create.sql | 2 +- 25 files changed, 1087 insertions(+), 26 deletions(-) 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;