sim-rs: simplify mempool behavior

Author: SupernaviX

sim-rs/implementations/LINEAR_LEIOS.md

@@ -11,13 +11,21 @@ Whenever a node creates an RB, it also creates an EB. The RB header contains a r
 
 RB headers are diffused separately from bodies. When a node receives an RB header, it checks whether that RB should be the new head of its chain. If so, it will request the RB body and the referenced EB (from the first peer which announces them).
 
-When a node receives an RB body, it immediately removes all referenced/conflicting transactions from its mempool. If the RB has an EB certificate, it also removes that EB’s transactions from its mempool (note that this should be redundant based on the procedure below).
-
 When a node receives an EB body, it runs lightweight validation and then propagates the body to peers. After this lightweight validation, it runs more expensive complete validation (presumably at the TX level) before voting.
 
 To detect equivocation, a node will wait until at least `3 * Δhdr` after an EB was generated before voting for it.
 
-When voting, a node runs a VRF lottery to decide how many times it can vote for that EB; if it has any votes, it will transmit them to all peers. If the EB has been certified after `L_vote` + `L_diff` slots have passed, the node removes all of its transactions from the mempool (under the assumption that the EB will make it on-chain).
+When voting, a node runs a VRF lottery to decide how many times it can vote for that EB; if it has any votes, it will transmit them to all peers.
+
+## Mempool behavior
+
+When a node creates an RB, it will follow these steps in order:
+ 1. Try to produce a cert for the parent RB's EB.
+   1. If this succeeds, remove all of this EB's transactions from its mempool.
+ 2. Fill the RB body with transactions from its mempool
+ 3. Create a new EB, filled with transactions from its mempool WITHOUT removing those transactions from the mempool.
+
+When a node receives an RB body, it immediately removes all referenced/conflicting transactions from its mempool. If the RB has an EB certificate, it also removes that EB’s transactions from its mempool.
 
 ## New parameters
 

sim-rs/sim-core/src/sim/linear_leios.rs

