From 121c1ac1dd631df845954602b1c0e41a3b62679d Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Tue, 22 Apr 2025 09:34:37 +0200
Subject: [PATCH] Upgrade supports cancelling

---
 .../src/scheduler/process_upgrade/mod.rs              |  8 +++++++-
 crates/milli/src/update/upgrade/mod.rs                | 11 +++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/crates/index-scheduler/src/scheduler/process_upgrade/mod.rs b/crates/index-scheduler/src/scheduler/process_upgrade/mod.rs
index 4feebabc4..e270ad2a8 100644
--- a/crates/index-scheduler/src/scheduler/process_upgrade/mod.rs
+++ b/crates/index-scheduler/src/scheduler/process_upgrade/mod.rs
@@ -1,5 +1,5 @@
-use meilisearch_types::milli;
 use meilisearch_types::milli::progress::{Progress, VariableNameStep};
+use meilisearch_types::milli::{self};
 
 use crate::{Error, IndexScheduler, Result};
 
@@ -16,6 +16,11 @@ impl IndexScheduler {
         let indexes = self.index_names()?;
 
         for (i, uid) in indexes.iter().enumerate() {
+            let must_stop_processing = self.scheduler.must_stop_processing.clone();
+
+            if must_stop_processing.get() {
+                return Err(Error::AbortedTask);
+            }
             progress.update_progress(VariableNameStep::<UpgradeIndex>::new(
                 format!("Upgrading index `{uid}`"),
                 i as u32,
@@ -27,6 +32,7 @@ impl IndexScheduler {
                 &mut index_wtxn,
                 &index,
                 db_version,
+                || must_stop_processing.get(),
                 progress.clone(),
             )
             .map_err(|e| Error::from_milli(e, Some(uid.to_string())))?;
diff --git a/crates/milli/src/update/upgrade/mod.rs b/crates/milli/src/update/upgrade/mod.rs
index 7c8dcf64a..de88b33e0 100644
--- a/crates/milli/src/update/upgrade/mod.rs
+++ b/crates/milli/src/update/upgrade/mod.rs
@@ -23,12 +23,16 @@ trait UpgradeIndex {
 }
 
 /// Return true if the cached stats of the index must be regenerated
-pub fn upgrade(
+pub fn upgrade<MSP>(
     wtxn: &mut RwTxn,
     index: &Index,
     db_version: (u32, u32, u32),
+    must_stop_processing: MSP,
     progress: Progress,
-) -> Result<bool> {
+) -> Result<bool>
+where
+    MSP: Fn() -> bool + Sync,
+{
     let from = index.get_version(wtxn)?.unwrap_or(db_version);
     let upgrade_functions: &[&dyn UpgradeIndex] = &[
         &V1_12_To_V1_12_3 {},
@@ -56,6 +60,9 @@ pub fn upgrade(
     let mut current_version = from;
     let mut regenerate_stats = false;
     for (i, upgrade) in upgrade_path.iter().enumerate() {
+        if (must_stop_processing)() {
+            return Err(crate::Error::InternalError(InternalError::AbortedIndexation));
+        }
         let target = upgrade.target_version();
         progress.update_progress(VariableNameStep::<UpgradeVersion>::new(
             format!(