Merge pull request #14 from stackrox/feature/per-core-mode

Add per-core mode

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 {

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))
@@ -49,8 +48,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);
 
@@ -60,13 +67,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:?}");

src/worker/endpoints.rs

@@ -27,6 +27,8 @@ impl EndpointWorker {
         let WorkloadConfig {
             restart_interval,
             workload: _,
+            per_core: _,
+            workers: _,
         } = workload;
 
         EndpointWorker {