feat(batcher): add queue limit (#1875)

Co-authored-by: Marcos Nicolau <[email protected]>
Co-authored-by: JuArce <[email protected]>
Co-authored-by: MauroFab <[email protected]>

Author: 4 people

batcher/aligned-batcher/src/config/mod.rs

@@ -46,6 +46,7 @@ pub struct BatcherConfigFromYaml {
     pub max_proof_size: usize,
     pub max_batch_byte_size: usize,
     pub max_batch_proof_qty: usize,
+    pub max_queue_size: usize,
     pub pre_verification_is_enabled: bool,
     pub metrics_port: u16,
     pub telemetry_ip_port_address: String,

batcher/aligned-batcher/src/lib.rs

@@ -210,7 +210,7 @@ impl Batcher {
         .expect("Failed to get fallback Service Manager contract");
 
         let mut user_states = HashMap::new();
-        let mut batch_state = BatchState::new();
+        let mut batch_state = BatchState::new(config.batcher.max_queue_size);
         let non_paying_config = if let Some(non_paying_config) = config.batcher.non_paying {
             warn!("Non-paying address configuration detected. Will replace non-paying address {} with configured address.",
                 non_paying_config.address);
@@ -228,7 +228,8 @@ impl Batcher {
                 non_paying_user_state,
             );
 
-            batch_state = BatchState::new_with_user_states(user_states);
+            batch_state =
+                BatchState::new_with_user_states(user_states, config.batcher.max_queue_size);
             Some(non_paying_config)
         } else {
             None
@@ -702,7 +703,7 @@ impl Batcher {
         // This is needed because we need to query the user state to make validations and
         // finally add the proof to the batch queue.
 
-        let batch_state_lock = self.batch_state.lock().await;
+        let mut batch_state_lock = self.batch_state.lock().await;
 
         let msg_max_fee = nonced_verification_data.max_fee;
         let Some(user_last_max_fee_limit) =
@@ -782,6 +783,8 @@ impl Batcher {
             return Ok(());
         }
 
+        // We check this after replacement logic because if user wants to replace a proof, their
+        // new_max_fee must be greater or equal than old_max_fee
         if msg_max_fee > user_last_max_fee_limit {
             std::mem::drop(batch_state_lock);
             warn!("Invalid max fee for address {addr}, had fee limit of {user_last_max_fee_limit:?}, sent {msg_max_fee:?}");
@@ -794,6 +797,67 @@ impl Batcher {
             return Ok(());
         }
 
+        // * ---------------------------------------------------------------------*
+        // *        Perform validation over batcher queue                         *
+        // * ---------------------------------------------------------------------*
+
+        if batch_state_lock.is_queue_full() {
+            debug!("Batch queue is full. Evaluating if the incoming proof can replace a lower-priority entry.");
+
+            // This cannot panic, if the batch queue is full it has at least one item
+            let (lowest_priority_entry, _) = batch_state_lock
+                .batch_queue
+                .peek()
+                .expect("Batch queue was expected to be full, but somehow no item was inside");
+
+            let lowest_fee_in_queue = lowest_priority_entry.nonced_verification_data.max_fee;
+
+            let new_proof_fee = nonced_verification_data.max_fee;
+
+            // We will keep the proof with the highest fee
+            // Note: we previously checked that if it's a new proof from the same user the fee is the same or lower
+            // So this will never eject a proof of the same user with a lower nonce
+            // which is the expected behaviour
+            if new_proof_fee > lowest_fee_in_queue {
+                // This cannot panic, if the batch queue is full it has at least one item
+                let (removed_entry, _) = batch_state_lock
+                    .batch_queue
+                    .pop()
+                    .expect("Batch queue was expected to be full, but somehow no item was inside");
+
+                info!(
+                    "Incoming proof (nonce: {}, fee: {}) has higher fee. Replacing lowest fee proof from sender {} with nonce {}.",
+                    nonced_verification_data.nonce,
+                    nonced_verification_data.max_fee,
+                    removed_entry.sender,
+                    removed_entry.nonced_verification_data.nonce
+                );
+
+                batch_state_lock.update_user_state_on_entry_removal(&removed_entry);
+
+                if let Some(removed_entry_ws) = removed_entry.messaging_sink {
+                    send_message(
+                        removed_entry_ws,
+                        SubmitProofResponseMessage::UnderpricedProof,
+                    )
+                    .await;
+                };
+            } else {
+                info!(
+                    "Incoming proof (nonce: {}, fee: {}) has lower priority than all entries in the full queue. Rejecting submission.",
+                    nonced_verification_data.nonce,
+                    nonced_verification_data.max_fee
+                );
+                std::mem::drop(batch_state_lock);
+                send_message(
+                    ws_conn_sink.clone(),
+                    SubmitProofResponseMessage::UnderpricedProof,
+                )
+                .await;
+                return Ok(());
+            }
+        }
+
         // * ---------------------------------------------------------------------*
         // *        Add message data into the queue and update user state         *
         // * ---------------------------------------------------------------------*
@@ -1727,6 +1791,16 @@ impl Batcher {
 
         let batch_state_lock = self.batch_state.lock().await;
 
+        if batch_state_lock.is_queue_full() {
+            error!("Can't add new entry, the batcher queue is full");
+            send_message(
+                ws_sink.clone(),
+                SubmitProofResponseMessage::UnderpricedProof,
+            )
+            .await;
+            return Ok(());
+        }
+
         let nonced_verification_data = NoncedVerificationData::new(
             client_msg.verification_data.verification_data.clone(),
             client_msg.verification_data.nonce,

batcher/aligned-batcher/src/types/batch_state.rs

@@ -10,22 +10,28 @@ use log::debug;
 pub(crate) struct BatchState {
     pub(crate) batch_queue: BatchQueue,
     pub(crate) user_states: HashMap<Address, UserState>,
+    pub(crate) max_size: usize,
 }
 
 impl BatchState {
     // CONSTRUCTORS:
 
-    pub(crate) fn new() -> Self {
+    pub(crate) fn new(max_size: usize) -> Self {
         Self {
             batch_queue: BatchQueue::new(),
             user_states: HashMap::new(),
+            max_size,
         }
     }
 
-    pub(crate) fn new_with_user_states(user_states: HashMap<Address, UserState>) -> Self {
+    pub(crate) fn new_with_user_states(
+        user_states: HashMap<Address, UserState>,
+        max_size: usize,
+    ) -> Self {
         Self {
             batch_queue: BatchQueue::new(),
             user_states,
+            max_size,
         }
     }
 
@@ -214,4 +220,39 @@ impl BatchState {
                 && entry.nonced_verification_data.max_fee < replacement_max_fee
         })
     }
+
+    /// Updates or removes a user's state when their latest proof entry is removed from the batch queue.
+    ///
+    /// If the user has no other proofs remaining in the queue, their state is removed entirely.
+    /// Otherwise, the user's state is updated to reflect the next most recent entry in the queue.
+    ///
+    /// Note: The given `removed_entry` must be the most recent (latest or highest nonce) entry for the user in the queue.
+    pub(crate) fn update_user_state_on_entry_removal(&mut self, removed_entry: &BatchQueueEntry) {
+        let addr = removed_entry.sender;
+
+        let new_last_max_fee_limit = match self
+            .batch_queue
+            .iter()
+            .filter(|(e, _)| e.sender == addr)
+            .next_back()
+        {
+            Some((last_entry, _)) => last_entry.nonced_verification_data.max_fee,
+            None => {
+                self.user_states.remove(&addr);
+                return;
+            }
+        };
+
+        if let Entry::Occupied(mut user_state) = self.user_states.entry(addr) {
+            user_state.get_mut().proofs_in_batch -= 1;
+            user_state.get_mut().nonce -= U256::one();
+            user_state.get_mut().total_fees_in_queue -=
+                removed_entry.nonced_verification_data.max_fee;
+            user_state.get_mut().last_max_fee_limit = new_last_max_fee_limit;
+        }
+    }
+
+    pub(crate) fn is_queue_full(&self) -> bool {
+        self.batch_queue.len() >= self.max_size
+    }
 }

batcher/aligned-sdk/src/communication/messaging.rs

@@ -265,6 +265,10 @@ async fn handle_batcher_response(msg: Message) -> Result<BatchInclusionData, Sub
             );
             Err(SubmitError::GenericError(e))
         }
+        Ok(SubmitProofResponseMessage::UnderpricedProof) => {
+            error!("Batcher responded with error: queue limit has been exceeded. Funds have not been spent.");
+            Err(SubmitError::BatchQueueLimitExceededError)
+        }
         Err(e) => {
             error!(
                 "Error while deserializing batch inclusion data: {}. Funds have not been spent.",

batcher/aligned-sdk/src/core/errors.rs

@@ -96,6 +96,7 @@ pub enum SubmitError {
     AddToBatchError,
     InvalidProofInclusionData,
     GetNonceError(String),
+    BatchQueueLimitExceededError,
     GenericError(String),
 }
 
@@ -210,6 +211,10 @@ impl fmt::Display for SubmitError {
             SubmitError::InvalidProofInclusionData => {
                 write!(f, "Batcher responded with invalid batch inclusion data. Can't verify your proof was correctly included in the batch.")
             }
+            SubmitError::BatchQueueLimitExceededError => {
+                write!(f, "Error while adding entry to batch, queue limit exeeded.")
+            }
+
             SubmitError::GetNonceError(e) => write!(f, "Error while getting nonce {}", e),
         }
     }

batcher/aligned-sdk/src/core/types.rs

@@ -399,6 +399,7 @@ pub enum SubmitProofResponseMessage {
     AddToBatchError,
     EthRpcError,
     InvalidPaymentServiceAddress(Address, Address),
+    UnderpricedProof,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

config-files/config-batcher-docker.yaml

@@ -21,9 +21,10 @@ batcher:
   block_interval: 3
   batch_size_interval: 10
   transaction_wait_timeout: 96000 # 8 blocks
-  max_proof_size: 67108864 # 64 MiB
+  max_proof_size: 4194304 # 4 MiB
   max_batch_byte_size: 268435456 # 256 MiB
   max_batch_proof_qty: 3000 # 3000 proofs in a batch
+  max_queue_size: 10000
   pre_verification_is_enabled: true
   metrics_port: 9093
   telemetry_ip_port_address: localhost:4001

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

@@ -19,9 +19,10 @@ batcher:
   block_interval: 3
   batch_size_interval: 10
   transaction_wait_timeout: 36000 # 3 blocks
-  max_proof_size: 67108864 # 64 MiB
+  max_proof_size: 4194304 # 4 MiB
   max_batch_byte_size: 268435456 # 256 MiB
   max_batch_proof_qty: 3000 # 3000 proofs in a batch
+  max_queue_size: 10000
   pre_verification_is_enabled: true
   metrics_port: 9093
   telemetry_ip_port_address: localhost:4001

config-files/config-batcher.yaml

@@ -21,9 +21,10 @@ batcher:
   block_interval: 3
   batch_size_interval: 10
   transaction_wait_timeout: 96000 # 8 blocks
-  max_proof_size: 67108864 # 64 MiB
+  max_proof_size: 4194304 # 4 MiB
   max_batch_byte_size: 268435456 # 256 MiB
   max_batch_proof_qty: 3000 # 3000 proofs in a batch
+  max_queue_size: 10000
   pre_verification_is_enabled: true
   metrics_port: 9093
   telemetry_ip_port_address: localhost:4001