From 705af94fd74c64118bd5248b7ab6503c713648b7 Mon Sep 17 00:00:00 2001
From: Tamo <tamo@meilisearch.com>
Date: Wed, 7 Sep 2022 00:22:58 +0200
Subject: [PATCH] add the task to the index db in the register task

---
 index-scheduler/src/lib.rs   |  9 +++++
 index-scheduler/src/task.rs  | 69 +++++++++++++++++++++++++-----------
 index-scheduler/src/utils.rs | 39 +++++++++++++-------
 3 files changed, 84 insertions(+), 33 deletions(-)

diff --git a/index-scheduler/src/lib.rs b/index-scheduler/src/lib.rs
index 52b22b6f6..978898732 100644
--- a/index-scheduler/src/lib.rs
+++ b/index-scheduler/src/lib.rs
@@ -110,6 +110,15 @@ impl IndexScheduler {
         self.all_tasks
             .append(&mut wtxn, &BEU32::new(task_id), &task)?;
 
+        if let Some(indexes) = task.indexes() {
+            for index in indexes {
+                self.update_index(&mut wtxn, index, |mut bitmap| {
+                    bitmap.insert(task_id);
+                    bitmap
+                })?;
+            }
+        }
+
         self.update_status(&mut wtxn, Status::Enqueued, |mut bitmap| {
             bitmap.insert(task_id);
             bitmap
diff --git a/index-scheduler/src/task.rs b/index-scheduler/src/task.rs
index 06ee76dfd..7205bf849 100644
--- a/index-scheduler/src/task.rs
+++ b/index-scheduler/src/task.rs
@@ -29,13 +29,20 @@ pub struct Task {
 }
 
 impl Task {
+    /// Persist all the temp files associated with the task.
     pub fn persist(&self) -> Result<()> {
         self.kind.persist()
     }
 
+    /// Delete all the files associated with the task.
     pub fn remove_data(&self) -> Result<()> {
         self.kind.remove_data()
     }
+
+    /// Return the list of indexes updated by this tasks.
+    pub fn indexes(&self) -> Option<Vec<&str>> {
+        self.kind.indexes()
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -98,8 +105,10 @@ impl KindWithContent {
     }
 
     pub fn persist(&self) -> Result<()> {
+        use KindWithContent::*;
+
         match self {
-            KindWithContent::DocumentAddition {
+            DocumentAddition {
                 index_name: _,
                 content_file: _,
             } => {
@@ -108,21 +117,23 @@ impl KindWithContent {
                 Ok(())
             }
             // There is nothing to persist for all these tasks
-            KindWithContent::DumpExport { .. }
-            | KindWithContent::DocumentDeletion { .. }
-            | KindWithContent::ClearAllDocuments { .. }
-            | KindWithContent::RenameIndex { .. }
-            | KindWithContent::CreateIndex { .. }
-            | KindWithContent::DeleteIndex { .. }
-            | KindWithContent::SwapIndex { .. }
-            | KindWithContent::CancelTask { .. }
-            | KindWithContent::Snapshot => Ok(()),
+            DumpExport { .. }
+            | DocumentDeletion { .. }
+            | ClearAllDocuments { .. }
+            | RenameIndex { .. }
+            | CreateIndex { .. }
+            | DeleteIndex { .. }
+            | SwapIndex { .. }
+            | CancelTask { .. }
+            | Snapshot => Ok(()),
         }
     }
 
     pub fn remove_data(&self) -> Result<()> {
+        use KindWithContent::*;
+
         match self {
-            KindWithContent::DocumentAddition {
+            DocumentAddition {
                 index_name: _,
                 content_file: _,
             } => {
@@ -131,15 +142,33 @@ impl KindWithContent {
                 Ok(())
             }
             // There is no data associated with all these tasks
-            KindWithContent::DumpExport { .. }
-            | KindWithContent::DocumentDeletion { .. }
-            | KindWithContent::ClearAllDocuments { .. }
-            | KindWithContent::RenameIndex { .. }
-            | KindWithContent::CreateIndex { .. }
-            | KindWithContent::DeleteIndex { .. }
-            | KindWithContent::SwapIndex { .. }
-            | KindWithContent::CancelTask { .. }
-            | KindWithContent::Snapshot => Ok(()),
+            DumpExport { .. }
+            | DocumentDeletion { .. }
+            | ClearAllDocuments { .. }
+            | RenameIndex { .. }
+            | CreateIndex { .. }
+            | DeleteIndex { .. }
+            | SwapIndex { .. }
+            | CancelTask { .. }
+            | Snapshot => Ok(()),
+        }
+    }
+
+    pub fn indexes(&self) -> Option<Vec<&str>> {
+        use KindWithContent::*;
+
+        match self {
+            DumpExport { .. } | Snapshot | CancelTask { .. } => None,
+            DocumentAddition { index_name, .. }
+            | DocumentDeletion { index_name, .. }
+            | ClearAllDocuments { index_name }
+            | CreateIndex { index_name, .. }
+            | DeleteIndex { index_name } => Some(vec![index_name]),
+            RenameIndex {
+                index_name: lhs,
+                new_name: rhs,
+            }
+            | SwapIndex { lhs, rhs } => Some(vec![lhs, rhs]),
         }
     }
 }
diff --git a/index-scheduler/src/utils.rs b/index-scheduler/src/utils.rs
index 2d27bd111..c3c4c7350 100644
--- a/index-scheduler/src/utils.rs
+++ b/index-scheduler/src/utils.rs
@@ -45,6 +45,19 @@ impl IndexScheduler {
         Ok(self.index_tasks.put(wtxn, index, bitmap)?)
     }
 
+    pub(crate) fn update_index(
+        &self,
+        wtxn: &mut RwTxn,
+        index: &str,
+        f: impl Fn(RoaringBitmap) -> RoaringBitmap,
+    ) -> Result<()> {
+        let tasks = self.get_index(&wtxn, index)?;
+        let tasks = f(tasks);
+        self.put_index(wtxn, index, &tasks)?;
+
+        Ok(())
+    }
+
     pub(crate) fn get_status(&self, rtxn: &RoTxn, status: Status) -> Result<RoaringBitmap> {
         Ok(self.status.get(&rtxn, &status)?.unwrap_or_default())
     }
@@ -58,19 +71,6 @@ impl IndexScheduler {
         Ok(self.status.put(wtxn, &status, bitmap)?)
     }
 
-    pub(crate) fn get_kind(&self, rtxn: &RoTxn, kind: Kind) -> Result<RoaringBitmap> {
-        Ok(self.kind.get(&rtxn, &kind)?.unwrap_or_default())
-    }
-
-    pub(crate) fn put_kind(
-        &self,
-        wtxn: &mut RwTxn,
-        kind: Kind,
-        bitmap: &RoaringBitmap,
-    ) -> Result<()> {
-        Ok(self.kind.put(wtxn, &kind, bitmap)?)
-    }
-
     pub(crate) fn update_status(
         &self,
         wtxn: &mut RwTxn,
@@ -84,6 +84,19 @@ impl IndexScheduler {
         Ok(())
     }
 
+    pub(crate) fn get_kind(&self, rtxn: &RoTxn, kind: Kind) -> Result<RoaringBitmap> {
+        Ok(self.kind.get(&rtxn, &kind)?.unwrap_or_default())
+    }
+
+    pub(crate) fn put_kind(
+        &self,
+        wtxn: &mut RwTxn,
+        kind: Kind,
+        bitmap: &RoaringBitmap,
+    ) -> Result<()> {
+        Ok(self.kind.put(wtxn, &kind, bitmap)?)
+    }
+
     pub(crate) fn update_kind(
         &self,
         wtxn: &mut RwTxn,