a3s_flow/worker/
runner.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use crate::engine::FlowEngine;
5use crate::error::Result;
6
7use super::{FlowTask, FlowTaskLease, FlowTaskOutcome, FlowTaskQueue, InMemoryFlowTaskQueue};
8
9/// Worker that handles queued workflow tasks against a [`FlowEngine`].
10#[derive(Clone)]
11pub struct FlowWorker {
12    engine: FlowEngine,
13    queue: Arc<dyn FlowTaskQueue>,
14    heartbeat_interval: Option<Duration>,
15}
16
17impl FlowWorker {
18    pub fn new(engine: FlowEngine, queue: Arc<dyn FlowTaskQueue>) -> Self {
19        Self {
20            engine,
21            queue,
22            heartbeat_interval: None,
23        }
24    }
25
26    pub fn in_memory(engine: FlowEngine) -> Self {
27        Self::new(engine, Arc::new(InMemoryFlowTaskQueue::new()))
28    }
29
30    pub fn engine(&self) -> &FlowEngine {
31        &self.engine
32    }
33
34    pub fn queue(&self) -> Arc<dyn FlowTaskQueue> {
35        Arc::clone(&self.queue)
36    }
37
38    /// Enables periodic lease heartbeats while a task is being handled.
39    ///
40    /// Every successful heartbeat rotates the lease fencing token. If a
41    /// heartbeat reports that the lease was lost, the in-progress handling
42    /// future is dropped and its outcome is not acknowledged.
43    pub fn with_heartbeat_interval(mut self, interval: Duration) -> Result<Self> {
44        if interval.is_zero() {
45            return Err(crate::FlowError::InvalidWorkerConfiguration(
46                "heartbeat interval must be greater than zero".to_string(),
47            ));
48        }
49        self.heartbeat_interval = Some(interval);
50        Ok(self)
51    }
52
53    pub fn heartbeat_interval(&self) -> Option<Duration> {
54        self.heartbeat_interval
55    }
56
57    pub async fn enqueue(&self, task: FlowTask) -> Result<()> {
58        self.queue.enqueue(task).await
59    }
60
61    pub async fn handle(&self, task: FlowTask) -> Result<FlowTaskOutcome> {
62        handle_flow_task(&self.engine, task).await
63    }
64
65    async fn handle_lease(&self, lease: FlowTaskLease) -> Result<FlowTaskOutcome> {
66        let mut lease_id = lease.lease_id;
67        let handling = self.handle(lease.task);
68        tokio::pin!(handling);
69
70        let outcome = if let Some(interval) = self.heartbeat_interval {
71            let first_heartbeat = tokio::time::Instant::now() + interval;
72            let mut heartbeats = tokio::time::interval_at(first_heartbeat, interval);
73            heartbeats.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
74            loop {
75                tokio::select! {
76                    biased;
77                    _ = heartbeats.tick() => {
78                        lease_id = self.queue.heartbeat(&lease_id).await?;
79                    }
80                    result = &mut handling => break result?,
81                }
82            }
83        } else {
84            handling.await?
85        };
86
87        self.queue.ack(&lease_id).await?;
88        Ok(outcome)
89    }
90
91    pub async fn run_once(&self) -> Result<Option<FlowTaskOutcome>> {
92        let Some(lease) = self.queue.lease().await? else {
93            return Ok(None);
94        };
95        let outcome = self.handle_lease(lease).await?;
96        Ok(Some(outcome))
97    }
98
99    pub async fn run_until_idle(&self) -> Result<Vec<FlowTaskOutcome>> {
100        let mut outcomes = Vec::new();
101        while let Some(outcome) = self.run_once().await? {
102            outcomes.push(outcome);
103        }
104        Ok(outcomes)
105    }
106}
107
108pub(super) async fn handle_flow_task(
109    engine: &FlowEngine,
110    task: FlowTask,
111) -> Result<FlowTaskOutcome> {
112    let mut outcome = FlowTaskOutcome::new(task.clone());
113    match task {
114        FlowTask::DriveRun { run_id } => {
115            engine.drive(&run_id).await?;
116            outcome.run_ids.push(run_id);
117        }
118        FlowTask::ResumeWait { run_id, wait_id } => {
119            engine.resume_wait(&run_id, &wait_id).await?;
120            outcome.run_ids.push(run_id.clone());
121            outcome.resumed_waits.push((run_id, wait_id));
122        }
123        FlowTask::ResumeHook {
124            run_id,
125            hook_id,
126            payload,
127        } => {
128            engine.resume_hook(&run_id, &hook_id, payload).await?;
129            outcome.run_ids.push(run_id.clone());
130            outcome.resumed_hook = Some((run_id, hook_id));
131        }
132        FlowTask::ResumeHookByToken { token, payload } => {
133            let (run_id, hook_id) = engine.resume_hook_by_token(&token, payload).await?;
134            outcome.run_ids.push(run_id.clone());
135            outcome.resumed_hook = Some((run_id, hook_id));
136        }
137        FlowTask::DisposeHook { run_id, hook_id } => {
138            engine.dispose_hook(&run_id, &hook_id).await?;
139            outcome.run_ids.push(run_id.clone());
140            outcome.disposed_hook = Some((run_id, hook_id));
141        }
142        FlowTask::DisposeHookByToken { token } => {
143            let (run_id, hook_id) = engine.dispose_hook_by_token(&token).await?;
144            outcome.run_ids.push(run_id.clone());
145            outcome.disposed_hook = Some((run_id, hook_id));
146        }
147        FlowTask::ResumeScheduledRun { run_id, now } => {
148            let resumed = engine.resume_scheduled_run(&run_id, now).await?;
149            outcome.run_ids.push(run_id);
150            for wakeup in resumed {
151                let target = (wakeup.run_id, wakeup.subject_id);
152                match wakeup.kind {
153                    crate::ScheduledWakeupKind::Wait => outcome.resumed_waits.push(target),
154                    crate::ScheduledWakeupKind::Retry => outcome.resumed_retries.push(target),
155                }
156            }
157        }
158        FlowTask::ResumeDueWaits { now } => {
159            let resumed = engine.resume_due_waits(now).await?;
160            for (run_id, _) in &resumed {
161                if !outcome.run_ids.contains(run_id) {
162                    outcome.run_ids.push(run_id.clone());
163                }
164            }
165            outcome.resumed_waits = resumed;
166        }
167        FlowTask::ResumeDueRetries { now } => {
168            let resumed = engine.resume_due_retries(now).await?;
169            for (run_id, _) in &resumed {
170                if !outcome.run_ids.contains(run_id) {
171                    outcome.run_ids.push(run_id.clone());
172                }
173            }
174            outcome.resumed_retries = resumed;
175        }
176    }
177    Ok(outcome)
178}
a3s_flow/worker/runner.rs

a3s_flow/worker/
runner.rs
