From 3655d4bdca2aaa8d08d214e8d01fc78d408c4672 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Thu, 28 Sep 2023 10:10:12 +0200
Subject: [PATCH] Move the puffin file export logic into the run function

---
 index-scheduler/src/lib.rs | 60 +++++++++++++++++++++++---------------
 1 file changed, 37 insertions(+), 23 deletions(-)

diff --git a/index-scheduler/src/lib.rs b/index-scheduler/src/lib.rs
index c8a7f0c8b..0b3a5d58a 100644
--- a/index-scheduler/src/lib.rs
+++ b/index-scheduler/src/lib.rs
@@ -579,17 +579,52 @@ impl IndexScheduler {
                 run.wake_up.wait();
 
                 loop {
+                    let puffin_enabled = match run.features() {
+                        Ok(features) => features.check_puffin().is_ok(),
+                        Err(e) => {
+                            log::error!("{e}");
+                            continue;
+                        }
+                    };
+                    puffin::set_scopes_on(puffin_enabled);
+                    puffin::GlobalProfiler::lock().new_frame();
+
                     match run.tick() {
                         Ok(TickOutcome::TickAgain(_)) => (),
                         Ok(TickOutcome::WaitForSignal) => run.wake_up.wait(),
                         Err(e) => {
-                            log::error!("{}", e);
+                            log::error!("{e}");
                             // Wait one second when an irrecoverable error occurs.
                             if !e.is_recoverable() {
                                 std::thread::sleep(Duration::from_secs(1));
                             }
                         }
                     }
+
+                    // Let's write the previous frame to disk but only if
+                    // the user wanted to profile with puffin.
+                    if puffin_enabled {
+                        let mut frame_view = run.puffin_frame.lock();
+                        if !frame_view.is_empty() {
+                            let now = OffsetDateTime::now_utc();
+                            let mut file = match File::create(format!("{}.puffin", now)) {
+                                Ok(file) => file,
+                                Err(e) => {
+                                    log::error!("{e}");
+                                    continue;
+                                }
+                            };
+                            if let Err(e) = frame_view.save_to_writer(&mut file) {
+                                log::error!("{e}");
+                            }
+                            if let Err(e) = file.sync_all() {
+                                log::error!("{e}");
+                            }
+                            // We erase this frame view as it is no more useful. We want to
+                            // measure the new frames now that we exported the previous ones.
+                            *frame_view = FrameView::default();
+                        }
+                    }
                 }
             })
             .unwrap();
@@ -1069,34 +1104,13 @@ impl IndexScheduler {
             self.breakpoint(Breakpoint::Start);
         }
 
-        let puffin_enabled = self.features()?.check_puffin().is_ok();
-        puffin::set_scopes_on(puffin_enabled);
-        puffin::GlobalProfiler::lock().new_frame();
-
         self.cleanup_task_queue()?;
 
         let rtxn = self.env.read_txn().map_err(Error::HeedTransaction)?;
         let batch =
             match self.create_next_batch(&rtxn).map_err(|e| Error::CreateBatch(Box::new(e)))? {
                 Some(batch) => batch,
-                None => {
-                    // Let's write the previous frame to disk but only if
-                    // the user wanted to profile with puffin.
-                    if puffin_enabled {
-                        let mut frame_view = self.puffin_frame.lock();
-                        if !frame_view.is_empty() {
-                            let now = OffsetDateTime::now_utc();
-                            let mut file = File::create(format!("{}.puffin", now))?;
-                            frame_view.save_to_writer(&mut file)?;
-                            file.sync_all()?;
-                            // We erase this frame view as it is no more useful. We want to
-                            // measure the new frames now that we exported the previous ones.
-                            *frame_view = FrameView::default();
-                        }
-                    }
-
-                    return Ok(TickOutcome::WaitForSignal);
-                }
+                None => return Ok(TickOutcome::WaitForSignal),
             };
         let index_uid = batch.index_uid().map(ToOwned::to_owned);
         drop(rtxn);