Update tests to match reversion strategy

Author: Mirko-von-Leipzig

crates/block-producer/src/mempool/graph/dag.rs

@@ -279,6 +279,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashSet;
+
     use miden_protocol::block::BlockNumber;
 
     use super::*;
@@ -314,4 +316,35 @@ mod tests {
             graph.selection_candidates().keys().map(|id| **id).collect();
         assert_eq!(candidates_after_parent, vec![2]);
     }
+
+    #[test]
+    fn expired_and_unselected_returns_only_expired_unselected_nodes() {
+        let mut graph = Graph::<TestNode>::default();
+
+        graph
+            .append(
+                TestNode::new(1).with_output_notes([1]).with_expires_at(BlockNumber::from(5u32)),
+            )
+            .unwrap();
+        graph
+            .append(
+                TestNode::new(2).with_output_notes([2]).with_expires_at(BlockNumber::from(6u32)),
+            )
+            .unwrap();
+        graph
+            .append(
+                TestNode::new(3)
+                    .with_output_notes([3])
+                    .with_expires_at(BlockNumber::from(10u32)),
+            )
+            .unwrap();
+
+        graph.select_candidate(1);
+
+        let expired = graph.expired_and_unselected(BlockNumber::from(6u32));
+        let mut expected = HashSet::new();
+        expected.insert(2u32);
+
+        assert_eq!(expired, expected);
+    }
 }

crates/block-producer/src/mempool/mod.rs

