moirai_executor/hybrid/
manager.rs

1use std::sync::{Arc, Mutex};
2
3use moirai_core::{
4    error::{ExecutorError, ExecutorResult, TaskError},
5    executor::{TaskManager, TaskStats, TaskStatus},
6    task::TaskId,
7};
8
9use super::HybridExecutor;
10use crate::registry::{CancelOutcome, TaskRegistry};
11use crate::schedule::WorkScheduler;
12use crate::task::TaskMetadata;
13
14fn lock_registry(
15    registry: &Mutex<TaskRegistry>,
16) -> ExecutorResult<std::sync::MutexGuard<'_, TaskRegistry>> {
17    registry
18        .lock()
19        .map_err(|_| ExecutorError::ResourceExhausted("task registry lock poisoned".to_string()))
20}
21
22/// Derive the observable [`TaskStatus`] from registry metadata.
23fn status_of(metadata: &TaskMetadata) -> TaskStatus {
24    if metadata.cancelled {
25        TaskStatus::Cancelled
26    } else if metadata.completed_at.is_some() {
27        TaskStatus::Completed
28    } else if metadata.started_at.is_some() {
29        TaskStatus::Running
30    } else {
31        TaskStatus::Queued
32    }
33}
34
35impl<S: WorkScheduler> TaskManager for HybridExecutor<S> {
36    /// Cooperative cancellation.
37    ///
38    /// Contract: a queued task that has not started skips its body when a
39    /// worker dequeues it — the task completes with `TaskError::Cancelled` and
40    /// status `Cancelled`. A task that already started is **not** preempted; it
41    /// runs to completion and this call still returns `Ok(())` (the request is
42    /// recorded but has no effect). Cancelling an already-completed task is a
43    /// no-op `Ok(())`. An unknown task ID is an error.
44    fn cancel_task(&self, id: TaskId) -> ExecutorResult<()> {
45        let registry = lock_registry(&self.task_registry)?;
46        match registry.request_cancel(id.0) {
47            Some(CancelOutcome::Requested | CancelOutcome::AlreadyCompleted) => Ok(()),
48            None => Err(ExecutorError::SpawnFailed(TaskError::InvalidOperation)),
49        }
50    }
51
52    fn task_status(&self, id: TaskId) -> Option<TaskStatus> {
53        let registry = self.task_registry.lock().ok()?;
54        registry
55            .get_metadata(id.0)
56            .map(|metadata| status_of(&metadata))
57    }
58
59    /// Event-driven completion wait.
60    ///
61    /// The future registers a completion waker with the task registry and
62    /// returns `Pending`; `mark_completed`/`mark_cancelled` wake it exactly
63    /// (no polling loop and no thread-blocking sleep). The deadline is checked
64    /// on every poll — at creation, on each completion wake, and on any
65    /// external poll after expiry. No in-scope timer exists (the async timer
66    /// lives in `moirai-async`), so if the task never completes, observing the
67    /// expiry requires the caller to poll after the deadline (e.g. a
68    /// timeout-aware runtime); a completion wake always resolves promptly.
69    fn wait_for_task(
70        &self,
71        id: TaskId,
72        timeout: Option<core::time::Duration>,
73    ) -> impl core::future::Future<Output = ExecutorResult<()>> + Send {
74        let registry = Arc::clone(&self.task_registry);
75        let deadline = timeout.and_then(|timeout| std::time::Instant::now().checked_add(timeout));
76
77        std::future::poll_fn(move |context| {
78            let registry = match lock_registry(&registry) {
79                Ok(registry) => registry,
80                Err(error) => return std::task::Poll::Ready(Err(error)),
81            };
82
83            if registry.is_completed(id.0) {
84                return std::task::Poll::Ready(Ok(()));
85            }
86            if registry.get_metadata(id.0).is_none() {
87                return std::task::Poll::Ready(Err(ExecutorError::SpawnFailed(
88                    TaskError::InvalidOperation,
89                )));
90            }
91            if deadline.is_some_and(|deadline| std::time::Instant::now() >= deadline) {
92                return std::task::Poll::Ready(Err(ExecutorError::SpawnFailed(TaskError::Timeout)));
93            }
94
95            registry.register_waker(id.0, context.waker());
96            // Re-check after registration: completion publishes the timestamp
97            // before taking the waker, so a completion that raced ahead of the
98            // registration is visible here and must not be lost.
99            if registry.is_completed(id.0) {
100                return std::task::Poll::Ready(Ok(()));
101            }
102
103            std::task::Poll::Pending
104        })
105    }
106
107    /// Statistics limited to what the executor actually tracks.
108    ///
109    /// `priority` is the value recorded at spawn; timing fields come from the
110    /// registry's lifecycle timestamps.
111    fn task_stats(&self, id: TaskId) -> Option<TaskStats> {
112        let registry = self.task_registry.lock().ok()?;
113        registry.get_metadata(id.0).map(|metadata| TaskStats {
114            id,
115            priority: metadata.priority,
116            status: status_of(&metadata),
117            spawn_time: metadata.created_at,
118            start_time: metadata.started_at,
119            completion_time: metadata.completed_at,
120            cpu_time_ns: metadata
121                .execution_duration()
122                .map_or(0, |duration| duration.as_nanos() as u64),
123        })
124    }
125}
moirai_executor/hybrid/manager.rs

moirai_executor/hybrid/
manager.rs
