diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0b7c7645e00..0b7d36e9f2c 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -27,7 +27,7 @@ use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; -use std::time::Duration; +use std::time::{Duration, Instant}; use api::region::RegionResponse; use api::v1::meta::TopicStat; @@ -59,6 +59,7 @@ use datafusion_common::tree_node::TreeNode; use datatypes::schema::SchemaRef; use either::Either; use futures_util::future::try_join_all; +use futures_util::stream::FuturesUnordered; use futures_util::{Stream, StreamExt}; use metric_engine::engine::MetricEngine; use mito2::engine::{MITO_ENGINE_NAME, MitoEngine}; @@ -95,7 +96,7 @@ use store_api::region_request::{ RegionOpenRequest, RegionRequest, }; use store_api::storage::RegionId; -use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot}; use tokio::time::{self, timeout}; use tonic::{Request, Response, Result as TonicResult}; @@ -551,6 +552,10 @@ impl RegionServer { let requests = RegionRequest::try_from_request_body(request).context(BuildRegionRequestsSnafu)?; + let requests = match self.try_handle_mito_batch_puts(requests).await? { + Either::Left(response) => return Ok(response), + Either::Right(requests) => requests, + }; // Try to optimize batch Put requests for metric engine // Returns either Some(response) or None(requests_back) match self.try_handle_metric_batch_puts(requests).await? { @@ -704,24 +709,126 @@ impl RegionServer { .map_err(BoxedError::new) .context(HandleRegionRequestSnafu { region_id: first_region_id, + }) + .map(RegionResponse::new); + record_ingest_result(request_type, batch_size, &result); + result.map(Either::Left) + } + + /// Puts into ready mito regions together through + /// [`RegionEngine::handle_batch_put_requests`], so that each mito worker gets the puts of the + /// request to its regions at once and writes them in one batch. One ingest task awaits every + /// response at once and accounts each, for the regions it answers, as the per-region path + /// does; it hands the first error back at once and keeps accounting the rest, also when the + /// caller goes away. + /// + /// Returns Either::Right(original_requests) unless + /// [`mito_engine_for_batch_puts`](Self::mito_engine_for_batch_puts) accepts the requests. + async fn try_handle_mito_batch_puts( + &self, + requests: Vec<(RegionId, RegionRequest)>, + ) -> Result>> { + let Some(engine) = self.mito_engine_for_batch_puts(&requests) else { + return Ok(Either::Right(requests)); + }; + + let request_type = requests[0].1.request_type(); + let puts = requests + .into_iter() + .map(|(region_id, request)| match request { + RegionRequest::Put(put) => (region_id, put), + _ => unreachable!("Already checked all are Put"), + }) + .collect(); + let started = Instant::now(); + let tracing_context = TracingContext::from_current_span(); + let (result_tx, result_rx) = oneshot::channel(); + let batch = common_runtime::spawn_ingest(async move { + let mut pending = engine + .handle_batch_put_requests(puts) + .into_iter() + .map(|(region_ids, response)| { + let span = tracing_context.attach(info_span!( + "RegionServer::handle_region_request", + region_id = region_ids[0].to_string() + )); + async move { (region_ids, response.await) }.trace(span) + }) + .collect::>(); + let histogram = + crate::metrics::HANDLE_REGION_REQUEST_ELAPSED.with_label_values(&[request_type]); + let mut result_tx = Some(result_tx); + let mut affected_rows = 0; + let mut extensions = HashMap::new(); + while let Some((region_ids, result)) = pending.next().await { + let elapsed = started.elapsed().as_secs_f64(); + for _ in ®ion_ids { + histogram.observe(elapsed); + } + let result = result.context(HandleRegionRequestSnafu { + region_id: region_ids[0], + }); + record_ingest_result(request_type, region_ids.len(), &result); + match result { + Ok(response) => { + affected_rows += response.affected_rows; + extensions.extend(response.extensions); + } + Err(err) => { + if let Some(result_tx) = result_tx.take() { + let _ = result_tx.send(Err(err)); + } + } + } + } + if let Some(result_tx) = result_tx { + let _ = result_tx.send(Ok(RegionResponse { + affected_rows, + extensions, + metadata: Vec::new(), + })); + } }); - match result { - Ok(total_affected) => { - crate::metrics::REGION_CHANGED_ROW_COUNT - .with_label_values(&[request_type]) - .inc_by(total_affected as u64); - Ok(Either::Left(RegionResponse::new(total_affected))) - } - Err(err) => { - crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT - .with_label_values(&[request_type, err.status_code().as_ref()]) - .inc_by(batch_size as u64); - Err(err) + match result_rx.await { + Ok(result) => result.map(Either::Left), + // The batch task answers on every path it completes, so it panicked or was cancelled. + Err(_) => { + batch.await.context(RuntimeJoinSnafu { request_type })?; + UnexpectedSnafu { + violated: "the batch put task completed without answering", + } + .fail() } } } + /// Returns the mito engine when there are at least two requests, every one a put into a ready + /// mito region. + fn mito_engine_for_batch_puts( + &self, + requests: &[(RegionId, RegionRequest)], + ) -> Option { + if requests.len() < 2 { + return None; + } + let first_region_id = requests[0].0; + let all_puts_to_ready_mito_regions = requests.iter().all(|(region_id, request)| { + matches!(request, RegionRequest::Put(_)) + && matches!( + self.inner.region_map.get(region_id).as_deref(), + Some(RegionEngineWithStatus::Ready(engine)) if engine.name() == MITO_ENGINE_NAME + ) + }); + if !all_puts_to_ready_mito_regions { + return None; + } + match self.inner.region_map.get(&first_region_id).as_deref() { + Some(RegionEngineWithStatus::Ready(engine)) => Some(engine.clone()), + _ => None, + } + } + async fn handle_sync_region_request(&self, request: &SyncRequest) -> Result { let region_id = RegionId::from_u64(request.region_id); let manifest_info = request @@ -1698,18 +1805,15 @@ impl RegionServerInner { // Sets corresponding region status to registering/deregistering before the operation. self.set_region_status_not_ready(region_id, &engine, ®ion_change); - match engine + let result = engine .handle_request(region_id, request) .await - .with_context(|_| HandleRegionRequestSnafu { region_id }) - { + .with_context(|_| HandleRegionRequestSnafu { region_id }); + if matches!(region_change, RegionChange::Ingest) { + record_ingest_result(request_type, 1, &result); + } + match result { Ok(result) => { - // Update metrics - if matches!(region_change, RegionChange::Ingest) { - crate::metrics::REGION_CHANGED_ROW_COUNT - .with_label_values(&[request_type]) - .inc_by(result.affected_rows as u64); - } // Sets corresponding region status to ready. self.set_region_status_ready(region_id, engine.clone(), region_change) .await?; @@ -1721,11 +1825,6 @@ impl RegionServerInner { }) } Err(err) => { - if matches!(region_change, RegionChange::Ingest) { - crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT - .with_label_values(&[request_type, err.status_code().as_ref()]) - .inc(); - } // Removes the region status if the operation fails. self.unset_region_status(region_id, &engine, region_change); Err(err) @@ -2052,6 +2151,19 @@ fn record_region_open_failures(engine: &str, status_code: StatusCode, count: usi .inc_by(count as u64); } +/// Records the result of `count` ingest requests answered by one response: the rows they +/// changed, or their failures. +fn record_ingest_result(request_type: &str, count: usize, result: &Result) { + match result { + Ok(response) => crate::metrics::REGION_CHANGED_ROW_COUNT + .with_label_values(&[request_type]) + .inc_by(response.affected_rows as u64), + Err(err) => crate::metrics::REGION_SERVER_INSERT_FAIL_COUNT + .with_label_values(&[request_type, err.status_code().as_ref()]) + .inc_by(count as u64), + } +} + fn parse_region_attribute( engine: &str, options: &HashMap, @@ -2094,15 +2206,22 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use api::v1::region::{InsertRequest, InsertRequests}; use api::v1::{Rows, SemanticType}; use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; + use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; + use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions}; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; use futures_util::StreamExt; - use mito2::test_util::CreateRequestBuilder; + use log_store::ObjectStoreLogStore; + use mito2::config::MitoConfig; + use mito2::test_util::{CreateRequestBuilder, TestEnv, build_rows, rows_schema}; + use object_store::ObjectStore; + use object_store::services::Memory; use query::options::FLOW_RETURN_REGION_SEQ; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_engine::RegionEngine; @@ -3059,4 +3178,294 @@ mod tests { .await .unwrap_err(); } + + /// Returns a region server with a Mito engine on a durable object store WAL, the WAL, and the + /// WAL options of a region on it. + async fn region_server_on_object_store_wal( + env: &mut TestEnv, + config: MitoConfig, + ) -> (RegionServer, Arc, String) { + let wal_config = ObjectStoreWalConfig::default(); + let log_store = ObjectStoreLogStore::try_new( + ObjectStore::new(Memory::default()).unwrap(), + &wal_config, + 0, + STANDALONE_GENERATION, + ) + .await + .unwrap(); + let engine = env + .create_engine_with_log_store(config, log_store.clone()) + .await; + let mut region_server = mock_region_server(); + region_server.register_engine(Arc::new(engine)); + let wal_options = serde_json::to_string(&WalOptions::ObjectStore( + ObjectStoreWalOptions::new(wal_config.node_prefix(0, STANDALONE_GENERATION)), + )) + .unwrap(); + (region_server, log_store, wal_options) + } + + /// Creates a region with `wal_options` and returns an insert of three rows into it. + async fn create_region( + region_server: &RegionServer, + region_id: RegionId, + wal_options: &str, + ) -> InsertRequest { + let request = CreateRequestBuilder::new() + .insert_option(WAL_OPTIONS_KEY, wal_options) + .build(); + let rows = Rows { + schema: rows_schema(&request), + rows: build_rows(0, 3), + }; + region_server + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + InsertRequest { + region_id: region_id.as_u64(), + rows: Some(rows), + partition_expr_version: None, + skip_wal: false, + } + } + + #[tokio::test] + async fn test_mito_puts_of_one_request_share_a_wal_append() { + // One worker that takes one message at a time, so that only puts that + // arrive in one message share a WAL append. + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts").await; + let config = MitoConfig { + num_workers: 1, + worker_request_batch_size: 1, + ..Default::default() + }; + let (region_server, log_store, wal_options) = + region_server_on_object_store_wal(&mut env, config).await; + let mut requests = Vec::new(); + for region_number in 1..=2 { + let region_id = RegionId::new(1, region_number); + requests.push(create_region(®ion_server, region_id, &wal_options).await); + } + + let response = region_server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { requests })) + .await + .unwrap(); + + assert_eq!(6, response.affected_rows); + // The worker got both puts at once and appended them together. + log_store.wait_for_admitted_appends(1).await.unwrap(); + assert!( + tokio::time::timeout(Duration::ZERO, log_store.wait_for_admitted_appends(2)) + .await + .is_err() + ); + } + + #[tokio::test] + async fn test_mito_puts_to_one_region_share_a_wal_append() { + // As above: only puts that arrive in one message share a WAL append. + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-one-region").await; + let config = MitoConfig { + num_workers: 1, + worker_request_batch_size: 1, + ..Default::default() + }; + let (region_server, log_store, wal_options) = + region_server_on_object_store_wal(&mut env, config).await; + let insert = create_region(®ion_server, RegionId::new(1, 1), &wal_options).await; + let requests = vec![insert.clone(), insert]; + + let response = region_server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { requests })) + .await + .unwrap(); + + assert_eq!(6, response.affected_rows); + log_store.wait_for_admitted_appends(1).await.unwrap(); + assert!( + tokio::time::timeout(Duration::ZERO, log_store.wait_for_admitted_appends(2)) + .await + .is_err() + ); + } + + #[tokio::test] + async fn test_mito_batch_puts_fail_without_waiting_for_other_regions() { + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-error").await; + let config = MitoConfig { + num_workers: 2, + ..Default::default() + }; + let (region_server, log_store, wal_options) = + region_server_on_object_store_wal(&mut env, config).await; + // Regions 1 and 2 of table 1 belong to different workers. The put to region 1 waits for a + // WAL object whose create is held; Mito does not know region 2, so its put fails at once. + let held = RegionId::new(1, 1); + let unknown = RegionId::new(1, 2); + let held_insert = create_region(®ion_server, held, &wal_options).await; + let unknown_insert = InsertRequest { + region_id: unknown.as_u64(), + ..held_insert.clone() + }; + let mito = region_server + .inner + .region_map + .get(&held) + .unwrap() + .clone() + .into_engine(); + region_server + .inner + .region_map + .insert(unknown, RegionEngineWithStatus::Ready(mito)); + log_store.hold_creates(); + + let requests = vec![held_insert, unknown_insert]; + let err = tokio::time::timeout( + Duration::from_secs(10), + region_server.handle_requests_in_parallel(region_request::Body::Inserts( + InsertRequests { requests }, + )), + ) + .await + .unwrap() + .unwrap_err(); + + assert_eq!(StatusCode::RegionNotFound, err.status_code()); + log_store.release_creates(); + } + + #[tokio::test] + async fn test_mito_batch_puts_to_a_stopped_worker_fail_as_a_single_put_does() { + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-stopped").await; + let (region_server, _log_store, wal_options) = + region_server_on_object_store_wal(&mut env, MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let insert = create_region(®ion_server, region_id, &wal_options).await; + let mito = region_server + .inner + .region_map + .get(®ion_id) + .unwrap() + .clone() + .into_engine(); + mito.stop().await.unwrap(); + + // One put takes the per-region path and two take the batch path. + for requests in [vec![insert.clone()], vec![insert.clone(), insert]] { + let err = region_server + .handle_requests_in_parallel(region_request::Body::Inserts(InsertRequests { + requests, + })) + .await + .unwrap_err(); + assert!(format!("{err:?}").contains("is stopped"), "{err:?}"); + } + } + + #[tokio::test] + async fn test_default_batch_puts_handle_each_put_as_a_request() { + let failing = RegionId::new(1, 2); + let (engine, _receiver) = MockRegionEngine::with_mock_fn( + "mock", + Box::new(move |region_id, request| { + assert!(matches!(request, RegionRequest::Put(_))); + if region_id == failing { + return error::UnexpectedSnafu { + violated: "test".to_string(), + } + .fail(); + } + Ok(region_id.region_number() as usize) + }), + ); + let region_ids = [1, 2, 3].map(|n| RegionId::new(1, n)); + let puts = region_ids + .iter() + .map(|region_id| { + let request = RegionPutRequest { + rows: Rows { + schema: Vec::new(), + rows: Vec::new(), + }, + hint: None, + skip_wal: false, + partition_expr_version: None, + }; + (*region_id, request) + }) + .collect(); + + let (answered, responses): (Vec<_>, Vec<_>) = + engine.handle_batch_put_requests(puts).into_iter().unzip(); + let responses = futures_util::future::join_all(responses).await; + + let answered_regions = region_ids.map(|region_id| vec![region_id]); + assert_eq!(answered_regions.to_vec(), answered); + assert_eq!(1, responses[0].as_ref().unwrap().affected_rows); + assert!(responses[1].is_err()); + assert_eq!(3, responses[2].as_ref().unwrap().affected_rows); + } + + #[tokio::test] + async fn test_only_several_puts_into_ready_mito_regions_are_batched() { + let mut env = TestEnv::with_prefix("region-server-mito-batch-puts-guard").await; + let mito: RegionEngineRef = Arc::new(env.create_engine(MitoConfig::default()).await); + let (other, _receiver) = MockRegionEngine::new("mock"); + let region_server = mock_region_server(); + let [mito_a, mito_b, registering, other_region] = [1, 2, 3, 4].map(|n| RegionId::new(1, n)); + let region_map = ®ion_server.inner.region_map; + region_map.insert(mito_a, RegionEngineWithStatus::Ready(mito.clone())); + region_map.insert(mito_b, RegionEngineWithStatus::Ready(mito.clone())); + region_map.insert(registering, RegionEngineWithStatus::Registering(mito)); + region_map.insert(other_region, RegionEngineWithStatus::Ready(other)); + let rows = || Rows { + schema: Vec::new(), + rows: Vec::new(), + }; + let put = |region_id| { + let request = RegionPutRequest { + rows: rows(), + hint: None, + skip_wal: false, + partition_expr_version: None, + }; + (region_id, RegionRequest::Put(request)) + }; + let delete = |region_id| { + let request = RegionDeleteRequest { + rows: rows(), + hint: None, + partition_expr_version: None, + }; + (region_id, RegionRequest::Delete(request)) + }; + + for requests in [ + vec![put(mito_a), put(mito_b)], + vec![put(mito_a), put(mito_a)], + ] { + assert!( + region_server + .mito_engine_for_batch_puts(&requests) + .is_some() + ); + } + for requests in [ + vec![put(mito_a)], + vec![put(mito_a), delete(mito_b)], + vec![put(mito_a), put(registering)], + vec![put(mito_a), put(other_region)], + vec![put(other_region), put(mito_a)], + ] { + assert!( + region_server + .mito_engine_for_batch_puts(&requests) + .is_none() + ); + } + } } diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index 08600cd00e4..c41aa649111 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -38,7 +38,8 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. ## Write path -`MitoEngine::handle_request` (`engine.rs`) → worker loop +`MitoEngine::handle_request` (`engine.rs`), or `handle_batch_put_requests` for the puts +of one request to several regions (one message per worker) → worker loop (`worker/handle_write.rs`) → sequence + WAL assembly (`region_write_ctx.rs`) → `wal.rs` → memtable (`memtable/`) → when buffer pressure trips, `flush.rs` writes SSTs via `access_layer.rs` and appends a `RegionEdit` to the manifest diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index cdfc78ff0c4..715357a39b3 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -108,7 +108,8 @@ use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream}; use common_stat::get_total_memory_bytes; use common_telemetry::{debug, info, tracing, warn}; use common_wal::options::WalOptions; -use futures::future::{join_all, try_join_all}; +use futures::FutureExt; +use futures::future::{self, join_all, try_join_all}; use futures::stream::{self, Stream, StreamExt}; use object_store::manager::ObjectStoreManagerRef; use region_hook::RegionHookRef; @@ -122,14 +123,14 @@ use store_api::metric_engine_consts::{ MANIFEST_INFO_EXTENSION_KEY, TABLE_COLUMN_METADATA_EXTENSION_KEY, }; use store_api::region_engine::{ - BatchResponses, MitoCopyRegionFromRequest, MitoCopyRegionFromResponse, RegionEngine, - RegionManifestInfo, RegionRole, RegionScannerRef, RegionStatistic, RemapManifestsRequest, - RemapManifestsResponse, SetRegionRoleStateResponse, SettableRegionRoleState, - SyncRegionFromRequest, SyncRegionFromResponse, + BatchPutResponse, BatchResponses, MitoCopyRegionFromRequest, MitoCopyRegionFromResponse, + RegionEngine, RegionManifestInfo, RegionRole, RegionScannerRef, RegionStatistic, + RemapManifestsRequest, RemapManifestsResponse, SetRegionRoleStateResponse, + SettableRegionRoleState, SyncRegionFromRequest, SyncRegionFromResponse, }; use store_api::region_info::RegionInfoEntry; use store_api::region_request::{ - AffectedRows, RegionCatchupRequest, RegionOpenRequest, RegionRequest, + AffectedRows, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, }; use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry}; use store_api::storage::{FileId, FileRefsManifest, RegionId, ScanRequest, SequenceNumber}; @@ -1407,6 +1408,46 @@ impl RegionEngine for MitoEngine { .map_err(BoxedError::new) } + /// Sends the puts to regions of the same worker in one message, so that the worker writes + /// them in one batch and one WAL append. The messages are sent on their own task, so every + /// response can be awaited at once. + fn handle_batch_put_requests( + self: Arc, + requests: Vec<(RegionId, RegionPutRequest)>, + ) -> Vec { + let mut writes = Vec::with_capacity(requests.len()); + let mut responses = Vec::with_capacity(requests.len()); + for (region_id, request) in requests { + let request = RegionRequest::Put(request); + let timer = HANDLE_REQUEST_ELAPSED + .with_label_values(&[request.request_type()]) + .start_timer(); + let region_metadata = self.inner.get_metadata(region_id).ok(); + let response = + match WorkerRequest::try_from_region_request(region_id, request, region_metadata) { + Ok((WorkerRequest::Write(write), receiver)) => { + writes.push(write); + async move { + let _timer = timer; + receiver + .await + .context(RecvSnafu) + .flatten() + .map(RegionResponse::new) + .map_err(BoxedError::new) + } + .boxed() + } + Ok(_) => unreachable!("a put converts to a write"), + Err(e) => future::ready(Err(BoxedError::new(e))).boxed(), + }; + responses.push((vec![region_id], response)); + } + let inner = self.inner.clone(); + common_runtime::spawn_ingest(async move { inner.workers.submit_writes(writes).await }); + responses + } + #[tracing::instrument(skip_all)] async fn handle_request( &self, diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 4cf475c7278..d084fc82c7e 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -758,6 +758,9 @@ pub(crate) enum WorkerRequest { /// Write to a region. Write(SenderWriteRequest), + /// Writes of one request to regions of the same worker. + BatchWrite(Vec), + /// Ddl request to a region. Ddl(SenderDdlRequest), diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 763a068e106..556f4a0757f 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -44,12 +44,12 @@ use common_meta::key::SchemaMetadataManagerRef; use common_runtime::JoinHandle; use common_stat::get_total_memory_bytes; use common_telemetry::{error, info, warn}; -use futures::future::try_join_all; +use futures::future::{join_all, try_join_all}; use object_store::ObjectStore; use object_store::manager::ObjectStoreManagerRef; use prometheus::{Histogram, IntGauge}; use rand::{Rng, rng}; -use snafu::{ResultExt, ensure}; +use snafu::ResultExt; use store_api::logstore::LogStore; use store_api::region_engine::{ RegionRole, SetRegionRoleStateResponse, SetRegionRoleStateSuccess, SettableRegionRoleState, @@ -318,6 +318,40 @@ impl WorkerGroup { self.worker(region_id).submit_request(request).await } + /// Submits the writes of one request, those of each worker in one message, so that the + /// worker handles them in one batch. The messages are sent concurrently, so a full worker + /// queue holds back no other worker. A write whose worker has stopped receives the error a + /// write submitted on its own gets. + pub(crate) async fn submit_writes(&self, writes: Vec) { + let mut writes_by_worker: HashMap> = HashMap::new(); + for write in writes { + let index = region_id_to_index(write.request.region_id, self.workers.len()); + writes_by_worker.entry(index).or_default().push(write); + } + join_all( + writes_by_worker + .into_iter() + .map(|(index, writes)| async move { + let worker = &self.workers[index]; + let Err(request) = worker + .try_submit_request(WorkerRequest::BatchWrite(writes)) + .await + else { + return; + }; + let WorkerRequest::BatchWrite(writes) = request else { + unreachable!("the worker returns the request it was given"); + }; + for write in writes { + write + .sender + .send(WorkerStoppedSnafu { id: worker.id }.fail()); + } + }), + ) + .await; + } + /// Returns true if the specific region exists. pub(crate) fn is_region_exists(&self, region_id: RegionId) -> bool { self.worker(region_id).is_region_exists(region_id) @@ -738,9 +772,23 @@ pub(crate) struct RegionWorker { impl RegionWorker { /// Submits request to background worker thread. async fn submit_request(&self, request: WorkerRequest) -> Result<()> { - ensure!(self.is_running(), WorkerStoppedSnafu { id: self.id }); + self.try_submit_request(request) + .await + .map_err(|_| WorkerStoppedSnafu { id: self.id }.build()) + } + + /// Submits request to background worker thread, or returns it if the worker is stopped. + async fn try_submit_request( + &self, + request: WorkerRequest, + ) -> std::result::Result<(), WorkerRequest> { + if !self.is_running() { + return Err(request); + } let request_with_time = WorkerRequestWithTime::new(request); - if self.sender.send(request_with_time).await.is_err() { + if let Err(mpsc::error::SendError(request_with_time)) = + self.sender.send(request_with_time).await + { warn!( "Worker {} is already exited but the running flag is still true", self.id @@ -750,7 +798,7 @@ impl RegionWorker { if let Some(state) = &self.series_index_task_state { state.stop(); } - return WorkerStoppedSnafu { id: self.id }.fail(); + return Err(request_with_time.request); } Ok(()) @@ -1057,6 +1105,7 @@ impl RegionWorkerLoop { match request_with_time.request { WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req), + WorkerRequest::BatchWrite(sender_reqs) => write_req_buffer.extend(sender_reqs), WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req), WorkerRequest::BulkInserts(bulk_insert) => { bulk_insert_req_num += 1; @@ -1110,6 +1159,9 @@ impl RegionWorkerLoop { match request_with_time.request { WorkerRequest::Write(sender_req) => write_req_buffer.push(sender_req), + WorkerRequest::BatchWrite(sender_reqs) => { + write_req_buffer.extend(sender_reqs) + } WorkerRequest::Ddl(sender_req) => ddl_req_buffer.push(sender_req), WorkerRequest::BulkInserts(bulk_insert) => { bulk_insert_req_num += 1; @@ -1184,7 +1236,7 @@ impl RegionWorkerLoop { ) { for worker_req in general_requests.drain(..) { match worker_req { - WorkerRequest::Write(_) | WorkerRequest::Ddl(_) => { + WorkerRequest::Write(_) | WorkerRequest::BatchWrite(_) | WorkerRequest::Ddl(_) => { // These requests are categorized before dispatching general requests. continue; } diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs index 36b8a471394..c95bf6d3b21 100644 --- a/src/store-api/src/region_engine.rs +++ b/src/store-api/src/region_engine.rs @@ -30,14 +30,15 @@ use datafusion_execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PhysicalExpr}; use datatypes::schema::SchemaRef; -use futures::future::join_all; +use futures::FutureExt; +use futures::future::{BoxFuture, join_all}; use serde::{Deserialize, Serialize}; use tokio::sync::Semaphore; use crate::logstore::entry; use crate::metadata::RegionMetadataRef; use crate::region_request::{ - BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionRequest, + BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest, RegionRequest, }; use crate::storage::{FileId, RegionId, ScanRequest, SequenceNumber}; @@ -548,6 +549,13 @@ pub type RegionScannerRef = Box; pub type BatchResponses = Vec<(RegionId, Result)>; +/// A response of [`RegionEngine::handle_batch_put_requests`] and the region of each put it +/// answers for. +pub type BatchPutResponse = ( + Vec, + BoxFuture<'static, Result>, +); + /// Represents the statistics of a region. #[derive(Debug, Deserialize, Serialize, Default)] pub struct RegionStatistic { @@ -981,6 +989,35 @@ pub trait RegionEngine: Send + Sync { Ok(join_all(tasks).await) } + /// Handles the puts of one request into regions of this engine and returns the responses + /// to await on their own, each with the region of every put it answers for: two puts to one + /// region are two entries. + /// + /// The default handles each put with [`handle_request`](Self::handle_request) and answers + /// each put on its own. An engine that writes several puts together may answer them with + /// one response whose affected rows count all of them. + fn handle_batch_put_requests( + self: Arc, + requests: Vec<(RegionId, RegionPutRequest)>, + ) -> Vec + where + Self: 'static, + { + requests + .into_iter() + .map(|(region_id, request)| { + let engine = self.clone(); + let response = async move { + engine + .handle_request(region_id, RegionRequest::Put(request)) + .await + } + .boxed(); + (vec![region_id], response) + }) + .collect() + } + async fn handle_batch_ddl_requests( &self, request: BatchRegionDdlRequest,