Skip to main content

lustre_executor/
task.rs

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