Skip to main content

a3s_flow/
scheduler.rs

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/// Result of one scheduler scan.
12#[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/// Scheduler that scans durable state and enqueues due workflow work.
26#[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    /// Backward-compatible name for [`Self::dispatcher`].
46    #[deprecated(since = "0.4.4", note = "use dispatcher()")]
47    pub fn queue(&self) -> Arc<dyn FlowTaskDispatcher> {
48        self.dispatcher()
49    }
50
51    /// Return the earliest wait or delayed retry that can wake the scheduler.
52    pub async fn next_wakeup(&self, now: DateTime<Utc>) -> Result<Option<WorkflowRunSuspension>> {
53        self.engine.next_wakeup(now).await
54    }
55
56    /// Return how long the host can sleep before the next scheduled wake-up.
57    ///
58    /// Due or overdue wake-ups return `Duration::ZERO`. Active hooks are not
59    /// represented because external callbacks are pushed into the queue by the
60    /// callback router instead of time.
61    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}