a3s_flow/
scheduler.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
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/// Result of one scheduler scan and its targeted per-run dispatches.
13#[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/// Scheduler that scans durable state once and enqueues one task per due run.
27#[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    /// Backward-compatible name for [`Self::dispatcher`].
47    #[deprecated(since = "0.4.4", note = "use dispatcher()")]
48    pub fn queue(&self) -> Arc<dyn FlowTaskDispatcher> {
49        self.dispatcher()
50    }
51
52    /// Return the earliest wait or delayed retry that can wake the scheduler.
53    pub async fn next_wakeup(&self, now: DateTime<Utc>) -> Result<Option<WorkflowRunSuspension>> {
54        self.engine.next_wakeup(now).await
55    }
56
57    /// Return how long the host can sleep before the next scheduled wake-up.
58    ///
59    /// Due or overdue wake-ups return `Duration::ZERO`. Active hooks are not
60    /// represented because external callbacks are pushed into the queue by the
61    /// callback router instead of time.
62    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 mut targets = BTreeMap::new();
92        for wakeup in due {
93            match targets.entry(wakeup.run_id) {
94                std::collections::btree_map::Entry::Vacant(entry) => {
95                    entry.insert(wakeup.runtime_build_id);
96                }
97                std::collections::btree_map::Entry::Occupied(entry)
98                    if entry.get() != &wakeup.runtime_build_id =>
99                {
100                    return Err(crate::FlowError::Store(format!(
101                        "scheduled wakeups for run {} disagree on runtime build identity",
102                        entry.key()
103                    )));
104                }
105                std::collections::btree_map::Entry::Occupied(_) => {}
106            }
107        }
108        for required_build_id in targets.values() {
109            self.dispatcher
110                .ensure_runtime_build_route(required_build_id.as_ref())?;
111        }
112        for (run_id, required_build_id) in targets {
113            self.dispatcher
114                .dispatch_for_runtime_build(
115                    required_build_id.as_ref(),
116                    FlowTask::ResumeScheduledRun { run_id, now },
117                )
118                .await?;
119            enqueued_tasks += 1;
120        }
121
122        Ok(FlowSchedulerTick {
123            due_waits,
124            due_retries,
125            enqueued_tasks,
126        })
127    }
128}
a3s_flow/scheduler.rs

a3s_flow/
scheduler.rs
