fix: flush batcher queue only when needed

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)
+ }
  }
  }
 }