in_memory_network.rs - iuna - iuna - experimental mainnet-candidate protocol

in_memory_network.rs (3478B)
      1 use std::collections::BTreeMap;
      2 
      3 use anyhow::Result;
      4 
      5 use super::{GossipEnvelope, NodeCore};
      6 
      7 #[derive(Debug, Default)]
      8 pub struct InMemoryNetwork {
      9     nodes: BTreeMap<String, NodeCore>,
     10 }
     11 
     12 impl InMemoryNetwork {
     13     pub fn insert(&mut self, id: impl Into<String>, node: NodeCore) {
     14         self.nodes.insert(id.into(), node);
     15     }
     16 
     17     pub fn node(&self, id: &str) -> Option<&NodeCore> {
     18         self.nodes.get(id)
     19     }
     20 
     21     pub fn node_mut(&mut self, id: &str) -> Option<&mut NodeCore> {
     22         self.nodes.get_mut(id)
     23     }
     24 
     25     pub fn deliver_until_idle(&mut self) -> Result<()> {
     26         self.deliver_until_idle_filtered(|_, _, _| true)
     27     }
     28 
     29     pub fn deliver_until_idle_filtered(
     30         &mut self,
     31         mut allow: impl FnMut(&str, &str, &GossipEnvelope) -> bool,
     32     ) -> Result<()> {
     33         loop {
     34             let mut outbound = Vec::new();
     35             for (id, node) in &mut self.nodes {
     36                 for envelope in node.drain_outbox() {
     37                     outbound.push((id.clone(), envelope));
     38                 }
     39             }
     40 
     41             if outbound.is_empty() {
     42                 return Ok(());
     43             }
     44 
     45             for (from, envelope) in outbound {
     46                 for (id, node) in &mut self.nodes {
     47                     if *id != from && allow(&from, id, &envelope) {
     48                         receive_in_memory_envelope(node, envelope.clone())?;
     49                     }
     50                 }
     51             }
     52         }
     53     }
     54 
     55     pub fn gossip_mempools_once(&mut self) -> Result<()> {
     56         self.gossip_mempools_once_filtered(|_, _, _| true)
     57     }
     58 
     59     pub fn gossip_mempools_once_filtered(
     60         &mut self,
     61         mut allow: impl FnMut(&str, &str, &GossipEnvelope) -> bool,
     62     ) -> Result<()> {
     63         let mut outbound = Vec::new();
     64         for (id, node) in &mut self.nodes {
     65             for envelope in node.mempool_gossip() {
     66                 outbound.push((id.clone(), envelope));
     67             }
     68         }
     69 
     70         for (from, envelope) in outbound {
     71             for (id, node) in &mut self.nodes {
     72                 if *id != from && allow(&from, id, &envelope) {
     73                     receive_in_memory_envelope(node, envelope.clone())?;
     74                 }
     75             }
     76         }
     77         Ok(())
     78     }
     79 
     80     pub fn sync_node_from_peer(&mut self, from: &str, to: &str, limit: usize) -> Result<bool> {
     81         let from_height = self
     82             .nodes
     83             .get(to)
     84             .map(|node| node.chain_height() + 1)
     85             .ok_or_else(|| anyhow::anyhow!("missing sync target node {to}"))?;
     86         let blocks = self
     87             .nodes
     88             .get(from)
     89             .map(|node| node.blocks_from(from_height, limit))
     90             .ok_or_else(|| anyhow::anyhow!("missing sync source node {from}"))?;
     91         if blocks.is_empty() {
     92             return Ok(false);
     93         }
     94 
     95         self.nodes
     96             .get_mut(to)
     97             .expect("sync target exists")
     98             .receive(GossipEnvelope::Blocks { blocks })?;
     99         Ok(true)
    100     }
    101 }
    102 
    103 fn receive_in_memory_envelope(node: &mut NodeCore, envelope: GossipEnvelope) -> Result<()> {
    104     let transaction_like = matches!(
    105         envelope,
    106         GossipEnvelope::Transaction(_)
    107             | GossipEnvelope::Transactions { .. }
    108             | GossipEnvelope::TransactionV2 { .. }
    109             | GossipEnvelope::TransactionsV2 { .. }
    110     );
    111     match node.receive(envelope) {
    112         Ok(()) => Ok(()),
    113         Err(_) if transaction_like => Ok(()),
    114         Err(error) => Err(error),
    115     }
    116 }