searchengines-web/MeiliSearch
searchengines-web/MeiliSearch
1
0
Fork 0
mirror of https://github.com/meilisearch/MeiliSearch synced 2024-06-01 10:28:02 +02:00
Code Issues Releases Wiki Activity

split the scheduler into multiples files

Browse Source
This commit is contained in:
Tamo 2022-09-07 00:10:14 +02:00 committed by Clément Renault
parent 22d24dba56
commit ed745591e1
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4
3 changed files with 213 additions and 175 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

110
index-scheduler/src/batch.rs Normal file
View File

@ -0,0 +1,110 @@
use crate::{
task::{KindWithContent, Status},
Error, IndexScheduler, Result,
};
use milli::heed::RoTxn;
use crate::{task::Kind, Task};
pub enum Batch {
Cancel(Task),
Snapshot(Vec<Task>),
Dump(Vec<Task>),
Contiguous { tasks: Vec<Task>, kind: Kind },
One(Task),
Empty,
}
impl IndexScheduler {
/// Create the next batch to be processed;
/// 1. We get the *last* task to cancel.
/// 2. We get the *next* snapshot to process.
/// 3. We get the *next* dump to process.
/// 4. We get the *next* tasks to process for a specific index.
fn get_next_batch(&self, rtxn: &RoTxn) -> Result<Batch> {
let enqueued = &self.get_status(rtxn, Status::Enqueued)?;
let to_cancel = self.get_kind(rtxn, Kind::CancelTask)? & enqueued;
// 1. we get the last task to cancel.
if let Some(task_id) = to_cancel.max() {
return Ok(Batch::Cancel(
self.get_task(rtxn, task_id)?
.ok_or(Error::CorruptedTaskQueue)?,
));
}
// 2. we batch the snapshot.
let to_snapshot = self.get_kind(rtxn, Kind::Snapshot)? & enqueued;
if !to_snapshot.is_empty() {
return Ok(Batch::Snapshot(self.get_existing_tasks(rtxn, to_snapshot)?));
}
// 3. we batch the dumps.
let to_dump = self.get_kind(rtxn, Kind::DumpExport)? & enqueued;
if !to_dump.is_empty() {
return Ok(Batch::Dump(self.get_existing_tasks(rtxn, to_dump)?));
}
// 4. We take the next task and try to batch all the tasks associated with this index.
if let Some(task_id) = enqueued.min() {
let task = self
.get_task(rtxn, task_id)?
.ok_or(Error::CorruptedTaskQueue)?;
match task.kind {
// We can batch all the consecutive tasks coming next which
// have the kind `DocumentAddition`.
KindWithContent::DocumentAddition { index_name, .. } => {
return self.batch_contiguous_kind(rtxn, &index_name, Kind::DocumentAddition)
}
// We can batch all the consecutive tasks coming next which
// have the kind `DocumentDeletion`.
KindWithContent::DocumentDeletion { index_name, .. } => {
return self.batch_contiguous_kind(rtxn, &index_name, Kind::DocumentAddition)
}
// The following tasks can't be batched
KindWithContent::ClearAllDocuments { .. }
| KindWithContent::RenameIndex { .. }
| KindWithContent::CreateIndex { .. }
| KindWithContent::DeleteIndex { .. }
| KindWithContent::SwapIndex { .. } => return Ok(Batch::One(task)),
// The following tasks have already been batched and thus can't appear here.
KindWithContent::CancelTask { .. }
| KindWithContent::DumpExport { .. }
| KindWithContent::Snapshot => {
unreachable!()
}
}
}
// If we found no tasks then we were notified for something that got autobatched
// somehow and there is nothing to do.
Ok(Batch::Empty)
}
/// Batch all the consecutive tasks coming next that shares the same `Kind`
/// for a specific index. There *MUST* be at least ONE task of this kind.
fn batch_contiguous_kind(&self, rtxn: &RoTxn, index: &str, kind: Kind) -> Result<Batch> {
let enqueued = &self.get_status(rtxn, Status::Enqueued)?;
// [1, 2, 4, 5]
let index_tasks = self.get_index(rtxn, &index)? & enqueued;
// [1, 2, 5]
let tasks_kind = &index_tasks & self.get_kind(rtxn, kind)?;
// [4]
let not_kind = &index_tasks - &tasks_kind;
// [1, 2]
let mut to_process = tasks_kind.clone();
if let Some(max) = not_kind.max() {
// it's safe to unwrap since we already ensured there
// was AT LEAST one task with the document addition tasks_kind.
to_process.remove_range(tasks_kind.min().unwrap()..max);
}
Ok(Batch::Contiguous {
tasks: self.get_existing_tasks(rtxn, to_process)?,
kind,
})
}
}

179
index-scheduler/src/lib.rs
View File

