apalis_core/backend/
custom.rs

1//! # Custom Backend
2//!
3//! A highly customizable backend for task processing that allows integration with any persistence engine by providing custom fetcher and sink functions.
4//!
5//! ## Overview
6//!
7//! The [`CustomBackend`] struct enables you to define how tasks are fetched from and persisted to
8//! your storage engine.
9//!
10//! You can use the [`BackendBuilder`] to construct a [`CustomBackend`] by
11//! providing the required database, fetcher, sink, and optional configuration and codec.
12//!
13//! ## Usage
14//!
15//! Use [`BackendBuilder`] to configure and build your custom backend:
16//!
17//! ## Example: CustomBackend with Worker
18//!
19//! ```rust
20//! # use std::collections::VecDeque;
21//! # use std::sync::Arc;
22//! # use futures_util::{lock::Mutex, sink, stream, sink::SinkExt};
23//! # use apalis_core::backend::custom::{BackendBuilder, CustomBackend};
24//! # use apalis_core::task::Task;
25//! # use apalis_core::task::task_id::{RandomId,TaskId};
26//! # use apalis_core::task::builder::TaskBuilder;
27//! # use apalis_core::worker::builder::WorkerBuilder;
28//! # use apalis_core::worker::context::WorkerContext;
29//! # use apalis_core::error::BoxDynError;
30//! # use std::time::Duration;
31//! # use futures_util::StreamExt;
32//! # use futures_util::FutureExt;
33//! # use apalis_core::backend::TaskSink;
34//! #[tokio::main]
35//! async fn main() {
36//!     // Create a memory-backed VecDeque
37//!     let memory = Arc::new(Mutex::new(VecDeque::<Task<u32>>::new()));
38//!
39//!     // Build the custom backend
40//!     let mut backend = BackendBuilder::new()
41//!         .database(memory)
42//!         .fetcher(|memory, _, _| {
43//!             stream::unfold(memory.clone(), |p| async move {
44//!                 let mut memory = p.lock().await;
45//!                 let item = memory.pop_front();
46//!                 drop(memory);
47//!                 match item {
48//!                     Some(item) => Some((Ok::<_, BoxDynError>(Some(item)), p)),
49//!                     None => Some((Ok::<_, BoxDynError>(None), p)),
50//!                 }
51//!             })
52//!             .boxed()
53//!         })
54//!         .sink(|memory, _| {
55//!             sink::unfold(memory.clone(), move |p, item: Task<_>| {
56//!                 async move {
57//!                     let mut memory = p.lock().await;
58//!                     let item = item
59//!                            .into_builder()
60//!                            .task_id(TaskId::String(RandomId::default().to_string()))
61//!                            .build();
62//!                     memory.push_back(item);
63//!                     drop(memory);
64//!                     Ok::<_, BoxDynError>(p)
65//!                 }
66//!                 .boxed()
67//!             })
68//!         })
69//!         .build()
70//!         .unwrap();
71//!
72//!     // Add a task to the backend;
73//!     backend.push(42).await.unwrap();
74//!
75//!     // Define the task handler
76//!     async fn task(task: u32, worker: WorkerContext) -> Result<(), BoxDynError> {
77//!         tokio::time::sleep(Duration::from_secs(1)).await;
78//! #       worker.stop().unwrap();
79//!         Ok(())
80//!     }
81//!
82//!     // Build and run the worker
83//!     let worker = WorkerBuilder::new("custom-worker")
84//!         .backend(backend)
85//!         .build(task);
86//!     worker.run().await.unwrap();
87//! }
88//! ```
89//!
90//! ## Features
91//!
92//! - **Custom Fetcher**: Define how jobs are fetched from your storage.
93//! - **Custom Sink**: Define how jobs are persisted to your storage.
94//! - **Configurable**: Pass custom configuration to your backend.
95//!
96use futures_sink::Sink;
97use futures_util::SinkExt;
98use futures_util::{Stream, StreamExt};
99use std::pin::Pin;
100use std::sync::Arc;
101use std::task::{Context, Poll};
102use std::{fmt, marker::PhantomData};
103use thiserror::Error;
104use tower_layer::Identity;
105
106use crate::backend::BackendConfig;
107use crate::backend::finalize::Ephemeral;
108use crate::error::BoxDynError;
109use crate::task::task_id::RandomId;
110use crate::{backend::Backend, task::Task, worker::context::WorkerContext};
111
112type Fetcher<DB, Config, Fetch> =
113    Arc<Box<dyn Fn(&mut DB, &Config, &WorkerContext) -> Fetch + Send + Sync>>;
114
115type Sinker<DB, Config, Sink> = Arc<Box<dyn Fn(&mut DB, &Config) -> Sink + Send + Sync>>;
116
117/// A highly customizable backend for integration with any persistence engine
118///
119/// This backend allows you to define how tasks are fetched from and persisted to your storage,
120/// meaning you can use it to integrate with existing systems.
121///
122/// # Example
123/// ```rust,ignore
124/// let backend = BackendBuilder::new()
125///     .database(my_db)
126///     .fetcher(my_fetcher_fn)
127///     .sink(my_sink_fn)
128///     .build()
129///     .unwrap();
130/// ```
131#[doc = features_table! {
132    setup = "{ unreachable!() }",
133    TaskSink => supported("Ability to push new tasks", false),
134    Serialization => supported("Serialization support for arguments", false),
135    FetchById => not_supported("Allow fetching a task by its ID"),
136    RegisterWorker => not_implemented("Allow registering a worker with the backend"),
137    PipeExt => limited("Allow other backends to pipe to this backend", false), // Would require Clone,
138    BackendFactory => not_implemented("Share the same [`CustomBackend`] across multiple workers", false),
139    Workflow => not_implemented("Flexible enough to support workflows"),
140    WaitForCompletion => not_implemented("Wait for tasks to complete without blocking"), // Would require Clone
141    ResumeById => not_supported("Resume a task by its ID"),
142    ResumeAbandoned => not_supported("Resume abandoned tasks"),
143    ListWorkers => not_implemented("List all workers registered with the backend"),
144    ListTasks => not_implemented("List all tasks in the backend"),
145}]
146#[pin_project::pin_project]
147#[must_use = "Custom backends must be polled or used as a sink"]
148pub struct CustomBackend<Args, DB, Fetch, Sink, Config = ()> {
149    _marker: PhantomData<Args>,
150    db: DB,
151    fetcher: Fetcher<DB, Config, Fetch>,
152    sinker: Sinker<DB, Config, Sink>,
153    #[pin]
154    current_sink: Sink,
155    config: Config,
156
157    stream: Option<Fetch>,
158}
159
160impl<Args, DB, Fetch, Sink, Config> Clone for CustomBackend<Args, DB, Fetch, Sink, Config>
161where
162    DB: Clone,
163    Config: Clone,
164{
165    fn clone(&self) -> Self {
166        let mut db = self.db.clone();
167        let current_sink = (self.sinker)(&mut db, &self.config);
168        Self {
169            _marker: PhantomData,
170            db,
171            fetcher: Arc::clone(&self.fetcher),
172            sinker: Arc::clone(&self.sinker),
173            current_sink,
174            config: self.config.clone(),
175
176            stream: None,
177        }
178    }
179}
180
181impl<Args, DB, Fetch, Sink, Config> fmt::Debug for CustomBackend<Args, DB, Fetch, Sink, Config>
182where
183    DB: fmt::Debug,
184    Config: fmt::Debug,
185{
186    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187        f.debug_struct("CustomBackend")
188            .field(
189                "_marker",
190                &format_args!("PhantomData<({})>", std::any::type_name::<Args>()),
191            )
192            .field("db", &self.db)
193            .field("fetcher", &"Fn(&mut DB, &Config, &WorkerContext) -> Fetch")
194            .field("sink", &"Fn(&mut DB, &Config) -> Sink")
195            .field("config", &self.config)
196            .finish()
197    }
198}
199
200type FetcherBuilder<DB, Config, Fetch> =
201    Box<dyn Fn(&mut DB, &Config, &WorkerContext) -> Fetch + Send + Sync + 'static>;
202
203type SinkerBuilder<DB, Config, Sink> =
204    Box<dyn Fn(&mut DB, &Config) -> Sink + Send + Sync + 'static>;
205
206/// Builder for [`CustomBackend`]
207///
208/// Lets you set the database, fetcher, sink, codec, and config
209pub struct BackendBuilder<Args, DB, Fetch, Sink, Config = ()> {
210    _marker: PhantomData<Args>,
211    database: Option<DB>,
212    fetcher: Option<FetcherBuilder<DB, Config, Fetch>>,
213    sink: Option<SinkerBuilder<DB, Config, Sink>>,
214    config: Option<Config>,
215}
216
217impl<Args, DB, Fetch, Sink, Config> fmt::Debug for BackendBuilder<Args, DB, Fetch, Sink, Config>
218where
219    DB: fmt::Debug,
220    Config: fmt::Debug,
221{
222    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
223        f.debug_struct("BackendBuilder")
224            .field(
225                "_marker",
226                &format_args!("PhantomData<({})>", std::any::type_name::<Args>(),),
227            )
228            .field("database", &self.database)
229            .field("fetcher", &self.fetcher.as_ref().map(|_| "Some(fn)"))
230            .field("sink", &self.sink.as_ref().map(|_| "Some(fn)"))
231            .field("config", &self.config)
232            .finish()
233    }
234}
235
236impl<Args, DB, Fetch, Sink, Config> Default for BackendBuilder<Args, DB, Fetch, Sink, Config> {
237    fn default() -> Self {
238        Self {
239            _marker: PhantomData,
240            database: None,
241            fetcher: None,
242            sink: None,
243            config: None,
244        }
245    }
246}
247
248impl<Args, DB, Fetch, Sink> BackendBuilder<Args, DB, Fetch, Sink, ()> {
249    /// Create a new `BackendBuilder` instance
250    #[must_use]
251    pub fn new() -> Self {
252        Self::new_with_cfg(())
253    }
254
255    /// Create a new `BackendBuilder` instance with custom configuration
256    pub fn new_with_cfg<Config>(config: Config) -> BackendBuilder<Args, DB, Fetch, Sink, Config> {
257        BackendBuilder {
258            config: Some(config),
259            ..Default::default()
260        }
261    }
262}
263
264impl<Args, DB, Fetch, Sink, Config> BackendBuilder<Args, DB, Fetch, Sink, Config> {
265    /// The custom backend persistence engine
266    #[must_use]
267    pub fn database(mut self, db: DB) -> Self {
268        self.database = Some(db);
269        self
270    }
271
272    /// The fetcher function to retrieve tasks from the database
273    #[must_use]
274    pub fn fetcher<F: Fn(&mut DB, &Config, &WorkerContext) -> Fetch + Send + Sync + 'static>(
275        mut self,
276        fetcher: F,
277    ) -> Self {
278        self.fetcher = Some(Box::new(fetcher));
279        self
280    }
281
282    /// The sink function to persist tasks to the database
283    #[must_use]
284    pub fn sink<F: Fn(&mut DB, &Config) -> Sink + Send + Sync + 'static>(
285        mut self,
286        sink: F,
287    ) -> Self {
288        self.sink = Some(Box::new(sink));
289        self
290    }
291
292    #[allow(clippy::type_complexity)]
293    /// Build the `CustomBackend` instance
294    pub fn build(self) -> Result<CustomBackend<Args, DB, Fetch, Sink, Config>, BuildError> {
295        let mut db = self.database.ok_or(BuildError::MissingDb)?;
296        let config = self.config.ok_or(BuildError::MissingConfig)?;
297        let sink_fn = self.sink.ok_or(BuildError::MissingSink)?;
298        let sink = sink_fn(&mut db, &config);
299
300        Ok(CustomBackend {
301            _marker: PhantomData,
302            db,
303            fetcher: self
304                .fetcher
305                .map(Arc::new)
306                .ok_or(BuildError::MissingFetcher)?,
307            current_sink: sink,
308            sinker: Arc::new(sink_fn),
309            config,
310            stream: None,
311        })
312    }
313}
314
315/// Errors encountered building a `CustomBackend`
316#[derive(Debug, Error)]
317#[non_exhaustive]
318pub enum BuildError {
319    /// Missing database db
320    #[error("Database db is required")]
321    MissingDb,
322    /// Missing fetcher function
323    #[error("Fetcher is required")]
324    MissingFetcher,
325    /// Missing sink function
326    #[error("Sink is required")]
327    MissingSink,
328    /// Missing configuration
329    #[error("Config is required")]
330    MissingConfig,
331}
332
333/// Errors encountered while using the `CustomBackend`
334#[derive(Debug, Error)]
335#[non_exhaustive]
336pub enum CustomBackendError {
337    /// Inner error
338    #[error("Inner error: {0}")]
339    Inner(#[from] BoxDynError),
340}
341
342impl<Args: 'static, DB, Fetch, S, E, Config> Backend for CustomBackend<Args, DB, Fetch, S, Config>
343where
344    Fetch: Stream<Item = Result<Option<Task<Args>>, E>> + Unpin + Send + 'static,
345    S: Sink<Task<Args>, Error = E> + Unpin + Send + 'static,
346    E: Into<BoxDynError>,
347{
348    type Task = Task<Args>;
349
350    type Error = CustomBackendError;
351
352    fn poll_ready(
353        &mut self,
354        _: &mut Context<'_>,
355        _: &WorkerContext,
356    ) -> Poll<Result<(), Self::Error>> {
357        Poll::Ready(Ok(()))
358    }
359
360    fn poll_next(
361        &mut self,
362        cx: &mut Context<'_>,
363        worker: &WorkerContext,
364    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
365        if self.stream.is_none() {
366            self.stream = Some((self.fetcher)(&mut self.db, &self.config, worker));
367        }
368        let stream = self.stream.as_mut().unwrap();
369        stream.poll_next_unpin(cx).map(|item| {
370            item.transpose()
371                .map(|res| res.flatten())
372                .map_err(|e| CustomBackendError::Inner(e.into()))
373                .transpose()
374        })
375    }
376
377    fn poll_close(
378        &mut self,
379        cx: &mut Context<'_>,
380        _: &WorkerContext,
381    ) -> Poll<Result<(), Self::Error>> {
382        self.current_sink
383            .poll_close_unpin(cx)
384            .map_err(|e| CustomBackendError::Inner(e.into()))
385    }
386}
387
388impl<Args, DB, Fetch, S, Config, E> BackendConfig for CustomBackend<Args, DB, Fetch, S, Config>
389where
390    Fetch: Stream<Item = Result<Option<Task<Args>>, E>> + Unpin + Send + 'static,
391    S: Sink<Task<Args>, Error = E> + Unpin + Send + 'static,
392    E: Into<BoxDynError>,
393    Args: 'static,
394{
395    type Id = RandomId;
396    type Args = Args;
397    type Kind = Ephemeral;
398
399    type Layer = Identity;
400
401    type Config = Config;
402
403    fn config(&self) -> &Self::Config {
404        &self.config
405    }
406
407    fn middleware(&mut self, _: &mut WorkerContext) -> Self::Layer {
408        Identity::new()
409    }
410}
411
412impl<Args, DB, Fetch, S, Config> Sink<Task<Args>> for CustomBackend<Args, DB, Fetch, S, Config>
413where
414    S: Sink<Task<Args>>,
415    S::Error: Into<BoxDynError>,
416{
417    type Error = CustomBackendError;
418
419    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
420        self.project()
421            .current_sink
422            .poll_ready_unpin(cx)
423            .map_err(|e| CustomBackendError::Inner(e.into()))
424    }
425
426    fn start_send(self: Pin<&mut Self>, item: Task<Args>) -> Result<(), Self::Error> {
427        self.project()
428            .current_sink
429            .start_send_unpin(item)
430            .map_err(|e| CustomBackendError::Inner(e.into()))
431    }
432
433    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
434        self.project()
435            .current_sink
436            .poll_flush_unpin(cx)
437            .map_err(|e| CustomBackendError::Inner(e.into()))
438    }
439
440    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
441        self.project()
442            .current_sink
443            .poll_close_unpin(cx)
444            .map_err(|e| CustomBackendError::Inner(e.into()))
445    }
446}
447
448#[cfg(test)]
449mod tests {
450    use std::{collections::VecDeque, time::Duration};
451
452    use futures_util::{FutureExt, lock::Mutex, sink, stream};
453
454    use crate::{
455        backend::TaskSink,
456        error::BoxDynError,
457        task::task_id::{RandomId, TaskId},
458        worker::{builder::WorkerBuilder, ext::event_listener::EventListenerExt},
459    };
460
461    use super::*;
462
463    const ITEMS: u32 = 10;
464
465    #[tokio::test]
466    async fn basic_custom_backend() {
467        let memory: Arc<Mutex<VecDeque<Task<u32>>>> = Arc::new(Mutex::new(VecDeque::new()));
468
469        let mut backend = BackendBuilder::new()
470            .database(memory)
471            .fetcher(|db, _, _| {
472                stream::unfold(db.clone(), |p| async move {
473                    tokio::time::sleep(Duration::from_millis(100)).await; // Debounce
474                    let mut db = p.try_lock().unwrap();
475                    let item = db.pop_front();
476                    drop(db);
477                    match item {
478                        Some(item) => Some((Ok::<_, CustomBackendError>(Some(item)), p)),
479                        None => Some((Ok::<_, CustomBackendError>(None), p)),
480                    }
481                })
482                .boxed()
483            })
484            .sink(|db, _| {
485                sink::unfold(db.clone(), move |p, item: Task<u32>| {
486                    async move {
487                        let mut db = p.try_lock().unwrap();
488                        let item = item
489                            .into_builder()
490                            .task_id(TaskId::String(RandomId::default().to_string()))
491                            .build();
492                        db.push_back(item);
493                        drop(db);
494                        Ok::<_, CustomBackendError>(p)
495                    }
496                    .boxed()
497                })
498            })
499            .build()
500            .unwrap();
501
502        for i in 0..ITEMS {
503            backend.push(i).await.unwrap();
504        }
505
506        async fn task(task: u32, worker: WorkerContext) -> Result<(), BoxDynError> {
507            tokio::time::sleep(Duration::from_secs(1)).await;
508            if task == ITEMS - 1 {
509                worker.stop().unwrap();
510                return Err("Worker stopped!")?;
511            }
512            Ok(())
513        }
514
515        let worker = WorkerBuilder::new("rango-tango")
516            .backend(backend)
517            .on_event(|worker, ev| {
518                println!("On Event = {ev:?} from {}", worker.name());
519            })
520            .build(task);
521        worker.run().await.unwrap();
522    }
523}
apalis_core/backend/custom.rs

apalis_core/backend/
custom.rs
