Skip to main content

lustre_executor/core/
task.rs

1//! Task abstraction with unique IDs and waker integration.
2
3use crate::core::id::{IdGenerator, TaskId};
4use futures::future::Future;
5use std::collections::VecDeque;
6use std::sync::{Arc, Mutex};
7use std::task::{Context, Poll as TaskPoll, Wake, Waker};
8
9/// A task wrapping an async future with a unique ID.
10pub struct Task {
11    id: TaskId<u64>,
12    future: std::pin::Pin<Box<dyn Future<Output = ()> + 'static>>,
13}
14
15impl Task {
16    /// Creates a new task from an async future.
17    pub fn new(future: impl Future<Output = ()> + 'static) -> Self {
18        Self {
19            id: TaskId::new(crate::core::id::AtomicU64Generator::generate()),
20            future: Box::pin(future),
21        }
22    }
23
24    pub(crate) fn id(&self) -> TaskId<u64> {
25        self.id
26    }
27
28    pub(crate) fn poll(&mut self, cx: &mut Context) -> TaskPoll<()> {
29        self.future.as_mut().poll(cx)
30    }
31}
32
33/// Custom waker that enqueues tasks for polling.
34pub(crate) struct TaskWaker {
35    task_id: TaskId<u64>,
36    task_queue: Arc<Mutex<VecDeque<TaskId<u64>>>>,
37}
38
39impl TaskWaker {
40    pub(crate) fn new(
41        task_id: TaskId<u64>,
42        task_queue: Arc<Mutex<VecDeque<TaskId<u64>>>>,
43    ) -> Waker {
44        Waker::from(Arc::new(Self {
45            task_id,
46            task_queue,
47        }))
48    }
49
50    fn wake_task(&self) {
51        self.task_queue.lock().unwrap().push_front(self.task_id);
52    }
53}
54
55impl Wake for TaskWaker {
56    fn wake(self: Arc<Self>) {
57        self.wake_task();
58    }
59
60    fn wake_by_ref(self: &Arc<Self>) {
61        self.wake_task();
62    }
63}