Change strategy to remove task instead of marking it succeeded

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

@@ -25,9 +25,50 @@ unsafe fn mark_tasks_as_succeeded(
     let env = env_options.max_dbs(TaskQueue::nb_db()).map_size(index_base_map_size).open(dst)?;
     let mut wtxn = env.write_txn()?;
     let task_queue = TaskQueue::new(&env, &mut wtxn)?;
-    for mut task in tasks.iter().cloned() {
-        task.status = Status::Succeeded;
-        task_queue.update_task(&mut wtxn, &task)?;
+
+    // Destructuring to ensure the code below gets updated if a database gets added in the future.
+    let TaskQueue {
+        all_tasks,
+        status,
+        kind,
+        index_tasks: _, // snapshot creation tasks are not index tasks
+        canceled_by,
+        enqueued_at,
+        started_at,
+        finished_at,
+    } = task_queue;
+
+    for task in tasks {
+        all_tasks.delete(&mut wtxn, &task.uid)?;
+
+        let mut tasks = status.get(&wtxn, &task.status)?.unwrap_or_default();
+        tasks.remove(task.uid);
+        status.put(&mut wtxn, &task.status, &tasks)?;
+
+        let mut tasks = kind.get(&wtxn, &task.kind.as_kind())?.unwrap_or_default();
+        tasks.remove(task.uid);
+        kind.put(&mut wtxn, &task.kind.as_kind(), &tasks)?;
+
+        canceled_by.delete(&mut wtxn, &task.uid)?;
+
+        let timestamp = task.enqueued_at.unix_timestamp_nanos();
+        let mut tasks = enqueued_at.get(&wtxn, &timestamp)?.unwrap_or_default();
+        tasks.remove(task.uid);
+        enqueued_at.put(&mut wtxn, &timestamp, &tasks)?;
+
+        if let Some(task_started_at) = task.started_at {
+            let timestamp = task_started_at.unix_timestamp_nanos();
+            let mut tasks = started_at.get(&wtxn, &timestamp)?.unwrap_or_default();
+            tasks.remove(task.uid);
+            started_at.put(&mut wtxn, &timestamp, &tasks)?;
+        }
+
+        if let Some(task_finished_at) = task.finished_at {
+            let timestamp = task_finished_at.unix_timestamp_nanos();
+            let mut tasks = finished_at.get(&wtxn, &timestamp)?.unwrap_or_default();
+            tasks.remove(task.uid);
+            finished_at.put(&mut wtxn, &timestamp, &tasks)?;
+        }
     }
     wtxn.commit()?;
     Ok(())
@@ -80,11 +121,7 @@ impl IndexScheduler {
         // This is safe because we open the env file we just created in a temporary directory.
         // We are sure it's not being used by any other process nor thread.
         unsafe {
-            mark_tasks_as_succeeded(
-                &tasks,
-                &dst,
-                self.index_mapper.index_base_map_size,
-            )?;
+            mark_tasks_as_succeeded(&tasks, &dst, self.index_mapper.index_base_map_size)?;
         }
 
         // 2.3 Create a read transaction on the index-scheduler