feat: add downstream peer limiting

For now this is a simple cmdline argument that just ignores peers beyond a certain count.
Proper handling will need to include idle timers etc. to kick out downstream peers that are just squatting.

Signed-Off-By: Roland Kuhn <rk@rkuhn.info>

Author: rkuhn

crates/amaru/src/bin/amaru/cmd/daemon.rs

@@ -54,6 +54,10 @@ pub struct Args {
     /// The address to listen on for incoming connections.
     #[arg(long, value_name = "LISTEN_ADDRESS", default_value = super::DEFAULT_LISTEN_ADDRESS)]
     listen_address: String,
+
+    /// The maximum number of downstream peers to connect to.
+    #[arg(long, value_name = "MAX_DOWNSTREAM_PEERS", default_value_t = 10)]
+    max_downstream_peers: usize,
 }
 
 pub async fn run(
@@ -112,5 +116,6 @@ fn parse_args(args: Args) -> Result<Config, Box<dyn std::error::Error>> {
         network: args.network,
         network_magic: args.network.to_network_magic(),
         listen_address: args.listen_address,
+        max_downstream_peers: args.max_downstream_peers,
     })
 }

crates/amaru/src/stages/consensus/chain_forward.rs

@@ -53,6 +53,7 @@ pub struct ForwardStage {
     pub runtime: AcTokio,
     pub listen_address: String,
     pub downstream: Option<ActoRef<ForwardEvent>>,
+    pub max_peers: usize,
 }
 
 #[derive(Debug, Clone)]
@@ -67,6 +68,7 @@ impl ForwardStage {
         store: Arc<Mutex<dyn ChainStore<Header>>>,
         network_magic: u64,
         listen_address: &str,
+        max_peers: usize,
     ) -> Self {
         #[allow(clippy::expect_used)]
         let runtime =
@@ -78,6 +80,7 @@ impl ForwardStage {
             runtime,
             listen_address: listen_address.to_string(),
             downstream,
+            max_peers,
         }
     }
 }
@@ -116,7 +119,7 @@ impl gasket::framework::Worker<ForwardStage> for Worker {
         let clients = stage
             .runtime
             .spawn_actor("chain_forward", |cell| {
-                client_supervisor(cell, stage.store.clone())
+                client_supervisor(cell, stage.store.clone(), stage.max_peers)
             })
             .me;
 
@@ -251,11 +254,17 @@ enum ClientMsg {
 async fn client_supervisor(
     mut cell: ActoCell<ClientMsg, impl ActoRuntime, Result<(), client_protocol::ClientError>>,
     store: Arc<Mutex<dyn ChainStore<Header>>>,
+    max_peers: usize,
 ) {
     let mut clients = HashMap::new();
     while let Some(msg) = cell.recv().await.has_senders() {
         match msg {
             ActoMsgSuper::Message(ClientMsg::Peer(peer, tip)) => {
+                if clients.len() >= max_peers {
+                    tracing::warn!(target: EVENT_TARGET, "max peers reached, dropping peer");
+                    continue;
+                }
+
                 let addr = peer
                     .accepted_address()
                     .map(|a| a.to_string())

crates/amaru/src/stages/consensus/chain_forward/tests.rs

@@ -37,15 +37,15 @@ impl ChainStore<Header> for TestStore {
     }
 
     fn get_nonces(&self, _header: &Hash<32>) -> Option<Nonces> {
-        todo!()
+        unimplemented!()
     }
 
     fn put_nonces(&mut self, _header: &Hash<32>, _nonces: Nonces) -> Result<(), StoreError> {
-        todo!()
+        unimplemented!()
     }
 
     fn era_history(&self) -> &slot_arithmetic::EraHistory {
-        todo!()
+        unimplemented!()
     }
 }
 
@@ -177,8 +177,10 @@ fn find_headers_between_tip_and_lost() {
     let tip = point(TIP_41_SLOT, TIP_41);
     let points = [point(LOST_41_SLOT, LOST_41)];
 
-    let result = find_headers_between(&store, &tip, &points);
-    assert!(result.is_none(), "{result:?}");
+    let result = find_headers_between(&store, &tip, &points).unwrap();
+    assert_eq!(result.0.len() as u64, TIP_41_HEIGHT);
+    assert_eq!(result.1 .0, Point::Origin);
+    assert_eq!(result.1 .1, 0);
 }
 
 #[test]
@@ -212,7 +214,7 @@ fn test_chain_sync() {
 
     // step 1: prepare the stage
     let (block_tx, block_rx) = mpsc::channel(1);
-    let mut stage = ForwardStage::new(Some(downstream), store, 42, "127.0.0.1:0");
+    let mut stage = ForwardStage::new(Some(downstream), store, 42, "127.0.0.1:0", 1);
     stage.upstream.connect(ChannelRecvAdapter::Mpsc(block_rx));
     let tether = spawn_stage(stage, Default::default());
 

crates/amaru/src/stages/mod.rs

@@ -49,6 +49,7 @@ pub struct Config {
     pub network: NetworkName,
     pub network_magic: u32,
     pub listen_address: String,
+    pub max_downstream_peers: usize,
 }
 
 /// A session with a peer, including the peer itself and a client to communicate with it.
@@ -102,6 +103,7 @@ pub fn bootstrap(
         chain_ref.clone(),
         config.network_magic as u64,
         &config.listen_address,
+        config.max_downstream_peers,
     );
 
     let (to_block_fetch, from_consensus_stage) = gasket::messaging::tokio::mpsc_channel(50);