Make sure that we do not batch tasks from different indexes

This commit is contained in:
Kerollmops 2022-10-13 11:07:36 +02:00 committed by Clément Renault
parent b029369653
commit e4d461ecba
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4
3 changed files with 7 additions and 6 deletions

View File

@ -413,16 +413,16 @@ impl IndexScheduler {
// matter. // matter.
let index_name = task.indexes().unwrap()[0]; 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 { let tasks_limit = if self.autobatching_enabled {
usize::MAX usize::MAX
} else { } else {
1 1
}; };
let enqueued = enqueued let enqueued = index_tasks
.into_iter() .into_iter()
.take(tasks_limit) .take(tasks_limit)
.map(|task_id| { .map(|task_id| {

View File

@ -276,7 +276,7 @@ impl IndexScheduler {
if let Some(index) = query.index_uid { if let Some(index) = query.index_uid {
let mut index_tasks = RoaringBitmap::new(); let mut index_tasks = RoaringBitmap::new();
for index in index { for index in index {
index_tasks |= self.get_index(&rtxn, &index)?; index_tasks |= self.index_tasks(&rtxn, &index)?;
} }
tasks &= index_tasks; tasks &= index_tasks;
} }

View File

@ -73,7 +73,8 @@ impl IndexScheduler {
Ok(()) 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()) Ok(self.index_tasks.get(rtxn, index)?.unwrap_or_default())
} }
@ -92,7 +93,7 @@ impl IndexScheduler {
index: &str, index: &str,
f: impl Fn(&mut RoaringBitmap), f: impl Fn(&mut RoaringBitmap),
) -> Result<()> { ) -> Result<()> {
let mut tasks = self.get_index(wtxn, index)?; let mut tasks = self.index_tasks(wtxn, index)?;
f(&mut tasks); f(&mut tasks);
self.put_index(wtxn, index, &tasks)?; self.put_index(wtxn, index, &tasks)?;