@ -1,10 +1,12 @@
mod batch;
pub mod error;
pub mod task;
mod utils;
use error::Error;
pub use error::Error;
use milli::heed::types::{DecodeIgnore, OwnedType, SerdeBincode, Str};
pub use task::Task;
use task::{Kind, KindWithContent, Status};
use task::{Kind, Status};
use std::collections::hash_map::Entry;
use std::sync::atomic::AtomicBool;
@ -134,181 +136,8 @@ impl IndexScheduler {
Ok(())
}
/// Create the next batch to be processed;
/// 1. We get the *last* task to cancel.
/// 2. We get the *next* snapshot to process.
/// 3. We get the *next* dump to process.
/// 4. We get the *next* tasks to process for a specific index.
fn get_next_batch(&self, rtxn: &RoTxn) -> Result<Batch> {
let enqueued = &self.get_status(rtxn, Status::Enqueued)?;
let to_cancel = self.get_kind(rtxn, Kind::CancelTask)? & enqueued;
// 1. we get the last task to cancel.
if let Some(task_id) = to_cancel.max() {
return Ok(Batch::Cancel(
self.get_task(rtxn, task_id)?
.ok_or(Error::CorruptedTaskQueue)?,
));
}
// 2. we batch the snapshot.
let to_snapshot = self.get_kind(rtxn, Kind::Snapshot)? & enqueued;
if !to_snapshot.is_empty() {
return Ok(Batch::Snapshot(self.get_existing_tasks(rtxn, to_snapshot)?));
}
// 3. we batch the dumps.
let to_dump = self.get_kind(rtxn, Kind::DumpExport)? & enqueued;
if !to_dump.is_empty() {
return Ok(Batch::Dump(self.get_existing_tasks(rtxn, to_dump)?));
}
// 4. We take the next task and try to batch all the tasks associated with this index.
if let Some(task_id) = enqueued.min() {
let task = self
.get_task(rtxn, task_id)?
.ok_or(Error::CorruptedTaskQueue)?;
match task.kind {
// We can batch all the consecutive tasks coming next which
// have the kind `DocumentAddition`.
KindWithContent::DocumentAddition { index_name, .. } => {
return self.batch_contiguous_kind(rtxn, &index_name, Kind::DocumentAddition)
}
// We can batch all the consecutive tasks coming next which
// have the kind `DocumentDeletion`.
KindWithContent::DocumentDeletion { index_name, .. } => {
return self.batch_contiguous_kind(rtxn, &index_name, Kind::DocumentAddition)
}
// The following tasks can't be batched
KindWithContent::ClearAllDocuments { .. }
| KindWithContent::RenameIndex { .. }
| KindWithContent::CreateIndex { .. }
| KindWithContent::DeleteIndex { .. }
| KindWithContent::SwapIndex { .. } => return Ok(Batch::One(task)),
// The following tasks have already been batched and thus can't appear here.
KindWithContent::CancelTask { .. }
| KindWithContent::DumpExport { .. }
| KindWithContent::Snapshot => {
unreachable!()
}
}
}
// If we found no tasks then we were notified for something that got autobatched
// somehow and there is nothing to do.
Ok(Batch::Empty)
}
/// Batch all the consecutive tasks coming next that shares the same `Kind`
/// for a specific index. There *MUST* be at least ONE task of this kind.
fn batch_contiguous_kind(&self, rtxn: &RoTxn, index: &str, kind: Kind) -> Result<Batch> {
let enqueued = &self.get_status(rtxn, Status::Enqueued)?;
// [1, 2, 4, 5]
let index_tasks = self.get_index(rtxn, &index)? & enqueued;
// [1, 2, 5]
let tasks_kind = &index_tasks & self.get_kind(rtxn, kind)?;
// [4]
let not_kind = &index_tasks - &tasks_kind;
// [1, 2]
let mut to_process = tasks_kind.clone();
if let Some(max) = not_kind.max() {
// it's safe to unwrap since we already ensured there
// was AT LEAST one task with the document addition tasks_kind.
to_process.remove_range(tasks_kind.min().unwrap()..max);
}
Ok(Batch::Contiguous {
tasks: self.get_existing_tasks(rtxn, to_process)?,
kind,
})
}
fn get_task(&self, rtxn: &RoTxn, task_id: TaskId) -> Result<Option<Task>> {
Ok(self.all_tasks.get(rtxn, &BEU32::new(task_id))?)
}
pub fn notify(&self) {
self.wake_up
.store(true, std::sync::atomic::Ordering::Relaxed);
}
// =========== Utility functions on the DBs
/// Convert an iterator to a `Vec` of tasks. The tasks MUST exist or a
// `CorruptedTaskQueue` error will be throwed.
fn get_existing_tasks(
&self,
rtxn: &RoTxn,
tasks: impl IntoIterator<Item = TaskId>,
) -> Result<Vec<Task>> {
tasks
.into_iter()
.map(|task_id| {
self.get_task(rtxn, task_id)
.and_then(|task| task.ok_or(Error::CorruptedTaskQueue))
})
.collect::<Result<_>>()
}
fn get_index(&self, rtxn: &RoTxn, index: &str) -> Result<RoaringBitmap> {
Ok(self.index_tasks.get(&rtxn, index)?.unwrap_or_default())
}
fn put_index(&self, wtxn: &mut RwTxn, index: &str, bitmap: &RoaringBitmap) -> Result<()> {
Ok(self.index_tasks.put(wtxn, index, bitmap)?)
}
fn get_status(&self, rtxn: &RoTxn, status: Status) -> Result<RoaringBitmap> {
Ok(self.status.get(&rtxn, &status)?.unwrap_or_default())
}
fn put_status(&self, wtxn: &mut RwTxn, status: Status, bitmap: &RoaringBitmap) -> Result<()> {
Ok(self.status.put(wtxn, &status, bitmap)?)
}
fn get_kind(&self, rtxn: &RoTxn, kind: Kind) -> Result<RoaringBitmap> {
Ok(self.kind.get(&rtxn, &kind)?.unwrap_or_default())
}
fn put_kind(&self, wtxn: &mut RwTxn, kind: Kind, bitmap: &RoaringBitmap) -> Result<()> {
Ok(self.kind.put(wtxn, &kind, bitmap)?)
}
fn update_status(
&self,
wtxn: &mut RwTxn,
status: Status,
f: impl Fn(RoaringBitmap) -> RoaringBitmap,
) -> Result<()> {
let tasks = self.get_status(&wtxn, status)?;
let tasks = f(tasks);
self.put_status(wtxn, status, &tasks)?;
Ok(())
}
fn update_kind(
&self,
wtxn: &mut RwTxn,
kind: Kind,
f: impl Fn(RoaringBitmap) -> RoaringBitmap,
) -> Result<()> {
let tasks = self.get_kind(&wtxn, kind)?;
let tasks = f(tasks);
self.put_kind(wtxn, kind, &tasks)?;
Ok(())
}
}
enum Batch {
Cancel(Task),
Snapshot(Vec<Task>),
Dump(Vec<Task>),
Contiguous { tasks: Vec<Task>, kind: Kind },
One(Task),
Empty,
}

