a3s_flow/worker/
runner.rs

1use std::sync::Arc;
2
3use crate::engine::FlowEngine;
4use crate::error::Result;
5
6use super::{FlowTask, FlowTaskOutcome, FlowTaskQueue, InMemoryFlowTaskQueue};
7
8/// Worker that handles queued workflow tasks against a [`FlowEngine`].
9#[derive(Clone)]
10pub struct FlowWorker {
11    engine: FlowEngine,
12    queue: Arc<dyn FlowTaskQueue>,
13}
14
15impl FlowWorker {
16    pub fn new(engine: FlowEngine, queue: Arc<dyn FlowTaskQueue>) -> Self {
17        Self { engine, queue }
18    }
19
20    pub fn in_memory(engine: FlowEngine) -> Self {
21        Self::new(engine, Arc::new(InMemoryFlowTaskQueue::new()))
22    }
23
24    pub fn engine(&self) -> &FlowEngine {
25        &self.engine
26    }
27
28    pub fn queue(&self) -> Arc<dyn FlowTaskQueue> {
29        Arc::clone(&self.queue)
30    }
31
32    pub async fn enqueue(&self, task: FlowTask) -> Result<()> {
33        self.queue.enqueue(task).await
34    }
35
36    pub async fn handle(&self, task: FlowTask) -> Result<FlowTaskOutcome> {
37        let mut outcome = FlowTaskOutcome::new(task.clone());
38        match task {
39            FlowTask::DriveRun { run_id } => {
40                self.engine.drive(&run_id).await?;
41                outcome.run_ids.push(run_id);
42            }
43            FlowTask::ResumeWait { run_id, wait_id } => {
44                self.engine.resume_wait(&run_id, &wait_id).await?;
45                outcome.run_ids.push(run_id.clone());
46                outcome.resumed_waits.push((run_id, wait_id));
47            }
48            FlowTask::ResumeHook {
49                run_id,
50                hook_id,
51                payload,
52            } => {
53                self.engine.resume_hook(&run_id, &hook_id, payload).await?;
54                outcome.run_ids.push(run_id.clone());
55                outcome.resumed_hook = Some((run_id, hook_id));
56            }
57            FlowTask::ResumeHookByToken { token, payload } => {
58                let (run_id, hook_id) = self.engine.resume_hook_by_token(&token, payload).await?;
59                outcome.run_ids.push(run_id.clone());
60                outcome.resumed_hook = Some((run_id, hook_id));
61            }
62            FlowTask::DisposeHook { run_id, hook_id } => {
63                self.engine.dispose_hook(&run_id, &hook_id).await?;
64                outcome.run_ids.push(run_id.clone());
65                outcome.disposed_hook = Some((run_id, hook_id));
66            }
67            FlowTask::DisposeHookByToken { token } => {
68                let (run_id, hook_id) = self.engine.dispose_hook_by_token(&token).await?;
69                outcome.run_ids.push(run_id.clone());
70                outcome.disposed_hook = Some((run_id, hook_id));
71            }
72            FlowTask::ResumeDueWaits { now } => {
73                let resumed = self.engine.resume_due_waits(now).await?;
74                for (run_id, _) in &resumed {
75                    if !outcome.run_ids.contains(run_id) {
76                        outcome.run_ids.push(run_id.clone());
77                    }
78                }
79                outcome.resumed_waits = resumed;
80            }
81            FlowTask::ResumeDueRetries { now } => {
82                let resumed = self.engine.resume_due_retries(now).await?;
83                for (run_id, _) in &resumed {
84                    if !outcome.run_ids.contains(run_id) {
85                        outcome.run_ids.push(run_id.clone());
86                    }
87                }
88                outcome.resumed_retries = resumed;
89            }
90        }
91        Ok(outcome)
92    }
93
94    pub async fn run_once(&self) -> Result<Option<FlowTaskOutcome>> {
95        let Some(lease) = self.queue.lease().await? else {
96            return Ok(None);
97        };
98        let outcome = self.handle(lease.task).await?;
99        self.queue.ack(&lease.lease_id).await?;
100        Ok(Some(outcome))
101    }
102
103    pub async fn run_until_idle(&self) -> Result<Vec<FlowTaskOutcome>> {
104        let mut outcomes = Vec::new();
105        while let Some(outcome) = self.run_once().await? {
106            outcomes.push(outcome);
107        }
108        Ok(outcomes)
109    }
110}
a3s_flow/worker/runner.rs

a3s_flow/worker/
runner.rs
