apalis_redis/
persist.rs

1use apalis_core::{
2    backend::{
3        WorkerFilter,
4        persistence::{Persistence, TaskEvent},
5    },
6    task::Task,
7    worker::context::WorkerContext,
8};
9use redis::aio::{ConnectionLike, ConnectionManager};
10use serde_json::Value;
11
12use crate::{
13    Config, RedisTask,
14    error::Error,
15    queries::{
16        self, keep_alive, push_tasks, reenqueue_orphaned, register_worker, release_leases,
17        renew_leases,
18    },
19};
20
21/// The default persistence implementation
22#[derive(Debug)]
23pub struct RedisPersistence<Conn = ConnectionManager> {
24    pub(crate) conn: Conn,
25    pub(crate) config: Config,
26}
27
28impl<C: Clone> Clone for RedisPersistence<C> {
29    fn clone(&self) -> Self {
30        RedisPersistence {
31            conn: self.conn.clone(),
32            config: self.config.clone(),
33        }
34    }
35}
36
37impl<C> Persistence for RedisPersistence<C>
38where
39    C: ConnectionLike + Send + Sync + 'static + Clone,
40{
41    type Compact = Vec<u8>;
42    type Error = Error;
43    type Response = Value;
44    async fn register(&mut self, worker: &WorkerContext) -> Result<(), Error> {
45        let count = reenqueue_orphaned(
46            &mut self.conn,
47            &self.config,
48            &WorkerFilter::Only(worker.name().to_owned()),
49        )
50        .await?;
51        if count > 0 {
52            tracing::debug!("{count} re-enqueued orphaned tasks",);
53        }
54        register_worker(&mut self.conn, worker, &self.config).await?;
55        tracing::debug!("registered worker successfully");
56        Ok(())
57    }
58    async fn heartbeat(&mut self, worker: &WorkerContext) -> Result<(), Error> {
59        let config = &self.config;
60        keep_alive(&mut self.conn, worker, config).await?;
61        if self.config.lock_tasks {
62            let count = renew_leases(&mut self.conn, worker, config).await?;
63            tracing::debug!("renewed leases for {count} tasks");
64        }
65        let count = reenqueue_orphaned(
66            &mut self.conn,
67            &self.config,
68            &WorkerFilter::AllExcept(worker.name().to_owned()),
69        )
70        .await?;
71
72        if count > 0 {
73            tracing::debug!("re-enqueued {count} orphaned task(s)");
74        }
75        Ok(())
76    }
77    async fn fetch_next(&mut self, worker: &WorkerContext) -> Result<Vec<RedisTask>, Error> {
78        queries::fetch_next(&mut self.conn, worker, &self.config).await
79    }
80
81    async fn handle_events(
82        &mut self,
83        messages: Vec<TaskEvent<Self::Response>>,
84        worker: &WorkerContext,
85    ) -> Result<(), Error> {
86        let lock_ids = messages
87            .iter()
88            .filter_map(|msg| {
89                if let TaskEvent::Lock { task_id, .. } = msg {
90                    Some(task_id.to_string())
91                } else {
92                    None
93                }
94            })
95            .collect::<Vec<_>>();
96
97        let ack_payloads = messages
98            .iter()
99            .filter_map(|msg| {
100                if let TaskEvent::Complete(payload) = msg {
101                    Some(payload)
102                } else {
103                    None
104                }
105            })
106            .collect::<Vec<_>>();
107        if lock_ids.is_empty() && ack_payloads.is_empty() {
108            return Ok(());
109        }
110
111        tracing::debug!(
112            "Processing {} messages ({} locks, {} acks)",
113            messages.len(),
114            lock_ids.len(),
115            ack_payloads.len()
116        );
117
118        if !lock_ids.is_empty() {
119            // TODO: Better error handling rather than stopping everything
120            queries::acquire_leases(&mut self.conn, worker, &self.config, &lock_ids).await?;
121        }
122        if !ack_payloads.is_empty() {
123            queries::handle_results(&mut self.conn, &ack_payloads, worker, &self.config).await?;
124        }
125
126        Ok(())
127    }
128
129    async fn reenqueue_abandoned(
130        &mut self,
131        tasks: Vec<RedisTask>,
132        worker: &WorkerContext,
133    ) -> Result<u64, Error> {
134        let config = &self.config;
135        let count = release_leases(&mut self.conn, config, worker, &tasks).await?;
136        if self.config.lock_tasks && count as usize != tasks.len() {
137            return Err(Error::ReenqueueMismatch {
138                queued: tasks.len(),
139                abandoned: count as usize,
140            });
141        }
142        Ok(count)
143    }
144    async fn push_tasks(&mut self, tasks: Vec<Task<Self::Compact>>) -> Result<(), Self::Error> {
145        let config = &self.config;
146        push_tasks(&mut self.conn, config, &tasks).await?;
147        Ok(())
148    }
149}
apalis_redis/persist.rs

apalis_redis/
persist.rs
