From a1e3da4aabbb966a154e1900d157a21f34894ec4 Mon Sep 17 00:00:00 2001
From: Noel Kwan <noelkwan1998@gmail.com>
Date: Fri, 9 Jun 2023 16:48:00 +0800
Subject: [PATCH] ensure ReadVersion is created with replication options

---
 .../hummock/event_handler/hummock_event_handler.rs    |  6 +++++-
 src/storage/src/hummock/event_handler/mod.rs          |  7 ++++++-
 src/storage/src/hummock/mod.rs                        |  1 +
 src/storage/src/hummock/store/version.rs              | 11 +++++------
 4 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/src/storage/src/hummock/event_handler/hummock_event_handler.rs b/src/storage/src/hummock/event_handler/hummock_event_handler.rs
index fc68d47f18455..00a884a830406 100644
--- a/src/storage/src/hummock/event_handler/hummock_event_handler.rs
+++ b/src/storage/src/hummock/event_handler/hummock_event_handler.rs
@@ -558,10 +558,14 @@ impl HummockEventHandler {
                         HummockEvent::RegisterReadVersion {
                             table_id,
                             new_read_version_sender,
+                            is_replicated,
                         } => {
                             let pinned_version = self.pinned_version.load();
                             let basic_read_version = Arc::new(RwLock::new(
-                                HummockReadVersion::new((**pinned_version).clone()),
+                                HummockReadVersion::new_with_replication_option(
+                                    (**pinned_version).clone(),
+                                    is_replicated,
+                                ),
                             ));
 
                             let instance_id = self.generate_instance_id();
diff --git a/src/storage/src/hummock/event_handler/mod.rs b/src/storage/src/hummock/event_handler/mod.rs
index 1a192192cab98..20ffe0cea0783 100644
--- a/src/storage/src/hummock/event_handler/mod.rs
+++ b/src/storage/src/hummock/event_handler/mod.rs
@@ -76,6 +76,7 @@ pub enum HummockEvent {
         table_id: TableId,
         new_read_version_sender:
             oneshot::Sender<(Arc<RwLock<HummockReadVersion>>, LocalInstanceGuard)>,
+        is_replicated: bool,
     },
 
     DestroyReadVersion {
@@ -116,7 +117,11 @@ impl HummockEvent {
             HummockEvent::RegisterReadVersion {
                 table_id,
                 new_read_version_sender: _,
-            } => format!("RegisterReadVersion table_id {:?}", table_id,),
+                is_replicated,
+            } => format!(
+                "RegisterReadVersion table_id {:?}, is_replicated: {:?}",
+                table_id, is_replicated
+            ),
 
             HummockEvent::DestroyReadVersion {
                 table_id,
diff --git a/src/storage/src/hummock/mod.rs b/src/storage/src/hummock/mod.rs
index 61a94c2914c52..8be31daede4bd 100644
--- a/src/storage/src/hummock/mod.rs
+++ b/src/storage/src/hummock/mod.rs
@@ -234,6 +234,7 @@ impl HummockStorage {
             .send(HummockEvent::RegisterReadVersion {
                 table_id: option.table_id,
                 new_read_version_sender: tx,
+                is_replicated: option.is_replicated,
             })
             .unwrap();
 
diff --git a/src/storage/src/hummock/store/version.rs b/src/storage/src/hummock/store/version.rs
index 690baf27b9024..d819bbe533b2a 100644
--- a/src/storage/src/hummock/store/version.rs
+++ b/src/storage/src/hummock/store/version.rs
@@ -215,7 +215,10 @@ pub struct HummockReadVersion {
 }
 
 impl HummockReadVersion {
-    fn new_inner(committed_version: CommittedVersion, is_replicated: bool) -> Self {
+    pub fn new_with_replication_option(
+        committed_version: CommittedVersion,
+        is_replicated: bool,
+    ) -> Self {
         // before build `HummockReadVersion`, we need to get the a initial version which obtained
         // from meta. want this initialization after version is initialized (now with
         // notification), so add a assert condition to guarantee correct initialization order
@@ -234,11 +237,7 @@ impl HummockReadVersion {
     }
 
     pub fn new(committed_version: CommittedVersion) -> Self {
-        Self::new_inner(committed_version, false)
-    }
-
-    pub fn new_replicated(committed_version: CommittedVersion) -> Self {
-        Self::new_inner(committed_version, true)
+        Self::new_with_replication_option(committed_version, false)
     }
 
     /// Updates the read version with `VersionUpdate`.