feat(mito2): WAL replay (#2264)

* feat: replay memtable when opening table

* test: region replay

* refactor: save logstore in TestEnv

* fix: some cr comments

* chore: rebase develop

* chore: update last entry id during replay

src/mito2/src/engine.rs

@@ -83,6 +83,11 @@ impl MitoEngine {
     fn handle_query(&self, region_id: RegionId, request: ScanRequest) -> Result<Scanner> {
         self.inner.handle_query(region_id, request)
     }
+
+    #[cfg(test)]
+    pub(crate) fn get_region(&self, id: RegionId) -> Option<crate::region::MitoRegionRef> {
+        self.inner.workers.get_region(id)
+    }
 }
 
 /// Inner struct of [MitoEngine].

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<Row> {
-    (0..num_rows)
+fn build_rows(start: usize, end: usize) -> Vec<Row> {
+    (start..end)
         .map(|i| api::v1::Row {
             values: vec![
                 api::v1::Value {
@@ -256,7 +257,7 @@ fn build_rows(num_rows: usize) -> Vec<Row> {
 
 #[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::<Vec<_>>();
+    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::<usize>());
+
+    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 }))

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)]

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<MitoRegion> {
+    pub(crate) async fn open<S: LogStore>(
+        self,
+        config: &MitoConfig,
+        wal: &Wal<S>,
+    ) -> Result<MitoRegion> {
         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<S: LogStore>(
+    wal: &Wal<S>,
+    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")

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,
         }
     }
 }

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<Arc<Error>>,
+    /// Sender to send write result to the waiter for this mutation.
+    sender: Option<Sender<Result<Output>>>,
+    /// Number of rows to be written.
+    num_rows: usize,
+}
+
+impl WriteNotify {
+    /// Creates a new notify from the `sender`.
+    fn new(sender: Option<Sender<Result<Output>>>, 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<WriteNotify>,
+}
+
+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<Rows>,
+        tx: Option<Sender<Result<Output>>>,
+    ) {
+        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<S: LogStore>(
+        &mut self,
+        wal_writer: &mut WalWriter<S>,
+    ) -> 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<Error>) {
+        // 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);
+    }
+}

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<Arc<RaftEngineLogStore>>,
+    object_store: Option<ObjectStore>,
 }
 
 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<Arc<RaftEngineLogStore>> {
+        self.logstore.clone()
+    }
+
+    pub fn get_object_store(&self) -> Option<ObjectStore> {
+        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.

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<S> {
     /// The underlying log store.
     store: Arc<S>,

src/mito2/src/worker.rs

@@ -365,9 +365,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
 
         self.handle_ddl_requests(ddl_requests).await;
     }
-}
 
-impl<S> RegionWorkerLoop<S> {
     /// Takes and handles all ddl requests.
     async fn handle_ddl_requests(&mut self, ddl_tasks: Vec<RegionTask>) {
         if ddl_tasks.is_empty() {
@@ -392,7 +390,9 @@ impl<S> RegionWorkerLoop<S> {
             }
         }
     }
+}
 
+impl<S> RegionWorkerLoop<S> {
     // Clean up the worker.
     async fn clean(&self) {
         // Closes remaining regions.

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<S> RegionWorkerLoop<S> {
+impl<S: LogStore> RegionWorkerLoop<S> {
     pub(crate) async fn handle_create_request(
         &mut self,
         region_id: RegionId,

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<S> RegionWorkerLoop<S> {
+impl<S: LogStore> RegionWorkerLoop<S> {
     pub(crate) async fn handle_open_request(
         &mut self,
         region_id: RegionId,
@@ -44,7 +45,7 @@ impl<S> RegionWorkerLoop<S> {
             self.object_store.clone(),
         )
         .region_dir(&request.region_dir)
-        .open(&self.config)
+        .open(&self.config, &self.wal)
         .await?;
 
         info!("Region {} is opened", region_id);

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<S: LogStore> RegionWorkerLoop<S> {
@@ -84,7 +78,10 @@ impl<S> RegionWorkerLoop<S> {
                 };
 
                 // Initialize the context.
-                e.insert(RegionWriteCtx::new(region));
+                e.insert(RegionWriteCtx::new(
+                    region.region_id,
+                    &region.version_control,
+                ));
             }
 
             // Safety: Now we ensure the region exists.
@@ -92,7 +89,7 @@ impl<S> RegionWorkerLoop<S> {
 
             // Checks whether request schema is compatible with region schema.
             if let Err(e) =
-                maybe_fill_missing_columns(&mut sender_req.request, &region_ctx.version.metadata)
+                maybe_fill_missing_columns(&mut sender_req.request, &region_ctx.version().metadata)
             {
                 send_result(sender_req.sender, Err(e));
 
@@ -100,7 +97,11 @@ impl<S> RegionWorkerLoop<S> {
             }
 
             // 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<Sender<Result<Output>>>, res: Result<Output>) {
         let _ = sender.send(res);
     }
 }
-
-/// Notifier to notify write result on drop.
-struct WriteNotify {
-    /// Error to send to the waiter.
-    err: Option<Arc<Error>>,
-    /// Sender to send write result to the waiter for this mutation.
-    sender: Option<Sender<Result<Output>>>,
-    /// Number of rows to be written.
-    num_rows: usize,
-}
-
-impl WriteNotify {
-    /// Creates a new notify from the `sender`.
-    fn new(sender: Option<Sender<Result<Output>>>, 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<WriteNotify>,
-}
-
-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<S: LogStore>(&mut self, wal_writer: &mut WalWriter<S>) -> 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<Error>) {
-        // 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);
-    }
-}