@@ -325,7 +325,6 @@ impl NodeImpl for LinearLeiosNode {
     }
 
     fn handle_new_slot(&mut self, slot: u64) -> EventResult {
-        self.process_certified_ebs(slot);
         self.try_generate_rb(slot);
 
         std::mem::take(&mut self.queued)
@@ -482,6 +481,41 @@ impl LinearLeiosNode {
             return;
         };
 
+        let parent = self.latest_rb().map(|rb| rb.header.id);
+        let endorsement = parent.and_then(|rb_id| {
+            if rb_id.slot
+                + self.sim_config.linear_vote_stage_length
+                + self.sim_config.linear_diffuse_stage_length
+                > slot
+            {
+                // This RB was generated too quickly after another; hasn't been time to gather all the votes.
+                // No endorsement.
+                return None;
+            }
+
+            let eb_id = self.leios.ebs_by_rb.get(&rb_id)?;
+            let votes = self.leios.votes_by_eb.get(eb_id)?;
+            let total_votes = votes.values().copied().sum::<usize>();
+            if (total_votes as u64) < self.sim_config.vote_threshold {
+                // Not enough votes. No endorsement.
+                return None;
+            }
+
+            let endorsement = Endorsement {
+                eb: *eb_id,
+                size_bytes: self.sim_config.sizes.cert(votes.len()),
+                votes: votes.clone(),
+            };
+
+            // We're endorsing this EB, so consider all of its transactions part of the ledger state
+            let Some(EndorserBlockView::Received { eb }) = self.leios.ebs.get(eb_id) else {
+                return None;
+            };
+            self.remove_eb_txs_from_mempool(&eb.clone());
+
+            Some(endorsement)
+        });
+
         let mut rb_transactions = vec![];
         if self.sim_config.praos_fallback {
             if let TransactionConfig::Mock(config) = &self.sim_config.transactions {
@@ -490,7 +524,11 @@ impl LinearLeiosNode {
                 self.tracker.track_transaction_generated(&tx, self.id);
                 rb_transactions.push(Arc::new(tx));
             } else {
-                self.sample_from_mempool(&mut rb_transactions, self.sim_config.max_block_size);
+                self.sample_from_mempool(
+                    &mut rb_transactions,
+                    self.sim_config.max_block_size,
+                    true,
+                );
             }
         }
 
@@ -516,36 +554,9 @@ impl LinearLeiosNode {
                 eb_transactions.push(Arc::new(tx));
             }
         } else {
-            self.sample_from_mempool(&mut eb_transactions, self.sim_config.max_eb_size);
+            self.sample_from_mempool(&mut eb_transactions, self.sim_config.max_eb_size, false);
         }
 
-        let parent = self.latest_rb().map(|rb| rb.header.id);
-        let endorsement = parent.and_then(|rb_id| {
-            if rb_id.slot
-                + self.sim_config.linear_vote_stage_length
-                + self.sim_config.linear_diffuse_stage_length
-                > slot
-            {
-                // This RB was generated too quickly after another; hasn't been time to gather all the votes.
-                // No endorsement.
-                return None;
-            }
-
-            let eb_id = self.leios.ebs_by_rb.get(&rb_id)?;
-            let votes = self.leios.votes_by_eb.get(eb_id)?;
-            let total_votes = votes.values().copied().sum::<usize>();
-            if (total_votes as u64) < self.sim_config.vote_threshold {
-                // Not enough votes. No endorsement.
-                return None;
-            }
-
-            Some(Endorsement {
-                eb: *eb_id,
-                size_bytes: self.sim_config.sizes.cert(votes.len()),
-                votes: votes.clone(),
-            })
-        });
-
         let rb = RankingBlock {
             header: RankingBlockHeader {
                 id: BlockId {
@@ -887,32 +898,6 @@ impl LinearLeiosNode {
     fn finish_validating_eb(&mut self, eb: Arc<EndorserBlock>, seen: Timestamp) {
         self.vote_for_endorser_block(&eb, seen);
     }
-
-    fn process_certified_ebs(&mut self, slot: u64) {
-        let Some(eb_creation_slot) = slot.checked_sub(
-            self.sim_config.linear_vote_stage_length + self.sim_config.linear_diffuse_stage_length,
-        ) else {
-            return;
-        };
-
-        let Some(rb) = self.latest_rb() else {
-            return;
-        };
-
-        if rb.header.id.slot != eb_creation_slot {
-            return;
-        }
-        // The EB on the head of the chain has had enough time to get certified.
-
-        let eb_id = rb.header.eb_announcement;
-        if self.leios.certified_ebs.contains(&eb_id) {
-            // This certified EB is eventually going on-chain.
-            // Get its contents out of the mempool.
-            if let Some(EndorserBlockView::Received { eb }) = self.leios.ebs.get(&eb_id) {
-                self.remove_eb_txs_from_mempool(&eb.clone());
-            };
-        }
-    }
 }
 
 // EB withholding:
@@ -1168,7 +1153,12 @@ impl LinearLeiosNode {
 
 // Ledger/mempool operations
 impl LinearLeiosNode {
-    fn sample_from_mempool(&mut self, txs: &mut Vec<Arc<Transaction>>, max_size: u64) {
+    fn sample_from_mempool(
+        &mut self,
+        txs: &mut Vec<Arc<Transaction>>,
+        max_size: u64,
+        remove: bool,
+    ) {
         let mut size = txs.iter().map(|tx| tx.bytes).sum::<u64>();
         let mut candidates: Vec<_> = self.praos.mempool.keys().copied().collect();
         if matches!(
@@ -1187,7 +1177,11 @@ impl LinearLeiosNode {
                 break;
             }
             size += tx.bytes;
-            txs.push(self.praos.mempool.remove(&id).unwrap());
+            if remove {
+                txs.push(self.praos.mempool.remove(&id).unwrap());
+            } else {
+                txs.push(tx.clone());
+            }
         }
     }