Skip to main content

apalis_postgres/queries/
keep_alive.rs

1use apalis_core::worker::context::WorkerContext;
2use sqlx::Executor;
3
4use crate::error::Error;
5
6/// Heartbeat for denoting liveliness of workers
7pub async fn keep_alive<E>(conn: &mut E, queue: &str, worker: &WorkerContext) -> Result<(), Error>
8where
9    for<'e> &'e mut E: Executor<'e, Database = sqlx::Postgres>,
10{
11    if cfg!(feature = "task-tracking") {
12        keep_alive_tracked(conn, queue, worker).await
13    } else {
14        keep_alive_untracked(conn, queue, worker).await
15    }
16}
17
18async fn keep_alive_tracked<E>(
19    conn: &mut E,
20    queue: &str,
21    worker: &WorkerContext,
22) -> Result<(), Error>
23where
24    for<'e> &'e mut E: Executor<'e, Database = sqlx::Postgres>,
25{
26    #[cfg(feature = "task-tracking")]
27    let tasks = worker
28        .tasks()
29        .map_err(|_| Error::WorkerOutOfSync)?
30        .iter()
31        .map(|task| task.task_id().to_owned())
32        .collect::<Vec<_>>();
33
34    //We will never reach this point if the feature is not enabled
35    #[cfg(not(feature = "task-tracking"))]
36    let tasks: Vec<String> = vec![];
37
38    let worker = worker.name();
39
40    let res = sqlx::query_file!("queries/backend/keep_alive.sql", worker, queue, &tasks)
41        .execute(conn)
42        .await?;
43    if res.rows_affected() != 1 {
44        return Err(Error::WorkerOutOfSync);
45    }
46    Ok(())
47}
48
49async fn keep_alive_untracked<E>(
50    conn: &mut E,
51    queue: &str,
52    worker: &WorkerContext,
53) -> Result<(), Error>
54where
55    for<'e> &'e mut E: Executor<'e, Database = sqlx::Postgres>,
56{
57    let worker = worker.name();
58
59    let res = sqlx::query_file!("queries/backend/keep_alive_untracked.sql", worker, queue)
60        .execute(conn)
61        .await?;
62    if res.rows_affected() != 1 {
63        return Err(Error::WorkerOutOfSync);
64    }
65    Ok(())
66}