Skip to main content

scheduler/model/
task.rs

1use crate::execution_guard::ExecutionGuardScope;
2use crate::model::{
3    JobTimeWindow, MissedRunPolicy, OverlapPolicy, RunSkipReason, Schedule, TimeWindowAlignment,
4};
5use chrono::{DateTime, Utc};
6use chrono_tz::Tz;
7use std::any::type_name;
8use std::future::{Future, ready};
9use std::panic::resume_unwind;
10use std::pin::Pin;
11use std::sync::Arc;
12
13/// The task return type used by scheduled jobs.
14pub type JobResult = Result<(), String>;
15/// The boxed future returned by a scheduled job.
16pub type JobFuture = Pin<Box<dyn Future<Output = JobResult> + Send>>;
17pub(crate) type TaskHandler<D> = Arc<dyn Fn(TriggeredTaskContext<D>) -> JobFuture + Send + Sync>;
18
19#[derive(Clone)]
20pub struct Task<D> {
21    pub(crate) handler: TaskHandler<D>,
22}
23
24impl<D> Task<D>
25where
26    D: Send + Sync + 'static,
27{
28    fn from_handler(handler: TaskHandler<D>) -> Self {
29        Self { handler }
30    }
31
32    /// Create an async task from the full [`TaskContext`].
33    pub fn from_async<F, Fut>(task: F) -> Self
34    where
35        F: Fn(TaskContext<D>) -> Fut + Send + Sync + 'static,
36        Fut: Future<Output = JobResult> + Send + 'static,
37    {
38        Self::from_handler(wrap_async_handler(Arc::new(move |context: TriggeredTaskContext<D>| {
39            task(TaskContext {
40                run: context.run,
41                deps: context.deps,
42            })
43        })))
44    }
45
46    /// Create a lightweight synchronous task from the full [`TaskContext`].
47    pub fn from_sync<F>(task: F) -> Self
48    where
49        F: Fn(TaskContext<D>) -> JobResult + Send + Sync + 'static,
50    {
51        Self::from_handler(wrap_sync_handler(Arc::new(move |context: TriggeredTaskContext<D>| {
52            task(TaskContext {
53                run: context.run,
54                deps: context.deps,
55            })
56        })))
57    }
58
59    /// Create a blocking synchronous task from the full [`TaskContext`].
60    pub fn from_blocking<F>(task: F) -> Self
61    where
62        F: Fn(TaskContext<D>) -> JobResult + Send + Sync + 'static,
63    {
64        Self::from_handler(wrap_blocking_handler(Arc::new(move |context: TriggeredTaskContext<D>| {
65            task(TaskContext {
66                run: context.run,
67                deps: context.deps,
68            })
69        })))
70    }
71
72    /// Create an async task from the full [`TriggeredTaskContext`].
73    pub fn from_async_with_trigger<F, Fut>(task: F) -> Self
74    where
75        F: Fn(TriggeredTaskContext<D>) -> Fut + Send + Sync + 'static,
76        Fut: Future<Output = JobResult> + Send + 'static,
77    {
78        Self::from_handler(wrap_async_handler(Arc::new(task)))
79    }
80
81    /// Create a lightweight synchronous task from the full [`TriggeredTaskContext`].
82    pub fn from_sync_with_trigger<F>(task: F) -> Self
83    where
84        F: Fn(TriggeredTaskContext<D>) -> JobResult + Send + Sync + 'static,
85    {
86        Self::from_handler(wrap_sync_handler(Arc::new(task)))
87    }
88
89    /// Create a blocking synchronous task from the full [`TriggeredTaskContext`].
90    pub fn from_blocking_with_trigger<F>(task: F) -> Self
91    where
92        F: Fn(TriggeredTaskContext<D>) -> JobResult + Send + Sync + 'static,
93    {
94        Self::from_handler(wrap_blocking_handler(Arc::new(task)))
95    }
96}
97
98#[derive(Clone)]
99pub struct Job<D = ()> {
100    pub job_id: String,
101    pub execution_resource_id: String,
102    pub guard_scope: ExecutionGuardScope,
103    pub schedule: Schedule,
104    pub time_window: Option<JobTimeWindow>,
105    pub time_window_alignment: TimeWindowAlignment,
106    pub max_runs: Option<u32>,
107    pub missed_run_policy: MissedRunPolicy,
108    pub overlap_policy: OverlapPolicy,
109    pub(crate) task: TaskHandler<D>,
110    pub(crate) deps: Arc<D>,
111}
112
113impl Job<()> {
114    /// Create a job that uses no injected dependencies.
115    pub fn without_deps(job_id: impl Into<String>, schedule: Schedule, task: Task<()>) -> Self {
116        Self::from_parts(job_id.into(), schedule, Arc::new(()), task)
117    }
118}
119
120impl<D> Job<D>
121where
122    D: Send + Sync + 'static,
123{
124    /// Create a job from explicit dependencies and a task handler.
125    pub fn new(
126        job_id: impl Into<String>,
127        schedule: Schedule,
128        deps: impl Into<Arc<D>>,
129        task: Task<D>,
130    ) -> Self {
131        Self::from_parts(job_id.into(), schedule, deps.into(), task)
132    }
133}
134
135impl<D> Job<D> {
136    fn default_policies() -> (MissedRunPolicy, OverlapPolicy) {
137        (MissedRunPolicy::CatchUpOnce, OverlapPolicy::Forbid)
138    }
139
140    fn from_parts(job_id: String, schedule: Schedule, deps: Arc<D>, task: Task<D>) -> Self {
141        let (missed_run_policy, overlap_policy) = Self::default_policies();
142        Self {
143            execution_resource_id: job_id.clone(),
144            job_id,
145            guard_scope: ExecutionGuardScope::Occurrence,
146            schedule,
147            time_window: None,
148            time_window_alignment: TimeWindowAlignment::Disabled,
149            max_runs: None,
150            missed_run_policy,
151            overlap_policy,
152            task: task.handler,
153            deps,
154        }
155    }
156
157    /// Limit how many triggers this job can consume before it exits.
158    ///
159    /// This applies to [`Schedule::Interval`], [`Schedule::StaggeredInterval`],
160    /// [`Schedule::GroupedInterval`], [`Schedule::WindowedInterval`],
161    /// [`Schedule::AtTimes`], and [`Schedule::Cron`].
162    /// A value of `0` makes the job exit immediately without running.
163    pub fn with_max_runs(mut self, max_runs: u32) -> Self {
164        self.max_runs = Some(max_runs);
165        self
166    }
167
168    pub fn with_missed_run_policy(mut self, policy: MissedRunPolicy) -> Self {
169        self.missed_run_policy = policy;
170        self
171    }
172
173    pub fn with_overlap_policy(mut self, policy: OverlapPolicy) -> Self {
174        self.overlap_policy = policy;
175        self
176    }
177
178    pub fn with_time_window(mut self, time_window: JobTimeWindow) -> Self {
179        self.time_window = Some(time_window);
180        self
181    }
182
183    pub fn with_time_window_alignment(mut self) -> Self {
184        self.time_window_alignment = TimeWindowAlignment::AlignToNextWindow;
185        self
186    }
187
188    pub fn with_execution_resource_id(mut self, resource_id: impl Into<String>) -> Self {
189        self.execution_resource_id = resource_id.into();
190        self
191    }
192
193    pub fn with_guard_scope(mut self, scope: ExecutionGuardScope) -> Self {
194        self.guard_scope = scope;
195        self
196    }
197
198    pub(crate) fn skip_reason_at(
199        &self,
200        now: DateTime<Utc>,
201        fallback_timezone: Tz,
202    ) -> Option<RunSkipReason> {
203        let window = self.time_window.as_ref()?;
204        if window.matches(now, fallback_timezone) {
205            None
206        } else {
207            Some(RunSkipReason::OutsideTimeWindow)
208        }
209    }
210}
211
212impl<D> std::fmt::Debug for Job<D> {
213    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214        f.debug_struct("Job")
215            .field("job_id", &self.job_id)
216            .field("execution_resource_id", &self.execution_resource_id)
217            .field("guard_scope", &self.guard_scope)
218            .field("schedule", &self.schedule)
219            .field("time_window", &self.time_window)
220            .field("time_window_alignment", &self.time_window_alignment)
221            .field("max_runs", &self.max_runs)
222            .field("missed_run_policy", &self.missed_run_policy)
223            .field("overlap_policy", &self.overlap_policy)
224            .field("deps", &type_name::<D>())
225            .finish_non_exhaustive()
226    }
227}
228
229fn wrap_async_handler<D, F, Fut>(task: Arc<F>) -> TaskHandler<D>
230where
231    D: Send + Sync + 'static,
232    F: Fn(TriggeredTaskContext<D>) -> Fut + Send + Sync + 'static,
233    Fut: Future<Output = JobResult> + Send + 'static,
234{
235    Arc::new(move |context| Box::pin((*task)(context)))
236}
237
238fn wrap_sync_handler<D, F>(task: Arc<F>) -> TaskHandler<D>
239where
240    D: Send + Sync + 'static,
241    F: Fn(TriggeredTaskContext<D>) -> JobResult + Send + Sync + 'static,
242{
243    Arc::new(move |context| Box::pin(ready((*task)(context))))
244}
245
246fn wrap_blocking_handler<D, F>(task: Arc<F>) -> TaskHandler<D>
247where
248    D: Send + Sync + 'static,
249    F: Fn(TriggeredTaskContext<D>) -> JobResult + Send + Sync + 'static,
250{
251    Arc::new(move |context| {
252        let task = task.clone();
253        Box::pin(async move { await_blocking(move || (*task)(context)).await })
254    })
255}
256
257#[derive(Debug, Clone)]
258pub struct RunContext {
259    pub job_id: String,
260    pub scheduled_at: DateTime<Utc>,
261    pub catch_up: bool,
262    /// The scheduler-configured timezone for downstream task logic.
263    pub timezone: Tz,
264}
265
266#[derive(Debug, Clone, Copy, PartialEq, Eq)]
267pub enum TriggerSource {
268    Scheduled,
269    Manual,
270}
271
272#[derive(Clone)]
273pub struct TaskContext<D> {
274    pub run: RunContext,
275    pub deps: Arc<D>,
276}
277
278impl<D> std::fmt::Debug for TaskContext<D> {
279    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
280        f.debug_struct("TaskContext")
281            .field("run", &self.run)
282            .field("deps", &type_name::<D>())
283            .finish()
284    }
285}
286
287#[derive(Clone)]
288pub struct TriggeredTaskContext<D> {
289    pub run: RunContext,
290    pub deps: Arc<D>,
291    pub source: TriggerSource,
292}
293
294impl<D> std::fmt::Debug for TriggeredTaskContext<D> {
295    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296        f.debug_struct("TriggeredTaskContext")
297            .field("run", &self.run)
298            .field("deps", &type_name::<D>())
299            .field("source", &self.source)
300            .finish()
301    }
302}
303
304async fn await_blocking<F>(task: F) -> JobResult
305where
306    F: FnOnce() -> JobResult + Send + 'static,
307{
308    match tokio::task::spawn_blocking(task).await {
309        Ok(result) => result,
310        Err(error) if error.is_panic() => resume_unwind(error.into_panic()),
311        Err(error) => panic!("blocking task failed to join: {error}"),
312    }
313}