store.rs 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. use crate::services::tasks::task::Task;
  2. use crate::services::tasks::{TaskId, TaskStatus};
  3. use std::collections::HashMap;
  4. use std::mem;
  5. use std::sync::atomic::AtomicU32;
  6. use std::sync::atomic::Ordering::SeqCst;
  7. pub(crate) struct GridTaskStore {
  8. tasks: HashMap<TaskId, Task>,
  9. task_id_counter: AtomicU32,
  10. }
  11. impl GridTaskStore {
  12. pub fn new() -> Self {
  13. Self {
  14. tasks: HashMap::new(),
  15. task_id_counter: AtomicU32::new(0),
  16. }
  17. }
  18. pub(crate) fn insert_task(&mut self, task: Task) {
  19. self.tasks.insert(task.id, task);
  20. }
  21. pub(crate) fn remove_task(&mut self, task_id: &TaskId) -> Option<Task> {
  22. self.tasks.remove(task_id)
  23. }
  24. #[allow(dead_code)]
  25. pub(crate) fn clear(&mut self) {
  26. let tasks = mem::take(&mut self.tasks);
  27. tasks.into_values().for_each(|mut task| {
  28. if task.ret.is_some() {
  29. let ret = task.ret.take().unwrap();
  30. task.set_status(TaskStatus::Cancel);
  31. let _ = ret.send(task.into());
  32. }
  33. });
  34. }
  35. pub(crate) fn next_task_id(&self) -> TaskId {
  36. let _ = self.task_id_counter.fetch_add(1, SeqCst);
  37. self.task_id_counter.load(SeqCst)
  38. }
  39. }