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