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

use crate::{Task, TaskId, TaskState};
use std::collections::HashMap;
use std::mem;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering::SeqCst;

pub(crate) struct TaskStore {
    tasks: HashMap<TaskId, Task>,
    task_id_counter: AtomicU32,
}

impl TaskStore {
    pub fn new() -> Self {
        Self {
            tasks: HashMap::new(),
            task_id_counter: AtomicU32::new(0),
        }
    }

    pub(crate) fn insert_task(&mut self, task: Task) {
        self.tasks.insert(task.id, task);
    }

    pub(crate) fn remove_task(&mut self, task_id: &TaskId) -> Option<Task> {
        self.tasks.remove(task_id)
    }

    pub(crate) fn mut_task(&mut self, task_id: &TaskId) -> Option<&mut Task> {
        self.tasks.get_mut(task_id)
    }

    pub(crate) fn read_task(&self, task_id: &TaskId) -> Option<&Task> {
        self.tasks.get(task_id)
    }

    pub(crate) fn clear(&mut self) {
        let tasks = mem::take(&mut self.tasks);
        tasks.into_values().for_each(|mut task| {
            if task.ret.is_some() {
                let ret = task.ret.take().unwrap();
                task.set_state(TaskState::Cancel);
                let _ = ret.send(task.into());
            }
        });
    }

    pub(crate) fn next_task_id(&self) -> TaskId {
        let _ = self.task_id_counter.fetch_add(1, SeqCst);
        self.task_id_counter.load(SeqCst)
    }
}