write a test ensuring the index-scheduler is effectively down when the upgrade task fail and try to process it when it restarts. There is a bug when deleting this task

crates/index-scheduler/src/scheduler/test_failure.rs

@@ -7,8 +7,10 @@ use meilisearch_types::milli::obkv_to_json;
 use meilisearch_types::milli::update::IndexDocumentsMethod::*;
 use meilisearch_types::milli::update::Setting;
 use meilisearch_types::tasks::KindWithContent;
+use roaring::RoaringBitmap;
 
 use crate::insta_snapshot::snapshot_index_scheduler;
+use crate::test_utils::Breakpoint;
 use crate::test_utils::Breakpoint::*;
 use crate::test_utils::{index_creation_task, read_json, FailureLocation};
 use crate::IndexScheduler;
@@ -249,3 +251,68 @@ fn panic_in_process_batch_for_index_creation() {
     // No matter what happens in process_batch, the index_scheduler should be internally consistent
     snapshot!(snapshot_index_scheduler(&index_scheduler), name: "index_creation_failed");
 }
+
+#[test]
+fn upgrade_failure() {
+    let (index_scheduler, mut handle) =
+        IndexScheduler::test(true, vec![(1, FailureLocation::ProcessUpgrade)]);
+
+    let kind = index_creation_task("catto", "mouse");
+    let _task = index_scheduler.register(kind, None, false).unwrap();
+    let upgrade_database_task = index_scheduler
+        .register(KindWithContent::UpgradeDatabase { from: (1, 12, 0) }, None, false)
+        .unwrap();
+    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "registered_a_task_and_upgrade_task");
+
+    handle.advance_one_failed_batch();
+    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "upgrade_task_failed");
+
+    // We can still register tasks
+    let kind = index_creation_task("doggo", "bone");
+    let _task = index_scheduler.register(kind, None, false).unwrap();
+
+    // But the scheduler is down and won't process anything ever again
+    handle.scheduler_is_down();
+
+    // =====> After a restart is it still working as expected?
+    let (index_scheduler, mut handle) =
+        handle.restart(index_scheduler, true, vec![(1, FailureLocation::ProcessUpgrade)]);
+
+    handle.advance_one_failed_batch();
+    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "upgrade_task_failed_again");
+    // We can still register tasks
+    let kind = index_creation_task("doggo", "bone");
+    let _task = index_scheduler.register(kind, None, false).unwrap();
+    // And the scheduler is still down and won't process anything ever again
+    handle.scheduler_is_down();
+
+    // =====> After a rerestart and without failure can we upgrade the indexes and process the tasks
+    let (index_scheduler, mut handle) = handle.restart(index_scheduler, true, vec![]);
+
+    handle.advance_one_successful_batch();
+    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "upgrade_task_succeeded");
+    // We can still register tasks
+    let kind = index_creation_task("girafo", "leaves");
+    let _task = index_scheduler.register(kind, None, false).unwrap();
+    // The scheduler is up and running
+    handle.advance_one_successful_batch();
+    handle.advance_one_successful_batch();
+    handle.advance_one_failed_batch(); // doggo already exists
+    handle.advance_one_successful_batch();
+    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "after_processing_everything");
+
+    // When deleting the single upgrade task it should remove the three associated batches
+    let _task = index_scheduler
+        .register(
+            KindWithContent::TaskDeletion {
+                query: String::from("test"),
+                tasks: RoaringBitmap::from_iter([upgrade_database_task.uid]),
+            },
+            None,
+            false,
+        )
+        .unwrap();
+
+    handle.advance_one_successful_batch();
+    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "after_removing_the_upgrade");
+}