1use std::future::Future;
10use std::pin::Pin;
11use std::time::Duration;
12
13use axum::extract::FromRequestParts;
14use serde::{Serialize, de::DeserializeOwned};
15
16use crate::state::AppState;
17use crate::{AutumnError, AutumnResult};
18
19pub type TaskHandler =
21 fn(AppState) -> Pin<Box<dyn Future<Output = crate::AutumnResult<()>> + Send>>;
22
23pub type OneOffTaskHandler =
25 fn(AppState, Vec<String>) -> Pin<Box<dyn Future<Output = AutumnResult<()>> + Send>>;
26
27pub struct OneOffTaskInfo {
29 pub name: String,
31 pub description: String,
33 pub handler: OneOffTaskHandler,
35}
36
37#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
39pub struct OneOffTaskListing {
40 pub name: String,
42 pub description: String,
44}
45
46pub struct TaskArgs<T>(pub T);
52
53pub trait TaskExtractor: Sized {
55 fn from_task_parts<'a>(
57 parts: &'a mut http::request::Parts,
58 state: &'a AppState,
59 ) -> Pin<Box<dyn Future<Output = AutumnResult<Self>> + Send + 'a>>;
60}
61
62impl<T> TaskExtractor for T
63where
64 T: FromRequestParts<AppState> + Send,
65 T::Rejection: Into<AutumnError> + Send,
66{
67 fn from_task_parts<'a>(
68 parts: &'a mut http::request::Parts,
69 state: &'a AppState,
70 ) -> Pin<Box<dyn Future<Output = AutumnResult<Self>> + Send + 'a>> {
71 Box::pin(async move {
72 T::from_request_parts(parts, state)
73 .await
74 .map_err(Into::into)
75 })
76 }
77}
78
79impl<T, S> FromRequestParts<S> for TaskArgs<T>
80where
81 T: DeserializeOwned + Send,
82 S: Send + Sync,
83{
84 type Rejection = AutumnError;
85
86 async fn from_request_parts(
87 parts: &mut http::request::Parts,
88 _state: &S,
89 ) -> Result<Self, Self::Rejection> {
90 let query = parts.uri.query().unwrap_or_default();
91 serde_urlencoded::from_str(query)
92 .map(Self)
93 .map_err(|error| AutumnError::bad_request_msg(format!("invalid task args: {error}")))
94 }
95}
96
97pub struct TaskInfo {
99 pub name: String,
101 pub schedule: Schedule,
103 pub coordination: TaskCoordination,
105 pub handler: TaskHandler,
107}
108
109#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, serde::Deserialize)]
111#[serde(rename_all = "snake_case")]
112pub enum TaskCoordination {
113 #[default]
115 Fleet,
116 PerReplica,
118}
119
120impl std::fmt::Display for TaskCoordination {
121 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122 match self {
123 Self::Fleet => f.write_str("fleet"),
124 Self::PerReplica => f.write_str("per_replica"),
125 }
126 }
127}
128
129#[non_exhaustive]
131pub enum Schedule {
132 FixedDelay(Duration),
134 Cron {
136 expression: String,
138 timezone: Option<String>,
140 },
141}
142
143impl std::fmt::Display for Schedule {
144 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145 match self {
146 Self::FixedDelay(d) => write!(f, "every {}s", d.as_secs()),
147 Self::Cron { expression, .. } => write!(f, "cron {expression}"),
148 }
149 }
150}
151
152#[must_use]
160pub fn parse_duration(s: &str) -> Option<Duration> {
161 let mut total_secs = 0u64;
162 let mut current_num = String::new();
163
164 for ch in s.chars() {
165 if ch.is_ascii_digit() {
166 current_num.push(ch);
167 } else if ch.is_ascii_alphabetic() {
168 let num: u64 = current_num.parse().ok()?;
169 current_num.clear();
170 match ch {
171 's' => total_secs = total_secs.checked_add(num)?,
172 'm' => total_secs = total_secs.checked_add(num.checked_mul(60)?)?,
173 'h' => total_secs = total_secs.checked_add(num.checked_mul(3600)?)?,
174 'd' => total_secs = total_secs.checked_add(num.checked_mul(86400)?)?,
175 _ => return None,
176 }
177 } else if ch == ' ' {
178 } else {
180 return None;
181 }
182 }
183
184 if !current_num.is_empty() {
185 return None; }
187
188 if total_secs == 0 {
189 return None;
190 }
191
192 Some(Duration::from_secs(total_secs))
193}
194
195pub fn task_args_to_query(args: &[String]) -> AutumnResult<String> {
204 let mut serializer = url::form_urlencoded::Serializer::new(String::new());
205 let mut i = 0;
206 while i < args.len() {
207 let token = &args[i];
208 let Some(flag) = token.strip_prefix("--") else {
209 return Err(AutumnError::bad_request_msg(format!(
210 "unexpected positional argument {token:?}; task args must use --flag value syntax"
211 )));
212 };
213 if flag.is_empty() {
214 return Err(AutumnError::bad_request_msg(
215 "empty task argument flag is not allowed",
216 ));
217 }
218
219 let (key, value) = if let Some((key, value)) = flag.split_once('=') {
220 (key, value.to_owned())
221 } else if args.get(i + 1).is_some_and(|next| !next.starts_with("--")) {
222 i += 1;
223 (flag, args[i].clone())
224 } else {
225 (flag, "true".to_owned())
226 };
227
228 serializer.append_pair(&key.replace('-', "_"), &value);
229 i += 1;
230 }
231
232 Ok(serializer.finish())
233}
234
235pub fn request_parts_for_task_args(args: &[String]) -> AutumnResult<http::request::Parts> {
242 let query = task_args_to_query(args)?;
243 let uri = if query.is_empty() {
244 "/".to_owned()
245 } else {
246 format!("/?{query}")
247 };
248 let request = http::Request::builder()
249 .uri(uri)
250 .body(())
251 .map_err(AutumnError::internal_server_error)?;
252 Ok(request.into_parts().0)
253}
254
255#[must_use]
257pub fn list_one_off_tasks(tasks: &[OneOffTaskInfo]) -> Vec<OneOffTaskListing> {
258 let mut listing: Vec<_> = tasks
259 .iter()
260 .map(|task| OneOffTaskListing {
261 name: task.name.clone(),
262 description: task.description.clone(),
263 })
264 .collect();
265 listing.sort_by(|a, b| a.name.cmp(&b.name));
266 listing
267}
268
269pub fn validate_unique_one_off_task_names(tasks: &[OneOffTaskInfo]) -> Result<(), String> {
275 let mut names = std::collections::HashSet::new();
276 for task in tasks {
277 if !names.insert(task.name.as_str()) {
278 return Err(format!("duplicate task name '{}'", task.name));
279 }
280 }
281 Ok(())
282}
283
284pub fn validate_unique_scheduled_task_names(tasks: &[TaskInfo]) -> Result<(), String> {
301 validate_unique_task_names(tasks.iter().map(|task| task.name.as_str()))
302}
303
304pub fn validate_unique_task_names<'a>(
316 names: impl IntoIterator<Item = &'a str>,
317) -> Result<(), String> {
318 let mut seen = std::collections::HashSet::new();
319 for name in names {
320 if !seen.insert(name) {
321 return Err(format!(
322 "duplicate scheduled task name '{name}': two tasks — hand-declared via \
323 tasks![...], or auto-collected from a #[repository(..., retention(...))] \
324 policy — registered the same name. The scheduler and fleet coordinator key \
325 state by this name, so both would compete and overwrite each other's \
326 actuator state. Rename one of them"
327 ));
328 }
329 }
330 Ok(())
331}
332
333#[cfg(test)]
334mod tests {
335 use super::*;
336 use serde::Deserialize;
337
338 #[test]
339 fn parse_seconds() {
340 assert_eq!(parse_duration("5s"), Some(Duration::from_secs(5)));
341 }
342
343 #[test]
344 fn parse_minutes() {
345 assert_eq!(parse_duration("5m"), Some(Duration::from_secs(300)));
346 }
347
348 #[test]
349 fn parse_hours() {
350 assert_eq!(parse_duration("2h"), Some(Duration::from_secs(7200)));
351 }
352
353 #[test]
354 fn parse_compound() {
355 assert_eq!(parse_duration("1h 30m"), Some(Duration::from_secs(5400)));
356 }
357
358 #[test]
359 fn parse_day() {
360 assert_eq!(parse_duration("1d"), Some(Duration::from_secs(86400)));
361 }
362
363 #[test]
364 fn invalid_unit() {
365 assert!(parse_duration("5x").is_none());
366 }
367
368 #[test]
369 fn trailing_number() {
370 assert!(parse_duration("5").is_none());
371 }
372
373 #[test]
374 fn empty() {
375 assert!(parse_duration("").is_none());
376 }
377
378 #[test]
379 fn zero_duration() {
380 assert!(parse_duration("0s").is_none());
381 assert!(parse_duration("0m").is_none());
382 }
383
384 #[test]
385 fn invalid_characters() {
386 assert!(parse_duration("1h_30m").is_none());
387 assert!(parse_duration("1h-30m").is_none());
388 }
389
390 #[test]
391 fn multiple_spaces() {
392 assert_eq!(parse_duration("1h 30m"), Some(Duration::from_secs(5400)));
393 }
394
395 #[test]
396 fn compound_trailing_number() {
397 assert!(parse_duration("1h 30").is_none());
398 }
399
400 #[test]
401 fn task_args_to_query_converts_long_flags_to_snake_case_fields() {
402 let args = vec![
403 "--user-id".to_string(),
404 "42".to_string(),
405 "--dry-run".to_string(),
406 ];
407
408 let query = task_args_to_query(&args).expect("task args should parse");
409
410 assert_eq!(query, "user_id=42&dry_run=true");
411 }
412
413 #[tokio::test]
414 async fn task_args_extractor_parses_struct_from_cli_style_args() {
415 #[derive(Debug, Deserialize, PartialEq, Eq)]
416 struct CleanupArgs {
417 user_id: i64,
418 dry_run: bool,
419 }
420
421 let raw = vec![
422 "--user-id".to_string(),
423 "42".to_string(),
424 "--dry-run".to_string(),
425 ];
426 let mut parts =
427 request_parts_for_task_args(&raw).expect("parts should be built from task args");
428 let state = AppState::for_test().with_profile("dev");
429
430 let TaskArgs(args) = <TaskArgs<CleanupArgs> as axum::extract::FromRequestParts<
431 AppState,
432 >>::from_request_parts(&mut parts, &state)
433 .await
434 .expect("task args should deserialize");
435
436 assert_eq!(
437 args,
438 CleanupArgs {
439 user_id: 42,
440 dry_run: true,
441 }
442 );
443 }
444
445 #[test]
446 fn task_args_to_query_rejects_values_without_a_flag_name() {
447 let error = task_args_to_query(&["42".to_string()])
448 .expect_err("bare positional values should be rejected");
449
450 assert!(
451 error.to_string().contains("unexpected positional argument"),
452 "unexpected error: {error}"
453 );
454 }
455
456 #[test]
457 fn list_one_off_tasks_sorts_by_name_and_keeps_descriptions() {
458 fn handler(
459 _state: AppState,
460 _args: Vec<String>,
461 ) -> Pin<Box<dyn Future<Output = AutumnResult<()>> + Send>> {
462 Box::pin(async { Ok(()) })
463 }
464
465 let tasks = vec![
466 OneOffTaskInfo {
467 name: "zeta".to_string(),
468 description: "Last task".to_string(),
469 handler,
470 },
471 OneOffTaskInfo {
472 name: "alpha".to_string(),
473 description: "First task".to_string(),
474 handler,
475 },
476 ];
477
478 let listing = list_one_off_tasks(&tasks);
479
480 assert_eq!(listing[0].name, "alpha");
481 assert_eq!(listing[0].description, "First task");
482 assert_eq!(listing[1].name, "zeta");
483 assert_eq!(listing[1].description, "Last task");
484 }
485
486 #[test]
487 fn validate_unique_one_off_task_names_rejects_duplicates() {
488 fn handler(
489 _state: AppState,
490 _args: Vec<String>,
491 ) -> Pin<Box<dyn Future<Output = AutumnResult<()>> + Send>> {
492 Box::pin(async { Ok(()) })
493 }
494
495 let tasks = vec![
496 OneOffTaskInfo {
497 name: "cleanup".to_string(),
498 description: String::new(),
499 handler,
500 },
501 OneOffTaskInfo {
502 name: "cleanup".to_string(),
503 description: String::new(),
504 handler,
505 },
506 ];
507
508 let error = validate_unique_one_off_task_names(&tasks)
509 .expect_err("duplicate task names should be rejected");
510
511 assert!(error.contains("duplicate task name 'cleanup'"));
512 }
513
514 #[test]
515 fn validate_unique_scheduled_task_names_rejects_duplicates() {
516 fn handler(_state: AppState) -> Pin<Box<dyn Future<Output = AutumnResult<()>> + Send>> {
523 Box::pin(async { Ok(()) })
524 }
525
526 let tasks = vec![
527 TaskInfo {
528 name: "retention-sweep-posts".to_string(),
529 schedule: Schedule::FixedDelay(Duration::from_secs(3600)),
530 coordination: TaskCoordination::Fleet,
531 handler,
532 },
533 TaskInfo {
534 name: "retention-sweep-posts".to_string(),
535 schedule: Schedule::FixedDelay(Duration::from_secs(60)),
536 coordination: TaskCoordination::PerReplica,
537 handler,
538 },
539 ];
540
541 let error = validate_unique_scheduled_task_names(&tasks)
542 .expect_err("duplicate scheduled task names should be rejected");
543
544 assert!(error.contains("duplicate scheduled task name 'retention-sweep-posts'"));
545 }
546
547 #[test]
548 fn validate_unique_scheduled_task_names_accepts_distinct_names() {
549 fn handler(_state: AppState) -> Pin<Box<dyn Future<Output = AutumnResult<()>> + Send>> {
550 Box::pin(async { Ok(()) })
551 }
552
553 let tasks = vec![
554 TaskInfo {
555 name: "retention-sweep-posts".to_string(),
556 schedule: Schedule::FixedDelay(Duration::from_secs(3600)),
557 coordination: TaskCoordination::Fleet,
558 handler,
559 },
560 TaskInfo {
561 name: "nightly-report".to_string(),
562 schedule: Schedule::FixedDelay(Duration::from_secs(3600)),
563 coordination: TaskCoordination::Fleet,
564 handler,
565 },
566 ];
567
568 assert!(validate_unique_scheduled_task_names(&tasks).is_ok());
569 }
570}
571
572#[cfg(test)]
573mod havoc_proptests {
574 use super::*;
575 use proptest::prelude::*;
576
577 proptest! {
578 #[test]
579 fn parse_duration_fuzz_panic(s in "[0-9]{15,30}[smhd]") {
580 let _ = parse_duration(&s);
581 }
582 }
583}