1#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/README.md"))]
2use std::{str::FromStr, sync::Arc};
3
4use apalis_board_types::ApiError;
5use apalis_core::{
6 backend::{
7 Backend, BackendConfig, FetchById, Filter, ListAllTasks, ListQueues, ListTasks,
8 ListWorkers, Metrics, QueueInfo, RunningWorker, Statistic, TaskSink, WireFormatBackend,
9 codec::Codec,
10 },
11 task::{ExecutionContext, Task, task_id::TaskId},
12};
13use serde::{
14 Deserialize, Deserializer, Serialize, Serializer, de::DeserializeOwned, ser::SerializeStruct,
15};
16use tokio::sync::RwLock;
17
18pub const PREFIX: &str = "x-task-";
20pub const META_PREFIX: &str = "x-task-meta-";
22
23pub mod framework;
25#[cfg(feature = "sse")]
27pub mod sse;
28#[cfg(feature = "ui")]
30pub mod ui;
31
32pub async fn push_task<Args, B>(
34 task: RawTask<Args>,
35 storage: Arc<RwLock<B>>,
36) -> Result<(), ApiError>
37where
38 B: BackendConfig,
39 Args: Serialize + DeserializeOwned + 'static,
40 B: TaskSink<Args, B::Kind> + Send,
41{
42 let res = storage.write().await.push_task(task.0).await;
43 match res {
44 Ok(_) => Ok(()),
45 Err(e) => Err(ApiError::BackendError(e.to_string())),
46 }
47}
48pub async fn stats_by_queue<S>(storage: Arc<RwLock<S>>) -> Result<Vec<Statistic>, ApiError>
50where
51 S::Error: std::error::Error,
52 S: Metrics,
53{
54 let storage = storage.read().await;
55 let stats = storage.fetch_by_queue().await;
56 match stats {
57 Ok(stats) => Ok(stats),
58 Err(e) => Err(ApiError::BackendError(e.to_string())),
59 }
60}
61
62pub async fn get_tasks<S, Args>(
64 storage: Arc<RwLock<S>>,
65 filter: Filter,
66) -> Result<Vec<Task<Args>>, ApiError>
67where
68 S: ListTasks + Send,
69 Args: Serialize,
70 S::Codec: Codec<Args, Compact = S::Compact>,
71 <S::Codec as Codec<Args>>::Error: std::error::Error,
72{
73 let storage = storage.read().await;
74
75 let tasks = storage.list_tasks(&filter).await;
76 let codec = storage.codec();
77 tasks
78 .map_err(|e| ApiError::BackendError(e.to_string()))?
79 .into_iter()
80 .map(|task| {
81 task.try_map_args(|compact| {
82 codec
83 .decode(&compact)
84 .map_err(|e| ApiError::CodecError(e.to_string()))
85 })
86 })
87 .collect()
88}
89
90pub async fn get_workers<S>(storage: Arc<RwLock<S>>) -> Result<Vec<RunningWorker>, ApiError>
92where
93 S: ListWorkers,
94 S::Error: std::error::Error,
95{
96 let storage = storage.read().await;
97 storage
98 .list_workers()
99 .await
100 .map_err(|e| ApiError::BackendError(e.to_string()))
101}
102
103pub async fn get_task_by_id<B, Args>(
105 task_id: String,
106 storage: Arc<RwLock<B>>,
107) -> Result<Option<Task<Args>>, ApiError>
108where
109 B: FetchById + WireFormatBackend + 'static,
110 B::Error: std::error::Error,
111 B: Backend<Task = Task<B::Compact>>,
112 B::Codec: Codec<Args, Compact = B::Compact>,
113 <B::Codec as Codec<Args>>::Error: std::error::Error,
114{
115 let task_id = TaskId::from_str(&task_id).map_err(|e| ApiError::BackendError(e.to_string()))?;
116
117 let mut storage = storage.write().await;
118 let task = storage
119 .fetch_by_id(&task_id)
120 .await
121 .map_err(|e| ApiError::BackendError(e.to_string()))?;
122 let codec = storage.codec();
123 task.map(|task| {
124 task.try_map_args(|compact| {
125 codec
126 .decode(&compact)
127 .map_err(|e| ApiError::CodecError(e.to_string()))
128 })
129 })
130 .transpose()
131}
132
133pub async fn get_all_tasks<S>(
135 storage: Arc<RwLock<S>>,
136 filter: Filter,
137) -> Result<Vec<Task<S::Compact>>, ApiError>
138where
139 S: ListAllTasks + Send,
140 S::Compact: Serialize,
141{
142 storage
143 .read()
144 .await
145 .list_all_tasks(&filter)
146 .await
147 .map_err(|e| ApiError::BackendError(e.to_string()))
148}
149
150pub async fn get_all_workers<S>(storage: Arc<RwLock<S>>) -> Result<Vec<RunningWorker>, ApiError>
152where
153 S: ListWorkers,
154 S::Error: std::error::Error,
155{
156 storage
157 .read()
158 .await
159 .list_all_workers()
160 .await
161 .map_err(|e| ApiError::BackendError(e.to_string()))
162}
163
164pub async fn fetch_queues<S>(storage: Arc<RwLock<S>>) -> Result<Vec<QueueInfo>, ApiError>
166where
167 S::Error: std::error::Error,
168 S: ListQueues,
169{
170 storage
171 .read()
172 .await
173 .list_queues()
174 .await
175 .map_err(|e| ApiError::BackendError(e.to_string()))
176}
177
178pub async fn overview<S>(storage: Arc<RwLock<S>>) -> Result<Vec<Statistic>, ApiError>
180where
181 S::Error: std::error::Error,
182 S: Metrics,
183{
184 let overview = storage.read().await.global().await;
185 match overview {
186 Ok(overview) => Ok(overview),
187 Err(e) => Err(ApiError::BackendError(e.to_string())),
188 }
189}
190
191#[derive(Debug, Clone)]
193pub struct RawTask<Args>(pub Task<Args>);
194
195impl<Args> Serialize for RawTask<Args>
196where
197 Args: Serialize,
198{
199 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
200 where
201 S: Serializer,
202 {
203 let mut state = serializer.serialize_struct("Task", 2)?;
204 state.serialize_field("args", &self.0.args)?;
205 state.serialize_field("ctx", self.0.ctx().as_ref())?;
206 state.end()
207 }
208}
209
210impl<'de, Args> Deserialize<'de> for RawTask<Args>
211where
212 Args: Deserialize<'de>,
213{
214 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
215 where
216 D: Deserializer<'de>,
217 {
218 #[derive(Deserialize)]
219 struct TaskData<Args> {
220 args: Args,
221 ctx: ExecutionContext,
222 }
223
224 let data = TaskData::deserialize(deserializer)?;
225
226 Ok(Self(Task::new(data.args).map_context(|_| data.ctx.into())))
227 }
228}