Skip to main content

moirai_executor/hybrid/
spawner.rs

1use std::panic::{AssertUnwindSafe, catch_unwind};
2use std::sync::Arc;
3
4use moirai_core::{
5    Priority,
6    error::ExecutorResult,
7    executor::TaskSpawner,
8    task::{Task, TaskHandle},
9};
10
11use super::{HybridExecutor, MetricsRef, async_state::AsyncFutureState};
12use crate::schedule::{BlockingTask, SyncTask, WorkScheduler};
13
14impl<S: WorkScheduler> HybridExecutor<S> {
15    /// `Task`-typed adapter over the canonical closure-based
16    /// [`HybridExecutor::spawn_result`] path, shared by `spawn` and
17    /// `spawn_with_priority`: the task executes inside the same
18    /// cancellation-aware scheduled job with panic containment, lifecycle
19    /// timing, and result publication.
20    fn spawn_task_job<T>(
21        &self,
22        task: T,
23        priority: Priority,
24        locality_hint: Option<usize>,
25    ) -> ExecutorResult<TaskHandle<T::Output>>
26    where
27        T: Task + Send + 'static,
28        T::Output: Send + 'static,
29    {
30        self.spawn_result::<SyncTask, _>(priority, locality_hint, move || task.execute())
31    }
32}
33
34impl<S: WorkScheduler> TaskSpawner for HybridExecutor<S> {
35    fn spawn<T>(&self, task: T) -> ExecutorResult<TaskHandle<T::Output>>
36    where
37        T: Task + Send + 'static,
38        T::Output: Send + 'static,
39    {
40        let priority = task.context().priority;
41        self.spawn_task_job(task, priority, None)
42    }
43
44    fn spawn_async<F>(&self, future: F) -> ExecutorResult<TaskHandle<F::Output>>
45    where
46        F: core::future::Future + Send + 'static,
47        F::Output: Send + 'static,
48    {
49        let (task_id, lifecycle) = self.register_scheduled_task(Priority::Normal)?;
50
51        let (handle, result_sender) = TaskHandle::new_pending(task_id);
52        let state = AsyncFutureState::new(
53            self.scheduler.clone(),
54            future,
55            lifecycle,
56            result_sender,
57            Arc::clone(&self.metrics),
58        );
59        Arc::clone(&state).schedule()?;
60
61        self.metrics.record_task_spawned();
62        Ok(handle)
63    }
64
65    fn spawn_blocking<F, R>(&self, func: F) -> ExecutorResult<TaskHandle<R>>
66    where
67        F: FnOnce() -> R + Send + 'static,
68        R: Send + 'static,
69    {
70        // No caller-supplied priority exists on this surface, so the spawn
71        // records the default priority rather than a fabricated one.
72        self.spawn_result::<BlockingTask, _>(Priority::Normal, None, func)
73    }
74
75    fn spawn_detached<F>(&self, func: F) -> ExecutorResult<()>
76    where
77        F: FnOnce() + Send + 'static,
78    {
79        // No result is collected, so no `TaskHandle::new_pending` and therefore
80        // no `Arc<TaskResultSlot>` heap allocation or atomic refcount — the win
81        // over routing through `spawn_result`. Lifecycle tracking and metrics
82        // are preserved so shutdown drain and counters stay accurate.
83        let (_task_id, lifecycle) = self.register_scheduled_task(Priority::Normal)?;
84        let metrics = MetricsRef::new(&self.metrics);
85
86        self.scheduler
87            .schedule::<BlockingTask, _>(Priority::Normal, None, move |worker_id| {
88                let Some(running) = lifecycle.start_unless_cancelled(worker_id) else {
89                    metrics.get().record_task_cancelled();
90                    return;
91                };
92                // Catch here (not only at the job level) so `complete()` runs and
93                // the executor-level completed/failed metric is recorded, matching
94                // `send_task_result`.
95                match catch_unwind(AssertUnwindSafe(func)) {
96                    Ok(()) => {
97                        metrics.get().record_task_completed(running.complete());
98                    }
99                    Err(_) => {
100                        running.complete();
101                        metrics.get().record_task_failed();
102                    }
103                }
104            })?;
105
106        self.metrics.record_task_spawned();
107        Ok(())
108    }
109
110    fn spawn_with_priority<T>(
111        &self,
112        task: T,
113        priority: Priority,
114        locality_hint: Option<usize>,
115    ) -> ExecutorResult<TaskHandle<T::Output>>
116    where
117        T: Task + Send + 'static,
118        T::Output: Send + 'static,
119    {
120        self.spawn_task_job(task, priority, locality_hint)
121    }
122}