hataori_runtime/
scope.rs

1use crate::{Action, Place, RemoteFuture, Runtime, RuntimeError, WireValue};
2use hataori_runtime_foundation::protocol::RequestId;
3use std::{
4    cell::RefCell,
5    collections::HashSet,
6    future::Future,
7    pin::Pin,
8    task::{Context, Poll},
9};
10
11pub struct RuntimeScope<'runtime> {
12    runtime: RefCell<&'runtime mut Runtime>,
13    tracked: RefCell<HashSet<RequestId>>,
14}
15
16impl std::fmt::Debug for RuntimeScope<'_> {
17    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18        formatter
19            .debug_struct("RuntimeScope")
20            .field("tracked", &self.tracked.borrow().len())
21            .finish_non_exhaustive()
22    }
23}
24
25impl RuntimeScope<'_> {
26    pub fn spawn_on<'scope, A: Action>(
27        &'scope self,
28        place: Place,
29        action: A,
30    ) -> Result<ScopedFuture<'scope, A::Output>, RuntimeError> {
31        let future = self.runtime.borrow().spawn_on(place, action)?;
32        let request = future.request_id();
33        self.tracked.borrow_mut().insert(request);
34        Ok(ScopedFuture {
35            inner: future,
36            tracked: &self.tracked,
37        })
38    }
39
40    pub fn block_on<T: WireValue>(&self, future: ScopedFuture<'_, T>) -> Result<T, RuntimeError> {
41        self.runtime.borrow_mut().block_on(future)
42    }
43}
44
45#[must_use = "a scoped future must be joined or dropped before the scope exits"]
46pub struct ScopedFuture<'scope, T: WireValue> {
47    inner: RemoteFuture<T>,
48    tracked: &'scope RefCell<HashSet<RequestId>>,
49}
50
51impl<T: WireValue> Future for ScopedFuture<'_, T> {
52    type Output = Result<T, RuntimeError>;
53
54    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
55        match Pin::new(&mut self.inner).poll(context) {
56            Poll::Ready(result) => {
57                self.tracked.borrow_mut().remove(&self.inner.request_id());
58                Poll::Ready(result)
59            }
60            Poll::Pending => Poll::Pending,
61        }
62    }
63}
64
65impl<T: WireValue> Drop for ScopedFuture<'_, T> {
66    fn drop(&mut self) {
67        self.tracked.borrow_mut().remove(&self.inner.request_id());
68    }
69}
70
71impl Runtime {
72    pub fn scope<R>(
73        &mut self,
74        function: impl for<'scope> FnOnce(&'scope RuntimeScope<'_>) -> Result<R, RuntimeError>,
75    ) -> Result<R, RuntimeError> {
76        let scope = RuntimeScope {
77            runtime: RefCell::new(self),
78            tracked: RefCell::new(HashSet::new()),
79        };
80        let result = function(&scope);
81        let requests: Vec<_> = scope.tracked.borrow().iter().copied().collect();
82        let client = scope.runtime.borrow().client();
83        for request in &requests {
84            client.cancel(*request);
85        }
86        if result.is_ok() && !requests.is_empty() {
87            scope.runtime.borrow_mut().progress(requests.len())?;
88        }
89        result
90    }
91}
hataori_runtime/scope.rs

hataori_runtime/
scope.rs
