Skip to main content

reifydb_sub_task/
handle.rs

1use tokio::sync::mpsc;
2
3use crate::{
4	coordinator::CoordinatorMessage,
5	registry::{TaskInfo, TaskRegistry},
6	task::{ScheduledTask, TaskId},
7};
8
9/// Handle for interacting with the task scheduler at runtime
10#[derive(Clone)]
11pub struct TaskHandle {
12	registry: TaskRegistry,
13	coordinator_tx: mpsc::Sender<CoordinatorMessage>,
14}
15
16impl TaskHandle {
17	pub(crate) fn new(registry: TaskRegistry, coordinator_tx: mpsc::Sender<CoordinatorMessage>) -> Self {
18		Self {
19			registry,
20			coordinator_tx,
21		}
22	}
23
24	pub async fn register_task(&self, task: ScheduledTask) -> Result<TaskId, String> {
25		let task_id = task.id;
26
27		self.coordinator_tx
28			.send(CoordinatorMessage::Register(task))
29			.await
30			.map_err(|e| format!("Failed to register task: {}", e))?;
31
32		Ok(task_id)
33	}
34
35	pub async fn unregister_task(&self, task_id: TaskId) -> Result<(), String> {
36		self.coordinator_tx
37			.send(CoordinatorMessage::Unregister(task_id))
38			.await
39			.map_err(|e| format!("Failed to unregister task: {}", e))?;
40
41		Ok(())
42	}
43
44	pub fn list_tasks(&self) -> Vec<TaskInfo> {
45		self.registry.iter().map(|entry| TaskInfo::from_entry(*entry.key(), entry.value())).collect()
46	}
47
48	pub fn get_task_info(&self, task_id: TaskId) -> Option<TaskInfo> {
49		self.registry.get(&task_id).map(|entry| TaskInfo::from_entry(task_id, entry.value()))
50	}
51}