From cb060616c663ab25804934bf9e7b21e0d4bb79a9 Mon Sep 17 00:00:00 2001
From: Dmitrii Dolgov <9erthalion6@gmail.com>
Date: Tue, 2 Apr 2024 16:48:17 +0200
Subject: [PATCH] Add per-core mode

Allow to configure number of workers to run, either total or in per-core
mode. It's neccessary because some of the workers will require only one
active process due to exclusive resources access.
---
 src/lib.rs              | 23 +++++++++++++++++++++++
 src/main.rs             | 25 ++++++++++++++++---------
 src/worker/endpoints.rs |  2 ++
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 22c0730..bce3e01 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -11,10 +11,29 @@ pub struct WorkloadConfig {
     /// An amount of time for workload payload to run before restarting.
     pub restart_interval: u64,
 
+    /// Controls per-core mode to handle number of workers. If per-core mode
+    /// is enabled, `workers` will be treated as a number of workers per CPU
+    /// core. Otherwise it will be treated as a total number of workers.
+    #[serde(default = "default_per_core")]
+    pub per_core: bool,
+
+    /// How many workers to spin, depending on `per_core` in either per-core
+    /// or total mode.
+    #[serde(default = "default_workers")]
+    pub workers: usize,
+
     /// Custom workload configuration.
     pub workload: Workload,
 }
 
+fn default_workers() -> usize {
+    1
+}
+
+fn default_per_core() -> bool {
+    true
+}
+
 /// Workload specific configuration, contains one enum value for each
 /// workload type.
 #[derive(Debug, Copy, Clone, Deserialize)]
@@ -114,6 +133,7 @@ mod tests {
         let WorkloadConfig {
             restart_interval,
             workload,
+            ..
         } = config;
         assert_eq!(restart_interval, 10);
         if let Workload::Processes {
@@ -152,6 +172,7 @@ mod tests {
         let WorkloadConfig {
             restart_interval,
             workload,
+            ..
         } = config;
         assert_eq!(restart_interval, 10);
 
@@ -189,6 +210,7 @@ mod tests {
         let WorkloadConfig {
             restart_interval,
             workload,
+            ..
         } = config;
         assert_eq!(restart_interval, 10);
 
@@ -224,6 +246,7 @@ mod tests {
         let WorkloadConfig {
             restart_interval,
             workload,
+            ..
         } = config;
         assert_eq!(restart_interval, 10);
         if let Workload::Syscalls { arrival_rate } = workload {
diff --git a/src/main.rs b/src/main.rs
index 8fd1ec0..4bdf92d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,6 +18,7 @@ extern crate log;
 extern crate core_affinity;
 
 use config::Config;
+use core_affinity::CoreId;
 use fork::{fork, Fork};
 use itertools::iproduct;
 use nix::sys::wait::waitpid;
@@ -26,8 +27,6 @@ use nix::unistd::Pid;
 use berserker::{worker::new_worker, WorkloadConfig};
 
 fn main() {
-    // Retrieve the IDs of all active CPU cores.
-    let core_ids = core_affinity::get_core_ids().unwrap();
     let config = Config::builder()
         // Add in `./Settings.toml`
         .add_source(config::File::with_name("/etc/berserker/workload.toml").required(false))
@@ -45,8 +44,16 @@ fn main() {
 
     env_logger::init();
 
-    // Create processes for each active CPU core.
-    let handles: Vec<_> = iproduct!(core_ids.into_iter(), 0..9)
+    info!("Config: {:?}", config);
+
+    let core_ids: Vec<CoreId> = if config.per_core {
+        // Retrieve the IDs of all active CPU cores.
+        core_affinity::get_core_ids().unwrap()
+    } else {
+        vec![CoreId { id: 0 }]
+    };
+
+    let handles: Vec<_> = iproduct!(core_ids.into_iter(), 0..config.workers)
         .map(|(cpu, process)| {
             let worker = new_worker(config, cpu, process, &mut lower, &mut upper);
 
@@ -56,13 +63,13 @@ fn main() {
                     Some(child)
                 }
                 Ok(Fork::Child) => {
-                    if core_affinity::set_for_current(cpu) {
-                        loop {
-                            worker.run_payload().unwrap();
-                        }
+                    if config.per_core {
+                        core_affinity::set_for_current(cpu);
                     }
 
-                    None
+                    loop {
+                        worker.run_payload().unwrap();
+                    }
                 }
                 Err(e) => {
                     warn!("Failed: {e:?}");
diff --git a/src/worker/endpoints.rs b/src/worker/endpoints.rs
index b437461..db7eecf 100644
--- a/src/worker/endpoints.rs
+++ b/src/worker/endpoints.rs
@@ -27,6 +27,8 @@ impl EndpointWorker {
         let WorkloadConfig {
             restart_interval,
             workload: _,
+            per_core: _,
+            workers: _,
         } = workload;
 
         EndpointWorker {