fix: profile server in bench_runner_server

Author: michaelmior

src/bin/argusdb.rs

@@ -57,6 +57,10 @@ struct Args {
     #[arg(long, default_value_t = false)]
     no_log: bool,
 
+    /// Profile output file
+    #[arg(long)]
+    profile: Option<String>,
+
     /// Print help
     #[arg(long, action = clap::ArgAction::Help)]
     help: Option<bool>,
@@ -296,6 +300,8 @@ async fn main() {
     tracing::subscriber::set_global_default(subscriber).unwrap();
     let settings: Settings = builder.build().unwrap().try_deserialize().unwrap();
 
+    let profile_guard = argusdb::bench_utils::start_profiling(&args.profile);
+
     let log_threshold = if settings.no_log {
         None
     } else {
@@ -316,12 +322,30 @@ async fn main() {
     let listener = TcpListener::bind(&server_addr).await.unwrap();
     info!("ArgusDB server listening on {}", server_addr);
 
+    let mut sigterm =
+        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()).unwrap();
+    let mut sigint =
+        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt()).unwrap();
+
     loop {
-        let (socket, _) = listener.accept().await.unwrap();
         let processor = processor.clone();
-
-        tokio::spawn(async move {
-            let _ = process_socket(socket, None, processor).await;
-        });
+        tokio::select! {
+            res = listener.accept() => {
+                let (socket, _) = res.unwrap();
+                tokio::spawn(async move {
+                    let _ = process_socket(socket, None, processor).await;
+                });
+            }
+            _ = sigterm.recv() => {
+                info!("Received SIGTERM, shutting down");
+                break;
+            }
+            _ = sigint.recv() => {
+                info!("Received SIGINT, shutting down");
+                break;
+            }
+        }
     }
+
+    argusdb::bench_utils::save_profile(profile_guard, &args.profile);
 }

src/bin/bench_runner_server.rs

@@ -1,9 +1,7 @@
 #[global_allocator]
 static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
-use argusdb::bench_utils::{
-    Args, Query, load_queries, run_measurement, save_profile, start_profiling,
-};
+use argusdb::bench_utils::{Args, Query, load_queries, run_measurement};
 use clap::Parser;
 use rand::prelude::IndexedRandom;
 use std::fs;
@@ -26,7 +24,12 @@ struct ServerGuard {
 
 impl Drop for ServerGuard {
     fn drop(&mut self) {
-        let _ = self.child.kill();
+        let pid = self.child.id();
+        let _ = Command::new("kill")
+            .arg("-s")
+            .arg("TERM")
+            .arg(pid.to_string())
+            .status();
         let _ = self.child.wait();
     }
 }
@@ -104,6 +107,10 @@ async fn main() {
         command.arg("--no-log");
     }
 
+    if let Some(profile) = &args.profile {
+        command.arg("--profile").arg(profile);
+    }
+
     let child = command
         .stdout(Stdio::inherit())
         .stderr(Stdio::inherit())
@@ -193,7 +200,6 @@ async fn main() {
     .await;
 
     println!("Starting measurement for {} seconds...", args.duration);
-    let guard = start_profiling(&args.profile);
 
     let results = run_measurement(
         args.concurrency,
@@ -205,8 +211,6 @@ async fn main() {
     )
     .await;
 
-    save_profile(guard, &args.profile);
-
     println!("Results:");
     for (name, (count, total_time)) in results {
         let avg = if count > 0 {