Merge #5152

5152: Make xtasks be able to use the specified binary r=dureuill a=Kerollmops

Makes it possible to specify the binary to run. It is useful to run PGO optimized binaries.

Co-authored-by: Kerollmops <clement@meilisearch.com>
Co-authored-by: Clément Renault <clement@meilisearch.com>

crates/meilisearch/src/main.rs

@@ -129,6 +129,11 @@ async fn try_main() -> anyhow::Result<()> {
 
     print_launch_resume(&opt, analytics.clone(), config_read_from);
 
+    tokio::spawn(async move {
+        tokio::signal::ctrl_c().await.unwrap();
+        std::process::exit(130);
+    });
+
     run_http(index_scheduler, auth_controller, opt, log_handle, Arc::new(analytics)).await?;
 
     Ok(())

crates/xtask/src/bench/meili_process.rs

@@ -1,12 +1,43 @@
 use std::collections::BTreeMap;
+use std::time::Duration;
 
 use anyhow::{bail, Context as _};
+use tokio::process::Command;
+use tokio::time;
 
 use super::assets::Asset;
 use super::client::Client;
 use super::workload::Workload;
 
 pub async fn kill(mut meilisearch: tokio::process::Child) {
+    let Some(id) = meilisearch.id() else { return };
+
+    match Command::new("kill").args(["--signal=TERM", &id.to_string()]).spawn() {
+        Ok(mut cmd) => {
+            let Err(error) = cmd.wait().await else { return };
+            tracing::warn!(
+                error = &error as &dyn std::error::Error,
+                "while awaiting the Meilisearch server kill"
+            );
+        }
+        Err(error) => {
+            tracing::warn!(
+                error = &error as &dyn std::error::Error,
+                "while terminating Meilisearch server with a kill -s TERM"
+            );
+            if let Err(error) = meilisearch.kill().await {
+                tracing::warn!(
+                    error = &error as &dyn std::error::Error,
+                    "while terminating Meilisearch server"
+                )
+            }
+            return;
+        }
+    };
+
+    match time::timeout(Duration::from_secs(5), meilisearch.wait()).await {
+        Ok(_) => (),
+        Err(_) => {
             if let Err(error) = meilisearch.kill().await {
                 tracing::warn!(
                     error = &error as &dyn std::error::Error,
@@ -14,10 +45,12 @@ pub async fn kill(mut meilisearch: tokio::process::Child) {
                 )
             }
         }
+    }
+}
 
 #[tracing::instrument]
 pub async fn build() -> anyhow::Result<()> {
-    let mut command = tokio::process::Command::new("cargo");
+    let mut command = Command::new("cargo");
     command.arg("build").arg("--release").arg("-p").arg("meilisearch");
 
     command.kill_on_drop(true);
@@ -37,17 +70,8 @@ pub async fn start(
     master_key: Option<&str>,
     workload: &Workload,
     asset_folder: &str,
+    mut command: Command,
 ) -> anyhow::Result<tokio::process::Child> {
-    let mut command = tokio::process::Command::new("cargo");
-    command
-        .arg("run")
-        .arg("--release")
-        .arg("-p")
-        .arg("meilisearch")
-        .arg("--bin")
-        .arg("meilisearch")
-        .arg("--");
-
     command.arg("--db-path").arg("./_xtask_benchmark.ms");
     if let Some(master_key) = master_key {
         command.arg("--master-key").arg(master_key);
@@ -86,7 +110,7 @@ async fn wait_for_health(
 
             return Ok(());
         }
-        tokio::time::sleep(std::time::Duration::from_millis(500)).await;
+        time::sleep(Duration::from_millis(500)).await;
         // check whether the Meilisearch instance exited early (cut the wait)
         if let Some(exit_code) =
             meilisearch.try_wait().context("cannot check Meilisearch server process status")?

crates/xtask/src/bench/mod.rs

@@ -86,6 +86,12 @@ pub struct BenchDeriveArgs {
     /// The maximum time in seconds we allow for fetching the task queue before timing out.
     #[arg(long, default_value_t = 60)]
     tasks_queue_timeout_secs: u64,
+
+    /// The path to the binary to run.
+    ///
+    /// If unspecified, runs `cargo run` after building Meilisearch with `cargo build`.
+    #[arg(long)]
+    binary_path: Option<PathBuf>,
 }
 
 pub fn run(args: BenchDeriveArgs) -> anyhow::Result<()> {
@@ -170,6 +176,7 @@ pub fn run(args: BenchDeriveArgs) -> anyhow::Result<()> {
                     args.master_key.as_deref(),
                     workload,
                     &args,
+                    args.binary_path.as_deref(),
                 )
                 .await?;
 

crates/xtask/src/bench/workload.rs

@@ -1,6 +1,7 @@
 use std::collections::BTreeMap;
 use std::fs::File;
 use std::io::{Seek as _, Write as _};
+use std::path::Path;
 
 use anyhow::{bail, Context as _};
 use futures_util::TryStreamExt as _;
@@ -85,13 +86,13 @@ pub async fn execute(
     master_key: Option<&str>,
     workload: Workload,
     args: &BenchDeriveArgs,
+    binary_path: Option<&Path>,
 ) -> anyhow::Result<()> {
     assets::fetch_assets(assets_client, &workload.assets, &args.asset_folder).await?;
 
     let workload_uuid = dashboard_client.create_workload(invocation_uuid, &workload).await?;
 
     let mut tasks = Vec::new();
-
     for i in 0..workload.run_count {
         tasks.push(
             execute_run(
@@ -102,6 +103,7 @@ pub async fn execute(
                 master_key,
                 &workload,
                 args,
+                binary_path,
                 i,
             )
             .await?,
@@ -109,7 +111,6 @@ pub async fn execute(
     }
 
     let mut reports = Vec::with_capacity(workload.run_count as usize);
-
     for task in tasks {
         reports.push(
             task.await
@@ -133,13 +134,31 @@ async fn execute_run(
     master_key: Option<&str>,
     workload: &Workload,
     args: &BenchDeriveArgs,
+    binary_path: Option<&Path>,
     run_number: u16,
 ) -> anyhow::Result<tokio::task::JoinHandle<anyhow::Result<std::fs::File>>> {
     meili_process::delete_db();
 
+    let run_command = match binary_path {
+        Some(binary_path) => tokio::process::Command::new(binary_path),
+        None => {
             meili_process::build().await?;
+            let mut command = tokio::process::Command::new("cargo");
+            command
+                .arg("run")
+                .arg("--release")
+                .arg("-p")
+                .arg("meilisearch")
+                .arg("--bin")
+                .arg("meilisearch")
+                .arg("--");
+            command
+        }
+    };
+
     let meilisearch =
-        meili_process::start(meili_client, master_key, workload, &args.asset_folder).await?;
+        meili_process::start(meili_client, master_key, workload, &args.asset_folder, run_command)
+            .await?;
 
     let processor = run_commands(
         dashboard_client,