@@ -266,14 +266,19 @@ impl Mempool {
     #[instrument(target = COMPONENT, name = "mempool.rollback_batch", skip_all)]
     pub fn rollback_batch(&mut self, batch: BatchId) {
         let reverted_batches = self.batches.revert_batch_and_descendants(batch);
-        let mut txs = Vec::new();
-        for reverted in reverted_batches {
-            self.transactions.requeue_transactions(&reverted);
-            txs.extend(reverted.transactions().iter().map(|tx| tx.id()));
+        for reverted in &reverted_batches {
+            self.transactions.requeue_transactions(reverted);
         }
 
-        let reverted_txs = self.transactions.increment_failure_count(txs.into_iter());
-        self.subscription.txs_reverted(reverted_txs);
+        // Find rolled back batch to mark the txs as failed.
+        //
+        // Note that its possible it doesn't exist, since this batch could have already been
+        // reverted as part of a separate rollback.
+        if let Some(batch) = reverted_batches.iter().find(|reverted| reverted.id() == batch) {
+            let failed_txs = batch.transactions().iter().map(|tx| tx.id());
+            let reverted_txs = self.transactions.increment_failure_count(failed_txs);
+            self.subscription.txs_reverted(reverted_txs);
+        }
 
         self.inject_telemetry();
     }
@@ -378,18 +383,18 @@ impl Mempool {
         // Revert the batches, and requeue the transactions for batch selection.
         //
         // Transactions which have failed excessively are also reverted.
-        let (_, batches) = self.pending_block.take().expect("we just checked it is some");
-        let mut txs = Vec::new();
-        for batch in batches {
+        let (_, block) = self.pending_block.take().expect("we just checked it is some");
+        for batch in &block {
             let reverted = self.batches.revert_batch_and_descendants(batch.id());
 
             for batch in reverted {
                 self.transactions.requeue_transactions(&batch);
-                txs.extend(batch.transactions().iter().map(|tx| tx.id()));
             }
         }
-
-        let reverted_txs = self.transactions.increment_failure_count(txs.into_iter());
+        let failed_txs = block
+            .iter()
+            .flat_map(|batch| batch.transactions().as_slice().iter().map(TransactionHeader::id));
+        let reverted_txs = self.transactions.increment_failure_count(failed_txs);
 
         self.subscription.txs_reverted(reverted_txs);
         self.inject_telemetry();

crates/block-producer/src/mempool/tests.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use miden_protocol::Word;
 use miden_protocol::block::{BlockHeader, BlockNumber};
+use miden_protocol::transaction::TransactionHeader;
 use pretty_assertions::assert_eq;
 use serial_test::serial;
 
@@ -114,6 +115,9 @@ fn failed_batch_transactions_are_requeued() {
     reference.select_batch().unwrap();
     reference.add_transaction(txs[1].clone()).unwrap();
     reference.add_transaction(txs[2].clone()).unwrap();
+    reference
+        .transactions
+        .increment_failure_count(failed_batch.txs().iter().map(|tx| tx.id()));
 
     assert_eq!(uut, reference);
 }
@@ -185,11 +189,9 @@ fn cannot_have_multiple_inflight_blocks() {
 // BLOCK FAILED TESTS
 // ================================================================================================
 
-/// A failed block should have all of its transactions reverted.
 #[test]
-fn block_failure_reverts_its_transactions() {
-    // We will revert everything so the reference should be the empty mempool.
-    let (mut uut, reference) = Mempool::for_tests();
+fn block_failure_increments_tx_failures() {
+    let (mut uut, mut reference) = Mempool::for_tests();
 
     let reverted_txs = MockProvenTxBuilder::sequential();
 
@@ -200,20 +202,51 @@ fn block_failure_reverts_its_transactions() {
     ])));
 
     // Block 1 will contain just the first batch.
-    let (number, _batches) = uut.select_block();
+    let (number, block) = uut.select_block();
 
     // Create another dependent batch.
     uut.add_transaction(reverted_txs[1].clone()).unwrap();
     uut.select_batch();
     // Create another dependent transaction.
     uut.add_transaction(reverted_txs[2].clone()).unwrap();
 
-    // Fail the block which should result in everything reverting.
     uut.rollback_block(number);
 
+    // Reference should contain all transactions, no batches, with tx failure from just that block.
+    reference.add_transaction(reverted_txs[0].clone()).unwrap();
+    reference.add_transaction(reverted_txs[1].clone()).unwrap();
+    reference.add_transaction(reverted_txs[2].clone()).unwrap();
+
+    reference.transactions.increment_failure_count(
+        block
+            .iter()
+            .flat_map(|batch| batch.transactions().as_slice().iter().map(TransactionHeader::id)),
+    );
+
     assert_eq!(uut, reference);
 }
 
+#[test]
+fn transactions_exceeding_failure_limit_are_removed() {
+    let (mut uut, _) = Mempool::for_tests();
+
+    let failing_tx = MockProvenTxBuilder::with_account_index(0).build();
+    let failing_tx = Arc::new(AuthenticatedTransaction::from_inner(failing_tx));
+    let tx_id = failing_tx.id();
+
+    uut.add_transaction(failing_tx).unwrap();
+
+    for _ in 0..3 {
+        let reverted = uut.transactions.increment_failure_count(std::iter::once(tx_id));
+        assert!(reverted.is_empty());
+        assert_eq!(uut.unbatched_transactions_count(), 1);
+    }
+
+    let reverted = uut.transactions.increment_failure_count(std::iter::once(tx_id));
+    assert!(reverted.contains(&tx_id));
+    assert_eq!(uut.unbatched_transactions_count(), 0);
+}
+
 /// We've decided that transactions from a rolled back batch should be requeued.
 ///
 /// This test checks this at a basic level by ensuring that rolling back a batch is the same as
@@ -244,6 +277,9 @@ fn transactions_from_reverted_batches_are_requeued() {
     reference.add_transaction(tx_set_a[1].clone()).unwrap();
     reference.add_transaction(tx_set_b[2].clone()).unwrap();
     reference.add_transaction(tx_set_a[2].clone()).unwrap();
+    reference
+        .transactions
+        .increment_failure_count([tx_set_a[1].id(), tx_set_b[1].id()].into_iter());
 
     assert_eq!(uut, reference);
 }
@@ -339,6 +375,9 @@ fn pass_through_txs_with_note_dependencies() {
     uut.rollback_batch(batch_a.id());
     reference.add_transaction(tx_pass_through_a).unwrap();
     reference.add_transaction(tx_pass_through_b).unwrap();
+    reference
+        .transactions
+        .increment_failure_count(batch_a.txs().iter().map(|tx| tx.id()));
 
     assert_eq!(uut, reference);
 }