Improve inference performance and reenable socket communciation

1 files changed, 8 insertions, 7 deletions

diff --git a/src/main.rs b/src/main.rs
index c56a41c..f0353e5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,5 @@
 mod bpf_skel;
-use benchmark::BenchmarkScheduler;
+use benchmark::{BenchmarkScheduler, Mode};
 pub use bpf_skel::*;
 pub mod bpf_intf;
 
@@ -36,9 +36,9 @@ fn main() -> Result<()> {
             Arg::new("benchmark")
                 .short('b')
                 .long("benchmark")
-                .help("Use this flag to enable benckmarking mode")
-                .action(ArgAction::SetTrue)
-                .required(false),
+                .help("Enable benchmarking mode. Use \"e\" for e cores and \"p\" for p cores")
+                .required(false)
+                .value_name("mode"),
         )
         .arg(
             Arg::new("power_cap")
@@ -51,13 +51,14 @@ fn main() -> Result<()> {
 
     let power_cap = *matches.get_one::<u64>("power_cap").unwrap_or(&u64::MAX);
     let use_mocking = matches.get_flag("mock");
-    let benchmark = matches.get_flag("benchmark");
+    let benchmark = matches.get_one::<String>("benchmark");
 
     // Initialize and load the scheduler.
     let mut open_object = MaybeUninit::uninit();
     let log_path = "/tmp/logs.csv";
-    if benchmark {
-        let mut sched = BenchmarkScheduler::init(&mut open_object, log_path)?;
+    if let Some(mode) = benchmark {
+        let mode = mode.trim().chars().next().unwrap();
+        let mut sched = BenchmarkScheduler::init(&mut open_object, log_path, Mode::from(mode))?;
         sched.run()?;
         return Ok(());
     }