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);