1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeSet;
4use std::sync::Arc;
5use std::time::Duration;
6
7use crate::engine::FlowEngine;
8use crate::error::Result;
9use crate::model::{ScheduledWakeupKind, WorkflowRunSuspension};
10use crate::worker::{FlowTask, FlowTaskDispatcher};
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
14pub struct FlowSchedulerTick {
15 pub due_waits: Vec<(String, String)>,
16 pub due_retries: Vec<(String, String)>,
17 pub enqueued_tasks: usize,
18}
19
20impl FlowSchedulerTick {
21 pub fn has_due_work(&self) -> bool {
22 !self.due_waits.is_empty() || !self.due_retries.is_empty()
23 }
24}
25
26#[derive(Clone)]
28pub struct FlowScheduler {
29 engine: FlowEngine,
30 dispatcher: Arc<dyn FlowTaskDispatcher>,
31}
32
33impl FlowScheduler {
34 pub fn new(engine: FlowEngine, dispatcher: Arc<dyn FlowTaskDispatcher>) -> Self {
35 Self { engine, dispatcher }
36 }
37
38 pub fn engine(&self) -> &FlowEngine {
39 &self.engine
40 }
41
42 pub fn dispatcher(&self) -> Arc<dyn FlowTaskDispatcher> {
43 Arc::clone(&self.dispatcher)
44 }
45
46 #[deprecated(since = "0.4.4", note = "use dispatcher()")]
48 pub fn queue(&self) -> Arc<dyn FlowTaskDispatcher> {
49 self.dispatcher()
50 }
51
52 pub async fn next_wakeup(&self, now: DateTime<Utc>) -> Result<Option<WorkflowRunSuspension>> {
54 self.engine.next_wakeup(now).await
55 }
56
57 pub async fn next_wakeup_delay(&self, now: DateTime<Utc>) -> Result<Option<Duration>> {
63 let Some(wakeup) = self.next_wakeup(now).await? else {
64 return Ok(None);
65 };
66 let Some(scheduled_at) = wakeup.scheduled_at() else {
67 return Ok(None);
68 };
69 Ok(Some(
70 scheduled_at
71 .signed_duration_since(now)
72 .to_std()
73 .unwrap_or(Duration::ZERO),
74 ))
75 }
76
77 pub async fn enqueue_due_work(&self, now: DateTime<Utc>) -> Result<FlowSchedulerTick> {
78 let due = self.engine.list_due_wakeups(now).await?;
79 let due_waits = due
80 .iter()
81 .filter(|wakeup| wakeup.kind == ScheduledWakeupKind::Wait)
82 .map(|wakeup| (wakeup.run_id.clone(), wakeup.subject_id.clone()))
83 .collect::<Vec<_>>();
84 let due_retries = due
85 .iter()
86 .filter(|wakeup| wakeup.kind == ScheduledWakeupKind::Retry)
87 .map(|wakeup| (wakeup.run_id.clone(), wakeup.subject_id.clone()))
88 .collect::<Vec<_>>();
89 let mut enqueued_tasks = 0usize;
90
91 let run_ids = due
92 .into_iter()
93 .map(|wakeup| wakeup.run_id)
94 .collect::<BTreeSet<_>>();
95 for run_id in run_ids {
96 self.dispatcher
97 .dispatch(FlowTask::ResumeScheduledRun { run_id, now })
98 .await?;
99 enqueued_tasks += 1;
100 }
101
102 Ok(FlowSchedulerTick {
103 due_waits,
104 due_retries,
105 enqueued_tasks,
106 })
107 }
108}