Delete a task's persisted data when appropriate

file-store/src/lib.rs

@@ -1,6 +1,8 @@
+use std::collections::BTreeSet;
 use std::fs::File as StdFile;
 use std::ops::{Deref, DerefMut};
 use std::path::{Path, PathBuf};
+use std::str::FromStr;
 
 use tempfile::NamedTempFile;
 use uuid::Uuid;
@@ -96,6 +98,20 @@ impl FileStore {
         std::fs::remove_file(path)?;
         Ok(())
     }
+
+    /// List the Uuids of the files in the FileStore
+    ///
+    /// This function is meant to be used by tests only.
+    #[doc(hidden)]
+    pub fn __all_uuids(&self) -> BTreeSet<Uuid> {
+        let mut uuids = BTreeSet::new();
+        for entry in self.path.read_dir().unwrap() {
+            let entry = entry.unwrap();
+            let uuid = Uuid::from_str(entry.file_name().to_str().unwrap()).unwrap();
+            uuids.insert(uuid);
+        }
+        uuids
+    }
 }
 
 pub struct File {

index-scheduler/src/batch.rs

@@ -953,15 +953,18 @@ impl IndexScheduler {
         let mut affected_kinds = HashSet::new();
 
         for task_id in to_delete_tasks.iter() {
-            // This should never fail, but there is no harm done if it does. The function
+            let task = self
-            // will still be 99% correct (the number of deleted tasks will be slightly incorrect).
+                .get_task(wtxn, task_id)?
-            if let Some(task) = self.get_task(wtxn, task_id)? {
+                .ok_or(Error::CorruptedTaskQueue)?;
-                if let Some(task_indexes) = task.indexes() {
+            if let Some(task_indexes) = task.indexes() {
-                    affected_indexes.extend(task_indexes.into_iter().map(|x| x.to_owned()));
+                affected_indexes.extend(task_indexes.into_iter().map(|x| x.to_owned()));
-                }
-                affected_statuses.insert(task.status);
-                affected_kinds.insert(task.kind.as_kind());
             }
+            affected_statuses.insert(task.status);
+            affected_kinds.insert(task.kind.as_kind());
+            // Note: don't delete the persisted task data since
+            // we can only delete succeeded, failed, and canceled tasks.
+            // In each of those cases, the persisted data is supposed to
+            // have been deleted already.
         }
         for index in affected_indexes {
             self.update_index(wtxn, &index, |bitmap| {

index-scheduler/src/lib.rs

@@ -381,7 +381,7 @@ impl IndexScheduler {
         match wtxn.commit() {
             Ok(()) => (),
             _e @ Err(_) => {
-                todo!("remove the data associated with the task");
+                self.delete_persisted_task_data(&task)?;
                 // _e?;
             }
         }
@@ -577,8 +577,8 @@ impl IndexScheduler {
                 for mut task in tasks {
                     task.started_at = Some(started_at);
                     task.finished_at = Some(finished_at);
-                    // TODO the info field should've been set by the process_batch function
                     self.update_task(&mut wtxn, &task)?;
+                    self.delete_persisted_task_data(&task)?;
                 }
                 log::info!("A batch of tasks was successfully completed.");
             }
@@ -606,6 +606,25 @@ impl IndexScheduler {
 
         Ok(processed_tasks)
     }
+
+    pub(crate) fn delete_persisted_task_data(&self, task: &Task) -> Result<()> {
+        match &task.kind {
+            KindWithContent::DocumentImport { content_file, .. } => {
+                self.delete_update_file(*content_file)
+            }
+            KindWithContent::DocumentDeletion { .. }
+            | KindWithContent::DocumentClear { .. }
+            | KindWithContent::Settings { .. }
+            | KindWithContent::IndexDeletion { .. }
+            | KindWithContent::IndexCreation { .. }
+            | KindWithContent::IndexUpdate { .. }
+            | KindWithContent::IndexSwap { .. }
+            | KindWithContent::CancelTask { .. }
+            | KindWithContent::TaskDeletion { .. }
+            | KindWithContent::DumpExport { .. }
+            | KindWithContent::Snapshot => Ok(()),
+        }
+    }
 }
 
 #[cfg(test)]

index-scheduler/src/snapshot.rs

@@ -15,7 +15,7 @@ pub fn snapshot_index_scheduler(scheduler: &IndexScheduler) -> String {
     let IndexScheduler {
         autobatching_enabled,
         processing_tasks,
-        file_store: _,
+        file_store,
         env,
         all_tasks,
         status,
@@ -59,11 +59,23 @@ pub fn snapshot_index_scheduler(scheduler: &IndexScheduler) -> String {
     snap.push_str(&snapshot_index_mapper(&rtxn, index_mapper));
     snap.push_str("\n----------------------------------------------------------------------\n");
 
+    snap.push_str("### File Store:\n");
+    snap.push_str(&snapshot_file_store(file_store));
+    snap.push_str("\n----------------------------------------------------------------------\n");
+
     rtxn.commit().unwrap();
 
     snap
 }
 
+fn snapshot_file_store(file_store: &file_store::FileStore) -> String {
+    let mut snap = String::new();
+    for uuid in file_store.__all_uuids() {
+        snap.push_str(&format!("{uuid}\n"));
+    }
+    snap
+}
+
 fn snapshot_bitmap(r: &RoaringBitmap) -> String {
     let mut snap = String::new();
     snap.push('[');

index-scheduler/src/snapshots/lib.rs/document_addition/1.snap

@@ -20,4 +20,8 @@ doggos [0,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+00000000-0000-0000-0000-000000000000
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/document_addition/2.snap

@@ -20,4 +20,8 @@ doggos [0,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+00000000-0000-0000-0000-000000000000
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/document_addition/3.snap

@@ -21,4 +21,7 @@ doggos [0,]
 ### Index Mapper:
 ["doggos"]
 ----------------------------------------------------------------------
+### File Store:
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/insert_task_while_another_task_is_processing/1.snap

@@ -23,4 +23,7 @@ doggos [2,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/register/1.snap

@@ -27,4 +27,7 @@ doggo [4,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_delete_same_task_twice/initial_tasks_enqueued.snap

@@ -22,4 +22,9 @@ doggo [1,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+00000000-0000-0000-0000-000000000000
+00000000-0000-0000-0000-000000000001
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_delete_same_task_twice/initial_tasks_processed.snap

@@ -23,4 +23,8 @@ doggo [1,]
 ### Index Mapper:
 ["catto"]
 ----------------------------------------------------------------------
+### File Store:
+00000000-0000-0000-0000-000000000001
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_delete_same_task_twice/task_deletion_processed.snap

@@ -25,4 +25,8 @@ doggo [1,]
 ### Index Mapper:
 ["catto"]
 ----------------------------------------------------------------------
+### File Store:
+00000000-0000-0000-0000-000000000001
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_deleteable/initial_tasks_enqueued.snap

@@ -22,4 +22,9 @@ doggo [1,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+00000000-0000-0000-0000-000000000000
+00000000-0000-0000-0000-000000000001
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_deleteable/initial_tasks_processed.snap

@@ -23,4 +23,8 @@ doggo [1,]
 ### Index Mapper:
 ["catto"]
 ----------------------------------------------------------------------
+### File Store:
+00000000-0000-0000-0000-000000000001
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_deleteable/task_deletion_processed.snap

@@ -24,4 +24,8 @@ doggo [1,]
 ### Index Mapper:
 ["catto"]
 ----------------------------------------------------------------------
+### File Store:
+00000000-0000-0000-0000-000000000001
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_undeleteable/initial_tasks_enqueued.snap

@@ -24,4 +24,7 @@ doggo [2,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_undeleteable/task_deletion_done.snap

@@ -27,4 +27,7 @@ doggo [2,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_undeleteable/task_deletion_enqueued.snap

@@ -26,4 +26,7 @@ doggo [2,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+
+----------------------------------------------------------------------
 

index-scheduler/src/snapshots/lib.rs/task_deletion_undeleteable/task_deletion_processing.snap

@@ -26,4 +26,7 @@ doggo [2,]
 ### Index Mapper:
 []
 ----------------------------------------------------------------------
+### File Store:
+
+----------------------------------------------------------------------