apalis_redis/
lib.rs

1#![warn(
2    missing_debug_implementations,
3    missing_docs,
4    rust_2018_idioms,
5    unreachable_pub
6)]
7#![cfg_attr(docsrs, feature(doc_cfg))]
8#![doc = include_str!("../README.md")]
9use std::{
10    io,
11    marker::PhantomData,
12    pin::Pin,
13    task::{Context, Poll, ready},
14};
15
16use apalis_codec::json::JsonCodec;
17use apalis_core::{
18    backend::{
19        Backend, BackendConfig, WireFormatBackend,
20        finalize::Durable,
21        future::BoxSyncFuture,
22        persistence::{Persisted, TaskPersistLayer},
23    },
24    features_table,
25    task::Task,
26    worker::context::WorkerContext,
27};
28
29use futures::{FutureExt, Sink};
30
31pub use redis::{
32    Client, IntoConnectionInfo, RedisError, aio::ConnectionLike, aio::ConnectionManager,
33};
34
35mod config;
36
37/// Raw queries used to make low level redis calls
38pub mod queries;
39
40/// Factory utilities for building [`RedisStorage`] instances.
41pub mod factory;
42
43use serde_json::Value;
44use ulid::Ulid;
45
46pub use crate::config::Config;
47
48pub use crate::{error::Error, persist::RedisPersistence};
49
50mod persist;
51
52mod error;
53
54/// A Redis task type alias
55pub type RedisTask<Args = Vec<u8>> = Task<Args>;
56
57/// Represents a [Backend] that uses Redis for storage.
58///
59#[doc = "# Feature Support\n"]
60#[doc = features_table! {
61    setup = r#"
62    # {
63    #    use apalis_redis::{RedisStorage, Config};
64    #    use std::env;
65    #    let redis_url = env::var("REDIS_URL").expect("REDIS_URL must be set");
66    #    let conn = apalis_redis::connect(redis_url).await.expect("Could not connect");
67    #    const MOD_PATH: &str = module_path!();
68    #    let config = Config::default().queue(MOD_PATH);
69    #    RedisStorage::<u32>::new(conn).with_config(config)
70    # };
71    "#,
72    TaskSink => supported("Ability to push new tasks", true),
73    MakeShared => supported("Share the same connection across multiple workers", false),
74    Workflow => supported("Supports workflows and orchestration", true),
75    WebUI => supported("Supports `apalis-board` for monitoring and managing tasks", true),
76    WaitForCompletion => supported("Wait for tasks to complete without blocking", true),
77    Serialization => supported("Supports multiple serialization formats such as JSON and MessagePack", false),
78    RegisterWorker => supported("Allow registering a worker with the backend", false),
79    ResumeAbandoned => supported("Resume abandoned tasks", false),
80}]
81#[derive(Debug)]
82pub struct RedisStorage<Args, Conn = ConnectionManager>
83where
84    Conn: ConnectionLike + Send + Sync + 'static + Clone,
85{
86    job_type: PhantomData<Args>,
87    persist: Persisted<RedisPersistence<Conn>>,
88    codec: JsonCodec,
89    cleanup: Option<BoxSyncFuture<Result<(), Error>>>,
90}
91
92impl<Args, Conn> Clone for RedisStorage<Args, Conn>
93where
94    Conn: ConnectionLike + Send + Sync + Clone + 'static,
95{
96    fn clone(&self) -> Self {
97        Self {
98            job_type: PhantomData,
99            persist: self.persist.clone(),
100            codec: self.codec.clone(),
101            cleanup: None,
102        }
103    }
104}
105
106impl<T, Conn> RedisStorage<T, Conn>
107where
108    Conn: ConnectionLike + Send + Sync + 'static + Clone,
109{
110    /// Start a new connection
111    pub fn new(conn: Conn) -> RedisStorage<T, Conn> {
112        let config = Config::default().queue(std::any::type_name::<T>());
113        RedisStorage {
114            job_type: PhantomData,
115            persist: Persisted::new(RedisPersistence { config, conn }),
116            codec: JsonCodec::default(),
117            cleanup: None,
118        }
119    }
120
121    /// Customize the backend
122    pub fn with_config(mut self, config: Config) -> RedisStorage<T, Conn> {
123        self.persist.config = config;
124        RedisStorage {
125            job_type: PhantomData,
126            persist: self.persist,
127            codec: self.codec,
128            cleanup: None,
129        }
130    }
131
132    /// Get current connection
133    pub fn get_connection(&mut self) -> &mut Conn {
134        &mut self.persist.conn
135    }
136
137    /// Get the config used by the storage
138    pub fn get_config(&self) -> &Config {
139        &self.persist.config
140    }
141}
142
143impl<Args, Conn> Backend for RedisStorage<Args, Conn>
144where
145    Conn: ConnectionLike + Send + Sync + 'static + Clone,
146{
147    type Task = RedisTask;
148
149    type Error = Error;
150
151    fn poll_ready(
152        &mut self,
153        cx: &mut Context<'_>,
154        worker: &WorkerContext,
155    ) -> Poll<Result<(), Self::Error>> {
156        let heartbeat_interval = self.persist.config.heartbeat_interval;
157        self.persist.poll_ready(cx, worker, heartbeat_interval)
158    }
159
160    fn poll_next(
161        &mut self,
162        cx: &mut Context<'_>,
163        worker: &WorkerContext,
164    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
165        self.persist.poll_next(cx, worker)
166    }
167
168    fn poll_close(
169        &mut self,
170        cx: &mut Context<'_>,
171        worker: &WorkerContext,
172    ) -> Poll<Result<(), Self::Error>> {
173        if self.cleanup.is_none() {
174            ready!(self.persist.poll_close(cx, worker))?;
175            let worker = worker.clone();
176            let config = self.config().clone();
177            let mut conn = self.persist.conn.clone();
178            self.cleanup = Some(
179                async move {
180                    redis::cmd("ZREM")
181                        .arg(config.workers_set())
182                        .arg(config.inflight_set_for(&worker))
183                        .query_async::<u32>(&mut conn)
184                        .await?;
185                    Ok(())
186                }
187                .boxed()
188                .into(),
189            );
190        };
191
192        self.cleanup.as_mut().unwrap().poll_unpin(cx)
193    }
194}
195
196impl<Args, Conn> BackendConfig for RedisStorage<Args, Conn>
197where
198    Conn: ConnectionLike + Send + Sync + 'static + Clone,
199{
200    type Args = Args;
201
202    type Id = Ulid;
203
204    type Kind = Durable;
205
206    type Config = Config;
207
208    type Layer = TaskPersistLayer<JsonCodec<Value>, Value>;
209
210    fn config(&self) -> &Self::Config {
211        &self.persist.config
212    }
213
214    fn middleware(&mut self, _worker: &mut WorkerContext) -> Self::Layer {
215        self.persist
216            .layer(JsonCodec::<Value>::default(), self.config().batch_size)
217            .persist_results(self.config().persist_results)
218            .lock_tasks(self.config().lock_tasks)
219    }
220}
221
222impl<Args, Conn> WireFormatBackend for RedisStorage<Args, Conn>
223where
224    Conn: ConnectionLike + Send + Sync + 'static + Clone,
225{
226    type Codec = JsonCodec;
227
228    type Compact = Vec<u8>;
229
230    fn codec(&self) -> &Self::Codec {
231        &self.codec
232    }
233}
234
235impl<Args, Conn> Sink<RedisTask> for RedisStorage<Args, Conn>
236where
237    Args: Unpin + Send + Sync + 'static,
238    Conn: ConnectionLike + Send + Sync + 'static + Clone + Unpin,
239{
240    type Error = Error;
241
242    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
243        Sink::poll_ready(Pin::new(&mut self.get_mut().persist), cx)
244    }
245
246    fn start_send(self: Pin<&mut Self>, item: RedisTask) -> Result<(), Self::Error> {
247        Sink::start_send(Pin::new(&mut self.get_mut().persist), item)
248    }
249
250    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
251        Sink::poll_flush(Pin::new(&mut self.get_mut().persist), cx)
252    }
253
254    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
255        Sink::poll_close(Pin::new(&mut self.get_mut().persist), cx)
256    }
257}
258
259/// Shorthand to create a client and connect
260pub async fn connect<S: IntoConnectionInfo>(redis: S) -> Result<ConnectionManager, RedisError> {
261    let client = Client::open(redis.into_connection_info()?)?;
262    let conn = client.get_connection_manager().await?;
263    Ok(conn)
264}
265
266fn build_error(message: &str) -> RedisError {
267    RedisError::from(io::Error::new(io::ErrorKind::InvalidData, message))
268}
269
270#[cfg(test)]
271mod tests {
272    use apalis_codec::msgpack::MsgPackCodec;
273    use apalis_core::{
274        backend::{ext::BackendExt, factory::BackendFactory},
275        error::BoxDynError,
276        task::{builder::TaskBuilder, context::TaskContext},
277        worker::ext::parallelize::ParallelizeExt,
278    };
279    use apalis_workflow::{SteppedFlow, WorkflowSink};
280
281    use redis::Client;
282    use std::{env, time::Duration};
283
284    use apalis_core::{
285        backend::TaskSink,
286        worker::{builder::WorkerBuilder, ext::event_listener::EventListenerExt},
287    };
288
289    use crate::factory::RedisStorageFactory;
290
291    use super::*;
292
293    const ITEMS: u32 = 10;
294
295    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
296    async fn basic_worker() {
297        let client = Client::open(env::var("REDIS_URL").unwrap()).unwrap();
298        let conn = client.get_connection_manager().await.unwrap();
299        let config = Config::default().queue("basic:queue").batch_size(100);
300        let mut backend = RedisStorage::new(conn).with_config(config);
301        for i in 0..ITEMS {
302            backend.push(i).await.unwrap();
303        }
304
305        async fn task(task: u32, ctx: TaskContext, wrk: WorkerContext) -> Result<(), BoxDynError> {
306            let handle = std::thread::current();
307            println!("{task:?}, {ctx:?}, Thread: {:?}", handle.id());
308            if task == ITEMS - 1 {
309                wrk.stop().unwrap();
310            }
311            Ok(())
312        }
313
314        let worker = WorkerBuilder::new("rango-tango-basic")
315            .backend(backend)
316            .on_event(|ctx, ev| {
317                println!("CTX {:?}, On Event = {:?}", ctx.name(), ev);
318            })
319            .build(task);
320        worker.run().await.unwrap();
321    }
322
323    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
324    async fn basic_worker_msgpack() {
325        let client = Client::open(env::var("REDIS_URL").unwrap()).unwrap();
326        let conn = client.get_connection_manager().await.unwrap();
327        let config = Config::default().queue("msgpack:queue").batch_size(100);
328        let mut backend = RedisStorage::new(conn)
329            .with_config(config)
330            .with_codec(MsgPackCodec::default());
331
332        for i in 0..ITEMS {
333            let req = TaskBuilder::new(i).build();
334            backend.push_task(req).await.unwrap();
335        }
336
337        async fn task(
338            task: u32,
339            meta: TaskContext,
340            wrk: WorkerContext,
341        ) -> Result<String, BoxDynError> {
342            let handle = std::thread::current();
343            println!("{task:?}, {meta:?}, Thread: {:?}", handle.id());
344            if task == ITEMS - 1 {
345                wrk.stop().unwrap();
346                return Err("Worker stopped!")?;
347            }
348            Ok("Worker".to_owned())
349        }
350
351        let worker = WorkerBuilder::new("rango-tango-msgpack")
352            .backend(backend)
353            .parallelize(tokio::spawn)
354            .on_event(|ctx, ev| {
355                println!("CTX {:?}, On Event = {:?}", ctx.name(), ev);
356            })
357            .build(task);
358        worker.run().await.unwrap();
359    }
360
361    #[tokio::test]
362    async fn shared_workers() {
363        let client = Client::open(env::var("REDIS_URL").unwrap()).unwrap();
364        let mut store = RedisStorageFactory::new(client).await.unwrap();
365
366        let mut string_store = store.create().unwrap();
367        let mut int_store = store.create().unwrap();
368
369        for i in 0..ITEMS {
370            string_store.push(format!("ITEM: {i}")).await.unwrap();
371            int_store.push(i).await.unwrap();
372        }
373
374        async fn task(job: u32, ctx: WorkerContext) -> Result<usize, BoxDynError> {
375            tokio::time::sleep(Duration::from_millis(2)).await;
376            if job == ITEMS - 1 {
377                ctx.stop().unwrap();
378                return Err("Worker stopped!")?;
379            }
380            Ok(job as usize)
381        }
382
383        let int_worker = WorkerBuilder::new("rango-tango-int")
384            .backend(int_store)
385            .on_event(|ctx, ev| {
386                println!("CTX {:?}, On Event = {:?}", ctx.name(), ev);
387            })
388            .build(task)
389            .run();
390
391        let string_worker = WorkerBuilder::new("rango-tango-string")
392            .backend(string_store)
393            .on_event(|ctx, ev| {
394                println!("CTX {:?}, On Event = {:?}", ctx.name(), ev);
395            })
396            .build(|req: String, ctx: WorkerContext| async move {
397                tokio::time::sleep(Duration::from_millis(3)).await;
398                println!("{req}");
399                if req.ends_with(&(ITEMS - 1).to_string()) {
400                    ctx.stop().unwrap();
401                }
402            })
403            .run();
404        let _ = futures::future::try_join(int_worker, string_worker)
405            .await
406            .unwrap();
407    }
408
409    #[tokio::test]
410    async fn workflow() {
411        async fn task1(job: u32) -> Result<Vec<u32>, BoxDynError> {
412            Ok((job..2).collect())
413        }
414
415        async fn task2(_: Vec<u32>) -> Result<usize, BoxDynError> {
416            Ok(42)
417        }
418
419        async fn task3(_: usize, wrk: WorkerContext, _: TaskContext) -> Result<(), io::Error> {
420            wrk.stop().unwrap();
421            Ok(())
422        }
423
424        let work_flow = SteppedFlow::new("sample-workflow")
425            .and_then(task1)
426            .delay_for(Duration::from_millis(1000))
427            .and_then(task2)
428            .and_then(task3);
429
430        let client = Client::open(env::var("REDIS_URL").unwrap()).unwrap();
431        let conn = client.get_connection_manager().await.unwrap();
432        let config = Config::default().queue("workflow:sample");
433        let mut backend = RedisStorage::new(conn)
434            .with_config(config)
435            .poll_with_interval(Duration::from_millis(1000));
436
437        backend.push_start(0u32).await.unwrap();
438
439        let worker = WorkerBuilder::new("rango-tango-workflow")
440            .backend(backend)
441            .on_event(|ctx, ev| {
442                println!("Worker {:?}, On Event = {:?}", ctx.name(), ev);
443            })
444            .build(work_flow);
445        worker.run().await.unwrap();
446    }
447}
apalis_redis/lib.rs

apalis_redis/
lib.rs