99
index-scheduler/src/utils.rs Normal file
View File

@ -0,0 +1,99 @@
//! Utility functions on the DBs. Mainly getter and setters.
use milli::{
heed::{RoTxn, RwTxn},
BEU32,
};
use roaring::RoaringBitmap;
use crate::{
task::{Kind, Status},
Error, IndexScheduler, Result, Task, TaskId,
};
impl IndexScheduler {
pub(crate) fn get_task(&self, rtxn: &RoTxn, task_id: TaskId) -> Result<Option<Task>> {
Ok(self.all_tasks.get(rtxn, &BEU32::new(task_id))?)
}
/// Convert an iterator to a `Vec` of tasks. The tasks MUST exist or a
/// `CorruptedTaskQueue` error will be throwed.
pub(crate) fn get_existing_tasks(
&self,
rtxn: &RoTxn,
tasks: impl IntoIterator<Item = TaskId>,
) -> Result<Vec<Task>> {
tasks
.into_iter()
.map(|task_id| {
self.get_task(rtxn, task_id)
.and_then(|task| task.ok_or(Error::CorruptedTaskQueue))
})
.collect::<Result<_>>()
}
pub(crate) fn get_index(&self, rtxn: &RoTxn, index: &str) -> Result<RoaringBitmap> {
Ok(self.index_tasks.get(&rtxn, index)?.unwrap_or_default())
}
pub(crate) fn put_index(
&self,
wtxn: &mut RwTxn,
index: &str,
bitmap: &RoaringBitmap,
) -> Result<()> {
Ok(self.index_tasks.put(wtxn, index, bitmap)?)
}
pub(crate) fn get_status(&self, rtxn: &RoTxn, status: Status) -> Result<RoaringBitmap> {
Ok(self.status.get(&rtxn, &status)?.unwrap_or_default())
}
pub(crate) fn put_status(
&self,
wtxn: &mut RwTxn,
status: Status,
bitmap: &RoaringBitmap,
) -> Result<()> {
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,
status: Status,
f: impl Fn(RoaringBitmap) -> RoaringBitmap,
) -> Result<()> {
let tasks = self.get_status(&wtxn, status)?;
let tasks = f(tasks);
self.put_status(wtxn, status, &tasks)?;
Ok(())
}
pub(crate) fn update_kind(
&self,
wtxn: &mut RwTxn,
kind: Kind,
f: impl Fn(RoaringBitmap) -> RoaringBitmap,
) -> Result<()> {
let tasks = self.get_kind(&wtxn, kind)?;
let tasks = f(tasks);
self.put_kind(wtxn, kind, &tasks)?;
Ok(())
}
}