diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 736c2a41a2..ef1e8f3cff 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -83,6 +83,11 @@ impl MitoEngine { fn handle_query(&self, region_id: RegionId, request: ScanRequest) -> Result { self.inner.handle_query(region_id, request) } + + #[cfg(test)] + pub(crate) fn get_region(&self, id: RegionId) -> Option { + self.inner.workers.get_region(id) + } } /// Inner struct of [MitoEngine]. diff --git a/src/mito2/src/engine/tests.rs b/src/mito2/src/engine/tests.rs index 2beb978740..d5da4f43ae 100644 --- a/src/mito2/src/engine/tests.rs +++ b/src/mito2/src/engine/tests.rs @@ -26,11 +26,12 @@ use store_api::storage::RegionId; use super::*; use crate::error::Error; +use crate::region::version::VersionControlData; use crate::test_util::{CreateRequestBuilder, TestEnv}; #[tokio::test] async fn test_engine_new_stop() { - let env = TestEnv::with_prefix("engine-stop"); + let mut env = TestEnv::with_prefix("engine-stop"); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -57,7 +58,7 @@ async fn test_engine_new_stop() { #[tokio::test] async fn test_engine_create_new_region() { - let env = TestEnv::with_prefix("new-region"); + let mut env = TestEnv::with_prefix("new-region"); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -72,7 +73,7 @@ async fn test_engine_create_new_region() { #[tokio::test] async fn test_engine_create_region_if_not_exists() { - let env = TestEnv::with_prefix("create-not-exists"); + let mut env = TestEnv::with_prefix("create-not-exists"); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -91,7 +92,7 @@ async fn test_engine_create_region_if_not_exists() { #[tokio::test] async fn test_engine_create_existing_region() { - let env = TestEnv::with_prefix("create-existing"); + let mut env = TestEnv::with_prefix("create-existing"); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -114,7 +115,7 @@ async fn test_engine_create_existing_region() { #[tokio::test] async fn test_engine_open_empty() { - let env = TestEnv::with_prefix("open-empty"); + let mut env = TestEnv::with_prefix("open-empty"); let engine = env.create_engine(MitoConfig::default()).await; let err = engine @@ -136,7 +137,7 @@ async fn test_engine_open_empty() { #[tokio::test] async fn test_engine_open_existing() { - let env = TestEnv::with_prefix("open-exiting"); + let mut env = TestEnv::with_prefix("open-exiting"); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -162,7 +163,7 @@ async fn test_engine_open_existing() { #[tokio::test] async fn test_engine_close_region() { - let env = TestEnv::with_prefix("close"); + let mut env = TestEnv::with_prefix("close"); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -194,7 +195,7 @@ async fn test_engine_close_region() { #[tokio::test] async fn test_engine_reopen_region() { - let env = TestEnv::with_prefix("reopen-region"); + let mut env = TestEnv::with_prefix("reopen-region"); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -236,8 +237,8 @@ fn column_metadata_to_column_schema(metadata: &ColumnMetadata) -> api::v1::Colum } } -fn build_rows(num_rows: usize) -> Vec { - (0..num_rows) +fn build_rows(start: usize, end: usize) -> Vec { + (start..end) .map(|i| api::v1::Row { values: vec![ api::v1::Value { @@ -256,7 +257,7 @@ fn build_rows(num_rows: usize) -> Vec { #[tokio::test] async fn test_write_to_region() { - let env = TestEnv::with_prefix("write-to-region"); + let mut env = TestEnv::with_prefix("write-to-region"); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -275,7 +276,7 @@ async fn test_write_to_region() { let num_rows = 42; let rows = Rows { schema: column_schemas, - rows: build_rows(num_rows), + rows: build_rows(0, num_rows), }; let output = engine .handle_request(region_id, RegionRequest::Put(RegionPutRequest { rows })) @@ -287,11 +288,100 @@ async fn test_write_to_region() { assert_eq!(num_rows, rows_inserted); } +#[tokio::test] +async fn test_region_replay() { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::with_prefix("region-replay"); + let engine = env.create_engine(MitoConfig::default()).await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let region_dir = request.region_dir.clone(); + + let column_schemas = request + .column_metadatas + .iter() + .map(column_metadata_to_column_schema) + .collect::>(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 20), + }; + let output = engine + .handle_request(region_id, RegionRequest::Put(RegionPutRequest { rows })) + .await + .unwrap(); + let Output::AffectedRows(rows_inserted) = output else { + unreachable!() + }; + assert_eq!(20, rows_inserted); + + let rows = Rows { + schema: column_schemas, + rows: build_rows(20, 42), + }; + let output = engine + .handle_request(region_id, RegionRequest::Put(RegionPutRequest { rows })) + .await + .unwrap(); + let Output::AffectedRows(rows_inserted) = output else { + unreachable!() + }; + assert_eq!(22, rows_inserted); + + engine.stop().await.unwrap(); + + let engine = MitoEngine::new( + MitoConfig::default(), + env.get_logstore().unwrap(), + env.get_object_store().unwrap(), + ); + + let open_region = engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + region_dir, + options: HashMap::default(), + }), + ) + .await + .unwrap(); + let Output::AffectedRows(rows) = open_region else { + unreachable!() + }; + assert_eq!(0, rows); + + let request = ScanRequest::default(); + let scanner = engine.handle_query(region_id, request).unwrap(); + let stream = scanner.scan().await.unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(42, batches.iter().map(|b| b.num_rows()).sum::()); + + let region = engine.get_region(region_id).unwrap(); + let VersionControlData { + committed_sequence, + last_entry_id, + .. + } = region.version_control.current(); + + assert_eq!(42, committed_sequence); + assert_eq!(2, last_entry_id); + + engine.stop().await.unwrap(); +} + // TODO(yingwen): build_rows() only generate one point for each series. We need to add tests // for series with multiple points and other cases. #[tokio::test] async fn test_write_query_region() { - let env = TestEnv::new(); + let mut env = TestEnv::new(); let engine = env.create_engine(MitoConfig::default()).await; let region_id = RegionId::new(1, 1); @@ -309,7 +399,7 @@ async fn test_write_query_region() { let rows = Rows { schema: column_schemas, - rows: build_rows(3), + rows: build_rows(0, 3), }; engine .handle_request(region_id, RegionRequest::Put(RegionPutRequest { rows })) diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index fcb6556ebc..afe5b3325e 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -33,6 +33,7 @@ pub mod memtable; pub mod read; #[allow(dead_code)] mod region; +mod region_write_ctx; #[allow(dead_code)] pub mod request; #[allow(dead_code)] diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 996bd147e5..64187f9eea 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -16,9 +16,12 @@ use std::sync::Arc; +use common_telemetry::info; +use futures::StreamExt; use object_store::util::join_dir; use object_store::ObjectStore; use snafu::{ensure, OptionExt}; +use store_api::logstore::LogStore; use store_api::metadata::RegionMetadata; use store_api::storage::RegionId; @@ -26,8 +29,10 @@ use crate::config::MitoConfig; use crate::error::{RegionCorruptedSnafu, RegionNotFoundSnafu, Result}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; use crate::memtable::MemtableBuilderRef; -use crate::region::version::{VersionBuilder, VersionControl}; +use crate::region::version::{VersionBuilder, VersionControl, VersionControlRef}; use crate::region::MitoRegion; +use crate::region_write_ctx::RegionWriteCtx; +use crate::wal::{EntryId, Wal}; /// Builder to create a new [MitoRegion] or open an existing one. pub(crate) struct RegionOpener { @@ -100,7 +105,11 @@ impl RegionOpener { /// Opens an existing region. /// /// Returns error if the region doesn't exist. - pub(crate) async fn open(self, config: &MitoConfig) -> Result { + pub(crate) async fn open( + self, + config: &MitoConfig, + wal: &Wal, + ) -> Result { let options = RegionManifestOptions { manifest_dir: new_manifest_dir(&self.region_dir), object_store: self.object_store, @@ -125,21 +134,58 @@ impl RegionOpener { } ); + let region_id = metadata.region_id; let mutable = self.memtable_builder.build(&metadata); let version = VersionBuilder::new(metadata, mutable).build(); + let flushed_sequence = version.flushed_entry_id; let version_control = Arc::new(VersionControl::new(version)); + replay_memtable(wal, region_id, flushed_sequence, &version_control).await?; - // TODO(yingwen): Replay. - - Ok(MitoRegion { + let region = MitoRegion { region_id: self.region_id, version_control, region_dir: self.region_dir, manifest_manager, - }) + }; + Ok(region) } } +/// Replays the mutations from WAL and inserts mutations to memtable of given region. +async fn replay_memtable( + wal: &Wal, + region_id: RegionId, + flushed_entry_id: EntryId, + version_control: &VersionControlRef, +) -> Result<()> { + let mut rows_replayed = 0; + let mut last_entry_id = EntryId::MIN; + let mut region_write_ctx = RegionWriteCtx::new(region_id, version_control); + let mut wal_stream = wal.scan(region_id, flushed_entry_id)?; + while let Some(res) = wal_stream.next().await { + let (entry_id, entry) = res?; + last_entry_id = last_entry_id.max(entry_id); + for mutation in entry.mutations { + rows_replayed += mutation + .rows + .as_ref() + .map(|rows| rows.rows.len()) + .unwrap_or(0); + region_write_ctx.push_mutation(mutation.op_type, mutation.rows, None); + } + } + + // set next_entry_id and write to memtable. + region_write_ctx.set_next_entry_id(last_entry_id + 1); + region_write_ctx.write_memtable(); + + info!( + "Replay WAL for region: {}, rows recovered: {}, last entry id: {}", + region_id, rows_replayed, last_entry_id + ); + Ok(()) +} + /// Returns the directory to the manifest files. fn new_manifest_dir(region_dir: &str) -> String { join_dir(region_dir, "manifest") diff --git a/src/mito2/src/region/version.rs b/src/mito2/src/region/version.rs index 8f0ff5a3b0..44d4a6d27b 100644 --- a/src/mito2/src/region/version.rs +++ b/src/mito2/src/region/version.rs @@ -94,8 +94,8 @@ pub(crate) struct Version { pub(crate) memtables: MemtableVersionRef, /// SSTs of the region. pub(crate) ssts: SstVersionRef, - /// Inclusive max sequence of flushed data. - pub(crate) flushed_sequence: SequenceNumber, + /// Inclusive max WAL entry id of flushed data. + pub(crate) flushed_entry_id: EntryId, // TODO(yingwen): RegionOptions. } @@ -120,7 +120,7 @@ impl VersionBuilder { metadata: self.metadata, memtables: Arc::new(MemtableVersion::new(self.mutable)), ssts: Arc::new(SstVersion::new()), - flushed_sequence: 0, + flushed_entry_id: 0, } } } diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs new file mode 100644 index 0000000000..b531f9f2ce --- /dev/null +++ b/src/mito2/src/region_write_ctx.rs @@ -0,0 +1,190 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::mem; +use std::sync::Arc; + +use api::v1::{Mutation, Rows, WalEntry}; +use common_query::Output; +use snafu::ResultExt; +use store_api::logstore::LogStore; +use store_api::storage::{RegionId, SequenceNumber}; +use tokio::sync::oneshot::Sender; + +use crate::error::{Error, Result, WriteGroupSnafu}; +use crate::memtable::KeyValues; +use crate::region::version::{VersionControlData, VersionControlRef, VersionRef}; +use crate::wal::{EntryId, WalWriter}; + +/// Context to keep region metadata and buffer write requests. + +/// Notifier to notify write result on drop. +struct WriteNotify { + /// Error to send to the waiter. + err: Option>, + /// Sender to send write result to the waiter for this mutation. + sender: Option>>, + /// Number of rows to be written. + num_rows: usize, +} + +impl WriteNotify { + /// Creates a new notify from the `sender`. + fn new(sender: Option>>, num_rows: usize) -> WriteNotify { + WriteNotify { + err: None, + sender, + num_rows, + } + } + + /// Send result to the waiter. + fn notify_result(&mut self) { + let Some(sender) = self.sender.take() else { + return; + }; + if let Some(err) = &self.err { + // Try to send the error to waiters. + let _ = sender.send(Err(err.clone()).context(WriteGroupSnafu)); + } else { + // Send success result. + let _ = sender.send(Ok(Output::AffectedRows(self.num_rows))); + } + } +} + +impl Drop for WriteNotify { + fn drop(&mut self) { + self.notify_result(); + } +} + +/// Context to keep region metadata and buffer write requests. +pub(crate) struct RegionWriteCtx { + /// Id of region to write. + region_id: RegionId, + /// Version of the region while creating the context. + version: VersionRef, + /// VersionControl of the region. + version_control: VersionControlRef, + /// Next sequence number to write. + /// + /// The context assigns a unique sequence number for each row. + next_sequence: SequenceNumber, + /// Next entry id of WAL to write. + next_entry_id: EntryId, + /// Valid WAL entry to write. + /// + /// We keep [WalEntry] instead of mutations to avoid taking mutations + /// out of the context to construct the wal entry when we write to the wal. + wal_entry: WalEntry, + /// Notifiers to send write results to waiters. + /// + /// The i-th notify is for i-th mutation. + notifiers: Vec, +} + +impl RegionWriteCtx { + /// Returns an empty context. + pub(crate) fn new(region_id: RegionId, version_control: &VersionControlRef) -> RegionWriteCtx { + let VersionControlData { + version, + committed_sequence, + last_entry_id, + } = version_control.current(); + + RegionWriteCtx { + region_id, + version, + version_control: version_control.clone(), + next_sequence: committed_sequence + 1, + next_entry_id: last_entry_id + 1, + wal_entry: WalEntry::default(), + notifiers: Vec::new(), + } + } + + /// Push [SenderWriteRequest] to the context. + pub(crate) fn push_mutation( + &mut self, + op_type: i32, + rows: Option, + tx: Option>>, + ) { + let num_rows = rows.as_ref().map(|rows| rows.rows.len()).unwrap_or(0); + self.wal_entry.mutations.push(Mutation { + op_type, + sequence: self.next_sequence, + rows, + }); + + let notify = WriteNotify::new(tx, num_rows); + // Notifiers are 1:1 map to mutations. + self.notifiers.push(notify); + + // Increase sequence number. + self.next_sequence += num_rows as u64; + } + + /// Encode and add WAL entry to the writer. + pub(crate) fn add_wal_entry( + &mut self, + wal_writer: &mut WalWriter, + ) -> Result<()> { + wal_writer.add_entry(self.region_id, self.next_entry_id, &self.wal_entry)?; + // We only call this method one time, but we still bump next entry id for consistency. + self.next_entry_id += 1; + Ok(()) + } + + pub(crate) fn version(&self) -> &VersionRef { + &self.version + } + + /// Sets error and marks all write operations are failed. + pub(crate) fn set_error(&mut self, err: Arc) { + // Set error for all notifiers + for notify in &mut self.notifiers { + notify.err = Some(err.clone()); + } + } + + /// Updates next entry id. + pub(crate) fn set_next_entry_id(&mut self, next_entry_id: EntryId) { + self.next_entry_id = next_entry_id + } + + /// Consumes mutations and writes them into mutable memtable. + pub(crate) fn write_memtable(&mut self) { + debug_assert_eq!(self.notifiers.len(), self.wal_entry.mutations.len()); + + let mutable = &self.version.memtables.mutable; + // Takes mutations from the wal entry. + let mutations = mem::take(&mut self.wal_entry.mutations); + for (mutation, notify) in mutations.into_iter().zip(&mut self.notifiers) { + // Write mutation to the memtable. + let Some(kvs) = KeyValues::new(&self.version.metadata, mutation) else { + continue; + }; + if let Err(e) = mutable.write(&kvs) { + notify.err = Some(Arc::new(e)); + } + } + + // Updates region sequence and entry id. Since we stores last sequence and entry id in region, we need + // to decrease `next_sequence` and `next_entry_id` by 1. + self.version_control + .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1); + } +} diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index f3b1a177e7..37d3f8410d 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -43,7 +43,8 @@ use crate::worker::WorkerGroup; pub struct TestEnv { /// Path to store data. data_home: TempDir, - // TODO(yingwen): Maybe provide a way to close the log store. + logstore: Option>, + object_store: Option, } impl Default for TestEnv { @@ -57,6 +58,8 @@ impl TestEnv { pub fn new() -> TestEnv { TestEnv { data_home: create_temp_dir(""), + logstore: None, + object_store: None, } } @@ -64,14 +67,27 @@ impl TestEnv { pub fn with_prefix(prefix: &str) -> TestEnv { TestEnv { data_home: create_temp_dir(prefix), + logstore: None, + object_store: None, } } + pub fn get_logstore(&self) -> Option> { + self.logstore.clone() + } + + pub fn get_object_store(&self) -> Option { + self.object_store.clone() + } + /// Creates a new engine with specific config under this env. - pub async fn create_engine(&self, config: MitoConfig) -> MitoEngine { + pub async fn create_engine(&mut self, config: MitoConfig) -> MitoEngine { let (log_store, object_store) = self.create_log_and_object_store().await; - MitoEngine::new(config, Arc::new(log_store), object_store) + let logstore = Arc::new(log_store); + self.logstore = Some(logstore.clone()); + self.object_store = Some(object_store.clone()); + MitoEngine::new(config, logstore, object_store) } /// Creates a new [WorkerGroup] with specific config under this env. diff --git a/src/mito2/src/wal.rs b/src/mito2/src/wal.rs index abad8055dd..010630653e 100644 --- a/src/mito2/src/wal.rs +++ b/src/mito2/src/wal.rs @@ -40,7 +40,7 @@ pub type WalEntryStream<'a> = BoxStream<'a, Result<(EntryId, WalEntry)>>; /// Write ahead log. /// /// All regions in the engine shares the same WAL instance. -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct Wal { /// The underlying log store. store: Arc, diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index cb3b988dbe..88e7f35319 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -365,9 +365,7 @@ impl RegionWorkerLoop { self.handle_ddl_requests(ddl_requests).await; } -} -impl RegionWorkerLoop { /// Takes and handles all ddl requests. async fn handle_ddl_requests(&mut self, ddl_tasks: Vec) { if ddl_tasks.is_empty() { @@ -392,7 +390,9 @@ impl RegionWorkerLoop { } } } +} +impl RegionWorkerLoop { // Clean up the worker. async fn clean(&self) { // Closes remaining regions. diff --git a/src/mito2/src/worker/handle_create.rs b/src/mito2/src/worker/handle_create.rs index d46f0138b6..3c11fe1cbc 100644 --- a/src/mito2/src/worker/handle_create.rs +++ b/src/mito2/src/worker/handle_create.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use common_query::Output; use common_telemetry::info; use snafu::{ensure, ResultExt}; +use store_api::logstore::LogStore; use store_api::metadata::RegionMetadataBuilder; use store_api::region_request::RegionCreateRequest; use store_api::storage::RegionId; @@ -27,7 +28,7 @@ use crate::error::{InvalidMetadataSnafu, RegionExistsSnafu, Result}; use crate::region::opener::RegionOpener; use crate::worker::RegionWorkerLoop; -impl RegionWorkerLoop { +impl RegionWorkerLoop { pub(crate) async fn handle_create_request( &mut self, region_id: RegionId, diff --git a/src/mito2/src/worker/handle_open.rs b/src/mito2/src/worker/handle_open.rs index 7848cabcee..c6b6c73fec 100644 --- a/src/mito2/src/worker/handle_open.rs +++ b/src/mito2/src/worker/handle_open.rs @@ -18,6 +18,7 @@ use std::sync::Arc; use common_query::Output; use common_telemetry::info; +use store_api::logstore::LogStore; use store_api::region_request::RegionOpenRequest; use store_api::storage::RegionId; @@ -25,7 +26,7 @@ use crate::error::Result; use crate::region::opener::RegionOpener; use crate::worker::RegionWorkerLoop; -impl RegionWorkerLoop { +impl RegionWorkerLoop { pub(crate) async fn handle_open_request( &mut self, region_id: RegionId, @@ -44,7 +45,7 @@ impl RegionWorkerLoop { self.object_store.clone(), ) .region_dir(&request.region_dir) - .open(&self.config) + .open(&self.config, &self.wal) .await?; info!("Region {} is opened", region_id); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index af4f04a4cd..c95d4316fe 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -15,23 +15,17 @@ //! Handling write requests. use std::collections::{hash_map, HashMap}; -use std::mem; use std::sync::Arc; -use api::v1::{Mutation, WalEntry}; use common_query::Output; -use snafu::ResultExt; use store_api::logstore::LogStore; use store_api::metadata::RegionMetadata; -use store_api::storage::{RegionId, SequenceNumber}; +use store_api::storage::RegionId; use tokio::sync::oneshot::Sender; -use crate::error::{Error, RegionNotFoundSnafu, Result, WriteGroupSnafu}; -use crate::memtable::KeyValues; -use crate::region::version::{VersionControlData, VersionRef}; -use crate::region::MitoRegionRef; +use crate::error::{RegionNotFoundSnafu, Result}; +use crate::region_write_ctx::RegionWriteCtx; use crate::request::{SenderWriteRequest, WriteRequest}; -use crate::wal::{EntryId, WalWriter}; use crate::worker::RegionWorkerLoop; impl RegionWorkerLoop { @@ -84,7 +78,10 @@ impl RegionWorkerLoop { }; // Initialize the context. - e.insert(RegionWriteCtx::new(region)); + e.insert(RegionWriteCtx::new( + region.region_id, + ®ion.version_control, + )); } // Safety: Now we ensure the region exists. @@ -92,7 +89,7 @@ impl RegionWorkerLoop { // Checks whether request schema is compatible with region schema. if let Err(e) = - maybe_fill_missing_columns(&mut sender_req.request, ®ion_ctx.version.metadata) + maybe_fill_missing_columns(&mut sender_req.request, ®ion_ctx.version().metadata) { send_result(sender_req.sender, Err(e)); @@ -100,7 +97,11 @@ impl RegionWorkerLoop { } // Collect requests by region. - region_ctx.push_sender_request(sender_req); + region_ctx.push_mutation( + sender_req.request.op_type as i32, + Some(sender_req.request.rows), + sender_req.sender, + ); } region_ctxs @@ -130,143 +131,3 @@ fn send_result(sender: Option>>, res: Result) { let _ = sender.send(res); } } - -/// Notifier to notify write result on drop. -struct WriteNotify { - /// Error to send to the waiter. - err: Option>, - /// Sender to send write result to the waiter for this mutation. - sender: Option>>, - /// Number of rows to be written. - num_rows: usize, -} - -impl WriteNotify { - /// Creates a new notify from the `sender`. - fn new(sender: Option>>, num_rows: usize) -> WriteNotify { - WriteNotify { - err: None, - sender, - num_rows, - } - } - - /// Send result to the waiter. - fn notify_result(&mut self) { - let Some(sender) = self.sender.take() else { - return; - }; - if let Some(err) = &self.err { - // Try to send the error to waiters. - let _ = sender.send(Err(err.clone()).context(WriteGroupSnafu)); - } else { - // Send success result. - let _ = sender.send(Ok(Output::AffectedRows(self.num_rows))); - } - } -} - -impl Drop for WriteNotify { - fn drop(&mut self) { - self.notify_result(); - } -} - -/// Context to keep region metadata and buffer write requests. -struct RegionWriteCtx { - /// Region to write. - region: MitoRegionRef, - /// Version of the region while creating the context. - version: VersionRef, - /// Next sequence number to write. - /// - /// The context assigns a unique sequence number for each row. - next_sequence: SequenceNumber, - /// Next entry id of WAL to write. - next_entry_id: EntryId, - /// Valid WAL entry to write. - /// - /// We keep [WalEntry] instead of mutations to avoid taking mutations - /// out of the context to construct the wal entry when we write to the wal. - wal_entry: WalEntry, - /// Notifiers to send write results to waiters. - /// - /// The i-th notify is for i-th mutation. - notifiers: Vec, -} - -impl RegionWriteCtx { - /// Returns an empty context. - fn new(region: MitoRegionRef) -> RegionWriteCtx { - let VersionControlData { - version, - committed_sequence, - last_entry_id, - } = region.version_control.current(); - RegionWriteCtx { - region, - version, - next_sequence: committed_sequence + 1, - next_entry_id: last_entry_id + 1, - wal_entry: WalEntry::default(), - notifiers: Vec::new(), - } - } - - /// Push [SenderWriteRequest] to the context. - fn push_sender_request(&mut self, sender_req: SenderWriteRequest) { - let num_rows = sender_req.request.rows.rows.len(); - - self.wal_entry.mutations.push(Mutation { - op_type: sender_req.request.op_type as i32, - sequence: self.next_sequence, - rows: Some(sender_req.request.rows), - }); - // Notifiers are 1:1 map to mutations. - self.notifiers - .push(WriteNotify::new(sender_req.sender, num_rows)); - - // Increase sequence number. - self.next_sequence += num_rows as u64; - } - - /// Encode and add WAL entry to the writer. - fn add_wal_entry(&mut self, wal_writer: &mut WalWriter) -> Result<()> { - wal_writer.add_entry(self.region.region_id, self.next_entry_id, &self.wal_entry)?; - // We only call this method one time, but we still bump next entry id for consistency. - self.next_entry_id += 1; - Ok(()) - } - - /// Sets error and marks all write operations are failed. - fn set_error(&mut self, err: Arc) { - // Set error for all notifiers - for notify in &mut self.notifiers { - notify.err = Some(err.clone()); - } - } - - /// Consumes mutations and writes them into mutable memtable. - fn write_memtable(&mut self) { - debug_assert_eq!(self.notifiers.len(), self.wal_entry.mutations.len()); - - let mutable = &self.version.memtables.mutable; - // Takes mutations from the wal entry. - let mutations = mem::take(&mut self.wal_entry.mutations); - for (mutation, notify) in mutations.into_iter().zip(&mut self.notifiers) { - // Write mutation to the memtable. - let Some(kvs) = KeyValues::new(&self.version.metadata, mutation) else { - continue; - }; - if let Err(e) = mutable.write(&kvs) { - notify.err = Some(Arc::new(e)); - } - } - - // Updates region sequence and entry id. Since we stores last sequence and entry id in region, we need - // to decrease `next_sequence` and `next_entry_id` by 1. - self.region - .version_control - .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1); - } -}