laterite_core/queue/
database.rs

1//! The `database` driver: jobs in the `jobs` table, the failed list in
2//! `failed_jobs`, claimed by one conditional update that only one worker
3//! can win.
4
5use std::time::Duration;
6
7use async_trait::async_trait;
8use chrono::{DateTime, Utc};
9use sea_query::{Expr, Iden, Order, Query};
10use sqlx::Row;
11
12use super::{time_text, FailedJob, QueueCount, QueueDriver, Queued};
13use crate::error::{CoreError, CoreResult};
14use crate::query::{bind_values, build, insert_returning_id, AnyRowExt};
15use crate::Db;
16
17#[derive(Iden)]
18pub(crate) enum Jobs {
19    Table,
20    Id,
21    Queue,
22    Name,
23    Payload,
24    Attempts,
25    AvailableAt,
26    ReservedAt,
27    ReservedBy,
28    UniqueKey,
29    CreatedAt,
30}
31
32#[derive(Iden)]
33pub(crate) enum FailedJobs {
34    Table,
35    Id,
36    Queue,
37    Name,
38    Payload,
39    Error,
40    FailedAt,
41}
42
43/// Jobs in the application's database.
44#[derive(Clone)]
45pub struct DatabaseQueue {
46    db: Db,
47}
48
49impl DatabaseQueue {
50    pub fn new(db: Db) -> Self {
51        Self { db }
52    }
53
54    fn due(now: &str, stale_before: &str) -> sea_query::SimpleExpr {
55        Expr::col(Jobs::ReservedAt)
56            .is_null()
57            .and(Expr::col(Jobs::AvailableAt).lte(now))
58            .or(Expr::col(Jobs::ReservedAt).lte(stale_before))
59    }
60
61    async fn find(&self, id: i64) -> CoreResult<Option<Queued>> {
62        let (sql, values) = build(
63            self.db.backend,
64            Query::select()
65                .columns([
66                    Jobs::Id,
67                    Jobs::Queue,
68                    Jobs::Name,
69                    Jobs::Payload,
70                    Jobs::Attempts,
71                    Jobs::AvailableAt,
72                    Jobs::ReservedAt,
73                    Jobs::ReservedBy,
74                    Jobs::CreatedAt,
75                ])
76                .from(Jobs::Table)
77                .and_where(Expr::col(Jobs::Id).eq(id))
78                .to_owned(),
79        );
80        let row = bind_values(sqlx::query(&sql), values)
81            .fetch_optional(&self.db.pool)
82            .await?;
83        row.map(|r| queued_of(&r)).transpose()
84    }
85}
86
87fn queued_of(r: &sqlx::any::AnyRow) -> CoreResult<Queued> {
88    let payload = r.get_text("payload")?;
89    Ok(Queued {
90        id: r.try_get::<i64, _>("id")?,
91        queue: r.get_text("queue")?,
92        name: r.get_text("name")?,
93        payload: serde_json::from_str(&payload)
94            .map_err(|e| CoreError::Internal(format!("job payload is not JSON: {e}")))?,
95        attempts: r.try_get::<i64, _>("attempts")?.max(0) as u32,
96        available_at: r.get_text("available_at")?,
97        reserved_at: r.get_text_opt("reserved_at")?,
98        reserved_by: r.get_text_opt("reserved_by")?,
99        created_at: r.get_text("created_at")?,
100    })
101}
102
103#[async_trait]
104impl QueueDriver for DatabaseQueue {
105    fn name(&self) -> &'static str {
106        "database"
107    }
108
109    async fn push(
110        &self,
111        queue: &str,
112        name: &str,
113        payload: &serde_json::Value,
114        available_at: DateTime<Utc>,
115        unique_key: Option<&str>,
116    ) -> CoreResult<Option<i64>> {
117        if let Some(key) = unique_key {
118            let (sql, values) = build(
119                self.db.backend,
120                Query::select()
121                    .expr(Expr::col(Jobs::Id).count())
122                    .from(Jobs::Table)
123                    .and_where(Expr::col(Jobs::UniqueKey).eq(key))
124                    .to_owned(),
125            );
126            let row = bind_values(sqlx::query(&sql), values)
127                .fetch_one(&self.db.pool)
128                .await?;
129            if row.try_get::<i64, _>(0)? > 0 {
130                return Ok(None);
131            }
132        }
133        let text = serde_json::to_string(payload)
134            .map_err(|e| CoreError::Internal(format!("job payload cannot serialize: {e}")))?;
135        let stmt = Query::insert()
136            .into_table(Jobs::Table)
137            .columns([
138                Jobs::Queue,
139                Jobs::Name,
140                Jobs::Payload,
141                Jobs::Attempts,
142                Jobs::AvailableAt,
143                Jobs::UniqueKey,
144                Jobs::CreatedAt,
145            ])
146            .values_panic([
147                queue.into(),
148                name.into(),
149                text.into(),
150                0i64.into(),
151                time_text(available_at).into(),
152                unique_key.into(),
153                time_text(Utc::now()).into(),
154            ])
155            .to_owned();
156        let id = insert_returning_id(&self.db, stmt, Jobs::Id).await?;
157        Ok(Some(id))
158    }
159
160    async fn claim(
161        &self,
162        queues: &[String],
163        retry_after: Duration,
164        worker: &str,
165    ) -> CoreResult<Option<Queued>> {
166        let stale = chrono::Duration::from_std(retry_after).unwrap_or(chrono::Duration::MAX);
167        // Two workers may pick the same candidate; the conditional update
168        // lets one of them have it and the other looks again.
169        for _ in 0..5 {
170            let now = time_text(Utc::now());
171            let stale_before = time_text(Utc::now() - stale);
172            let (sql, values) = build(
173                self.db.backend,
174                Query::select()
175                    .column(Jobs::Id)
176                    .from(Jobs::Table)
177                    .and_where(Expr::col(Jobs::Queue).is_in(queues.iter().cloned()))
178                    .and_where(Self::due(&now, &stale_before))
179                    .order_by(Jobs::Id, Order::Asc)
180                    .limit(1)
181                    .to_owned(),
182            );
183            let row = bind_values(sqlx::query(&sql), values)
184                .fetch_optional(&self.db.pool)
185                .await?;
186            let Some(row) = row else {
187                return Ok(None);
188            };
189            let id = row.try_get::<i64, _>("id")?;
190            let (sql, values) = build(
191                self.db.backend,
192                Query::update()
193                    .table(Jobs::Table)
194                    .value(Jobs::ReservedAt, now.clone())
195                    .value(Jobs::ReservedBy, worker)
196                    .value(Jobs::Attempts, Expr::col(Jobs::Attempts).add(1))
197                    .and_where(Expr::col(Jobs::Id).eq(id))
198                    .and_where(Self::due(&now, &stale_before))
199                    .to_owned(),
200            );
201            let done = bind_values(sqlx::query(&sql), values)
202                .execute(&self.db.pool)
203                .await?;
204            if done.rows_affected() == 1 {
205                return self.find(id).await;
206            }
207        }
208        Ok(None)
209    }
210
211    async fn release(&self, id: i64, available_at: DateTime<Utc>) -> CoreResult<()> {
212        let (sql, values) = build(
213            self.db.backend,
214            Query::update()
215                .table(Jobs::Table)
216                .value(Jobs::ReservedAt, Option::<String>::None)
217                .value(Jobs::ReservedBy, Option::<String>::None)
218                .value(Jobs::AvailableAt, time_text(available_at))
219                .and_where(Expr::col(Jobs::Id).eq(id))
220                .to_owned(),
221        );
222        bind_values(sqlx::query(&sql), values)
223            .execute(&self.db.pool)
224            .await?;
225        Ok(())
226    }
227
228    async fn delete(&self, id: i64) -> CoreResult<()> {
229        let (sql, values) = build(
230            self.db.backend,
231            Query::delete()
232                .from_table(Jobs::Table)
233                .and_where(Expr::col(Jobs::Id).eq(id))
234                .to_owned(),
235        );
236        bind_values(sqlx::query(&sql), values)
237            .execute(&self.db.pool)
238            .await?;
239        Ok(())
240    }
241
242    async fn fail(&self, job: &Queued, error: &str) -> CoreResult<i64> {
243        let text = serde_json::to_string(&job.payload)
244            .map_err(|e| CoreError::Internal(format!("job payload cannot serialize: {e}")))?;
245        let stmt = Query::insert()
246            .into_table(FailedJobs::Table)
247            .columns([
248                FailedJobs::Queue,
249                FailedJobs::Name,
250                FailedJobs::Payload,
251                FailedJobs::Error,
252                FailedJobs::FailedAt,
253            ])
254            .values_panic([
255                job.queue.as_str().into(),
256                job.name.as_str().into(),
257                text.into(),
258                error.into(),
259                time_text(Utc::now()).into(),
260            ])
261            .to_owned();
262        let failed_id = insert_returning_id(&self.db, stmt, FailedJobs::Id).await?;
263        self.delete(job.id).await?;
264        Ok(failed_id)
265    }
266
267    async fn counts(&self) -> CoreResult<Vec<QueueCount>> {
268        let (sql, values) = build(
269            self.db.backend,
270            Query::select()
271                .columns([Jobs::Queue, Jobs::ReservedAt])
272                .from(Jobs::Table)
273                .to_owned(),
274        );
275        let rows = bind_values(sqlx::query(&sql), values)
276            .fetch_all(&self.db.pool)
277            .await?;
278        let mut counts: std::collections::BTreeMap<String, QueueCount> = Default::default();
279        for r in rows {
280            let queue = r.get_text("queue")?;
281            let entry = counts.entry(queue.clone()).or_insert(QueueCount {
282                queue,
283                pending: 0,
284                reserved: 0,
285            });
286            if r.get_text_opt("reserved_at")?.is_some() {
287                entry.reserved += 1;
288            } else {
289                entry.pending += 1;
290            }
291        }
292        Ok(counts.into_values().collect())
293    }
294
295    async fn failed(&self) -> CoreResult<Vec<FailedJob>> {
296        let (sql, values) = build(
297            self.db.backend,
298            Query::select()
299                .columns([
300                    FailedJobs::Id,
301                    FailedJobs::Queue,
302                    FailedJobs::Name,
303                    FailedJobs::Payload,
304                    FailedJobs::Error,
305                    FailedJobs::FailedAt,
306                ])
307                .from(FailedJobs::Table)
308                .order_by(FailedJobs::Id, Order::Desc)
309                .to_owned(),
310        );
311        let rows = bind_values(sqlx::query(&sql), values)
312            .fetch_all(&self.db.pool)
313            .await?;
314        let mut out = Vec::with_capacity(rows.len());
315        for r in rows {
316            let payload = r.get_text("payload")?;
317            out.push(FailedJob {
318                id: r.try_get::<i64, _>("id")?,
319                queue: r.get_text("queue")?,
320                name: r.get_text("name")?,
321                payload: serde_json::from_str(&payload)
322                    .map_err(|e| CoreError::Internal(format!("job payload is not JSON: {e}")))?,
323                error: r.get_text("error")?,
324                failed_at: r.get_text("failed_at")?,
325            });
326        }
327        Ok(out)
328    }
329
330    async fn retry(&self, id: Option<i64>) -> CoreResult<usize> {
331        let failed = self.failed().await?;
332        let mut moved = 0;
333        for job in failed
334            .into_iter()
335            .filter(|f| id.is_none_or(|id| f.id == id))
336        {
337            self.push(&job.queue, &job.name, &job.payload, Utc::now(), None)
338                .await?;
339            self.forget(Some(job.id)).await?;
340            moved += 1;
341        }
342        Ok(moved)
343    }
344
345    async fn forget(&self, id: Option<i64>) -> CoreResult<usize> {
346        let (sql, values) = {
347            let mut stmt = Query::delete();
348            stmt.from_table(FailedJobs::Table);
349            if let Some(id) = id {
350                stmt.and_where(Expr::col(FailedJobs::Id).eq(id));
351            }
352            build(self.db.backend, stmt)
353        };
354        let done = bind_values(sqlx::query(&sql), values)
355            .execute(&self.db.pool)
356            .await?;
357        Ok(done.rows_affected() as usize)
358    }
359
360    async fn clear(&self, queue: Option<&str>) -> CoreResult<usize> {
361        let (sql, values) = {
362            let mut stmt = Query::delete();
363            stmt.from_table(Jobs::Table)
364                .and_where(Expr::col(Jobs::ReservedAt).is_null());
365            if let Some(queue) = queue {
366                stmt.and_where(Expr::col(Jobs::Queue).eq(queue));
367            }
368            build(self.db.backend, stmt)
369        };
370        let done = bind_values(sqlx::query(&sql), values)
371            .execute(&self.db.pool)
372            .await?;
373        Ok(done.rows_affected() as usize)
374    }
375}
laterite_core/queue/database.rs

laterite_core/queue/
database.rs
