refactor: remove_proofs_from_queue called after submit is OK

batcher/aligned-batcher/src/lib.rs

@@ -1115,12 +1115,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;
 
@@ -1152,7 +1153,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,
@@ -1172,7 +1173,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();
 
@@ -1188,17 +1208,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,
@@ -1256,6 +1292,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,
@@ -1295,6 +1332,12 @@ impl Batcher {
  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 a COPY of the current batch queue, , and applyies an 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,10 +295,7 @@ 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();
-
- assert!(resulting_batch_queue.is_empty());
+ let batch = try_build_batch(batch_queue, gas_price, 5000000, 50).unwrap();
 
  assert_eq!(batch[0].nonced_verification_data.max_fee, max_fee_3);
  assert_eq!(batch[1].nonced_verification_data.max_fee, max_fee_2);
@@ -404,13 +395,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();
+ let 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
+ // All entries from the batch queue should be in
  // the finalized batch.
- assert!(resulting_batch_queue.is_empty());
  assert_eq!(finalized_batch.len(), 3);
  assert_eq!(
  finalized_batch[0].nonced_verification_data.max_fee,
@@ -515,14 +503,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();
+ let 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.
+ // All but one entries from the batch queue should be in the finalized batch.
 
- assert_eq!(resulting_batch_queue.len(), 1);
  assert_eq!(finalized_batch.len(), 2);
  assert_eq!(
  finalized_batch[0].nonced_verification_data.max_fee,
@@ -628,10 +612,9 @@ 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) =
+ let finalized_batch =
  try_build_batch(batch_queue, gas_price, 5000000, max_batch_proof_qty).unwrap();
 
- assert_eq!(resulting_batch_queue.len(), 1);
  assert_eq!(finalized_batch.len(), 2);
  assert_eq!(
  finalized_batch[0].nonced_verification_data.max_fee,

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

@@ -20,6 +20,7 @@ pub enum BatcherError {
  BatchCostTooHigh,
  WsSinkEmpty,
  AddressNotFoundInUserStates(Address),
+ QueueRemoveError(String),
 }
 
 impl From<tungstenite::Error> for BatcherError {
@@ -98,6 +99,9 @@ impl fmt::Debug for BatcherError {
  reason
  )
  }
+ BatcherError::QueueRemoveError(e) => {
+ write!(f, "Error while removing entry from queue: {}", e)
+ }
  }
  }
 }