diff --git a/core/bin/zksync_server/src/node_builder.rs b/core/bin/zksync_server/src/node_builder.rs index 0403622f3f1..3af547ff2ee 100644 --- 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 index 5fed50e0f53..c1334f43959 100644 --- 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 index 889e61c9139..de2751474b2 100644 --- 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, + pub stop_node_on_completion: bool, } impl NodeStorageInitializer { - pub fn new(strategy: NodeInitializationStrategy, pool: ConnectionPool) -> Self { - Self { strategy, pool } + pub fn new( + strategy: NodeInitializationStrategy, + pool: ConnectionPool, + stop_node_on_completion: bool, + ) -> Self { + Self { + strategy, + pool, + stop_node_on_completion, + } } /// Returns the preferred kind of storage initialization.