mirror of
https://github.com/meilisearch/meilisearch.git
synced 2025-08-02 19:59:58 +00:00
Remove once for all the meilisearch-lib crate
This commit is contained in:
Clément Renault
@ -1,420 +0,0 @@
|
||||
mod store;
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::io::{BufWriter, Write};
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use log::debug;
|
||||
use milli::heed::{Env, RwTxn};
|
||||
use time::OffsetDateTime;
|
||||
|
||||
use super::batch::BatchContent;
|
||||
use super::error::TaskError;
|
||||
use super::scheduler::Processing;
|
||||
use super::task::{Task, TaskContent, TaskId};
|
||||
use super::Result;
|
||||
use crate::tasks::task::TaskEvent;
|
||||
use crate::update_file_store::UpdateFileStore;
|
||||
|
||||
#[cfg(test)]
|
||||
pub use store::test::MockStore as Store;
|
||||
#[cfg(not(test))]
|
||||
pub use store::Store;
|
||||
|
||||
type FilterFn = Box<dyn Fn(&Task) -> bool + Sync + Send + 'static>;
|
||||
|
||||
/// Defines constraints to be applied when querying for Tasks from the store.
|
||||
#[derive(Default)]
|
||||
pub struct TaskFilter {
|
||||
indexes: Option<HashSet<String>>,
|
||||
filter_fn: Option<FilterFn>,
|
||||
}
|
||||
|
||||
impl TaskFilter {
|
||||
fn pass(&self, task: &Task) -> bool {
|
||||
match task.index_uid() {
|
||||
Some(index_uid) => self
|
||||
.indexes
|
||||
.as_ref()
|
||||
.map_or(true, |indexes| indexes.contains(index_uid)),
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn filtered_indexes(&self) -> Option<&HashSet<String>> {
|
||||
self.indexes.as_ref()
|
||||
}
|
||||
|
||||
/// Adds an index to the filter, so the filter must match this index.
|
||||
pub fn filter_index(&mut self, index: String) {
|
||||
self.indexes
|
||||
.get_or_insert_with(Default::default)
|
||||
.insert(index);
|
||||
}
|
||||
|
||||
pub fn filter_fn(&mut self, f: FilterFn) {
|
||||
self.filter_fn.replace(f);
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TaskStore {
|
||||
store: Arc<Store>,
|
||||
}
|
||||
|
||||
impl Clone for TaskStore {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
store: self.store.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TaskStore {
|
||||
pub fn new(env: Arc<milli::heed::Env>) -> Result<Self> {
|
||||
let store = Arc::new(Store::new(env)?);
|
||||
Ok(Self { store })
|
||||
}
|
||||
|
||||
pub async fn register(&self, content: TaskContent) -> Result<Task> {
|
||||
debug!("registering update: {:?}", content);
|
||||
let store = self.store.clone();
|
||||
let task = tokio::task::spawn_blocking(move || -> Result<Task> {
|
||||
let mut txn = store.wtxn()?;
|
||||
let next_task_id = store.next_task_id(&mut txn)?;
|
||||
let created_at = TaskEvent::Created(OffsetDateTime::now_utc());
|
||||
let task = Task {
|
||||
id: next_task_id,
|
||||
content,
|
||||
events: vec![created_at],
|
||||
};
|
||||
|
||||
store.put(&mut txn, &task)?;
|
||||
txn.commit()?;
|
||||
|
||||
Ok(task)
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(task)
|
||||
}
|
||||
|
||||
pub fn register_raw_update(&self, wtxn: &mut RwTxn, task: &Task) -> Result<()> {
|
||||
self.store.put(wtxn, task)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_task(&self, id: TaskId, filter: Option<TaskFilter>) -> Result<Task> {
|
||||
let store = self.store.clone();
|
||||
let task = tokio::task::spawn_blocking(move || -> Result<_> {
|
||||
let txn = store.rtxn()?;
|
||||
let task = store.get(&txn, id)?;
|
||||
Ok(task)
|
||||
})
|
||||
.await??
|
||||
.ok_or(TaskError::UnexistingTask(id))?;
|
||||
|
||||
match filter {
|
||||
Some(filter) => filter
|
||||
.pass(&task)
|
||||
.then_some(task)
|
||||
.ok_or(TaskError::UnexistingTask(id)),
|
||||
None => Ok(task),
|
||||
}
|
||||
}
|
||||
|
||||
/// This methods takes a `Processing` which contains the next task ids to process, and returns
|
||||
/// the corresponding tasks along with the ownership to the passed processing.
|
||||
///
|
||||
/// We need get_processing_tasks to take ownership over `Processing` because we need it to be
|
||||
/// valid for 'static.
|
||||
pub async fn get_processing_tasks(
|
||||
&self,
|
||||
processing: Processing,
|
||||
) -> Result<(Processing, BatchContent)> {
|
||||
let store = self.store.clone();
|
||||
let tasks = tokio::task::spawn_blocking(move || -> Result<_> {
|
||||
let txn = store.rtxn()?;
|
||||
|
||||
let content = match processing {
|
||||
Processing::DocumentAdditions(ref ids) => {
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for id in ids.iter() {
|
||||
let task = store
|
||||
.get(&txn, *id)?
|
||||
.ok_or(TaskError::UnexistingTask(*id))?;
|
||||
tasks.push(task);
|
||||
}
|
||||
BatchContent::DocumentsAdditionBatch(tasks)
|
||||
}
|
||||
Processing::IndexUpdate(id) => {
|
||||
let task = store.get(&txn, id)?.ok_or(TaskError::UnexistingTask(id))?;
|
||||
BatchContent::IndexUpdate(task)
|
||||
}
|
||||
Processing::Dump(id) => {
|
||||
let task = store.get(&txn, id)?.ok_or(TaskError::UnexistingTask(id))?;
|
||||
debug_assert!(matches!(task.content, TaskContent::Dump { .. }));
|
||||
BatchContent::Dump(task)
|
||||
}
|
||||
Processing::Nothing => BatchContent::Empty,
|
||||
};
|
||||
|
||||
Ok((processing, content))
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
pub async fn update_tasks(&self, tasks: Vec<Task>) -> Result<Vec<Task>> {
|
||||
let store = self.store.clone();
|
||||
|
||||
let tasks = tokio::task::spawn_blocking(move || -> Result<_> {
|
||||
let mut txn = store.wtxn()?;
|
||||
|
||||
for task in &tasks {
|
||||
store.put(&mut txn, task)?;
|
||||
}
|
||||
|
||||
txn.commit()?;
|
||||
|
||||
Ok(tasks)
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
pub async fn fetch_unfinished_tasks(&self, offset: Option<TaskId>) -> Result<Vec<Task>> {
|
||||
let store = self.store.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let txn = store.rtxn()?;
|
||||
let tasks = store.fetch_unfinished_tasks(&txn, offset)?;
|
||||
Ok(tasks)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn list_tasks(
|
||||
&self,
|
||||
offset: Option<TaskId>,
|
||||
filter: Option<TaskFilter>,
|
||||
limit: Option<usize>,
|
||||
) -> Result<Vec<Task>> {
|
||||
let store = self.store.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let txn = store.rtxn()?;
|
||||
let tasks = store.list_tasks(&txn, offset, filter, limit)?;
|
||||
Ok(tasks)
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
pub async fn dump(
|
||||
env: Arc<Env>,
|
||||
dir_path: impl AsRef<Path>,
|
||||
update_file_store: UpdateFileStore,
|
||||
) -> Result<()> {
|
||||
let store = Self::new(env)?;
|
||||
let update_dir = dir_path.as_ref().join("updates");
|
||||
let updates_file = update_dir.join("data.jsonl");
|
||||
let tasks = store.list_tasks(None, None, None).await?;
|
||||
|
||||
let dir_path = dir_path.as_ref().to_path_buf();
|
||||
tokio::task::spawn_blocking(move || -> Result<()> {
|
||||
std::fs::create_dir(&update_dir)?;
|
||||
let updates_file = std::fs::File::create(updates_file)?;
|
||||
let mut updates_file = BufWriter::new(updates_file);
|
||||
|
||||
for task in tasks {
|
||||
serde_json::to_writer(&mut updates_file, &task)?;
|
||||
updates_file.write_all(b"\n")?;
|
||||
|
||||
if !task.is_finished() {
|
||||
if let Some(content_uuid) = task.get_content_uuid() {
|
||||
update_file_store.dump(content_uuid, &dir_path)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
updates_file.flush()?;
|
||||
Ok(())
|
||||
})
|
||||
.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn load_dump(src: impl AsRef<Path>, env: Arc<Env>) -> anyhow::Result<()> {
|
||||
// create a dummy update field store, since it is not needed right now.
|
||||
let store = Self::new(env.clone())?;
|
||||
|
||||
let src_update_path = src.as_ref().join("updates");
|
||||
let update_data = std::fs::File::open(&src_update_path.join("data.jsonl"))?;
|
||||
let update_data = std::io::BufReader::new(update_data);
|
||||
|
||||
let stream = serde_json::Deserializer::from_reader(update_data).into_iter::<Task>();
|
||||
|
||||
let mut wtxn = env.write_txn()?;
|
||||
for entry in stream {
|
||||
store.register_raw_update(&mut wtxn, &entry?)?;
|
||||
}
|
||||
wtxn.commit()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test {
|
||||
use crate::tasks::{scheduler::Processing, task_store::store::test::tmp_env};
|
||||
|
||||
use super::*;
|
||||
|
||||
use meilisearch_types::index_uid::IndexUid;
|
||||
use nelson::Mocker;
|
||||
use proptest::{
|
||||
strategy::Strategy,
|
||||
test_runner::{Config, TestRunner},
|
||||
};
|
||||
|
||||
pub enum MockTaskStore {
|
||||
Real(TaskStore),
|
||||
Mock(Arc<Mocker>),
|
||||
}
|
||||
|
||||
impl Clone for MockTaskStore {
|
||||
fn clone(&self) -> Self {
|
||||
match self {
|
||||
Self::Real(x) => Self::Real(x.clone()),
|
||||
Self::Mock(x) => Self::Mock(x.clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl MockTaskStore {
|
||||
pub fn new(env: Arc<milli::heed::Env>) -> Result<Self> {
|
||||
Ok(Self::Real(TaskStore::new(env)?))
|
||||
}
|
||||
|
||||
pub async fn dump(
|
||||
env: Arc<milli::heed::Env>,
|
||||
path: impl AsRef<Path>,
|
||||
update_file_store: UpdateFileStore,
|
||||
) -> Result<()> {
|
||||
TaskStore::dump(env, path, update_file_store).await
|
||||
}
|
||||
|
||||
pub fn mock(mocker: Mocker) -> Self {
|
||||
Self::Mock(Arc::new(mocker))
|
||||
}
|
||||
|
||||
pub async fn update_tasks(&self, tasks: Vec<Task>) -> Result<Vec<Task>> {
|
||||
match self {
|
||||
Self::Real(s) => s.update_tasks(tasks).await,
|
||||
Self::Mock(m) => unsafe {
|
||||
m.get::<_, Result<Vec<Task>>>("update_tasks").call(tasks)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_task(&self, id: TaskId, filter: Option<TaskFilter>) -> Result<Task> {
|
||||
match self {
|
||||
Self::Real(s) => s.get_task(id, filter).await,
|
||||
Self::Mock(m) => unsafe { m.get::<_, Result<Task>>("get_task").call((id, filter)) },
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_processing_tasks(
|
||||
&self,
|
||||
tasks: Processing,
|
||||
) -> Result<(Processing, BatchContent)> {
|
||||
match self {
|
||||
Self::Real(s) => s.get_processing_tasks(tasks).await,
|
||||
Self::Mock(m) => unsafe { m.get("get_pending_task").call(tasks) },
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn fetch_unfinished_tasks(&self, from: Option<TaskId>) -> Result<Vec<Task>> {
|
||||
match self {
|
||||
Self::Real(s) => s.fetch_unfinished_tasks(from).await,
|
||||
Self::Mock(m) => unsafe { m.get("fetch_unfinished_tasks").call(from) },
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_tasks(
|
||||
&self,
|
||||
from: Option<TaskId>,
|
||||
filter: Option<TaskFilter>,
|
||||
limit: Option<usize>,
|
||||
) -> Result<Vec<Task>> {
|
||||
match self {
|
||||
Self::Real(s) => s.list_tasks(from, filter, limit).await,
|
||||
Self::Mock(m) => unsafe { m.get("list_tasks").call((from, filter, limit)) },
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn register(&self, content: TaskContent) -> Result<Task> {
|
||||
match self {
|
||||
Self::Real(s) => s.register(content).await,
|
||||
Self::Mock(_m) => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn register_raw_update(&self, wtxn: &mut RwTxn, task: &Task) -> Result<()> {
|
||||
match self {
|
||||
Self::Real(s) => s.register_raw_update(wtxn, task),
|
||||
Self::Mock(_m) => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load_dump(path: impl AsRef<Path>, env: Arc<Env>) -> anyhow::Result<()> {
|
||||
TaskStore::load_dump(path, env)
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_increment_task_id() {
|
||||
let tmp = tmp_env();
|
||||
let store = Store::new(tmp.env()).unwrap();
|
||||
|
||||
let mut txn = store.wtxn().unwrap();
|
||||
assert_eq!(store.next_task_id(&mut txn).unwrap(), 0);
|
||||
txn.abort().unwrap();
|
||||
|
||||
let gen_task = |id: TaskId| Task {
|
||||
id,
|
||||
content: TaskContent::IndexCreation {
|
||||
primary_key: None,
|
||||
index_uid: IndexUid::new_unchecked("test"),
|
||||
},
|
||||
events: Vec::new(),
|
||||
};
|
||||
|
||||
let mut runner = TestRunner::new(Config::default());
|
||||
runner
|
||||
.run(&(0..100u32).prop_map(gen_task), |task| {
|
||||
let mut txn = store.wtxn().unwrap();
|
||||
let previous_id = store.next_task_id(&mut txn).unwrap();
|
||||
|
||||
store.put(&mut txn, &task).unwrap();
|
||||
|
||||
let next_id = store.next_task_id(&mut txn).unwrap();
|
||||
|
||||
// if we put a task whose task_id is less than the next_id, then the next_id remains
|
||||
// unchanged, otherwise it becomes task.id + 1
|
||||
if task.id < previous_id {
|
||||
assert_eq!(next_id, previous_id)
|
||||
} else {
|
||||
assert_eq!(next_id, task.id + 1);
|
||||
}
|
||||
|
||||
txn.commit().unwrap();
|
||||
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
@ -1,377 +0,0 @@
|
||||
#[allow(clippy::upper_case_acronyms)]
|
||||
|
||||
type BEU32 = milli::heed::zerocopy::U32<milli::heed::byteorder::BE>;
|
||||
|
||||
const INDEX_UIDS_TASK_IDS: &str = "index-uids-task-ids";
|
||||
const TASKS: &str = "tasks";
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::ops::Bound::{Excluded, Unbounded};
|
||||
use std::result::Result as StdResult;
|
||||
use std::sync::Arc;
|
||||
|
||||
use milli::heed::types::{OwnedType, SerdeJson, Str};
|
||||
use milli::heed::{Database, Env, RoTxn, RwTxn};
|
||||
use milli::heed_codec::RoaringBitmapCodec;
|
||||
use roaring::RoaringBitmap;
|
||||
|
||||
use crate::tasks::task::{Task, TaskId};
|
||||
|
||||
use super::super::Result;
|
||||
use super::TaskFilter;
|
||||
|
||||
pub struct Store {
|
||||
env: Arc<Env>,
|
||||
/// Maps an index uid to the set of tasks ids associated to it.
|
||||
index_uid_task_ids: Database<Str, RoaringBitmapCodec>,
|
||||
tasks: Database<OwnedType<BEU32>, SerdeJson<Task>>,
|
||||
}
|
||||
|
||||
impl Drop for Store {
|
||||
fn drop(&mut self) {
|
||||
if Arc::strong_count(&self.env) == 1 {
|
||||
self.env.as_ref().clone().prepare_for_closing();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Store {
|
||||
/// Create a new store from the specified `Path`.
|
||||
/// Be really cautious when calling this function, the returned `Store` may
|
||||
/// be in an invalid state, with dangling processing tasks.
|
||||
/// You want to patch all un-finished tasks and put them in your pending
|
||||
/// queue with the `reset_and_return_unfinished_update` method.
|
||||
pub fn new(env: Arc<milli::heed::Env>) -> Result<Self> {
|
||||
let index_uid_task_ids = env.create_database(Some(INDEX_UIDS_TASK_IDS))?;
|
||||
let tasks = env.create_database(Some(TASKS))?;
|
||||
|
||||
Ok(Self {
|
||||
env,
|
||||
index_uid_task_ids,
|
||||
tasks,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn wtxn(&self) -> Result<RwTxn> {
|
||||
Ok(self.env.write_txn()?)
|
||||
}
|
||||
|
||||
pub fn rtxn(&self) -> Result<RoTxn> {
|
||||
Ok(self.env.read_txn()?)
|
||||
}
|
||||
|
||||
/// Returns the id for the next task.
|
||||
///
|
||||
/// The required `mut txn` acts as a reservation system. It guarantees that as long as you commit
|
||||
/// the task to the store in the same transaction, no one else will have this task id.
|
||||
pub fn next_task_id(&self, txn: &mut RwTxn) -> Result<TaskId> {
|
||||
let id = self
|
||||
.tasks
|
||||
.lazily_decode_data()
|
||||
.last(txn)?
|
||||
.map(|(id, _)| id.get() + 1)
|
||||
.unwrap_or(0);
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
pub fn put(&self, txn: &mut RwTxn, task: &Task) -> Result<()> {
|
||||
self.tasks.put(txn, &BEU32::new(task.id), task)?;
|
||||
// only add the task to the indexes index if it has an index_uid
|
||||
if let Some(index_uid) = task.index_uid() {
|
||||
let mut tasks_set = self
|
||||
.index_uid_task_ids
|
||||
.get(txn, index_uid)?
|
||||
.unwrap_or_default();
|
||||
|
||||
tasks_set.insert(task.id);
|
||||
|
||||
self.index_uid_task_ids.put(txn, index_uid, &tasks_set)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn get(&self, txn: &RoTxn, id: TaskId) -> Result<Option<Task>> {
|
||||
let task = self.tasks.get(txn, &BEU32::new(id))?;
|
||||
Ok(task)
|
||||
}
|
||||
|
||||
/// Returns the unfinished tasks starting from the given taskId in ascending order.
|
||||
pub fn fetch_unfinished_tasks(&self, txn: &RoTxn, from: Option<TaskId>) -> Result<Vec<Task>> {
|
||||
// We must NEVER re-enqueue an already processed task! It's content uuid would point to an unexisting file.
|
||||
//
|
||||
// TODO(marin): This may create some latency when the first batch lazy loads the pending updates.
|
||||
let from = from.unwrap_or_default();
|
||||
|
||||
let result: StdResult<Vec<_>, milli::heed::Error> = self
|
||||
.tasks
|
||||
.range(txn, &(BEU32::new(from)..))?
|
||||
.map(|r| r.map(|(_, t)| t))
|
||||
.filter(|result| result.as_ref().map_or(true, |t| !t.is_finished()))
|
||||
.collect();
|
||||
|
||||
result.map_err(Into::into)
|
||||
}
|
||||
|
||||
/// Returns all the tasks starting from the given taskId and going in descending order.
|
||||
pub fn list_tasks(
|
||||
&self,
|
||||
txn: &RoTxn,
|
||||
from: Option<TaskId>,
|
||||
filter: Option<TaskFilter>,
|
||||
limit: Option<usize>,
|
||||
) -> Result<Vec<Task>> {
|
||||
let from = match from {
|
||||
Some(from) => from,
|
||||
None => self.tasks.last(txn)?.map_or(0, |(id, _)| id.get()),
|
||||
};
|
||||
|
||||
let filter_fn = |task: &Task| {
|
||||
filter
|
||||
.as_ref()
|
||||
.and_then(|f| f.filter_fn.as_ref())
|
||||
.map_or(true, |f| f(task))
|
||||
};
|
||||
|
||||
let result: Result<Vec<_>> = match filter.as_ref().and_then(|f| f.filtered_indexes()) {
|
||||
Some(indexes) => self
|
||||
.compute_candidates(txn, indexes, from)?
|
||||
.filter(|result| result.as_ref().map_or(true, filter_fn))
|
||||
.take(limit.unwrap_or(usize::MAX))
|
||||
.collect(),
|
||||
None => self
|
||||
.tasks
|
||||
.rev_range(txn, &(..=BEU32::new(from)))?
|
||||
.map(|r| r.map(|(_, t)| t).map_err(Into::into))
|
||||
.filter(|result| result.as_ref().map_or(true, filter_fn))
|
||||
.take(limit.unwrap_or(usize::MAX))
|
||||
.collect(),
|
||||
};
|
||||
|
||||
result.map_err(Into::into)
|
||||
}
|
||||
|
||||
fn compute_candidates<'a>(
|
||||
&'a self,
|
||||
txn: &'a RoTxn,
|
||||
indexes: &HashSet<String>,
|
||||
from: TaskId,
|
||||
) -> Result<impl Iterator<Item = Result<Task>> + 'a> {
|
||||
let mut candidates = RoaringBitmap::new();
|
||||
|
||||
for index_uid in indexes {
|
||||
if let Some(tasks_set) = self.index_uid_task_ids.get(txn, index_uid)? {
|
||||
candidates |= tasks_set;
|
||||
}
|
||||
}
|
||||
|
||||
candidates.remove_range((Excluded(from), Unbounded));
|
||||
|
||||
let iter = candidates
|
||||
.into_iter()
|
||||
.rev()
|
||||
.filter_map(|id| self.get(txn, id).transpose());
|
||||
|
||||
Ok(iter)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub mod test {
|
||||
use itertools::Itertools;
|
||||
use meilisearch_types::index_uid::IndexUid;
|
||||
use milli::heed::EnvOpenOptions;
|
||||
use nelson::Mocker;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use crate::tasks::task::TaskContent;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// TODO: use this mock to test the task store properly.
|
||||
#[allow(dead_code)]
|
||||
pub enum MockStore {
|
||||
Real(Store),
|
||||
Fake(Mocker),
|
||||
}
|
||||
|
||||
pub struct TmpEnv(TempDir, Arc<milli::heed::Env>);
|
||||
|
||||
impl TmpEnv {
|
||||
pub fn env(&self) -> Arc<milli::heed::Env> {
|
||||
self.1.clone()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tmp_env() -> TmpEnv {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
|
||||
let mut options = EnvOpenOptions::new();
|
||||
options.map_size(4096 * 100000);
|
||||
options.max_dbs(1000);
|
||||
let env = Arc::new(options.open(tmp.path()).unwrap());
|
||||
|
||||
TmpEnv(tmp, env)
|
||||
}
|
||||
|
||||
impl MockStore {
|
||||
pub fn new(env: Arc<milli::heed::Env>) -> Result<Self> {
|
||||
Ok(Self::Real(Store::new(env)?))
|
||||
}
|
||||
|
||||
pub fn wtxn(&self) -> Result<RwTxn> {
|
||||
match self {
|
||||
MockStore::Real(index) => index.wtxn(),
|
||||
MockStore::Fake(_) => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn rtxn(&self) -> Result<RoTxn> {
|
||||
match self {
|
||||
MockStore::Real(index) => index.rtxn(),
|
||||
MockStore::Fake(_) => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn next_task_id(&self, txn: &mut RwTxn) -> Result<TaskId> {
|
||||
match self {
|
||||
MockStore::Real(index) => index.next_task_id(txn),
|
||||
MockStore::Fake(_) => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn put(&self, txn: &mut RwTxn, task: &Task) -> Result<()> {
|
||||
match self {
|
||||
MockStore::Real(index) => index.put(txn, task),
|
||||
MockStore::Fake(_) => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get(&self, txn: &RoTxn, id: TaskId) -> Result<Option<Task>> {
|
||||
match self {
|
||||
MockStore::Real(index) => index.get(txn, id),
|
||||
MockStore::Fake(_) => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn fetch_unfinished_tasks(
|
||||
&self,
|
||||
txn: &RoTxn,
|
||||
from: Option<TaskId>,
|
||||
) -> Result<Vec<Task>> {
|
||||
match self {
|
||||
MockStore::Real(index) => index.fetch_unfinished_tasks(txn, from),
|
||||
MockStore::Fake(_) => todo!(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn list_tasks(
|
||||
&self,
|
||||
txn: &RoTxn,
|
||||
from: Option<TaskId>,
|
||||
filter: Option<TaskFilter>,
|
||||
limit: Option<usize>,
|
||||
) -> Result<Vec<Task>> {
|
||||
match self {
|
||||
MockStore::Real(index) => index.list_tasks(txn, from, filter, limit),
|
||||
MockStore::Fake(_) => todo!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ordered_filtered_updates() {
|
||||
let tmp = tmp_env();
|
||||
let store = Store::new(tmp.env()).unwrap();
|
||||
|
||||
let tasks = (0..100)
|
||||
.map(|_| Task {
|
||||
id: rand::random(),
|
||||
content: TaskContent::IndexDeletion {
|
||||
index_uid: IndexUid::new_unchecked("test"),
|
||||
},
|
||||
events: vec![],
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let mut txn = store.env.write_txn().unwrap();
|
||||
tasks
|
||||
.iter()
|
||||
.try_for_each(|t| store.put(&mut txn, t))
|
||||
.unwrap();
|
||||
|
||||
let mut filter = TaskFilter::default();
|
||||
filter.filter_index("test".into());
|
||||
|
||||
let tasks = store.list_tasks(&txn, None, Some(filter), None).unwrap();
|
||||
|
||||
assert!(tasks
|
||||
.iter()
|
||||
.map(|t| t.id)
|
||||
.tuple_windows()
|
||||
.all(|(a, b)| a > b));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_same_index_prefix() {
|
||||
let tmp = tmp_env();
|
||||
let store = Store::new(tmp.env()).unwrap();
|
||||
|
||||
let task_1 = Task {
|
||||
id: 1,
|
||||
content: TaskContent::IndexDeletion {
|
||||
index_uid: IndexUid::new_unchecked("test"),
|
||||
},
|
||||
events: vec![],
|
||||
};
|
||||
|
||||
let task_2 = Task {
|
||||
id: 0,
|
||||
content: TaskContent::IndexDeletion {
|
||||
index_uid: IndexUid::new_unchecked("test1"),
|
||||
},
|
||||
events: vec![],
|
||||
};
|
||||
|
||||
let mut txn = store.wtxn().unwrap();
|
||||
store.put(&mut txn, &task_1).unwrap();
|
||||
store.put(&mut txn, &task_2).unwrap();
|
||||
|
||||
let mut filter = TaskFilter::default();
|
||||
filter.filter_index("test".into());
|
||||
|
||||
let tasks = store.list_tasks(&txn, None, Some(filter), None).unwrap();
|
||||
|
||||
txn.abort().unwrap();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks.first().as_ref().unwrap().index_uid().unwrap(), "test");
|
||||
|
||||
// same thing but invert the ids
|
||||
let task_1 = Task {
|
||||
id: 0,
|
||||
content: TaskContent::IndexDeletion {
|
||||
index_uid: IndexUid::new_unchecked("test"),
|
||||
},
|
||||
events: vec![],
|
||||
};
|
||||
let task_2 = Task {
|
||||
id: 1,
|
||||
content: TaskContent::IndexDeletion {
|
||||
index_uid: IndexUid::new_unchecked("test1"),
|
||||
},
|
||||
events: vec![],
|
||||
};
|
||||
|
||||
let mut txn = store.wtxn().unwrap();
|
||||
store.put(&mut txn, &task_1).unwrap();
|
||||
store.put(&mut txn, &task_2).unwrap();
|
||||
|
||||
let mut filter = TaskFilter::default();
|
||||
filter.filter_index("test".into());
|
||||
|
||||
let tasks = store.list_tasks(&txn, None, Some(filter), None).unwrap();
|
||||
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks.first().as_ref().unwrap().index_uid().unwrap(), "test");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user