From 7dc37d282fe158ccd2eece5a786b68c1921576cc Mon Sep 17 00:00:00 2001
From: Vitaly Terekhov <va.tehb@gmail.com>
Date: Sun, 23 Jun 2024 15:09:54 +0000
Subject: [PATCH] feature(collator): moved pre-accept blocks under feature

---
 cli/Cargo.toml               |  2 +-
 collator/Cargo.toml          |  3 ++-
 collator/src/collator/mod.rs | 13 +++++++++++--
 collator/src/manager/mod.rs  |  3 ++-
 collator/src/state_node.rs   |  8 +++++++-
 5 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/cli/Cargo.toml b/cli/Cargo.toml
index b76201363..da81d8598 100644
--- a/cli/Cargo.toml
+++ b/cli/Cargo.toml
@@ -56,7 +56,7 @@ tycho-storage = { workspace = true }
 tycho-util = { workspace = true }
 
 [dev-dependencies]
-tycho-collator = { workspace = true, features = ["test"] }
+tycho-collator = { workspace = true, features = ["test", "pre-accept-blocks"] }
 tycho-storage = { workspace = true, features = ["test"] }
 
 [build-dependencies]
diff --git a/collator/Cargo.toml b/collator/Cargo.toml
index d9f89a290..917a9bcae 100644
--- a/collator/Cargo.toml
+++ b/collator/Cargo.toml
@@ -57,9 +57,10 @@ tycho-storage = { workspace = true, features = ["test"] }
 tycho-util = { workspace = true, features = ["test"] }
 
 [features]
-default = []
+default = ["pre-accept-blocks"]
 test = []
 block-creator-stats = []
+pre-accept-blocks = []
 
 [lints]
 workspace = true
diff --git a/collator/src/collator/mod.rs b/collator/src/collator/mod.rs
index b234933f7..3ef8c4751 100644
--- a/collator/src/collator/mod.rs
+++ b/collator/src/collator/mod.rs
@@ -383,9 +383,18 @@ impl CollatorStdImpl {
         let _histogram =
             HistogramGuardWithLabels::begin("tycho_collator_update_mc_data_time", &labels);
 
-        let prev_states_loaded_from_storage = self.working_state().prev_states_loaded_from_storage;
+        #[cfg(feature = "pre-accept-blocks")]
+        let new_mc_state = self
+            .state_node_adapter
+            .load_state(new_mc_state.block_id())
+            .await?;
+
+        #[cfg(feature = "pre-accept-blocks")]
+        let load_prev_states_from_storage = !self.working_state().prev_states_loaded_from_storage;
+        #[cfg(not(feature = "pre-accept-blocks"))]
+        let load_prev_states_from_storage = false;
 
-        if prev_states_loaded_from_storage {
+        if !load_prev_states_from_storage {
             let working_state_mut = self
                 .working_state
                 .as_mut()
diff --git a/collator/src/manager/mod.rs b/collator/src/manager/mod.rs
index 911a3be3b..124a2568b 100644
--- a/collator/src/manager/mod.rs
+++ b/collator/src/manager/mod.rs
@@ -900,7 +900,8 @@ where
             .clone();
 
         // pre-accept shard block to store new shard state
-        if !candidate_id.shard.is_masterchain() {
+        #[cfg(feature = "pre-accept-blocks")]
+        {
             let histogram = HistogramGuard::begin_with_labels(
                 "tycho_collator_process_collated_block_candidate_pre_accept_total_time",
                 labels,
diff --git a/collator/src/state_node.rs b/collator/src/state_node.rs
index 7bc733ae2..4e83ce868 100644
--- a/collator/src/state_node.rs
+++ b/collator/src/state_node.rs
@@ -57,6 +57,7 @@ pub trait StateNodeAdapter: Send + Sync + 'static {
     async fn load_block_handle(&self, block_id: &BlockId) -> Result<Option<BlockHandle>>;
 
     /// Just store updated state
+    #[cfg(feature = "pre-accept-blocks")]
     async fn pre_accept_block(
         &self,
         block: BlockStuffForSync,
@@ -133,6 +134,7 @@ impl StateNodeAdapter for StateNodeAdapterStdImpl {
         Ok(self.storage.block_handle_storage().load_handle(block_id))
     }
 
+    #[cfg(feature = "pre-accept-blocks")]
     async fn pre_accept_block(
         &self,
         block: BlockStuffForSync,
@@ -145,7 +147,11 @@ impl StateNodeAdapter for StateNodeAdapterStdImpl {
         let block_stuff = &block.block_stuff_aug.data;
         let archive_data = &block.block_stuff_aug.archive_data;
         let info = block_stuff.load_info()?;
-        let mc_ref_seqno = info.min_ref_mc_seqno + 1;
+        let mc_ref_seqno = if block.block_id.shard.is_masterchain() {
+            block.block_id.seqno
+        } else {
+            info.min_ref_mc_seqno + 1
+        };
         let store_res = block_storage
             .store_block_data(block_stuff, archive_data, BlockMetaData {
                 is_key_block: info.key_block,