feat(batcher): poll for unlock events and remove users with unlocked funds (#2123)

Co-authored-by: MauroFab <[email protected]>
Co-authored-by: Mauro Toscano <[email protected]>

Author: JuArce

config-files/config-batcher-docker.yaml

@@ -37,4 +37,6 @@ batcher:
   # When replacing the message, how much higher should the max fee in comparison to the original one
   # The calculation is replacement_max_fee >= original_max_fee + original_max_fee * min_bump_percentage / 100
   min_bump_percentage: 10
+  # How often to poll for BalanceUnlocked events in seconds (default: 600 seconds = 10 minutes)
+  balance_unlock_polling_interval_seconds: 600
 

config-files/config-batcher-ethereum-package.yaml

@@ -35,4 +35,6 @@ batcher:
   # When replacing the message, how much higher should the max fee in comparison to the original one
   # The calculation is replacement_max_fee >= original_max_fee + original_max_fee * min_bump_percentage / 100
   min_bump_percentage: 10
+  # How often to poll for BalanceUnlocked events in seconds (default: 600 seconds = 10 minutes)
+  balance_unlock_polling_interval_seconds: 600
 

config-files/config-batcher.yaml

@@ -37,3 +37,5 @@ batcher:
   # When replacing the message, how much higher should the max fee in comparison to the original one
   # The calculation is replacement_max_fee >= original_max_fee + original_max_fee * min_bump_percentage / 100
   min_bump_percentage: 10
+  # How often to poll for BalanceUnlocked events in seconds (default: 600 seconds = 10 minutes)
+  balance_unlock_polling_interval_seconds: 600

crates/batcher/src/config/mod.rs

@@ -53,6 +53,7 @@ pub struct BatcherConfigFromYaml {
     pub non_paying: Option<NonPayingConfigFromYaml>,
     pub amount_of_proofs_for_min_max_fee: usize,
     pub min_bump_percentage: u64,
+    pub balance_unlock_polling_interval_seconds: u64,
 }
 
 #[derive(Debug, Deserialize)]

crates/batcher/src/lib.rs

@@ -7,9 +7,10 @@ use eth::utils::{calculate_bumped_gas_price, get_batcher_signer, get_gas_price};
 use ethers::contract::ContractError;
 use ethers::signers::Signer;
 use retry::batcher_retryables::{
-    cancel_create_new_task_retryable, create_new_task_retryable, get_user_balance_retryable,
-    get_user_nonce_from_ethereum_retryable, simulate_create_new_task_retryable,
-    user_balance_is_unlocked_retryable,
+    cancel_create_new_task_retryable, create_new_task_retryable,
+    get_current_block_number_retryable, get_user_balance_retryable,
+    get_user_nonce_from_ethereum_retryable, query_balance_unlocked_events_retryable,
+    simulate_create_new_task_retryable, user_balance_is_unlocked_retryable,
 };
 use retry::{retry_function, RetryError};
 use tokio::time::{timeout, Instant};
@@ -39,8 +40,8 @@ use aligned_sdk::common::types::{
 
 use aws_sdk_s3::client::Client as S3Client;
 use eth::payment_service::{BatcherPaymentService, CreateNewTaskFeeParams, SignerMiddlewareT};
-use ethers::prelude::{Middleware, Provider};
-use ethers::types::{Address, Signature, TransactionReceipt, U256};
+use ethers::prelude::{Http, Middleware, Provider};
+use ethers::types::{Address, Signature, TransactionReceipt, U256, U64};
 use futures_util::{future, join, SinkExt, StreamExt, TryStreamExt};
 use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
 use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend;
@@ -50,6 +51,7 @@ use tokio::sync::{Mutex, MutexGuard, RwLock};
 
 // Message handler lock timeout
 const MESSAGE_HANDLER_LOCK_TIMEOUT: Duration = Duration::from_secs(10);
+const POLLING_EVENTS_LOCK_TIMEOUT: Duration = Duration::from_secs(300);
 use tokio_tungstenite::tungstenite::{Error, Message};
 use types::batch_queue::{self, BatchQueueEntry, BatchQueueEntryPriority};
 use types::errors::{BatcherError, TransactionSendError};
@@ -86,6 +88,8 @@ pub struct Batcher {
     eth_ws_url_fallback: String,
     batcher_signer: Arc<SignerMiddlewareT>,
     batcher_signer_fallback: Arc<SignerMiddlewareT>,
+    eth_http_provider: Provider<Http>,
+    eth_http_provider_fallback: Provider<Http>,
     chain_id: U256,
     payment_service: BatcherPaymentService,
     payment_service_fallback: BatcherPaymentService,
@@ -103,6 +107,7 @@ pub struct Batcher {
     current_min_max_fee: RwLock<U256>,
     amount_of_proofs_for_min_max_fee: usize,
     min_bump_percentage: U256,
+    balance_unlock_polling_interval_seconds: u64,
 
     // Shared state access:
     // Two kinds of threads interact with the shared state:
@@ -315,6 +320,8 @@ impl Batcher {
             eth_ws_url_fallback: config.eth_ws_url_fallback,
             batcher_signer,
             batcher_signer_fallback,
+            eth_http_provider,
+            eth_http_provider_fallback,
             chain_id,
             payment_service,
             payment_service_fallback,
@@ -327,6 +334,9 @@ impl Batcher {
             max_batch_proof_qty: config.batcher.max_batch_proof_qty,
             amount_of_proofs_for_min_max_fee: config.batcher.amount_of_proofs_for_min_max_fee,
             min_bump_percentage: U256::from(config.batcher.min_bump_percentage),
+            balance_unlock_polling_interval_seconds: config
+                .batcher
+                .balance_unlock_polling_interval_seconds,
             last_uploaded_batch_block: Mutex::new(last_uploaded_batch_block),
             pre_verification_is_enabled: config.batcher.pre_verification_is_enabled,
             non_paying_config,
@@ -436,12 +446,16 @@ impl Batcher {
         }
     }
 
-    /// Helper to apply 15-second timeout to batch lock acquisition with consistent logging and metrics
-    async fn try_batch_lock_with_timeout<F, T>(&self, lock_future: F) -> Option<T>
+    /// Helper to apply `duration` timeout to batch lock acquisition with consistent logging and metrics
+    async fn try_batch_lock_with_timeout<F, T>(
+        &self,
+        lock_future: F,
+        duration: Duration,
+    ) -> Option<T>
     where
         F: std::future::Future<Output = T>,
     {
-        match timeout(MESSAGE_HANDLER_LOCK_TIMEOUT, lock_future).await {
+        match timeout(duration, lock_future).await {
             Ok(result) => Some(result),
             Err(_) => {
                 warn!("Batch lock acquisition timed out");
@@ -491,6 +505,204 @@ impl Batcher {
         .map_err(|e| e.inner())
     }
 
+    /// Poll for BalanceUnlocked events from BatcherPaymentService contract.
+    /// Runs at configurable intervals and checks recent blocks for events (2x the polling interval).
+    /// When an event is detected, removes user's proofs from queue and resets UserState.
+    pub async fn poll_balance_unlocked_events(self: Arc<Self>) -> Result<(), BatcherError> {
+        let mut interval = tokio::time::interval(Duration::from_secs(
+            self.balance_unlock_polling_interval_seconds,
+        ));
+        let mut from_block = self.get_current_block_number().await.map_err(|e| {
+            BatcherError::EthereumProviderError(format!(
+                "Failed to get current block number: {:?}",
+                e
+            ))
+        })?;
+
+        loop {
+            interval.tick().await;
+
+            match self.process_balance_unlocked_events(from_block).await {
+                Ok(current_block) => {
+                    from_block = current_block;
+                }
+                Err(e) => {
+                    error!("Error processing BalanceUnlocked events: {:?}", e);
+                    // On error, keep from_block unchanged to retry the same range next time
+                }
+            }
+        }
+    }
+
+    async fn process_balance_unlocked_events(&self, from_block: U64) -> Result<U64, BatcherError> {
+        // Get current block number using HTTP providers
+        let current_block = self.get_current_block_number().await.map_err(|e| {
+            BatcherError::EthereumProviderError(format!(
+                "Failed to get current block number: {:?}",
+                e
+            ))
+        })?;
+
+        // Query events with retry logic
+        let events = self
+            .query_balance_unlocked_events(from_block, current_block)
+            .await
+            .map_err(|e| {
+                BatcherError::EthereumProviderError(format!(
+                    "Failed to query BalanceUnlocked events: {:?}",
+                    e
+                ))
+            })?;
+
+        info!(
+            "Found {} BalanceUnlocked events in blocks {} to {}",
+            events.len(),
+            from_block,
+            current_block
+        );
+
+        // Process each event
+        for event in events {
+            let user_address = event.user;
+            debug!(
+                "Processing BalanceUnlocked event for user: {:?}",
+                user_address
+            );
+
+            // Check if user has proofs in queue
+            //
+            // Double-check that funds are still unlocked by calling the contract
+            // This is necessary because we query events over a block range, and the
+            // user’s state may have changed (e.g., funds could be locked again) after
+            // the event was emitted. Verifying on-chain ensures we don’t act on stale data.
+            //
+            // There is a brief period between the checks and the removal during which the user's
+            // proofs could be sent. This is acceptable, as the removal will not fail;
+            // it will simply clear the user's state.
+            if self.user_has_proofs_in_queue(user_address).await
+                && self.user_balance_is_unlocked(&user_address).await
+            {
+                info!(
+                    "User {:?} has proofs in queue and funds are unlocked, proceeding to remove proofs and resetting UserState",
+                    user_address
+                );
+                self.remove_user_proofs_and_reset_state(user_address).await;
+            }
+        }
+
+        Ok(current_block)
+    }
+
+    /// Gets the current block number from Ethereum.
+    /// Retries on recoverable errors using exponential backoff up to `ETHEREUM_CALL_MAX_RETRIES` times:
+    /// (0,5 secs - 1 secs - 2 secs - 4 secs - 8 secs).
+    async fn get_current_block_number(&self) -> Result<U64, RetryError<String>> {
+        retry_function(
+            || {
+                get_current_block_number_retryable(
+                    &self.eth_http_provider,
+                    &self.eth_http_provider_fallback,
+                )
+            },
+            ETHEREUM_CALL_MIN_RETRY_DELAY,
+            ETHEREUM_CALL_BACKOFF_FACTOR,
+            ETHEREUM_CALL_MAX_RETRIES,
+            ETHEREUM_CALL_MAX_RETRY_DELAY,
+        )
+        .await
+    }
+
+    /// Queries BalanceUnlocked events from the BatcherPaymentService contract.
+    /// Retries on recoverable errors using exponential backoff up to `ETHEREUM_CALL_MAX_RETRIES` times:
+    /// (0,5 secs - 1 secs - 2 secs - 4 secs - 8 secs).
+    async fn query_balance_unlocked_events(
+        &self,
+        from_block: U64,
+        to_block: U64,
+    ) -> Result<
+        Vec<aligned_sdk::eth::batcher_payment_service::BalanceUnlockedFilter>,
+        RetryError<String>,
+    > {
+        retry_function(
+            || {
+                query_balance_unlocked_events_retryable(
+                    &self.payment_service,
+                    &self.payment_service_fallback,
+                    from_block,
+                    to_block,
+                )
+            },
+            ETHEREUM_CALL_MIN_RETRY_DELAY,
+            ETHEREUM_CALL_BACKOFF_FACTOR,
+            ETHEREUM_CALL_MAX_RETRIES,
+            ETHEREUM_CALL_MAX_RETRY_DELAY,
+        )
+        .await
+    }
+
+    async fn user_has_proofs_in_queue(&self, user_address: Address) -> bool {
+        let user_states = self.user_states.read().await;
+        let Some(user_state) = user_states.get(&user_address) else {
+            return false;
+        };
+
+        let Some(user_state_guard) = self
+            .try_user_lock_with_timeout(user_address, user_state.lock())
+            .await
+        else {
+            return false;
+        };
+
+        user_state_guard.proofs_in_batch > 0
+    }
+
+    async fn remove_user_proofs_and_reset_state(&self, user_address: Address) {
+        let mut user_states = self.user_states.write().await;
+
+        let mut batch_state_guard = match self
+            .try_batch_lock_with_timeout(self.batch_state.lock(), POLLING_EVENTS_LOCK_TIMEOUT)
+            .await
+        {
+            Some(guard) => guard,
+            None => {
+                error!(
+                    "Failed to acquire batch lock when trying to remove proofs from user {:?}, skipping removal",
+                    user_address
+                );
+                self.metrics.inc_unlocked_event_polling_batch_lock_timeout();
+                return;
+            }
+        };
+
+        let removed_entries = batch_state_guard
+            .batch_queue
+            .extract_if(|entry, _| entry.sender == user_address);
+
+        // Notify user via websocket
+        for (entry, _) in removed_entries {
+            if let Some(ws_sink) = entry.messaging_sink {
+                let ws_sink_clone = ws_sink.clone();
+                tokio::spawn(async move {
+                    send_message(
+                        ws_sink_clone.clone(),
+                        SubmitProofResponseMessage::UserFundsUnlocked,
+                    )
+                    .await;
+                });
+            }
+            info!(
+                "Removed proof with nonce {} for user {:?} from batch queue",
+                entry.nonced_verification_data.nonce, user_address
+            );
+        }
+
+        user_states.remove(&user_address);
+        info!(
+            "Removed UserState entry for user {:?} after processing BalanceUnlocked event",
+            user_address
+        );
+    }
+
     pub async fn listen_new_blocks_retryable(
         self: Arc<Self>,
     ) -> Result<(), RetryError<BatcherError>> {
@@ -1052,7 +1264,7 @@ impl Batcher {
         // * ---------------------------------------------------------------------*
 
         let Some(mut batch_state_lock) = self
-            .try_batch_lock_with_timeout(self.batch_state.lock())
+            .try_batch_lock_with_timeout(self.batch_state.lock(), MESSAGE_HANDLER_LOCK_TIMEOUT)
             .await
         else {
             send_message(ws_conn_sink.clone(), SubmitProofResponseMessage::ServerBusy).await;
@@ -1222,7 +1434,7 @@ impl Batcher {
         let replacement_max_fee = nonced_verification_data.max_fee;
         let nonce = nonced_verification_data.nonce;
         let Some(mut batch_state_guard) = self
-            .try_batch_lock_with_timeout(self.batch_state.lock())
+            .try_batch_lock_with_timeout(self.batch_state.lock(), MESSAGE_HANDLER_LOCK_TIMEOUT)
             .await
         else {
             drop(user_state_guard);
@@ -1817,14 +2029,16 @@ impl Batcher {
                 warn!("Failed to send task status to telemetry: {:?}", e);
             }
 
-            // decide if i want to flush the queue:
             match e {
+                // This should never happen, there is a task that regularly cleans up
+                // user proofs with unlocked states
+                // (and it runs more frequently than the 1H the user needs to withdraw funds)
                 BatcherError::TransactionSendError(
                     TransactionSendError::SubmissionInsufficientBalance(address),
                 ) => {
+                    // In the future we could do a more granular recovery
                     warn!("User {:?} has insufficient balance, flushing entire queue as safety measure", address);
-                    // TODO: In the future, we should re-add the failed batch back to the queue
-                    // For now, we flush everything as a safety measure
+
                     self.flush_queue_and_clear_nonce_cache().await;
                 }
                 _ => {

crates/batcher/src/main.rs

@@ -55,6 +55,16 @@ async fn main() -> Result<(), BatcherError> {
         }
     });
 
+    // spawn task to poll for BalanceUnlocked events
+    tokio::spawn({
+        let app = batcher.clone();
+        async move {
+            app.poll_balance_unlocked_events()
+                .await
+                .expect("Error polling BalanceUnlocked events")
+        }
+    });
+
     batcher.metrics.inc_batcher_restart();
 
     batcher.listen_connections(&address).await?;