fix: flush batcher queue only when needed (#1512)

Co-authored-by: Julian Ventura <[email protected]>
Co-authored-by: avilagaston9 <[email protected]>
Co-authored-by: Marcos Nicolau <[email protected]>
Co-authored-by: Marcos Nicolau <[email protected]>
Co-authored-by: Avila Gastón <[email protected]>

Author: 4 people

batcher/aligned-batcher/src/lib.rs

@@ -48,7 +48,7 @@ use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::{Mutex, MutexGuard, RwLock};
 use tokio_tungstenite::tungstenite::{Error, Message};
 use types::batch_queue::{self, BatchQueueEntry, BatchQueueEntryPriority};
-use types::errors::BatcherError;
+use types::errors::{BatcherError, TransactionSendError};
 
 use crate::config::{ConfigFromYaml, ContractDeploymentOutput};
 use crate::telemetry::sender::TelemetrySender;
@@ -1116,12 +1116,13 @@ impl Batcher {
     /// an empty batch, even if the block interval has been reached.
     /// Once the batch meets the conditions for submission, the finalized batch is then passed to the
     /// `finalize_batch` function.
+    /// This function doesn't remove the proofs from the queue.
     async fn is_batch_ready(
         &self,
         block_number: u64,
         gas_price: U256,
     ) -> Option<Vec<BatchQueueEntry>> {
-        let mut batch_state_lock = self.batch_state.lock().await;
+        let batch_state_lock = self.batch_state.lock().await;
         let current_batch_len = batch_state_lock.batch_queue.len();
         let last_uploaded_batch_block_lock = self.last_uploaded_batch_block.lock().await;
 
@@ -1153,7 +1154,7 @@ impl Batcher {
         // Set the batch posting flag to true
         *batch_posting = true;
         let batch_queue_copy = batch_state_lock.batch_queue.clone();
-        let (resulting_batch_queue, finalized_batch) = batch_queue::try_build_batch(
+        let finalized_batch = batch_queue::try_build_batch(
             batch_queue_copy,
             gas_price,
             self.max_batch_byte_size,
@@ -1173,7 +1174,26 @@ impl Batcher {
         })
         .ok()?;
 
-        batch_state_lock.batch_queue = resulting_batch_queue;
+        Some(finalized_batch)
+    }
+
+    /// Takes the submitted proofs and removes them from the queue.
+    /// This function should be called only AFTER the submission was confirmed onchain
+    async fn remove_proofs_from_queue(
+        &self,
+        finalized_batch: Vec<BatchQueueEntry>,
+    ) -> Result<(), BatcherError> {
+        info!("Removing proofs from queue...");
+        let mut batch_state_lock = self.batch_state.lock().await;
+
+        finalized_batch.iter().for_each(|entry| {
+            if batch_state_lock.batch_queue.remove(entry).is_none() {
+                // If this happens, we have a bug in our code
+                error!("Some proofs were not found in the queue. This should not happen.");
+            }
+        });
+
+        // now we calculate the new user_states
         let new_user_states = // proofs, max_fee_limit, total_fees_in_queue
             batch_state_lock.calculate_new_user_states_data();
 
@@ -1189,17 +1209,33 @@ impl Batcher {
             // informative error.
 
             // Now we update the user states related to the batch (proof count in batch and min fee in batch)
-            batch_state_lock.update_user_proof_count(addr, *proof_count)?;
-            batch_state_lock.update_user_max_fee_limit(addr, *max_fee_limit)?;
-            batch_state_lock.update_user_total_fees_in_queue(addr, *total_fees_in_queue)?;
+            batch_state_lock
+                .update_user_proof_count(addr, *proof_count)
+                .ok_or(BatcherError::QueueRemoveError(
+                    "Could not update_user_proof_count".into(),
+                ))?;
+            batch_state_lock
+                .update_user_max_fee_limit(addr, *max_fee_limit)
+                .ok_or(BatcherError::QueueRemoveError(
+                    "Could not update_user_max_fee_limit".into(),
+                ))?;
+            batch_state_lock
+                .update_user_total_fees_in_queue(addr, *total_fees_in_queue)
+                .ok_or(BatcherError::QueueRemoveError(
+                    "Could not update_user_total_fees_in_queue".into(),
+                ))?;
         }
 
-        Some(finalized_batch)
+        Ok(())
     }
 
-    /// Takes the finalized batch as input and builds the merkle tree, posts verification data batch
-    /// to s3, creates new task in Aligned contract and sends responses to all clients that added proofs
-    /// to the batch. The last uploaded batch block is updated once the task is created in Aligned.
+    /// Takes the finalized batch as input and:
+    ///     builds the merkle tree
+    ///     posts verification data batch to s3
+    ///     creates new task in Aligned contract
+    ///     removes the proofs from the queue, once they are succesfully submitted on-chain
+    ///     sends responses to all clients that added proofs to the batch.
+    /// The last uploaded batch block is updated once the task is created in Aligned.
     async fn finalize_batch(
         &self,
         block_number: u64,
@@ -1257,6 +1293,7 @@ impl Batcher {
             warn!("Failed to initialize task trace on telemetry: {:?}", e);
         }
 
+        // Here we submit the batch on-chain
         if let Err(e) = self
             .submit_batch(
                 &batch_bytes,
@@ -1275,27 +1312,30 @@ impl Batcher {
             {
                 warn!("Failed to send task status to telemetry: {:?}", e);
             }
-            for entry in finalized_batch.into_iter() {
-                if let Some(ws_sink) = entry.messaging_sink {
-                    let merkle_root = hex::encode(batch_merkle_tree.root);
-                    send_message(
-                        ws_sink.clone(),
-                        SubmitProofResponseMessage::CreateNewTaskError(
-                            merkle_root,
-                            format!("{:?}", e),
-                        ),
-                    )
-                    .await
-                } else {
-                    warn!("Websocket sink was found empty. This should only happen in tests");
+
+            // decide if i want to flush the queue:
+            match e {
+                BatcherError::TransactionSendError(
+                    TransactionSendError::SubmissionInsufficientBalance,
+                ) => {
+                    // TODO calling remove_proofs_from_queue here is a better solution, flushing only the failed batch
+                    // this would also need a message sent to the clients
+                    self.flush_queue_and_clear_nonce_cache().await;
+                }
+                _ => {
+                    // Add more cases here if we want in the future
                 }
             }
 
-            self.flush_queue_and_clear_nonce_cache().await;
-
             return Err(e);
         };
 
+        // Once the submit is succesfull, we remove the submitted proofs from the queue
+        // TODO handle error case:
+        if let Err(e) = self.remove_proofs_from_queue(finalized_batch.clone()).await {
+            error!("Unexpected error while updating queue: {:?}", e);
+        }
+
         connection::send_batch_inclusion_data_responses(finalized_batch, &batch_merkle_tree).await
     }
 

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

@@ -134,27 +134,25 @@ pub(crate) fn calculate_batch_size(batch_queue: &BatchQueue) -> Result<usize, Ba
 }
 
 /// This function tries to build a batch to be submitted to Aligned.
-/// Given a copy of the current batch queue, , and applyies an algorithm to find the biggest batch
+/// Given the current batch queue applies the following algorithm to find the biggest batch
 /// of proofs from users that are willing to pay for it:
 /// 1. Traverse each batch priority queue, starting from the one with minimum max fee.
 /// 2. Calculate the `fee_per_proof` for the whole batch and compare with the `max_fee` of the entry.
 /// 3. If `fee_per_proof` is less than the `max_fee` of the current entry, submit the batch. If not, pop this entry
-///     from the queue and push it to `resulting_priority_queue`, then repeat step 1.
+///     from the queue. then repeat step 1.
 ///
-/// `resulting_priority_queue` will be the batch queue composed of all entries that were not willing to pay for the batch.
-/// This is outputted in along with the finalized batch.
+/// Returns the finalized batch.
 pub(crate) fn try_build_batch(
     batch_queue: BatchQueue,
     gas_price: U256,
     max_batch_byte_size: usize,
     max_batch_proof_qty: usize,
-) -> Result<(BatchQueue, Vec<BatchQueueEntry>), BatcherError> {
-    let mut batch_queue = batch_queue;
-    let mut batch_size = calculate_batch_size(&batch_queue)?;
-    let mut resulting_priority_queue = BatchQueue::new();
+) -> Result<Vec<BatchQueueEntry>, BatcherError> {
+    let mut finalized_batch = batch_queue;
+    let mut batch_size = calculate_batch_size(&finalized_batch)?;
 
-    while let Some((entry, _)) = batch_queue.peek() {
-        let batch_len = batch_queue.len();
+    while let Some((entry, _)) = finalized_batch.peek() {
+        let batch_len = finalized_batch.len();
         let fee_per_proof = calculate_fee_per_proof(batch_len, gas_price);
 
         if batch_size > max_batch_byte_size
@@ -173,8 +171,7 @@ pub(crate) fn try_build_batch(
                     .len();
             batch_size -= verification_data_size;
 
-            let (not_working_entry, not_working_priority) = batch_queue.pop().unwrap();
-            resulting_priority_queue.push(not_working_entry, not_working_priority);
+            finalized_batch.pop();
 
             continue;
         }
@@ -183,16 +180,13 @@ pub(crate) fn try_build_batch(
         break;
     }
 
-    // If `batch_queue_copy` is empty, this means that all the batch queue was traversed and we didn't find
+    // If `finalized_batch` is empty, this means that all the batch queue was traversed and we didn't find
     // any user willing to pay fot the fee per proof.
-    if batch_queue.is_empty() {
+    if finalized_batch.is_empty() {
         return Err(BatcherError::BatchCostTooHigh);
     }
 
-    Ok((
-        resulting_priority_queue,
-        batch_queue.clone().into_sorted_vec(),
-    ))
+    Ok(finalized_batch.clone().into_sorted_vec())
 }
 
 fn calculate_fee_per_proof(batch_len: usize, gas_price: U256) -> U256 {
@@ -301,14 +295,20 @@ mod test {
         batch_queue.push(entry_3, batch_priority_3);
 
         let gas_price = U256::from(1);
-        let (resulting_batch_queue, batch) =
-            try_build_batch(batch_queue, gas_price, 5000000, 50).unwrap();
+        let finalized_batch = try_build_batch(batch_queue, gas_price, 5000000, 50).unwrap();
 
-        assert!(resulting_batch_queue.is_empty());
-
-        assert_eq!(batch[0].nonced_verification_data.max_fee, max_fee_3);
-        assert_eq!(batch[1].nonced_verification_data.max_fee, max_fee_2);
-        assert_eq!(batch[2].nonced_verification_data.max_fee, max_fee_1);
+        assert_eq!(
+            finalized_batch[0].nonced_verification_data.max_fee,
+            max_fee_3
+        );
+        assert_eq!(
+            finalized_batch[1].nonced_verification_data.max_fee,
+            max_fee_2
+        );
+        assert_eq!(
+            finalized_batch[2].nonced_verification_data.max_fee,
+            max_fee_1
+        );
     }
 
     #[test]
@@ -404,13 +404,11 @@ mod test {
         batch_queue.push(entry_3, batch_priority_3);
 
         let gas_price = U256::from(1);
-        let (resulting_batch_queue, finalized_batch) =
-            try_build_batch(batch_queue, gas_price, 5000000, 50).unwrap();
+        let finalized_batch = try_build_batch(batch_queue.clone(), gas_price, 5000000, 50).unwrap();
 
-        // The resulting batch queue (entries from the old batch queue that were not willing to pay
-        // in this batch), should be empty and hence, all entries from the batch queue should be in
+        // All entries from the batch queue should be in
         // the finalized batch.
-        assert!(resulting_batch_queue.is_empty());
+        assert_eq!(batch_queue.len(), 3);
         assert_eq!(finalized_batch.len(), 3);
         assert_eq!(
             finalized_batch[0].nonced_verification_data.max_fee,
@@ -515,14 +513,10 @@ mod test {
         batch_queue.push(entry_3, batch_priority_3);
 
         let gas_price = U256::from(1);
-        let (resulting_batch_queue, finalized_batch) =
-            try_build_batch(batch_queue, gas_price, 5000000, 50).unwrap();
-
-        // The resulting batch queue (entries from the old batch queue that were not willing to pay
-        // in this batch), should be empty and hence, all entries from the batch queue should be in
-        // the finalized batch.
+        let finalized_batch = try_build_batch(batch_queue.clone(), gas_price, 5000000, 50).unwrap();
 
-        assert_eq!(resulting_batch_queue.len(), 1);
+        // All but one entries from the batch queue should be in the finalized batch.
+        assert_eq!(batch_queue.len(), 3);
         assert_eq!(finalized_batch.len(), 2);
         assert_eq!(
             finalized_batch[0].nonced_verification_data.max_fee,
@@ -628,10 +622,10 @@ mod test {
         // The max batch len is 2, so the algorithm should stop at the second entry.
         let max_batch_proof_qty = 2;
 
-        let (resulting_batch_queue, finalized_batch) =
-            try_build_batch(batch_queue, gas_price, 5000000, max_batch_proof_qty).unwrap();
+        let finalized_batch =
+            try_build_batch(batch_queue.clone(), gas_price, 5000000, max_batch_proof_qty).unwrap();
 
-        assert_eq!(resulting_batch_queue.len(), 1);
+        assert_eq!(batch_queue.len(), 3);
         assert_eq!(finalized_batch.len(), 2);
         assert_eq!(
             finalized_batch[0].nonced_verification_data.max_fee,

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

@@ -27,12 +27,13 @@ impl From<Bytes> for TransactionSendError {
             "0xa3a8658a" => TransactionSendError::NoFeePerProof,     // can't happen, don't flush
             "0x7899ec71" => TransactionSendError::InsufficientFeeForAggregator, // shouldn't happen, don't flush
             // returning the proofs and retrying later may help
-            "0x4f779ceb" => TransactionSendError::SubmissionInsufficientBalance, // shouldn't happen,
-            // flush can help if something went wrong
             "0x3102f10c" => TransactionSendError::BatchAlreadySubmitted, // can happen, don't flush
             "0x5c54305e" => TransactionSendError::InsufficientFunds, // shouldn't happen, don't flush
             "0x152bc288" => TransactionSendError::OnlyBatcherAllowed, // won't happen, don't flush
+            "0x4f779ceb" => TransactionSendError::SubmissionInsufficientBalance, // shouldn't happen,
+            // flush can help if something went wrong
             _ => {
+                // flush because unkown error
                 TransactionSendError::Generic(format!("Unknown bytestring error: {}", byte_string))
             }
         }
@@ -56,6 +57,7 @@ pub enum BatcherError {
     BatchCostTooHigh,
     WsSinkEmpty,
     AddressNotFoundInUserStates(Address),
+    QueueRemoveError(String),
 }
 
 impl From<tungstenite::Error> for BatcherError {
@@ -134,6 +136,9 @@ impl fmt::Debug for BatcherError {
                     reason
                 )
             }
+            BatcherError::QueueRemoveError(e) => {
+                write!(f, "Error while removing entry from queue: {}", e)
+            }
         }
     }
 }