Merge pull request #5219 from stacks-network/feat/mempool-opt

Feat: optimize mempool iteration by skipping last invocation
stacks-network · Sep 20, 2024 · b5250c6 · b5250c6
2 parents 7cac88a + cb1a47c
commit b5250c6
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 12 deletions.
diff --git a/stackslib/src/chainstate/stacks/miner.rs b/stackslib/src/chainstate/stacks/miner.rs
@@ -2228,10 +2228,10 @@ impl StacksBlockBuilder {
 
         debug!("Block transaction selection begins (parent height = {tip_height})");
         let result = {
-            let mut intermediate_result: Result<_, Error> = Ok(0);
+            let mut loop_result = Ok(());
             while block_limit_hit != BlockLimitFunction::LIMIT_REACHED {
                 let mut num_considered = 0;
-                intermediate_result = mempool.iterate_candidates(
+                let intermediate_result = mempool.iterate_candidates(
                     epoch_tx,
                     &mut tx_events,
                     mempool_settings.clone(),
@@ -2390,16 +2390,27 @@ impl StacksBlockBuilder {
                     let _ = mempool.drop_and_blacklist_txs(&to_drop_and_blacklist);
                 }
 
-                if intermediate_result.is_err() {
-                    break;
+                match intermediate_result {
+                    Err(e) => {
+                        loop_result = Err(e);
+                        break;
+                    }
+                    Ok((_txs_considered, stop_reason)) => {
+                        match stop_reason {
+                            MempoolIterationStopReason::NoMoreCandidates => break,
+                            MempoolIterationStopReason::DeadlineReached => break,
+                            // if the iterator function exited, let the loop tick: it checks the block limits
+                            MempoolIterationStopReason::IteratorExited => {}
+                        }
+                    }
                 }
 
                 if num_considered == 0 {
                     break;
                 }
             }
             debug!("Block transaction selection finished (parent height {}): {} transactions selected ({} considered)", &tip_height, num_txs, considered.len());
-            intermediate_result
+            loop_result
         };
 
         mempool.drop_txs(&invalidated_txs)?;

diff --git a/stackslib/src/chainstate/stacks/tests/block_construction.rs b/stackslib/src/chainstate/stacks/tests/block_construction.rs
@@ -5072,6 +5072,7 @@ fn paramaterized_mempool_walk_test(
                         },
                     )
                     .unwrap()
+                    .0
                     == 0
                 {
                     break;

diff --git a/stackslib/src/core/mempool.rs b/stackslib/src/core/mempool.rs
@@ -144,6 +144,14 @@ pub enum MemPoolSyncData {
     TxTags([u8; 32], Vec<TxTag>),
 }
 
+pub enum MempoolIterationStopReason {
+    NoMoreCandidates,
+    DeadlineReached,
+    /// If the iteration function supplied to mempool iteration exited
+    ///  (i.e., the transaction evaluator returned an early exit command)
+    IteratorExited,
+}
+
 impl StacksMessageCodec for MemPoolSyncData {
     fn consensus_serialize<W: Write>(&self, fd: &mut W) -> Result<(), codec_error> {
         match *self {
@@ -1592,7 +1600,7 @@ impl MemPoolDB {
         output_events: &mut Vec<TransactionEvent>,
         settings: MemPoolWalkSettings,
         mut todo: F,
-    ) -> Result<u64, E>
+    ) -> Result<(u64, MempoolIterationStopReason), E>
     where
         C: ClarityConnection,
         F: FnMut(
@@ -1643,11 +1651,11 @@ impl MemPoolDB {
             .query(NO_PARAMS)
             .map_err(|err| Error::SqliteError(err))?;
 
-        loop {
+        let stop_reason = loop {
             if start_time.elapsed().as_millis() > settings.max_walk_time_ms as u128 {
                 debug!("Mempool iteration deadline exceeded";
                        "deadline_ms" => settings.max_walk_time_ms);
-                break;
+                break MempoolIterationStopReason::DeadlineReached;
             }
 
             let start_with_no_estimate =
@@ -1687,7 +1695,7 @@ impl MemPoolDB {
                                 ),
                                 None => {
                                     debug!("No more transactions to consider in mempool");
-                                    break;
+                                    break MempoolIterationStopReason::NoMoreCandidates;
                                 }
                             }
                         }
@@ -1875,7 +1883,7 @@ impl MemPoolDB {
                 }
                 None => {
                     debug!("Mempool iteration early exit from iterator");
-                    break;
+                    break MempoolIterationStopReason::IteratorExited;
                 }
             }
 
@@ -1885,7 +1893,7 @@ impl MemPoolDB {
                 candidate_cache.len()
             );
             candidate_cache.reset();
-        }
+        };
 
         // drop these rusqlite statements and queries, since their existence as immutable borrows on the
         // connection prevents us from beginning a transaction below (which requires a mutable
@@ -1908,7 +1916,7 @@ impl MemPoolDB {
             "considered_txs" => u128::from(total_considered),
             "elapsed_ms" => start_time.elapsed().as_millis()
         );
-        Ok(total_considered)
+        Ok((total_considered, stop_reason))
     }
 
     pub fn conn(&self) -> &DBConn {