fix

matter-labs · Dec 19, 2024 · 2a0e964 · 2a0e964
1 parent 4a60208
commit 2a0e964
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 6 deletions.
diff --git a/core/bin/zksync_server/src/node_builder.rs b/core/bin/zksync_server/src/node_builder.rs
@@ -680,6 +680,9 @@ impl MainNodeBuilder {
             contracts: self.contracts_config.clone(),
         });
         let mut layer = NodeStorageInitializerLayer::new();
+        if matches!(kind, LayerKind::TaskEndingExecution) {
+            layer = layer.stop_node_on_completion();
+        }
         if matches!(kind, LayerKind::Precondition) {
             layer = layer.as_precondition();
         }
@@ -696,7 +699,7 @@ impl MainNodeBuilder {
             .add_blob_client_layer(l1_chain_id)?
             .add_query_eth_client_layer()?
             .add_healthcheck_layer()?
-            .add_storage_initialization_layer(LayerKind::Task, false)?;
+            .add_storage_initialization_layer(LayerKind::TaskEndingExecution, false)?;
 
         Ok(self.node.build())
     }
@@ -709,7 +712,7 @@ impl MainNodeBuilder {
             .add_blob_client_layer(l1_chain_id)?
             .add_query_eth_client_layer()?
             .add_healthcheck_layer()?
-            .add_storage_initialization_layer(LayerKind::Task, true)?;
+            .add_storage_initialization_layer(LayerKind::TaskEndingExecution, true)?;
 
         Ok(self.node.build())
     }
@@ -840,4 +843,5 @@ impl MainNodeBuilder {
 enum LayerKind {
     Task,
     Precondition,
+    TaskEndingExecution,
 }
diff --git a/core/node/node_framework/src/implementations/layers/node_storage_init/mod.rs b/core/node/node_framework/src/implementations/layers/node_storage_init/mod.rs
@@ -25,6 +25,7 @@ pub mod main_node_strategy;
 #[derive(Debug, Default)]
 pub struct NodeStorageInitializerLayer {
     as_precondition: bool,
+    stop_node_on_completion: bool,
 }
 
 impl NodeStorageInitializerLayer {
@@ -37,6 +38,11 @@ impl NodeStorageInitializerLayer {
         self.as_precondition = true;
         self
     }
+
+    pub fn stop_node_on_completion(mut self) -> Self {
+        self.stop_node_on_completion = true;
+        self
+    }
 }
 
 #[derive(Debug, FromContext)]
@@ -87,7 +93,7 @@ impl WiringLayer for NodeStorageInitializerLayer {
         let pool = input.master_pool.get().await?;
         let NodeInitializationStrategyResource(strategy) = input.strategy;
 
-        let initializer = NodeStorageInitializer::new(strategy, pool);
+        let initializer = NodeStorageInitializer::new(strategy, pool, self.stop_node_on_completion);
 
         // Insert either task or precondition.
         let output = if self.as_precondition {
@@ -103,7 +109,11 @@ impl WiringLayer for NodeStorageInitializerLayer {
 #[async_trait::async_trait]
 impl Task for NodeStorageInitializer {
     fn kind(&self) -> TaskKind {
-        TaskKind::UnconstrainedOneshotTask
+        if self.stop_node_on_completion {
+            TaskKind::UnconstrainedTask
+        } else {
+            TaskKind::UnconstrainedOneshotTask
+        }
     }
 
     fn id(&self) -> TaskId {

diff --git a/core/node/node_storage_init/src/lib.rs b/core/node/node_storage_init/src/lib.rs
@@ -48,11 +48,20 @@ pub struct NodeInitializationStrategy {
 pub struct NodeStorageInitializer {
     strategy: NodeInitializationStrategy,
     pool: ConnectionPool<Core>,
+    pub stop_node_on_completion: bool,
 }
 
 impl NodeStorageInitializer {
-    pub fn new(strategy: NodeInitializationStrategy, pool: ConnectionPool<Core>) -> Self {
-        Self { strategy, pool }
+    pub fn new(
+        strategy: NodeInitializationStrategy,
+        pool: ConnectionPool<Core>,
+        stop_node_on_completion: bool,
+    ) -> Self {
+        Self {
+            strategy,
+            pool,
+            stop_node_on_completion,
+        }
     }
 
     /// Returns the preferred kind of storage initialization.