Skip to main content

apalis_core/backend/ext/
mod.rs

1//! Extension traits and combinators for [`Backend`] implementations.
2//!
3//! This module provides additional functionality for [`Backend`] implementations,
4//! including:
5//!
6//! - [`BackendExt`]: Extension methods for transforming, composing, and managing
7//!   backends.
8//! - [`InspectErr`]: A wrapper that allows inspection of errors produced by a backend.
9//! - [`MapErr`]: A wrapper that maps backend errors from one type to another.
10//! - [`Pipe`]: A utility for piping tasks from one backend to another.
11//! - [`BeforeStart`]: A lifecycle wrapper that runs an action before the backend starts.
12//! - [`AfterStart`]: A lifecycle wrapper that runs an action after the backend starts.
13//! - [`BeforeStop`]: A lifecycle wrapper that runs an action before the backend stops.
14//! - [`AfterStop`]: A lifecycle wrapper that runs an action after the backend stops.
15use std::{
16    task::{Context, Poll},
17    time::Duration,
18};
19
20use futures_core::Stream;
21
22#[cfg(feature = "tracing")]
23use crate::backend::ext::instrument::Instrumented;
24use crate::{
25    backend::{
26        Backend, BackendConfig, WireFormatBackend,
27        codec::Codec,
28        ext::{
29            inspect_err::InspectErr,
30            interleave::Interleave,
31            lifecycle::{AfterStart, AfterStop, BeforeStart, BeforeStop},
32            map_err::MapErr,
33            pipe::Pipe,
34            poll_strategy::{PollStrategy, PollWith, StreamStrategy},
35            shared::Shared,
36            wake_on_push::WakeOnPush,
37            with_codec::WithCodec,
38        },
39    },
40    error::BoxDynError,
41    task::Task,
42    worker::context::WorkerContext,
43};
44
45#[cfg(feature = "sleep")]
46use crate::backend::ext::poll_strategy::{BackoffConfig, BackoffStrategy, IntervalStrategy};
47
48#[macro_use]
49pub mod delegate;
50/// A wrapper that allows inspecting errors produced by a backend.
51pub mod inspect_err;
52
53/// Extension allowing backends to be instrumented with a [tracing::Span].
54#[cfg(feature = "tracing")]
55pub mod instrument;
56/// A wrapper that allows merging a backend with a stream
57pub mod interleave;
58
59/// A wrapper that allows mapping the error type of a backend from `Self::Error` to another error type `E2`.
60pub mod map_err;
61pub mod pipe;
62pub mod poll_strategy;
63/// A wrapper that wakes the worker when a new item is fetched.
64pub mod wake_on_push;
65pub mod with_codec;
66
67pub mod lifecycle;
68
69/// A wrapper that makes a backend clonable
70pub mod shared;
71
72/// A wrapper that allows a backend to be used as a stream of tasks, without needing to know the concrete backend type at compile time.
73#[derive(Debug, thiserror::Error)]
74#[non_exhaustive]
75pub enum PollNextArgsError<B: Backend> {
76    /// The backend produced an error while polling for the next task.
77    #[error("backend error: {0}")]
78    BackendError(B::Error),
79    /// The backend produced a task, but the task's arguments could not be decoded.
80    #[error("failed to decode task args: {0}")]
81    DecodeError(BoxDynError),
82}
83
84/// Extension trait for `Backend` that provides additional combinators and utilities.
85pub trait BackendExt: Backend {
86    /// A convenience method for calling `poll_next` and decoding the `Args` in one step,
87    /// returning a `Task<Self::Args, ..>` instead of `Task<Self::Compact, ..>`.
88    #[allow(clippy::type_complexity)]
89    fn poll_next_args(
90        &mut self,
91        cx: &mut Context<'_>,
92        worker: &WorkerContext,
93    ) -> Poll<Option<Result<Task<Self::Args>, PollNextArgsError<Self>>>>
94    where
95        Self: Sized + BackendConfig + WireFormatBackend + Backend<Task = Task<Self::Compact>>,
96        Self::Codec: Codec<Self::Args, Compact = Self::Compact>,
97        <Self::Codec as Codec<Self::Args>>::Error: std::error::Error + Send + Sync + 'static,
98    {
99        let next = self.poll_next(cx, worker);
100        let codec = self.codec();
101        next.map(move |item| match item {
102            Some(Ok(task)) => {
103                let task = task.try_map_args(|compact| codec.decode(&compact));
104                Some(task.map_err(|e| PollNextArgsError::DecodeError(e.into())))
105            }
106            Some(Err(e)) => Some(Err(PollNextArgsError::BackendError(e))),
107            None => None,
108        })
109    }
110
111    /// Pipes every task polled from this backend into `sink`
112    ///
113    /// Useful for bridging two backend implementations — e.g. draining an
114    /// ephemeral/legacy queue into a durable one, or fanning a lightweight
115    /// source into a shared sink that multiple producers write into.
116    fn pipe_to<Dst>(self, backend: Dst) -> Pipe<Dst, Self>
117    where
118        Self: Sized,
119    {
120        Pipe::new(self, backend)
121    }
122
123    /// Attaches a callback `F` to be run on each error produced while polling the backend.
124    fn inspect_err<F>(self, f: F) -> InspectErr<Self, F>
125    where
126        Self: Sized,
127        F: Fn(&Self::Error),
128    {
129        InspectErr { backend: self, f }
130    }
131
132    /// Maps errors produced by the backend from `Self::Error` into `E2`, useful for
133    /// heterogeneous composed backends.
134    fn map_err<F, E2>(self, f: F) -> MapErr<Self, F>
135    where
136        Self: Sized,
137        F: Fn(Self::Error) -> E2,
138    {
139        MapErr { backend: self, f }
140    }
141
142    /// Swaps out the backend's serialization codec entirely (JSON,
143    /// MessagePack, Protobuf, ...) without touching storage logic.
144    fn with_codec<NewCodec>(self, codec: NewCodec) -> WithCodec<Self, NewCodec>
145    where
146        Self: Sized + BackendConfig,
147        NewCodec: Codec<Self::Args>,
148    {
149        WithCodec::new(self, codec)
150    }
151
152    /// Wake the worker when a stream receives a new item
153    fn poll_with_stream<S>(self, stream: S) -> PollWith<Self, StreamStrategy<S>>
154    where
155        Self: Sized,
156        S: Stream + Unpin + Send + 'static,
157    {
158        let strategy = StreamStrategy::new(stream);
159        PollWith::new(self, strategy)
160    }
161
162    /// Wake the worker periodically
163    #[cfg(feature = "sleep")]
164    fn poll_with_interval(self, duration: Duration) -> PollWith<Self, IntervalStrategy>
165    where
166        Self: Sized,
167    {
168        let strategy = IntervalStrategy::new(duration);
169        PollWith::new(self, strategy)
170    }
171
172    /// Wake the worker periodically with a backoff
173    #[cfg(feature = "sleep")]
174    fn poll_with_backoff(
175        self,
176        interval: Duration,
177        config: BackoffConfig,
178    ) -> PollWith<Self, BackoffStrategy>
179    where
180        Self: Sized,
181    {
182        let strategy = IntervalStrategy::new(interval).with_backoff(config);
183        PollWith::new(self, strategy)
184    }
185
186    /// Wake the worker with a custom strategy
187    fn poll_with_strategy<S>(self, strategy: S) -> PollWith<Self, S>
188    where
189        Self: Sized,
190        S: PollStrategy,
191    {
192        PollWith::new(self, strategy)
193    }
194
195    #[cfg(feature = "tracing")]
196    /// Provides a span to decorate emitted events
197    fn instrumented(self, span: tracing::Span) -> Instrumented<Self>
198    where
199        Self: Sized,
200    {
201        Instrumented::new(self, span)
202    }
203
204    /// Runs an async callback once, before the backend's first `poll_ready` is delegated.
205    fn before_start<F, Fut>(self, f: F) -> BeforeStart<Self, Self::Error>
206    where
207        Self: Sized,
208        F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
209        Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
210    {
211        BeforeStart::new(self, f)
212    }
213
214    /// Runs an async callback once, before the backend's poll_close is called.
215    fn before_stop<F, Fut>(self, f: F) -> BeforeStop<Self, Self::Error>
216    where
217        Self: Sized,
218        F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
219        Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
220    {
221        BeforeStop::new(self, f)
222    }
223
224    /// Runs an async callback once, after the backend's first `poll_ready` is successful.
225    fn after_start<F, Fut>(self, f: F) -> AfterStart<Self, Self::Error>
226    where
227        Self: Sized,
228        for<'c> F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
229        Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
230    {
231        AfterStart::new(self, f)
232    }
233
234    /// Runs an async callback once, after the worker has stopped and backend has cleaned up.
235    fn after_stop<F, Fut>(self, f: F) -> AfterStop<Self, Self::Error>
236    where
237        Self: Sized,
238        F: Fn(&mut Self) -> Fut + Send + Sync + 'static,
239        Fut: Future<Output = Result<(), Self::Error>> + Send + 'static,
240    {
241        AfterStop::new(self, f)
242    }
243
244    /// Interleaves the external stream with the backend.
245    fn interleave<S>(self, stream: S) -> Interleave<Self, S>
246    where
247        Self: Sized,
248        S: Stream<Item = Result<Self::Task, Self::Error>> + Unpin,
249    {
250        Interleave::new(self, stream)
251    }
252
253    /// Wakes the worker when a new item is pushed
254    fn wake_on_push(self) -> WakeOnPush<Self>
255    where
256        Self: Sized,
257    {
258        WakeOnPush::new(self)
259    }
260
261    /// Create a cloneable handle to the inner backend where all handles are clone.
262    fn shared(self) -> Shared<Self>
263    where
264        Self: WireFormatBackend + Send,
265        Self::Codec: Clone,
266    {
267        Shared::new(self)
268    }
269}
270
271impl<B: Backend> BackendExt for B {}