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(TaskContext<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(task)))
39    }
40
41    /// Create a lightweight synchronous task from the full [`TaskContext`].
42    pub fn from_sync<F>(task: F) -> Self
43    where
44        F: Fn(TaskContext<D>) -> JobResult + Send + Sync + 'static,
45    {
46        Self::from_handler(wrap_sync_handler(Arc::new(task)))
47    }
48
49    /// Create a blocking synchronous task from the full [`TaskContext`].
50    pub fn from_blocking<F>(task: F) -> Self
51    where
52        F: Fn(TaskContext<D>) -> JobResult + Send + Sync + 'static,
53    {
54        Self::from_handler(wrap_blocking_handler(Arc::new(task)))
55    }
56}
57
58#[derive(Clone)]
59pub struct Job<D = ()> {
60    pub job_id: String,
61    pub execution_resource_id: String,
62    pub guard_scope: ExecutionGuardScope,
63    pub schedule: Schedule,
64    pub time_window: Option<JobTimeWindow>,
65    pub time_window_alignment: TimeWindowAlignment,
66    pub max_runs: Option<u32>,
67    pub missed_run_policy: MissedRunPolicy,
68    pub overlap_policy: OverlapPolicy,
69    pub(crate) task: TaskHandler<D>,
70    pub(crate) deps: Arc<D>,
71}
72
73impl Job<()> {
74    /// Create a job that uses no injected dependencies.
75    pub fn without_deps(job_id: impl Into<String>, schedule: Schedule, task: Task<()>) -> Self {
76        Self::from_parts(job_id.into(), schedule, Arc::new(()), task)
77    }
78}
79
80impl<D> Job<D>
81where
82    D: Send + Sync + 'static,
83{
84    /// Create a job from explicit dependencies and a task handler.
85    pub fn new(
86        job_id: impl Into<String>,
87        schedule: Schedule,
88        deps: impl Into<Arc<D>>,
89        task: Task<D>,
90    ) -> Self {
91        Self::from_parts(job_id.into(), schedule, deps.into(), task)
92    }
93}
94
95impl<D> Job<D> {
96    fn default_policies() -> (MissedRunPolicy, OverlapPolicy) {
97        (MissedRunPolicy::CatchUpOnce, OverlapPolicy::Forbid)
98    }
99
100    fn from_parts(job_id: String, schedule: Schedule, deps: Arc<D>, task: Task<D>) -> Self {
101        let (missed_run_policy, overlap_policy) = Self::default_policies();
102        Self {
103            execution_resource_id: job_id.clone(),
104            job_id,
105            guard_scope: ExecutionGuardScope::Occurrence,
106            schedule,
107            time_window: None,
108            time_window_alignment: TimeWindowAlignment::Disabled,
109            max_runs: None,
110            missed_run_policy,
111            overlap_policy,
112            task: task.handler,
113            deps,
114        }
115    }
116
117    /// Limit how many triggers this job can consume before it exits.
118    ///
119    /// This applies to [`Schedule::Interval`], [`Schedule::StaggeredInterval`],
120    /// [`Schedule::GroupedInterval`], [`Schedule::WindowedInterval`],
121    /// [`Schedule::AtTimes`], and [`Schedule::Cron`].
122    /// A value of `0` makes the job exit immediately without running.
123    pub fn with_max_runs(mut self, max_runs: u32) -> Self {
124        self.max_runs = Some(max_runs);
125        self
126    }
127
128    pub fn with_missed_run_policy(mut self, policy: MissedRunPolicy) -> Self {
129        self.missed_run_policy = policy;
130        self
131    }
132
133    pub fn with_overlap_policy(mut self, policy: OverlapPolicy) -> Self {
134        self.overlap_policy = policy;
135        self
136    }
137
138    pub fn with_time_window(mut self, time_window: JobTimeWindow) -> Self {
139        self.time_window = Some(time_window);
140        self
141    }
142
143    pub fn with_time_window_alignment(mut self) -> Self {
144        self.time_window_alignment = TimeWindowAlignment::AlignToNextWindow;
145        self
146    }
147
148    pub fn with_execution_resource_id(mut self, resource_id: impl Into<String>) -> Self {
149        self.execution_resource_id = resource_id.into();
150        self
151    }
152
153    pub fn with_guard_scope(mut self, scope: ExecutionGuardScope) -> Self {
154        self.guard_scope = scope;
155        self
156    }
157
158    pub(crate) fn skip_reason_at(
159        &self,
160        now: DateTime<Utc>,
161        fallback_timezone: Tz,
162    ) -> Option<RunSkipReason> {
163        let window = self.time_window.as_ref()?;
164        if window.matches(now, fallback_timezone) {
165            None
166        } else {
167            Some(RunSkipReason::OutsideTimeWindow)
168        }
169    }
170}
171
172impl<D> std::fmt::Debug for Job<D> {
173    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174        f.debug_struct("Job")
175            .field("job_id", &self.job_id)
176            .field("execution_resource_id", &self.execution_resource_id)
177            .field("guard_scope", &self.guard_scope)
178            .field("schedule", &self.schedule)
179            .field("time_window", &self.time_window)
180            .field("time_window_alignment", &self.time_window_alignment)
181            .field("max_runs", &self.max_runs)
182            .field("missed_run_policy", &self.missed_run_policy)
183            .field("overlap_policy", &self.overlap_policy)
184            .field("deps", &type_name::<D>())
185            .finish_non_exhaustive()
186    }
187}
188
189fn wrap_async_handler<D, F, Fut>(task: Arc<F>) -> TaskHandler<D>
190where
191    D: Send + Sync + 'static,
192    F: Fn(TaskContext<D>) -> Fut + Send + Sync + 'static,
193    Fut: Future<Output = JobResult> + Send + 'static,
194{
195    Arc::new(move |context| Box::pin((*task)(context)))
196}
197
198fn wrap_sync_handler<D, F>(task: Arc<F>) -> TaskHandler<D>
199where
200    D: Send + Sync + 'static,
201    F: Fn(TaskContext<D>) -> JobResult + Send + Sync + 'static,
202{
203    Arc::new(move |context| Box::pin(ready((*task)(context))))
204}
205
206fn wrap_blocking_handler<D, F>(task: Arc<F>) -> TaskHandler<D>
207where
208    D: Send + Sync + 'static,
209    F: Fn(TaskContext<D>) -> JobResult + Send + Sync + 'static,
210{
211    Arc::new(move |context| {
212        let task = task.clone();
213        Box::pin(async move { await_blocking(move || (*task)(context)).await })
214    })
215}
216
217#[derive(Debug, Clone)]
218pub struct RunContext {
219    pub job_id: String,
220    pub scheduled_at: DateTime<Utc>,
221    pub catch_up: bool,
222    /// The scheduler-configured timezone for downstream task logic.
223    pub timezone: Tz,
224}
225
226#[derive(Clone)]
227pub struct TaskContext<D> {
228    pub run: RunContext,
229    pub deps: Arc<D>,
230}
231
232impl<D> std::fmt::Debug for TaskContext<D> {
233    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
234        f.debug_struct("TaskContext")
235            .field("run", &self.run)
236            .field("deps", &type_name::<D>())
237            .finish()
238    }
239}
240
241async fn await_blocking<F>(task: F) -> JobResult
242where
243    F: FnOnce() -> JobResult + Send + 'static,
244{
245    match tokio::task::spawn_blocking(task).await {
246        Ok(result) => result,
247        Err(error) if error.is_panic() => resume_unwind(error.into_panic()),
248        Err(error) => panic!("blocking task failed to join: {error}"),
249    }
250}