From 058322fc9e0cf66588ac68ea9feeedb141314bb8 Mon Sep 17 00:00:00 2001
From: Danil <deniallugo@gmail.com>
Date: Fri, 21 Jun 2024 12:35:15 +0200
Subject: [PATCH] Split to experimental and normal config for tree

Signed-off-by: Danil <deniallugo@gmail.com>
---
 .../src/proto/config/experimental.proto       |  5 ++
 .../src/proto/config/snapshot_recovery.proto  |  5 +-
 .../protobuf_config/src/snapshot_recovery.rs  | 63 +++++++++++--------
 3 files changed, 45 insertions(+), 28 deletions(-)

diff --git a/core/lib/protobuf_config/src/proto/config/experimental.proto b/core/lib/protobuf_config/src/proto/config/experimental.proto
index ddcb4265319d..6f9ec426d8bb 100644
--- a/core/lib/protobuf_config/src/proto/config/experimental.proto
+++ b/core/lib/protobuf_config/src/proto/config/experimental.proto
@@ -12,3 +12,8 @@ message DB {
   optional uint64 processing_delay_ms = 4;
   optional bool include_indices_and_filters_in_block_cache = 5;
 }
+
+// Experimental part of the Snapshot recovery configuration.
+message SnapshotRecovery {
+  optional uint64 tree_recovery_parallel_persistence_buffer = 1;
+}
diff --git a/core/lib/protobuf_config/src/proto/config/snapshot_recovery.proto b/core/lib/protobuf_config/src/proto/config/snapshot_recovery.proto
index 5faf0399ba69..9eceda12ad86 100644
--- a/core/lib/protobuf_config/src/proto/config/snapshot_recovery.proto
+++ b/core/lib/protobuf_config/src/proto/config/snapshot_recovery.proto
@@ -1,10 +1,11 @@
 syntax = "proto3";
+import "zksync/config/object_store.proto";
+import "zksync/config/experimental.proto";
 
 package zksync.config.snapshot_recovery;
 
 message Tree {
   optional uint64 chunk_size = 1;
-  optional uint64 parallel_persistence_buffer = 2;
 }
 
 message Postgres {
@@ -16,4 +17,6 @@ message SnapshotRecovery {
   optional Postgres postgres = 2;
   optional Tree tree = 3;
   optional uint32 l1_batch = 4;
+  optional config.object_store.ObjectStore object_store = 5;
+  optional experimental.SnapshotRecovery experimental = 6;
 }
diff --git a/core/lib/protobuf_config/src/snapshot_recovery.rs b/core/lib/protobuf_config/src/snapshot_recovery.rs
index 237f6057d794..ea16d1b158b4 100644
--- a/core/lib/protobuf_config/src/snapshot_recovery.rs
+++ b/core/lib/protobuf_config/src/snapshot_recovery.rs
@@ -10,26 +10,6 @@ use zksync_protobuf::ProtoRepr;
 
 use crate::{proto::snapshot_recovery as proto, read_optional_repr};
 
-impl ProtoRepr for proto::Tree {
-    type Type = TreeRecoveryConfig;
-
-    fn read(&self) -> anyhow::Result<Self::Type> {
-        Ok(Self::Type {
-            chunk_size: self.chunk_size,
-            parallel_persistence_buffer: self
-                .parallel_persistence_buffer
-                .and_then(|a| NonZeroUsize::new(a as usize)),
-        })
-    }
-
-    fn build(this: &Self::Type) -> Self {
-        Self {
-            chunk_size: this.chunk_size,
-            parallel_persistence_buffer: this.parallel_persistence_buffer.map(|a| a.get() as u64),
-        }
-    }
-}
-
 impl ProtoRepr for proto::Postgres {
     type Type = PostgresRecoveryConfig;
 
@@ -52,11 +32,29 @@ impl ProtoRepr for proto::SnapshotRecovery {
     type Type = SnapshotRecoveryConfig;
 
     fn read(&self) -> anyhow::Result<Self::Type> {
+        let tree = self
+            .tree
+            .as_ref()
+            .map(|tree| {
+                let chunk_size = tree.chunk_size;
+                let parallel_persistence_buffer = self
+                    .experimental
+                    .as_ref()
+                    .and_then(|a| {
+                        a.tree_recovery_parallel_persistence_buffer
+                            .map(|a| NonZeroUsize::new(a as usize))
+                    })
+                    .flatten();
+                TreeRecoveryConfig {
+                    chunk_size,
+                    parallel_persistence_buffer,
+                }
+            })
+            .unwrap_or_default();
+
         Ok(Self::Type {
             enabled: self.enabled.unwrap_or_default(),
-            tree: read_optional_repr(&self.tree)
-                .context("tree")?
-                .unwrap_or_default(),
+            tree,
             postgres: read_optional_repr(&self.postgres)
                 .context("postgres")?
                 .unwrap_or_default(),
@@ -65,10 +63,20 @@ impl ProtoRepr for proto::SnapshotRecovery {
     }
 
     fn build(this: &Self::Type) -> Self {
-        let tree = if this.tree == TreeRecoveryConfig::default() {
-            None
+        let (tree, experimental) = if this.tree == TreeRecoveryConfig::default() {
+            (None, None)
         } else {
-            Some(this.tree.clone())
+            (
+                Some(proto::Tree {
+                    chunk_size: this.tree.chunk_size,
+                }),
+                Some(crate::proto::experimental::SnapshotRecovery {
+                    tree_recovery_parallel_persistence_buffer: this
+                        .tree
+                        .parallel_persistence_buffer
+                        .map(|a| a.get() as u64),
+                }),
+            )
         };
         let postgres = if this.postgres == PostgresRecoveryConfig::default() {
             None
@@ -78,7 +86,8 @@ impl ProtoRepr for proto::SnapshotRecovery {
         Self {
             enabled: Some(this.enabled),
             postgres: postgres.as_ref().map(ProtoRepr::build),
-            tree: tree.as_ref().map(ProtoRepr::build),
+            tree,
+            experimental,
             l1_batch: this.l1_batch.map(|a| a.0),
         }
     }