feat: handling graceful block closing with separate message

Author: Mohiiit

madara/crates/client/block_production/src/executor/mod.rs

@@ -49,7 +49,10 @@ pub enum ExecutorMessage {
         exec_ctx: BlockExecutionContext,
     },
     BatchExecuted(BatchExecutionResult),
+    /// Normal block closing (block time reached, block full, or explicit CloseBlock).
     EndBlock(Box<BlockExecutionSummary>),
+    /// Final block closing during graceful shutdown. Only sent when executor detects shutdown.
+    EndFinalBlock(Box<BlockExecutionSummary>),
 }
 
 #[derive(Default, Debug)]

madara/crates/client/block_production/src/executor/thread.rs

@@ -297,18 +297,18 @@ impl ExecutorThread {
                                 // This uses the executor's current state - no re-execution needed
                                 match execution_state.executor.finalize() {
                                     Ok(block_exec_summary) => {
-                                        // Send EndBlock message so main loop can close the block
+                                        // Send EndFinalBlock message so main loop can close the block during shutdown
                                         if self
                                             .replies_sender
-                                            .blocking_send(super::ExecutorMessage::EndBlock(Box::new(
+                                            .blocking_send(super::ExecutorMessage::EndFinalBlock(Box::new(
                                                 block_exec_summary,
                                             )))
                                             .is_err()
                                         {
                                             // Receiver closed - main loop already shut down
                                             // Block will remain preconfirmed and be handled on restart
                                             tracing::warn!(
-                                                "Could not send EndBlock during shutdown, block will remain preconfirmed"
+                                                "Could not send EndFinalBlock during shutdown, block will remain preconfirmed"
                                             );
                                         }
                                     }

madara/crates/client/block_production/src/lib.rs

@@ -58,6 +58,10 @@
 //!   receives this message it will proceed to finalize (seal) the pending block and store it to db
 //!   as a full block.
 //!
+//! - [`EndFinalBlock`]: this message is sent by the [`ExecutorThread`] only during graceful shutdown
+//!   when it detects the batch channel closure. It signals the final block closing and the executor
+//!   thread exits immediately after sending it.
+//!
 //! ## Pending Phase
 //!
 //! One important detail to note is that the [`PendingBlockState`] kept in the
@@ -76,8 +80,8 @@
 //! When a cancellation signal is received:
 //! 1. The batcher detects cancellation and exits gracefully, closing the `send_batch` channel
 //! 2. The executor detects the channel closure and finalizes any open block
-//! 3. The executor sends an `EndBlock` message and then completes
-//! 4. The main loop processes the `EndBlock`, closes the block, and exits when both tasks complete
+//! 3. The executor sends an `EndFinalBlock` message (shutdown-specific) and then completes
+//! 4. The main loop processes the `EndFinalBlock`, closes the block, and exits when both tasks complete
 //!
 //! ### Batcher Panic/Error
 //!
@@ -101,6 +105,7 @@
 //! [`StartNewBlock`]: ExecutorMessage::StartNewBlock
 //! [`BatchExecuted`]: ExecutorMessage::BatchExecuted
 //! [`EndBlock`]: ExecutorMessage::EndBlock
+//! [`EndFinalBlock`]: ExecutorMessage::EndFinalBlock
 //! [`ExecutorThreadHandle::send_batch`]: executor::ExecutorThreadHandle::send_batch
 //! [`ExecutorThread::incoming_batches`]: executor::thread::ExecutorThread::incoming_batches
 //! [`ExecutorThread`]: executor::thread::ExecutorThread
@@ -698,52 +703,61 @@ impl BlockProductionTask {
             }
             ExecutorMessage::EndBlock(block_exec_summary) => {
                 tracing::debug!("Received ExecutorMessage::EndBlock");
-                let current_state = self.current_state.take().context("No current state")?;
-                let TaskState::Executing(state) = current_state else {
-                    anyhow::bail!("Invalid executor state transition: expected current state to be Executing")
-                };
+                self.close_block(block_exec_summary).await?;
+            }
+            ExecutorMessage::EndFinalBlock(block_exec_summary) => {
+                tracing::debug!("Received ExecutorMessage::EndFinalBlock (shutdown)");
+                self.close_block(block_exec_summary).await?;
+            }
+        }
 
-                tracing::debug!("Close and save block block_n={}", state.block_number);
-                let start_time = Instant::now();
-
-                let n_txs = self
-                    .backend
-                    .block_view_on_preconfirmed()
-                    .context("No current pre-confirmed block")?
-                    .num_executed_transactions();
-
-                // Convert state_diff and close block using helper function
-                let state_diff: mp_state_update::StateDiff = block_exec_summary.state_diff.into();
-                Self::close_preconfirmed_block_with_state_diff(
-                    self.backend.clone(),
-                    state.block_number,
-                    state.consumed_core_contract_nonces,
-                    &block_exec_summary.bouncer_weights,
-                    state_diff,
-                )
-                .await
-                .context("Closing block")?;
-
-                let time_to_close = start_time.elapsed();
-                tracing::info!(
-                    "⛏️  Closed block #{} with {n_txs} transactions - {time_to_close:?}",
-                    state.block_number
-                );
+        Ok(())
+    }
+
+    /// Close and save a block using the execution summary.
+    /// Used for both normal block closing (EndBlock) and shutdown (EndFinalBlock).
+    async fn close_block(&mut self, block_exec_summary: Box<BlockExecutionSummary>) -> anyhow::Result<()> {
+        let current_state = self.current_state.take().context("No current state")?;
+        let TaskState::Executing(state) = current_state else {
+            anyhow::bail!("Invalid executor state transition: expected current state to be Executing")
+        };
 
-                // Record metrics
-                let attributes = [
-                    KeyValue::new("transactions_added", n_txs.to_string()),
-                    KeyValue::new("closing_time", time_to_close.as_secs_f32().to_string()),
-                ];
+        tracing::debug!("Close and save block block_n={}", state.block_number);
+        let start_time = Instant::now();
 
-                self.metrics.block_counter.add(1, &[]);
-                self.metrics.block_gauge.record(state.block_number, &attributes);
-                self.metrics.transaction_counter.add(n_txs as u64, &[]);
+        let n_txs = self
+            .backend
+            .block_view_on_preconfirmed()
+            .context("No current pre-confirmed block")?
+            .num_executed_transactions();
 
-                self.current_state = Some(TaskState::NotExecuting { latest_block_n: Some(state.block_number) });
-                self.send_state_notification(BlockProductionStateNotification::ClosedBlock);
-            }
-        }
+        // Convert state_diff and close block using helper function
+        let state_diff: mp_state_update::StateDiff = block_exec_summary.state_diff.into();
+        Self::close_preconfirmed_block_with_state_diff(
+            self.backend.clone(),
+            state.block_number,
+            state.consumed_core_contract_nonces,
+            &block_exec_summary.bouncer_weights,
+            state_diff,
+        )
+        .await
+        .context("Closing block")?;
+
+        let time_to_close = start_time.elapsed();
+        tracing::info!("⛏️  Closed block #{} with {n_txs} transactions - {time_to_close:?}", state.block_number);
+
+        // Record metrics
+        let attributes = [
+            KeyValue::new("transactions_added", n_txs.to_string()),
+            KeyValue::new("closing_time", time_to_close.as_secs_f32().to_string()),
+        ];
+
+        self.metrics.block_counter.add(1, &[]);
+        self.metrics.block_gauge.record(state.block_number, &attributes);
+        self.metrics.transaction_counter.add(n_txs as u64, &[]);
+
+        self.current_state = Some(TaskState::NotExecuting { latest_block_n: Some(state.block_number) });
+        self.send_state_notification(BlockProductionStateNotification::ClosedBlock);
 
         Ok(())
     }
@@ -789,7 +803,7 @@ impl BlockProductionTask {
         // Track shutdown state: both batcher and executor must complete before shutdown finishes.
         // Both tasks only complete during shutdown scenarios (cancellation, error, or panic).
         let mut batcher_completed = false;
-        let mut executor_completed = false;
+        let mut end_final_block_received = false; // Track if EndFinalBlock has been processed (executor completed with block)
         let mut batcher_error: Option<anyhow::Error> = None; // Store batcher error to return after graceful shutdown
 
         // Main loop: handles normal operation and graceful shutdown
@@ -813,36 +827,42 @@ impl BlockProductionTask {
                     }
                 }
 
-                // Path 2: Executor replies (EndBlock message during normal operation or shutdown)
+                // Path 2: Executor replies (EndBlock for normal operation, EndFinalBlock for shutdown)
                 Some(reply) = executor.replies.recv() => {
+                    let is_end_final_block = matches!(reply, ExecutorMessage::EndFinalBlock(_));
                     self.process_reply(reply).await.context("Processing reply from executor thread")?;
+                    // Mark executor as completed only after processing EndFinalBlock
+                    if is_end_final_block {
+                        end_final_block_received = true;
+                        tracing::debug!("EndFinalBlock processed, executor completed");
+                    }
                 }
 
                 // Path 3: Executor thread stopped (normal completion or panic)
+                // This fires when executor exits. If EndFinalBlock was already processed, executor is done.
+                // If not, executor exited without a block (no EndFinalBlock sent) - continue loop to check exit conditions.
                 res = executor.stop.recv() => {
-                    executor_completed = true;
                     res.context("In executor thread")?;
-                    tracing::debug!("Executor thread stopped");
                 }
             }
 
-            // Exit conditions:
-            // 1. Both completed → shutdown done (return error if any, otherwise success)
-            // 2. Batcher errored without preconfirmed block → exit immediately with error
-            if batcher_completed && executor_completed {
-                tracing::debug!("Shutdown complete: batcher completed, executor completed");
+            // Exit conditions (checked after each select iteration):
+            // Shutdown is complete when batcher completed AND:
+            // - EndFinalBlock was processed (block case), OR
+            // - No preconfirmed block exists (executor stopped without block)
+            if batcher_completed && (end_final_block_received || !self.backend.has_preconfirmed_block()) {
+                tracing::debug!(
+                    "Shutdown complete: batcher completed, end_final_block_received={}, has_preconfirmed_block={}",
+                    end_final_block_received,
+                    self.backend.has_preconfirmed_block()
+                );
                 return batcher_error
                     .map(|e| {
                         tracing::warn!("Shutdown completed but batcher had error: {e:?}");
                         Err(e)
                     })
                     .unwrap_or(Ok(()));
             }
-
-            if batcher_error.is_some() && !self.backend.has_preconfirmed_block() {
-                tracing::debug!("Batcher errored without preconfirmed block, exiting");
-                return Err(batcher_error.take().unwrap());
-            }
         }
     }
 }