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::GroupedCron`],
161    /// [`Schedule::WindowedInterval`],
162    /// [`Schedule::AtTimes`], and [`Schedule::Cron`].
163    /// A value of `0` makes the job exit immediately without running.
164    pub fn with_max_runs(mut self, max_runs: u32) -> Self {
165        self.max_runs = Some(max_runs);
166        self
167    }
168
169    pub fn with_missed_run_policy(mut self, policy: MissedRunPolicy) -> Self {
170        self.missed_run_policy = policy;
171        self
172    }
173
174    pub fn with_overlap_policy(mut self, policy: OverlapPolicy) -> Self {
175        self.overlap_policy = policy;
176        self
177    }
178
179    pub fn with_time_window(mut self, time_window: JobTimeWindow) -> Self {
180        self.time_window = Some(time_window);
181        self
182    }
183
184    pub fn with_time_window_alignment(mut self) -> Self {
185        self.time_window_alignment = TimeWindowAlignment::AlignToNextWindow;
186        self
187    }
188
189    pub fn with_execution_resource_id(mut self, resource_id: impl Into<String>) -> Self {
190        self.execution_resource_id = resource_id.into();
191        self
192    }
193
194    pub fn with_guard_scope(mut self, scope: ExecutionGuardScope) -> Self {
195        self.guard_scope = scope;
196        self
197    }
198
199    pub(crate) fn skip_reason_at(
200        &self,
201        now: DateTime<Utc>,
202        fallback_timezone: Tz,
203    ) -> Option<RunSkipReason> {
204        let window = self.time_window.as_ref()?;
205        if window.matches(now, fallback_timezone) {
206            None
207        } else {
208            Some(RunSkipReason::OutsideTimeWindow)
209        }
210    }
211}
212
213impl<D> std::fmt::Debug for Job<D> {
214    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
215        f.debug_struct("Job")
216            .field("job_id", &self.job_id)
217            .field("execution_resource_id", &self.execution_resource_id)
218            .field("guard_scope", &self.guard_scope)
219            .field("schedule", &self.schedule)
220            .field("time_window", &self.time_window)
221            .field("time_window_alignment", &self.time_window_alignment)
222            .field("max_runs", &self.max_runs)
223            .field("missed_run_policy", &self.missed_run_policy)
224            .field("overlap_policy", &self.overlap_policy)
225            .field("deps", &type_name::<D>())
226            .finish_non_exhaustive()
227    }
228}
229
230fn wrap_async_handler<D, F, Fut>(task: Arc<F>) -> TaskHandler<D>
231where
232    D: Send + Sync + 'static,
233    F: Fn(TriggeredTaskContext<D>) -> Fut + Send + Sync + 'static,
234    Fut: Future<Output = JobResult> + Send + 'static,
235{
236    Arc::new(move |context| Box::pin((*task)(context)))
237}
238
239fn wrap_sync_handler<D, F>(task: Arc<F>) -> TaskHandler<D>
240where
241    D: Send + Sync + 'static,
242    F: Fn(TriggeredTaskContext<D>) -> JobResult + Send + Sync + 'static,
243{
244    Arc::new(move |context| Box::pin(ready((*task)(context))))
245}
246
247fn wrap_blocking_handler<D, F>(task: Arc<F>) -> TaskHandler<D>
248where
249    D: Send + Sync + 'static,
250    F: Fn(TriggeredTaskContext<D>) -> JobResult + Send + Sync + 'static,
251{
252    Arc::new(move |context| {
253        let task = task.clone();
254        Box::pin(async move { await_blocking(move || (*task)(context)).await })
255    })
256}
257
258#[derive(Debug, Clone)]
259pub struct RunContext {
260    pub job_id: String,
261    pub scheduled_at: DateTime<Utc>,
262    pub catch_up: bool,
263    /// The scheduler-configured timezone for downstream task logic.
264    pub timezone: Tz,
265}
266
267#[derive(Debug, Clone, Copy, PartialEq, Eq)]
268pub enum TriggerSource {
269    Scheduled,
270    Manual,
271}
272
273#[derive(Clone)]
274pub struct TaskContext<D> {
275    pub run: RunContext,
276    pub deps: Arc<D>,
277}
278
279impl<D> std::fmt::Debug for TaskContext<D> {
280    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
281        f.debug_struct("TaskContext")
282            .field("run", &self.run)
283            .field("deps", &type_name::<D>())
284            .finish()
285    }
286}
287
288#[derive(Clone)]
289pub struct TriggeredTaskContext<D> {
290    pub run: RunContext,
291    pub deps: Arc<D>,
292    pub source: TriggerSource,
293}
294
295impl<D> std::fmt::Debug for TriggeredTaskContext<D> {
296    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
297        f.debug_struct("TriggeredTaskContext")
298            .field("run", &self.run)
299            .field("deps", &type_name::<D>())
300            .field("source", &self.source)
301            .finish()
302    }
303}
304
305async fn await_blocking<F>(task: F) -> JobResult
306where
307    F: FnOnce() -> JobResult + Send + 'static,
308{
309    match tokio::task::spawn_blocking(task).await {
310        Ok(result) => result,
311        Err(error) if error.is_panic() => resume_unwind(error.into_panic()),
312        Err(error) => panic!("blocking task failed to join: {error}"),
313    }
314}