AppFlowy/frontend/rust-lib/flowy-task/src/queue.rs

use crate::{PendingTask, Task};
use atomic_refcell::AtomicRefCell;
use std::cmp::Ordering;
use std::collections::hash_map::Entry;
use std::collections::{BinaryHeap, HashMap};
use std::ops::{Deref, DerefMut};
use std::sync::Arc;

#[derive(Default)]
pub(crate) struct TaskQueue {
  // index_tasks for quick access
  index_tasks: HashMap<TaskHandlerId, Arc<AtomicRefCell<TaskList>>>,
  queue: BinaryHeap<Arc<AtomicRefCell<TaskList>>>,
}

impl TaskQueue {
  pub(crate) fn new() -> Self {
    Self::default()
  }

  pub(crate) fn push(&mut self, task: &Task) {
    if task.content.is_none() {
      tracing::warn!(
        "The task:{} with empty content will be not executed",
        task.id
      );
      return;
    }

    let pending_task = PendingTask {
      qos: task.qos,
      id: task.id,
    };
    match self.index_tasks.entry(task.handler_id.clone()) {
      Entry::Occupied(entry) => {
        let mut list = entry.get().borrow_mut();
        assert!(list
          .peek()
          .map(|old_id| pending_task.id >= old_id.id)
          .unwrap_or(true));
        list.push(pending_task);
      },
      Entry::Vacant(entry) => {
        let mut task_list = TaskList::new(entry.key());
        task_list.push(pending_task);
        let task_list = Arc::new(AtomicRefCell::new(task_list));
        entry.insert(task_list.clone());
        self.queue.push(task_list);
      },
    }
  }

  #[allow(dead_code)]
  pub(crate) fn clear(&mut self) {
    self.queue.clear();
  }

  pub(crate) fn mut_head<T, F>(&mut self, mut f: F) -> Option<T>
  where
    F: FnMut(&mut TaskList) -> Option<T>,
  {
    let head = self.queue.pop()?;
    let result = {
      let mut ref_head = head.borrow_mut();
      f(&mut ref_head)
    };
    if !head.borrow().tasks.is_empty() {
      self.queue.push(head);
    } else {
      self.index_tasks.remove(&head.borrow().id);
    }
    result
  }
}

pub type TaskHandlerId = String;

#[derive(Debug)]
pub(crate) struct TaskList {
  pub(crate) id: TaskHandlerId,
  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(id: &str) -> Self {
    Self {
      id: id.to_owned(),
      tasks: BinaryHeap::new(),
    }
  }
}

impl PartialEq for TaskList {
  fn eq(&self, other: &Self) -> bool {
    self.id == other.id
  }
}

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))
  }
}