tinyio_core/runtime/
task.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    sync::{Arc, Mutex},
5    task::{Context, Wake, Waker},
6};
7
8use crate::util::signal::{Counter, Signal};
9
10use super::RUNTIME;
11
12pub struct Task {
13    future: Mutex<Option<Pin<Box<dyn Future<Output = ()> + 'static + Send>>>>,
14    mark: Signal,
15}
16
17impl Wake for Task {
18    fn wake_by_ref(self: &Arc<Self>) {
19        RUNTIME.get().unwrap().push(self.clone());
20    }
21
22    fn wake(self: Arc<Self>) {
23        RUNTIME.get().unwrap().push(self.clone());
24    }
25}
26
27impl Task {
28    pub fn new(future: impl Future<Output = ()> + 'static + Send) -> Self {
29        Self {
30            future: Mutex::new(Some(Box::pin(future))),
31            mark: Signal::new(false),
32        }
33    }
34
35    /// Increase the counter if the task is not marked, then mark the task.
36    pub fn mark(self: &Arc<Self>, counter: &Counter) {
37        if !self.mark.value() {
38            counter.add(1);
39            self.mark.set(true);
40        }
41    }
42
43    /// Poll a task, if pending, send it back.
44    pub fn try_poll(self: &Arc<Self>, counter: &Counter) {
45        let mut future_slot = self.future.lock().unwrap();
46        if let Some(mut future) = future_slot.take() {
47            let waker = Waker::from(self.clone());
48            let mut cx = Context::from_waker(&waker);
49            if future.as_mut().poll(&mut cx).is_pending() {
50                *future_slot = Some(future);
51            } else {
52                counter.sub(1);
53            }
54        }
55    }
56}
tinyio_core/runtime/task.rs

tinyio_core/runtime/
task.rs
