Expose an experimental flag to write the puffin reports to disk

index-scheduler/src/lib.rs

@@ -33,6 +33,7 @@ pub type Result<T> = std::result::Result<T, Error>;
 pub type TaskId = u32;
 
 use std::collections::{BTreeMap, HashMap};
+use std::fs::File;
 use std::ops::{Bound, RangeBounds};
 use std::path::{Path, PathBuf};
 use std::sync::atomic::AtomicBool;
@@ -256,6 +257,8 @@ pub struct IndexSchedulerOptions {
     /// The maximum number of tasks stored in the task queue before starting
     /// to auto schedule task deletions.
     pub max_number_of_tasks: usize,
+    /// Weither we need to export indexation puffin reports.
+    pub profile_with_puffin: bool,
     /// The experimental features enabled for this instance.
     pub instance_features: InstanceTogglableFeatures,
 }
@@ -314,6 +317,9 @@ pub struct IndexScheduler {
     /// the finished tasks automatically.
     pub(crate) max_number_of_tasks: usize,
 
+    /// Wether we must output indexation  profiling files to disk.
+    pub(crate) puffin_frame: Option<Arc<puffin::GlobalFrameView>>,
+
     /// The path used to create the dumps.
     pub(crate) dumps_path: PathBuf,
 
@@ -364,6 +370,7 @@ impl IndexScheduler {
             wake_up: self.wake_up.clone(),
             autobatching_enabled: self.autobatching_enabled,
             max_number_of_tasks: self.max_number_of_tasks,
+            puffin_frame: self.puffin_frame.clone(),
             snapshots_path: self.snapshots_path.clone(),
             dumps_path: self.dumps_path.clone(),
             auth_path: self.auth_path.clone(),
@@ -457,6 +464,9 @@ impl IndexScheduler {
             env,
             // we want to start the loop right away in case meilisearch was ctrl+Ced while processing things
             wake_up: Arc::new(SignalEvent::auto(true)),
+            puffin_frame: options
+                .profile_with_puffin
+                .then(|| Arc::new(puffin::GlobalFrameView::default())),
             autobatching_enabled: options.autobatching_enabled,
             max_number_of_tasks: options.max_number_of_tasks,
             dumps_path: options.dumps_path,
@@ -1070,7 +1080,21 @@ impl IndexScheduler {
         let batch =
             match self.create_next_batch(&rtxn).map_err(|e| Error::CreateBatch(Box::new(e)))? {
                 Some(batch) => batch,
-                None => return Ok(TickOutcome::WaitForSignal),
+                None => {
+                    // Let's write the previous save to disk but only if
+                    // the user wanted to profile with puffin.
+                    if let Some(global_frame_view) = &self.puffin_frame {
+                        let frame_view = global_frame_view.lock();
+                        if !frame_view.is_empty() {
+                            let mut file =
+                                File::create(format!("{}.puffin", OffsetDateTime::now_utc()))?;
+                            frame_view.save_to_writer(&mut file)?;
+                            file.sync_all()?;
+                        }
+                    }
+
+                    return Ok(TickOutcome::WaitForSignal);
+                }
             };
         let index_uid = batch.index_uid().map(ToOwned::to_owned);
         drop(rtxn);

meilisearch/src/analytics/segment_analytics.rs

@@ -285,6 +285,7 @@ impl From<Opt> for Infos {
             db_path,
             experimental_enable_metrics,
             experimental_reduce_indexing_memory_usage,
+            experimental_profile_with_puffin: _,
             http_addr,
             master_key: _,
             env,

meilisearch/src/lib.rs

@@ -237,6 +237,7 @@ fn open_or_create_database_unchecked(
             indexer_config: (&opt.indexer_options).try_into()?,
             autobatching_enabled: true,
             max_number_of_tasks: 1_000_000,
+            profile_with_puffin: opt.experimental_profile_with_puffin,
             index_growth_amount: byte_unit::Byte::from_str("10GiB").unwrap().get_bytes() as usize,
             index_count: DEFAULT_INDEX_COUNT,
             instance_features,

meilisearch/src/option.rs

@@ -51,6 +51,7 @@ const MEILI_LOG_LEVEL: &str = "MEILI_LOG_LEVEL";
 const MEILI_EXPERIMENTAL_ENABLE_METRICS: &str = "MEILI_EXPERIMENTAL_ENABLE_METRICS";
 const MEILI_EXPERIMENTAL_REDUCE_INDEXING_MEMORY_USAGE: &str =
     "MEILI_EXPERIMENTAL_REDUCE_INDEXING_MEMORY_USAGE";
+const MEILI_EXPERIMENTAL_PROFILE_WITH_PUFFIN: &str = "MEILI_EXPERIMENTAL_PROFILE_WITH_PUFFIN";
 
 const DEFAULT_CONFIG_FILE_PATH: &str = "./config.toml";
 const DEFAULT_DB_PATH: &str = "./data.ms";
@@ -301,6 +302,11 @@ pub struct Opt {
     #[serde(default)]
     pub experimental_reduce_indexing_memory_usage: bool,
 
+    /// Experimental flag to export puffin reports, see: <https://github.com/meilisearch/product/discussions/xxx>
+    #[clap(long, env = MEILI_EXPERIMENTAL_PROFILE_WITH_PUFFIN)]
+    #[serde(default)]
+    pub experimental_profile_with_puffin: bool,
+
     #[serde(flatten)]
     #[clap(flatten)]
     pub indexer_options: IndexerOpts,
@@ -394,6 +400,7 @@ impl Opt {
             no_analytics,
             experimental_enable_metrics: enable_metrics_route,
             experimental_reduce_indexing_memory_usage: reduce_indexing_memory_usage,
+            experimental_profile_with_puffin,
         } = self;
         export_to_env_if_not_present(MEILI_DB_PATH, db_path);
         export_to_env_if_not_present(MEILI_HTTP_ADDR, http_addr);
@@ -439,6 +446,10 @@ impl Opt {
             MEILI_EXPERIMENTAL_REDUCE_INDEXING_MEMORY_USAGE,
             reduce_indexing_memory_usage.to_string(),
         );
+        export_to_env_if_not_present(
+            MEILI_EXPERIMENTAL_PROFILE_WITH_PUFFIN,
+            experimental_profile_with_puffin.to_string(),
+        );
         indexer_options.export_to_env();
     }