screeps_async/
job.rs

1//! See [JobHandle]
2
3use async_task::Task;
4use std::cell::RefCell;
5use std::future::Future;
6use std::pin::Pin;
7use std::rc::Rc;
8use std::task::{Context, Poll};
9
10/// Reference to a [Future] that has been scheduled via [send](crate::ScreepsRuntime::spawn)
11///
12/// Dropping a [`JobHandle`] cancels it, which means its future won't be polled again.
13/// To drop the [`JobHandle`] handle without canceling it, use [detach()](JobHandle::detach) instead.
14/// To cancel a task gracefully and wait until it is fully destroyed, use the [cancel()](JobHandle::cancel) method
15///
16/// This type implements [Future] to allow awaiting on the result of the spawned task
17pub struct JobHandle<T> {
18    pub(crate) fut_res: Rc<RefCell<Option<T>>>,
19    task: Task<()>,
20    complete: bool,
21}
22
23impl<T> JobHandle<T> {
24    pub(crate) fn new(fut_res: Rc<RefCell<Option<T>>>, task: Task<()>) -> Self {
25        Self {
26            fut_res,
27            task,
28            complete: false,
29        }
30    }
31
32    /// Cancels the task and waits for it to stop running.
33    ///
34    /// Returns the task's output if it was completed just before it got canceled, or [`None`] if
35    /// it didn't complete.
36    ///
37    /// While it's possible to simply drop the [`JobHandle`] to cancel it, this is a cleaner way of
38    /// canceling because it also waits for the task to stop running.
39    pub async fn cancel(self) -> Option<T> {
40        let _ = self.task.cancel().await;
41        self.fut_res.take()
42    }
43
44    /// Detaches the task to let it keep running in the background.
45    pub fn detach(self) {
46        self.task.detach()
47    }
48
49    /// Check whether this job has completed
50    pub fn is_complete(&self) -> bool {
51        self.complete || self.fut_res.borrow().is_some()
52    }
53}
54
55impl<T> Future for JobHandle<T> {
56    type Output = T;
57
58    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
59        if self.complete {
60            panic!("Cannot await on a JobHandle that has already completed")
61        }
62        let res = self.fut_res.borrow_mut().take();
63
64        match res {
65            Some(res) => {
66                self.complete = true;
67                Poll::Ready(res)
68            }
69            None => {
70                cx.waker().wake_by_ref();
71                Poll::Pending
72            }
73        }
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use crate::spawn;
80    use crate::tests::init_test;
81
82    #[test]
83    fn test_cancel() {
84        init_test();
85
86        let result = crate::block_on(async move {
87            let task = spawn(async move { true });
88
89            task.cancel().await
90        })
91        .unwrap();
92
93        assert_eq!(None, result);
94    }
95
96    #[test]
97    fn test_await() {
98        init_test();
99
100        let result = crate::block_on(async move { spawn(async move { true }).await }).unwrap();
101
102        assert!(result, "Failed to await spawned future");
103    }
104}
screeps_async/job.rs

screeps_async/
job.rs
