1use async_trait::async_trait;
2
3use crate::error::Result;
4use crate::runtime_build::RuntimeBuildId;
5
6use super::{FlowTask, FlowTaskLease};
7
8#[async_trait]
10pub trait FlowTaskDispatcher: Send + Sync {
11 async fn dispatch(&self, task: FlowTask) -> Result<()>;
12
13 fn has_runtime_build_route(&self, required_build_id: Option<&RuntimeBuildId>) -> bool {
15 required_build_id.is_none()
16 }
17
18 fn ensure_runtime_build_route(&self, required_build_id: Option<&RuntimeBuildId>) -> Result<()> {
20 if self.has_runtime_build_route(required_build_id) {
21 return Ok(());
22 }
23 Err(crate::FlowError::RuntimeBuildRouteNotFound {
24 required_build_id: required_build_id.cloned(),
25 })
26 }
27
28 async fn dispatch_for_runtime_build(
35 &self,
36 required_build_id: Option<&RuntimeBuildId>,
37 task: FlowTask,
38 ) -> Result<()> {
39 self.ensure_runtime_build_route(required_build_id)?;
40 self.dispatch(task).await
41 }
42}
43
44#[async_trait]
46pub trait FlowTaskQueue: Send + Sync {
47 async fn enqueue(&self, task: FlowTask) -> Result<()>;
48
49 async fn lease(&self) -> Result<Option<FlowTaskLease>>;
50
51 async fn heartbeat(&self, lease_id: &str) -> Result<String>;
56
57 async fn ack(&self, lease_id: &str) -> Result<()>;
63
64 async fn requeue_inflight(&self) -> Result<usize> {
65 Ok(0)
66 }
67
68 async fn dequeue(&self) -> Result<Option<FlowTask>> {
69 let Some(lease) = self.lease().await? else {
70 return Ok(None);
71 };
72 let task = lease.task.clone();
73 self.ack(&lease.lease_id).await?;
74 Ok(Some(task))
75 }
76
77 async fn len(&self) -> Result<usize>;
78
79 async fn is_empty(&self) -> Result<bool> {
80 Ok(self.len().await? == 0)
81 }
82}
83
84#[async_trait]
85impl<T> FlowTaskDispatcher for T
86where
87 T: FlowTaskQueue + ?Sized,
88{
89 async fn dispatch(&self, task: FlowTask) -> Result<()> {
90 FlowTaskQueue::enqueue(self, task).await
91 }
92}