From ba3f8f4e892b9fe2a13d75b0ce5c28801b379d42 Mon Sep 17 00:00:00 2001
From: Lee Smet <lee.smet@hotmail.com>
Date: Fri, 2 Aug 2024 10:59:36 +0200
Subject: [PATCH] Make amount of update workers configurable

Signed-off-by: Lee Smet <lee.smet@hotmail.com>
---
 mycelium/src/lib.rs    |  7 +++++++
 mycelium/src/router.rs | 15 +++++++++------
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/mycelium/src/lib.rs b/mycelium/src/lib.rs
index be28e38..9e6c170 100644
--- a/mycelium/src/lib.rs
+++ b/mycelium/src/lib.rs
@@ -78,6 +78,12 @@ pub struct Config<M> {
     // the TUN's file descriptor to mycelium.
     #[cfg(any(target_os = "android", target_os = "ios"))]
     pub tun_fd: Option<i32>,
+
+    /// The maount of worker tasks spawned to process updates. Up to this amound of updates can be
+    /// processed in parallel. Because processing an update is a CPU bound task, it is pointless to
+    /// set this to a value which is higher than the amount of logical CPU cores available to the
+    /// system.
+    pub update_workers: usize,
 }
 
 /// The Node is the main structure in mycelium. It governs the entire data flow.
@@ -126,6 +132,7 @@ where
 
         // Creating a new Router instance
         let router = match router::Router::new(
+            config.update_workers,
             tun_tx,
             node_subnet,
             vec![node_subnet],
diff --git a/mycelium/src/router.rs b/mycelium/src/router.rs
index 75e06fb..0ef7573 100644
--- a/mycelium/src/router.rs
+++ b/mycelium/src/router.rs
@@ -80,6 +80,7 @@ pub struct Router<M> {
     /// Channel to notify the router of expired SourceKey's.
     expired_source_key_sink: mpsc::Sender<SourceKey>,
     seqno_cache: SeqnoCache,
+    update_workers: usize,
     metrics: M,
 }
 
@@ -88,6 +89,7 @@ where
     M: Metrics + Clone + Send + 'static,
 {
     pub fn new(
+        update_workers: usize,
         node_tun: UnboundedSender<DataPacket>,
         node_tun_subnet: Subnet,
         static_routes: Vec<Subnet>,
@@ -125,6 +127,7 @@ where
             expired_source_key_sink,
             seqno_cache,
             update_filters: Arc::new(update_filters),
+            update_workers,
             metrics,
         };
 
@@ -577,9 +580,8 @@ where
 
     /// Background task to process Update TLV's.
     async fn update_processor(self, mut update_rx: UnboundedReceiver<(Update, Peer)>) {
-        const WORKERS: usize = 5;
-
-        let senders = core::array::from_fn::<_, WORKERS, _>(|_| {
+        let mut senders = Vec::with_capacity(self.update_workers);
+        for _ in 0..self.update_workers {
             let router = self.clone();
             let (tx, mut rx) = mpsc::unbounded_channel::<(_, Peer)>();
 
@@ -602,13 +604,13 @@ where
                 warn!("Update processor task exitted");
             });
 
-            tx
-        });
+            senders.push(tx);
+        }
 
         while let Some(item) = update_rx.recv().await {
             let mut hasher = std::hash::DefaultHasher::new();
             item.0.subnet().network().hash(&mut hasher);
-            let slot = hasher.finish() as usize % WORKERS;
+            let slot = hasher.finish() as usize % self.update_workers;
 
             if senders[slot].send(item).is_err() {
                 break;
@@ -1754,6 +1756,7 @@ where
             dead_peer_sink: self.dead_peer_sink.clone(),
             expired_source_key_sink: self.expired_source_key_sink.clone(),
             seqno_cache: self.seqno_cache.clone(),
+            update_workers: self.update_workers,
             metrics: self.metrics.clone(),
         }
     }