Make get_best_tx_ids_by_score_and_ancestry non-recursive

Cargo.toml

@@ -127,7 +127,7 @@ utxo = { path = "utxo" }
 
 [workspace.package]
 edition = "2024"
-rust-version = "1.91"
+rust-version = "1.92"
 version = "1.3.0"
 license = "MIT"
 

mempool/src/pool/tx_pool/collect_txs.rs

@@ -359,6 +359,23 @@ pub fn get_best_tx_ids_by_score_and_ancestry<M>(
 mod get_best_tx_ids_by_score_and_ancestry_impl {
     use super::*;
 
+    // Stack item for collect_nearest_selected_ancestors.
+    struct StackItem<'a, I: Iterator<Item = &'a Id<Transaction>>> {
+        tx_entry: &'a TxMempoolEntry,
+        tx_ancestors: BTreeSet<Id<Transaction>>,
+        parents_iter: I,
+    }
+
+    fn new_stack_item<'a>(
+        tx_entry: &'a TxMempoolEntry,
+    ) -> StackItem<'a, impl Iterator<Item = &'a Id<Transaction>>> {
+        StackItem {
+            tx_entry,
+            tx_ancestors: BTreeSet::new(),
+            parents_iter: tx_entry.parents(),
+        }
+    }
+
     // For each ancestor path starting from the specified tx, take the nearest ancestor that is in `selected_tx_ids`
     // and put it into a set. E.g. in
     //  /--S1--**--S2
@@ -379,30 +396,41 @@ mod get_best_tx_ids_by_score_and_ancestry_impl {
         selected_tx_ids: &BTreeSet<Id<Transaction>>,
         ancestors_map: &mut BTreeMap<Id<Transaction>, BTreeSet<Id<Transaction>>>,
     ) -> Result<(), TxCollectionError> {
-        let tx_id = tx_entry.tx_id();
+        if ancestors_map.contains_key(tx_entry.tx_id()) {
+            return Ok(());
+        }
 
-        if !ancestors_map.contains_key(tx_id) {
-            let mut tx_ancestors = BTreeSet::new();
+        let mut stack = vec![new_stack_item(tx_entry)];
 
-            for parent_id in tx_entry.parents() {
+        'outer: loop {
+            // We start with a non-empty stack and then check if it's empty at the end of the outer loop.
+            let item = stack.last_mut().expect("known to be present");
+
+            for parent_id in item.parents_iter.by_ref() {
                 if selected_tx_ids.contains(parent_id) {
                     // The nearest selected ancestor has been found, no need to go deeper.
-                    tx_ancestors.insert(*parent_id);
+                    item.tx_ancestors.insert(*parent_id);
+                } else if let Some(existing_ancestors) = ancestors_map.get(parent_id) {
+                    item.tx_ancestors.extend(existing_ancestors);
                 } else {
                     let parent_tx_entry = mempool.store.get_existing_entry(parent_id)?;
-                    collect_nearest_selected_ancestors(
-                        mempool,
-                        parent_tx_entry,
-                        selected_tx_ids,
-                        ancestors_map,
-                    )?;
-
-                    let parent_ancestors = ancestors_map.get(parent_id).expect("must be present");
-                    tx_ancestors.extend(parent_ancestors);
+                    stack.push(new_stack_item(parent_tx_entry));
+                    continue 'outer;
                 }
             }
 
-            ancestors_map.insert(*tx_id, tx_ancestors);
+            // We've completely processed the current item, so put the collected `tx_ancestors`
+            // into `ancestors_map`. If there is another item on top of the stack, it's this item's
+            // child, whose `tx_ancestors` must be extended with this item's `tx_ancestors`.
+            // Otherwise we're done.
+            let item = stack.pop().expect("known to be present");
+            let entry = ancestors_map.entry(*item.tx_entry.tx_id()).insert_entry(item.tx_ancestors);
+
+            if let Some(child_item) = stack.last_mut() {
+                child_item.tx_ancestors.extend(entry.get());
+            } else {
+                break;
+            }
         }
 
         Ok(())

mempool/src/pool/tx_pool/tests/tx_ids_by_score_and_ancestry.rs

@@ -174,3 +174,87 @@ async fn tx_ids_by_score_and_ancestry(#[case] seed: Seed) {
         ]
     );
 }
+
+// While using a small stack and unlimited `max_cluster_tx_count`, create a long chain of
+// dependent txs and pass the first and the last of them to `get_best_tx_ids_by_score_and_ancestry`
+// (so that it has to check all the txs in between). The expected result is that the test doesn't
+// crash. This basically checks that `get_best_tx_ids_by_score_and_ancestry` doesn't use recursion.
+#[rstest]
+#[case(Seed::from_entropy())]
+#[trace]
+fn stack_overflow_on_misconfigured_cluster_size_limit(#[case] seed: Seed) {
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(2)
+        // Use a small stack, same logic as in the `stack_overflow_on_transaction_addition` test.
+        .thread_stack_size(256 * 1024)
+        .enable_all()
+        .build()
+        .unwrap();
+
+    let test_body = move || {
+        let mut rng = make_seedable_rng(seed);
+        let mut tf = TestFramework::builder(&mut rng).build();
+        let genesis_id = tf.genesis().get_id();
+
+        // This was enough to crash the original code using the given stack size both in
+        // debug and release builds (actually, 1000 was also enough, but we use a bigger
+        // value "just in case").
+        let chain_len = 1500;
+
+        let mut amount = 1_000_000;
+        let funding_tx = make_tx(
+            &mut rng,
+            &[(OutPointSourceId::BlockReward(genesis_id.into()), 0)],
+            &[amount],
+        );
+        let funding_tx_id = funding_tx.transaction().get_id();
+        tf.make_block_builder()
+            .add_transaction(funding_tx)
+            .build_and_process(&mut rng)
+            .unwrap();
+
+        let mut mempool = setup_with_chainstate_generic(
+            tf.chainstate(),
+            MempoolConfig {
+                min_tx_relay_fee_rate: FeeRate::from_amount_per_kb(Amount::ZERO).into(),
+                max_cluster_size_bytes: usize::MAX.into(),
+                max_cluster_tx_count: usize::MAX.into(),
+            },
+            Default::default(),
+        );
+
+        // Disable heavy checks because the test is already relatively slow.
+        mempool.disable_heavy_validity_checks();
+
+        let mut prev_source: OutPointSourceId = funding_tx_id.into();
+        let mut tx_ids = Vec::new();
+
+        for _ in 0..chain_len {
+            let tx = make_tx(&mut rng, &[(prev_source, 0)], &[amount - 1]);
+            amount -= 1;
+            let tx_id = tx.transaction().get_id();
+            prev_source = tx_id.into();
+
+            let result = mempool.add_transaction_test(tx);
+
+            assert_eq!(result, Ok(TxStatus::InMempool));
+            tx_ids.push(tx_id);
+        }
+
+        let selected_tx_ids = [*tx_ids.first().unwrap(), *tx_ids.last().unwrap()];
+        let result = mempool
+            .get_best_tx_ids_by_score_and_ancestry(&BTreeSet::from(selected_tx_ids), tx_ids.len())
+            .unwrap();
+        assert_eq!(result, selected_tx_ids);
+    };
+
+    runtime.block_on(async move {
+        // Same as in `stack_overflow_on_transaction_addition`, need to use `spawn` or `spawn_blocking`
+        // here.
+        tokio::task::spawn_blocking(move || {
+            test_body();
+        })
+        .await
+        .unwrap();
+    });
+}