From e4d461ecbac18a57b7d5aef253b74c10e93e990a Mon Sep 17 00:00:00 2001
From: Kerollmops <clement@meilisearch.com>
Date: Thu, 13 Oct 2022 11:07:36 +0200
Subject: [PATCH] Make sure that we do not batch tasks from different indexes

---
 index-scheduler/src/batch.rs | 6 +++---
 index-scheduler/src/lib.rs   | 2 +-
 index-scheduler/src/utils.rs | 5 +++--
 3 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/index-scheduler/src/batch.rs b/index-scheduler/src/batch.rs
index ef1a740f0..6f7f9c82d 100644
--- a/index-scheduler/src/batch.rs
+++ b/index-scheduler/src/batch.rs
@@ -413,16 +413,16 @@ impl IndexScheduler {
             // matter.
             let index_name = task.indexes().unwrap()[0];
 
-            let _index = self.get_index(rtxn, index_name)? & enqueued;
+            let index_tasks = self.index_tasks(rtxn, index_name)? & enqueued;
 
-            // If the autobatching is disabled we only take one task at a time.
+            // If autobatching is disabled we only take one task at a time.
             let tasks_limit = if self.autobatching_enabled {
                 usize::MAX
             } else {
                 1
             };
 
-            let enqueued = enqueued
+            let enqueued = index_tasks
                 .into_iter()
                 .take(tasks_limit)
                 .map(|task_id| {
diff --git a/index-scheduler/src/lib.rs b/index-scheduler/src/lib.rs
index 090e0cbb6..bf6181935 100644
--- a/index-scheduler/src/lib.rs
+++ b/index-scheduler/src/lib.rs
@@ -276,7 +276,7 @@ impl IndexScheduler {
         if let Some(index) = query.index_uid {
             let mut index_tasks = RoaringBitmap::new();
             for index in index {
-                index_tasks |= self.get_index(&rtxn, &index)?;
+                index_tasks |= self.index_tasks(&rtxn, &index)?;
             }
             tasks &= index_tasks;
         }
diff --git a/index-scheduler/src/utils.rs b/index-scheduler/src/utils.rs
index 640d43df8..5a39cc378 100644
--- a/index-scheduler/src/utils.rs
+++ b/index-scheduler/src/utils.rs
@@ -73,7 +73,8 @@ impl IndexScheduler {
         Ok(())
     }
 
-    pub(crate) fn get_index(&self, rtxn: &RoTxn, index: &str) -> Result<RoaringBitmap> {
+    /// Returns the whole set of tasks that belongs to this index.
+    pub(crate) fn index_tasks(&self, rtxn: &RoTxn, index: &str) -> Result<RoaringBitmap> {
         Ok(self.index_tasks.get(rtxn, index)?.unwrap_or_default())
     }
 
@@ -92,7 +93,7 @@ impl IndexScheduler {
         index: &str,
         f: impl Fn(&mut RoaringBitmap),
     ) -> Result<()> {
-        let mut tasks = self.get_index(wtxn, index)?;
+        let mut tasks = self.index_tasks(wtxn, index)?;
         f(&mut tasks);
         self.put_index(wtxn, index, &tasks)?;