Stop processing users while restoring the queue

Author: MauroToscano

crates/batcher/src/lib.rs

@@ -102,6 +102,11 @@ pub struct Batcher {
     /// - Batch creation needs to be able to change all the states, so all processing
     ///   needs to be stopped, and all user_states locks need to be taken
     batch_state: Mutex<BatchState>,
+    /// Flag to indicate when restoration is in progress
+    /// When true, message handlers will return ServerBusy responses
+    /// It's used a way to "lock" all the user_states at the same time
+    /// If one needed is taken in the handle message it will timeout
+    is_restoration_in_progress: RwLock<bool>,
     user_states: DashMap<Address, Arc<Mutex<UserState>>>,
 
     last_uploaded_batch_block: Mutex<u64>,
@@ -113,6 +118,8 @@ pub struct Batcher {
 
     disabled_verifiers: Mutex<U256>,
 
+
+
     // Observability and monitoring
     pub metrics: metrics::BatcherMetrics,
     pub telemetry: TelemetrySender,
@@ -292,6 +299,7 @@ impl Batcher {
             batch_state: Mutex::new(batch_state),
             user_states,
             disabled_verifiers: Mutex::new(disabled_verifiers),
+            is_restoration_in_progress: RwLock::new(false),
             metrics,
             telemetry,
         }
@@ -665,6 +673,17 @@ impl Batcher {
         mut address: Address,
         ws_conn_sink: WsMessageSink,
     ) -> Result<(), Error> {
+        // Check if restoration is in progress
+        if *self.is_restoration_in_progress.read().await {
+            warn!(
+                "Rejecting nonce request from {} during restoration",
+                address
+            );
+            let response = GetNonceResponseMessage::ServerBusy;
+            send_message(ws_conn_sink, response).await;
+            return Ok(());
+        }
+
         // If the address is not paying, we will return the nonce of the aligned_payment_address
         if !self.has_to_pay(&address) {
             info!("Handling nonpaying message");
@@ -749,6 +768,17 @@ impl Batcher {
         debug!("Received message with nonce: {msg_nonce:?}");
         self.metrics.received_proofs.inc();
 
+        // Check if restoration is in progress
+        if *self.is_restoration_in_progress.read().await {
+            warn!(
+                "Rejecting proof submission from {} during restoration (nonce: {})",
+                client_msg.verification_data.verification_data.proof_generator_addr, msg_nonce
+            );
+            let response = SubmitProofResponseMessage::ServerBusy;
+            send_message(ws_conn_sink, response).await;
+            return Ok(());
+        }
+
         // * ---------------------------------------------------*
         // *        Perform validations over the message        *
         // * ---------------------------------------------------*
@@ -1530,6 +1560,9 @@ impl Batcher {
             failed_batch.len()
         );
 
+        // Set restoration flag to stop handling new user messages
+        *self.is_restoration_in_progress.write().await = true;
+
         let mut batch_state_lock = self.batch_state.lock().await;
         let mut restored_entries = Vec::new();
 
@@ -1591,10 +1624,10 @@ impl Batcher {
         let users_with_restored_proofs: std::collections::HashSet<Address> =
             restored_entries.iter().map(|entry| entry.sender).collect();
 
-        drop(batch_state_lock); // Release batch lock before user state updates
-
         // Update user states for successfully restored proofs
         info!("Updating user states after proof restoration...");
+        // TODO: We may have ejected some users that didn't have restored proofs, 
+        // we should include in this list the ejected users
         if let Err(e) = self
             .update_user_states_from_queue_state(users_with_restored_proofs)
             .await
@@ -1604,6 +1637,10 @@ impl Batcher {
                 e
             );
         }
+
+        // Clear restoration flag to allow normal user message handling
+        *self.is_restoration_in_progress.write().await = false;
+        info!("Proof restoration completed, resuming normal operations");
     }
 
     /// Takes the finalized batch as input and: