feat(lib): auto-batching

meilisearch-lib/src/tasks/scheduler.rs

@@ -1,253 +1,526 @@
+use std::cmp::Ordering;
+use std::collections::{hash_map::Entry, BinaryHeap, HashMap, VecDeque};
+use std::ops::{Deref, DerefMut};
+use std::path::Path;
 use std::sync::Arc;
 use std::time::Duration;
 
+use atomic_refcell::AtomicRefCell;
 use chrono::Utc;
-use serde::{Deserialize, Serialize};
+use milli::update::IndexDocumentsMethod;
+use tokio::sync::{watch, RwLock};
+
+use crate::options::SchedulerConfig;
+use crate::update_file_store::UpdateFileStore;
 
 use super::batch::Batch;
 use super::error::Result;
-#[cfg(test)]
-use super::task_store::test::MockTaskStore as TaskStore;
-use super::task_store::Pending;
-#[cfg(not(test))]
-use super::task_store::TaskStore;
-use super::TaskPerformer;
-use crate::tasks::task::TaskEvent;
+use super::task::{Job, Task, TaskContent, TaskEvent, TaskId};
+use super::update_loop::UpdateLoop;
+use super::{TaskFilter, TaskPerformer, TaskStore};
 
-/// The scheduler roles is to perform batches of tasks one at a time. It will monitor the TaskStore
-/// for new tasks, put them in a batch, and process the batch as soon as possible.
-///
-/// When a batch is currently processing, the scheduler is just waiting.
-pub struct Scheduler<P: TaskPerformer> {
-    store: TaskStore,
-    performer: Arc<P>,
-
-    /// The interval at which the the `TaskStore` should be checked for new updates
-    task_store_check_interval: Duration,
+#[derive(Eq, Debug, Clone, Copy)]
+enum TaskType {
+    DocumentAddition { number: usize },
+    DocumentUpdate { number: usize },
+    Other,
 }
 
-impl<P> Scheduler<P>
-where
-    P: TaskPerformer + Send + Sync + 'static,
-    P::Error: Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static,
-{
-    pub fn new(store: TaskStore, performer: Arc<P>, task_store_check_interval: Duration) -> Self {
+/// Two tasks are equal if they have the same type.
+impl PartialEq for TaskType {
+    fn eq(&self, other: &Self) -> bool {
+        matches!(
+            (self, other),
+            (Self::DocumentAddition { .. }, Self::DocumentAddition { .. })
+                | (Self::DocumentUpdate { .. }, Self::DocumentUpdate { .. })
+        )
+    }
+}
+
+#[derive(Eq, Debug, Clone, Copy)]
+struct PendingTask {
+    kind: TaskType,
+    id: TaskId,
+}
+
+impl PartialEq for PendingTask {
+    fn eq(&self, other: &Self) -> bool {
+        self.id.eq(&other.id)
+    }
+}
+
+impl PartialOrd for PendingTask {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for PendingTask {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.id.cmp(&other.id).reverse()
+    }
+}
+
+#[derive(Debug)]
+struct TaskList {
+    index: String,
+    tasks: BinaryHeap<PendingTask>,
+}
+
+impl Deref for TaskList {
+    type Target = BinaryHeap<PendingTask>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.tasks
+    }
+}
+
+impl DerefMut for TaskList {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.tasks
+    }
+}
+
+impl TaskList {
+    fn new(index: String) -> Self {
         Self {
-            store,
-            performer,
-            task_store_check_interval,
+            index,
+            tasks: Default::default(),
         }
     }
+}
 
-    pub async fn run(self) {
-        loop {
-            if let Err(e) = self.process_next_batch().await {
-                log::error!("an error occured while processing an update batch: {}", e);
+impl PartialEq for TaskList {
+    fn eq(&self, other: &Self) -> bool {
+        self.index == other.index
+    }
+}
+
+impl Eq for TaskList {}
+
+impl Ord for TaskList {
+    fn cmp(&self, other: &Self) -> Ordering {
+        match (self.peek(), other.peek()) {
+            (None, None) => Ordering::Equal,
+            (None, Some(_)) => Ordering::Less,
+            (Some(_), None) => Ordering::Greater,
+            (Some(lhs), Some(rhs)) => lhs.cmp(rhs),
+        }
+    }
+}
+
+impl PartialOrd for TaskList {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+#[derive(Default)]
+struct TaskQueue {
+    /// Maps index uids to their TaskList, for quick access
+    index_tasks: HashMap<String, Arc<AtomicRefCell<TaskList>>>,
+    /// A queue that orders TaskList by the priority of their fist update
+    queue: BinaryHeap<Arc<AtomicRefCell<TaskList>>>,
+}
+
+impl TaskQueue {
+    fn insert(&mut self, task: Task) {
+        let uid = task.index_uid.into_inner();
+        let id = task.id;
+        let kind = match task.content {
+            TaskContent::DocumentAddition {
+                documents_count,
+                merge_strategy: IndexDocumentsMethod::ReplaceDocuments,
+                ..
+            } => TaskType::DocumentAddition {
+                number: documents_count,
+            },
+            TaskContent::DocumentAddition {
+                documents_count,
+                merge_strategy: IndexDocumentsMethod::UpdateDocuments,
+                ..
+            } => TaskType::DocumentUpdate {
+                number: documents_count,
+            },
+            _ => TaskType::Other,
+        };
+        let task = PendingTask { kind, id };
+
+        match self.index_tasks.entry(uid) {
+            Entry::Occupied(entry) => {
+                // A task list already exists for this index, all we have to to is to push the new
+                // update to the end of the list. This won't change the order since ids are
+                // monotically increasing.
+                let mut list = entry.get().borrow_mut();
+
+                // We only need the first element to be lower than the one we want to
+                // insert to preserve the order in the queue.
+                assert!(list.peek().map(|old_id| id >= old_id.id).unwrap_or(true));
+
+                list.push(task);
+            }
+            Entry::Vacant(entry) => {
+                let mut task_list = TaskList::new(entry.key().to_owned());
+                task_list.push(task);
+                let task_list = Arc::new(AtomicRefCell::new(task_list));
+                entry.insert(task_list.clone());
+                self.queue.push(task_list);
             }
         }
     }
 
-    async fn process_next_batch(&self) -> Result<()> {
-        match self.prepare_batch().await? {
-            Some(mut batch) => {
-                for task in &mut batch.tasks {
-                    match task {
-                        Pending::Task(task) => task.events.push(TaskEvent::Processing(Utc::now())),
-                        Pending::Job(_) => (),
+    /// Passes a context with a view to the task list of the next index to schedule. It is
+    /// guaranteed that the first id from task list will be the lowest pending task id.
+    fn head_mut<R>(&mut self, mut f: impl FnMut(&mut TaskList) -> R) -> Option<R> {
+        let head = self.queue.pop()?;
+        let result = {
+            let mut ref_head = head.borrow_mut();
+            f(&mut *ref_head)
+        };
+        if !head.borrow().tasks.is_empty() {
+            // After being mutated, the head is reinserted to the correct position.
+            self.queue.push(head);
+        } else {
+            self.index_tasks.remove(&head.borrow().index);
+        }
+
+        Some(result)
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.queue.is_empty() && self.index_tasks.is_empty()
+    }
+}
+
+pub struct Scheduler {
+    jobs: VecDeque<Job>,
+    tasks: TaskQueue,
+
+    store: TaskStore,
+    processing: Vec<TaskId>,
+    next_fetched_task_id: TaskId,
+    config: SchedulerConfig,
+    /// Notifies the update loop that a new task was received
+    notifier: watch::Sender<()>,
+}
+
+impl Scheduler {
+    pub fn new<P>(
+        store: TaskStore,
+        performer: Arc<P>,
+        mut config: SchedulerConfig,
+    ) -> Result<Arc<RwLock<Self>>>
+    where
+        P: TaskPerformer,
+    {
+        let (notifier, rcv) = watch::channel(());
+
+        let debounce_time = config.debounce_duration_sec;
+
+        // Disable autobatching
+        if !config.enable_autobatching {
+            config.max_batch_size = Some(1);
+        }
+
+        let this = Self {
+            jobs: VecDeque::new(),
+            tasks: TaskQueue::default(),
+
+            store,
+            processing: Vec::new(),
+            next_fetched_task_id: 0,
+            config,
+            notifier,
+        };
+
+        // Notify update loop to start processing pending updates immediately after startup.
+        this.notify();
+
+        let this = Arc::new(RwLock::new(this));
+
+        let update_loop = UpdateLoop::new(
+            this.clone(),
+            performer,
+            debounce_time.filter(|&v| v > 0).map(Duration::from_secs),
+            rcv,
+        );
+
+        tokio::task::spawn_local(update_loop.run());
+
+        Ok(this)
+    }
+
+    pub async fn dump(&self, path: &Path, file_store: UpdateFileStore) -> Result<()> {
+        self.store.dump(path, file_store).await
+    }
+
+    fn register_task(&mut self, task: Task) {
+        assert!(!task.is_finished());
+        self.tasks.insert(task);
+    }
+
+    /// Clears the processing list, this method should be called when the processing of a batch is finished.
+    pub fn finish(&mut self) {
+        self.processing.clear();
+    }
+
+    pub fn notify(&self) {
+        let _ = self.notifier.send(());
+    }
+
+    fn notify_if_not_empty(&self) {
+        if !self.jobs.is_empty() || !self.tasks.is_empty() {
+            self.notify();
+        }
+    }
+
+    pub async fn update_tasks(&self, tasks: Vec<Task>) -> Result<Vec<Task>> {
+        self.store.update_tasks(tasks).await
+    }
+
+    pub async fn get_task(&self, id: TaskId, filter: Option<TaskFilter>) -> Result<Task> {
+        self.store.get_task(id, filter).await
+    }
+
+    pub async fn list_tasks(
+        &self,
+        offset: Option<TaskId>,
+        filter: Option<TaskFilter>,
+        limit: Option<usize>,
+    ) -> Result<Vec<Task>> {
+        self.store.list_tasks(offset, filter, limit).await
+    }
+
+    pub async fn get_processing_tasks(&self) -> Result<Vec<Task>> {
+        let mut tasks = Vec::new();
+
+        for id in self.processing.iter() {
+            let task = self.store.get_task(*id, None).await?;
+            tasks.push(task);
+        }
+
+        Ok(tasks)
+    }
+
+    pub async fn schedule_job(&mut self, job: Job) {
+        self.jobs.push_back(job);
+        self.notify();
+    }
+
+    async fn fetch_pending_tasks(&mut self) -> Result<()> {
+        // 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 mut filter = TaskFilter::default();
+        filter.filter_fn(|task| !task.is_finished());
+
+        self.store
+            .list_tasks(Some(self.next_fetched_task_id), Some(filter), None)
+            .await?
+            .into_iter()
+            // The tasks arrive in reverse order, and we need to insert them in order.
+            .rev()
+            .for_each(|t| {
+                self.next_fetched_task_id = t.id + 1;
+                self.register_task(t);
+            });
+
+        Ok(())
+    }
+
+    /// Prepare the next batch, and set `processing` to the ids in that batch.
+    pub async fn prepare(&mut self) -> Result<Pending> {
+        // If there is a job to process, do it first.
+        if let Some(job) = self.jobs.pop_front() {
+            // There is more work to do, notify the update loop
+            self.notify_if_not_empty();
+            return Ok(Pending::Job(job));
+        }
+        // Try to fill the queue with pending tasks.
+        self.fetch_pending_tasks().await?;
+
+        make_batch(&mut self.tasks, &mut self.processing, &self.config);
+
+        log::debug!("prepared batch with {} tasks", self.processing.len());
+
+        if !self.processing.is_empty() {
+            let ids = std::mem::take(&mut self.processing);
+
+            let (ids, mut tasks) = self.store.get_pending_tasks(ids).await?;
+
+            // The batch id is the id of the first update it contains
+            let id = match tasks.first() {
+                Some(Task { id, .. }) => *id,
+                _ => panic!("invalid batch"),
+            };
+
+            tasks.iter_mut().for_each(|t| {
+                t.events.push(TaskEvent::Batched {
+                    batch_id: id,
+                    timestamp: Utc::now(),
+                })
+            });
+
+            self.processing = ids;
+
+            let batch = Batch {
+                id,
+                created_at: Utc::now(),
+                tasks,
+            };
+
+            // There is more work to do, notify the update loop
+            self.notify_if_not_empty();
+
+            Ok(Pending::Batch(batch))
+        } else {
+            Ok(Pending::Nothing)
+        }
+    }
+}
+
+#[derive(Debug)]
+pub enum Pending {
+    Batch(Batch),
+    Job(Job),
+    Nothing,
+}
+
+fn make_batch(tasks: &mut TaskQueue, processing: &mut Vec<TaskId>, config: &SchedulerConfig) {
+    processing.clear();
+
+    let mut doc_count = 0;
+    tasks.head_mut(|list| match list.peek().copied() {
+        Some(PendingTask {
+            kind: TaskType::Other,
+            id,
+        }) => {
+            processing.push(id);
+            list.pop();
+        }
+        Some(PendingTask { kind, .. }) => loop {
+            match list.peek() {
+                Some(pending) if pending.kind == kind => {
+                    // We always need to process at least one task for the scheduler to make progress.
+                    if processing.len() >= config.max_batch_size.unwrap_or(usize::MAX).max(1) {
+                        break;
+                    }
+                    let pending = list.pop().unwrap();
+                    processing.push(pending.id);
+
+                    // We add the number of documents to the count if we are scheduling document additions and
+                    // stop adding if we already have enough.
+                    //
+                    // We check that bound only after adding the current task to the batch, so that a batch contains at least one task.
+                    match pending.kind {
+                        TaskType::DocumentUpdate { number }
+                        | TaskType::DocumentAddition { number } => {
+                            doc_count += number;
+
+                            if doc_count >= config.max_documents_per_batch.unwrap_or(usize::MAX) {
+                                break;
+                            }
+                        }
+                        _ => (),
                     }
                 }
-
-                // the jobs are ignored
-                batch.tasks = self.store.update_tasks(batch.tasks).await?;
-
-                let performer = self.performer.clone();
-                let batch_result = performer.process(batch).await;
-                self.handle_batch_result(batch_result).await?;
+                _ => break,
             }
-            None => {
-                // No update found to create a batch we wait a bit before we retry.
-                tokio::time::sleep(self.task_store_check_interval).await;
-            }
-        }
-
-        Ok(())
-    }
-
-    /// Checks for pending tasks and groups them in a batch. If there are no pending update,
-    /// return Ok(None)
-    ///
-    /// Until batching is properly implemented, the batches contain only one task.
-    async fn prepare_batch(&self) -> Result<Option<Batch>> {
-        match self.store.peek_pending_task().await {
-            Some(Pending::Task(next_task_id)) => {
-                let mut task = self.store.get_task(next_task_id, None).await?;
-
-                task.events.push(TaskEvent::Batched {
-                    timestamp: Utc::now(),
-                    batch_id: 0,
-                });
-
-                let batch = Batch {
-                    id: 0,
-                    // index_uid: task.index_uid.clone(),
-                    created_at: Utc::now(),
-                    tasks: vec![Pending::Task(task)],
-                };
-                Ok(Some(batch))
-            }
-            Some(Pending::Job(job)) => Ok(Some(Batch {
-                id: 0,
-                created_at: Utc::now(),
-                tasks: vec![Pending::Job(job)],
-            })),
-            None => Ok(None),
-        }
-    }
-
-    /// Handles the result from a batch processing.
-    ///
-    /// When a task is processed, the result of the processing is pushed to its event list. The
-    /// handle batch result make sure that the new state is save into its store.
-    /// The tasks are then removed from the processing queue.
-    async fn handle_batch_result(&self, mut batch: Batch) -> Result<()> {
-        let tasks = self.store.update_tasks(batch.tasks).await?;
-        batch.tasks = tasks;
-        self.store.delete_pending(&batch.tasks[0]).await;
-        self.performer.finish(&batch).await;
-        Ok(())
-    }
+        },
+        None => (),
+    });
 }
 
 #[cfg(test)]
 mod test {
-    use nelson::Mocker;
+    use milli::update::IndexDocumentsMethod;
+    use uuid::Uuid;
 
-    use crate::index_resolver::IndexUid;
-    use crate::tasks::task::Task;
-    use crate::tasks::task_store::TaskFilter;
+    use crate::{index_resolver::IndexUid, tasks::task::TaskContent};
 
-    use super::super::task::{TaskContent, TaskEvent, TaskId, TaskResult};
-    use super::super::MockTaskPerformer;
     use super::*;
 
-    #[tokio::test]
-    async fn test_prepare_batch_full() {
-        let mocker = Mocker::default();
-
-        mocker
-            .when::<(TaskId, Option<TaskFilter>), Result<Option<Task>>>("get_task")
-            .once()
-            .then(|(id, _filter)| {
-                let task = Task {
-                    id,
-                    index_uid: IndexUid::new("Test".to_string()).unwrap(),
-                    content: TaskContent::IndexDeletion,
-                    events: vec![TaskEvent::Created(Utc::now())],
-                };
-                Ok(Some(task))
-            });
-
-        mocker
-            .when::<(), Option<Pending<TaskId>>>("peek_pending_task")
-            .then(|()| Some(Pending::Task(1)));
-
-        let store = TaskStore::mock(mocker);
-        let performer = Arc::new(MockTaskPerformer::new());
-
-        let scheduler = Scheduler {
-            store,
-            performer,
-            task_store_check_interval: Duration::from_millis(1),
-        };
-
-        let batch = scheduler.prepare_batch().await.unwrap().unwrap();
-
-        assert_eq!(batch.tasks.len(), 1);
-        assert!(
-            matches!(batch.tasks[0], Pending::Task(Task { id: 1, .. })),
-            "{:?}",
-            batch.tasks[0]
-        );
-    }
-
-    #[tokio::test]
-    async fn test_prepare_batch_empty() {
-        let mocker = Mocker::default();
-        mocker
-            .when::<(), Option<Pending<TaskId>>>("peek_pending_task")
-            .then(|()| None);
-
-        let store = TaskStore::mock(mocker);
-        let performer = Arc::new(MockTaskPerformer::new());
-
-        let scheduler = Scheduler {
-            store,
-            performer,
-            task_store_check_interval: Duration::from_millis(1),
-        };
-
-        assert!(scheduler.prepare_batch().await.unwrap().is_none());
-    }
-
-    #[tokio::test]
-    async fn test_loop_run_normal() {
-        let mocker = Mocker::default();
-        let mut id = Some(1);
-        mocker
-            .when::<(), Option<Pending<TaskId>>>("peek_pending_task")
-            .then(move |()| id.take().map(Pending::Task));
-        mocker
-            .when::<(TaskId, Option<TaskFilter>), Result<Task>>("get_task")
-            .once()
-            .then(|(id, _)| {
-                let task = Task {
-                    id,
-                    index_uid: IndexUid::new("Test".to_string()).unwrap(),
-                    content: TaskContent::IndexDeletion,
-                    events: vec![TaskEvent::Created(Utc::now())],
-                };
-                Ok(task)
-            });
-
-        mocker
-            .when::<Vec<Pending<Task>>, Result<Vec<Pending<Task>>>>("update_tasks")
-            .times(2)
-            .then(|tasks| {
-                assert_eq!(tasks.len(), 1);
-                Ok(tasks)
-            });
-
-        mocker.when::<(), ()>("delete_pending").once().then(|_| ());
-
-        let store = TaskStore::mock(mocker);
-
-        let mut performer = MockTaskPerformer::new();
-        performer.expect_process().once().returning(|mut batch| {
-            batch.tasks.iter_mut().for_each(|t| match t {
-                Pending::Task(Task { ref mut events, .. }) => events.push(TaskEvent::Succeded {
-                    result: TaskResult::Other,
-                    timestamp: Utc::now(),
-                }),
-                _ => panic!("expected a task, found a job"),
-            });
-
-            batch
-        });
-
-        performer.expect_finish().once().returning(|_| ());
-
-        let performer = Arc::new(performer);
-
-        let scheduler = Scheduler {
-            store,
-            performer,
-            task_store_check_interval: Duration::from_millis(1),
-        };
-
-        let handle = tokio::spawn(scheduler.run());
-
-        if let Ok(r) = tokio::time::timeout(Duration::from_millis(100), handle).await {
-            r.unwrap();
+    fn gen_task(id: TaskId, index_uid: &str, content: TaskContent) -> Task {
+        Task {
+            id,
+            index_uid: IndexUid::new_unchecked(index_uid.to_owned()),
+            content,
+            events: vec![],
         }
     }
+
+    #[test]
+    fn register_updates_multiples_indexes() {
+        let mut queue = TaskQueue::default();
+        queue.insert(gen_task(0, "test1", TaskContent::IndexDeletion));
+        queue.insert(gen_task(1, "test2", TaskContent::IndexDeletion));
+        queue.insert(gen_task(2, "test2", TaskContent::IndexDeletion));
+        queue.insert(gen_task(3, "test2", TaskContent::IndexDeletion));
+        queue.insert(gen_task(4, "test1", TaskContent::IndexDeletion));
+        queue.insert(gen_task(5, "test1", TaskContent::IndexDeletion));
+        queue.insert(gen_task(6, "test2", TaskContent::IndexDeletion));
+
+        let test1_tasks = queue
+            .head_mut(|tasks| tasks.drain().map(|t| t.id).collect::<Vec<_>>())
+            .unwrap();
+
+        assert_eq!(test1_tasks, &[0, 4, 5]);
+
+        let test2_tasks = queue
+            .head_mut(|tasks| tasks.drain().map(|t| t.id).collect::<Vec<_>>())
+            .unwrap();
+
+        assert_eq!(test2_tasks, &[1, 2, 3, 6]);
+
+        assert!(queue.index_tasks.is_empty());
+        assert!(queue.queue.is_empty());
+    }
+
+    #[test]
+    fn test_make_batch() {
+        let mut queue = TaskQueue::default();
+        let content = TaskContent::DocumentAddition {
+            content_uuid: Uuid::new_v4(),
+            merge_strategy: IndexDocumentsMethod::ReplaceDocuments,
+            primary_key: Some("test".to_string()),
+            documents_count: 0,
+            allow_index_creation: true,
+        };
+        queue.insert(gen_task(0, "test1", content.clone()));
+        queue.insert(gen_task(1, "test2", content.clone()));
+        queue.insert(gen_task(2, "test2", TaskContent::IndexDeletion));
+        queue.insert(gen_task(3, "test2", content.clone()));
+        queue.insert(gen_task(4, "test1", content.clone()));
+        queue.insert(gen_task(5, "test1", TaskContent::IndexDeletion));
+        queue.insert(gen_task(6, "test2", content.clone()));
+        queue.insert(gen_task(7, "test1", content));
+
+        let mut batch = Vec::new();
+
+        let config = SchedulerConfig::default();
+        make_batch(&mut queue, &mut batch, &config);
+        assert_eq!(batch, &[0, 4]);
+
+        batch.clear();
+        make_batch(&mut queue, &mut batch, &config);
+        assert_eq!(batch, &[1]);
+
+        batch.clear();
+        make_batch(&mut queue, &mut batch, &config);
+        assert_eq!(batch, &[2]);
+
+        batch.clear();
+        make_batch(&mut queue, &mut batch, &config);
+        assert_eq!(batch, &[3, 6]);
+
+        batch.clear();
+        make_batch(&mut queue, &mut batch, &config);
+        assert_eq!(batch, &[5]);
+
+        batch.clear();
+        make_batch(&mut queue, &mut batch, &config);
+        assert_eq!(batch, &[7]);
+
+        assert!(queue.is_empty());
+    }
 }