apalis_core/backend/impls/
memory.rs

1//! # In-memory backend based on channels
2//!
3//! An in-memory backend suitable for testing, prototyping, or lightweight task processing scenarios where persistence is not required.
4//!
5//! ## Features
6//! - Generic in-memory queue for any task type.
7//! - Implements [`Backend`] for integration with workers.
8//! - Sink support: Ability to push new tasks.
9//!
10//! A detailed feature list can be found in the [capabilities](crate::backend::memory::MemoryStorage#capabilities) section.
11//!
12//! ## Example
13//!
14//! ```rust
15//! # use apalis_core::backend::memory::MemoryStorage;
16//! # use apalis_core::worker::context::WorkerContext;
17//! # use apalis_core::worker::builder::WorkerBuilder;
18//! # use apalis_core::backend::TaskSink;
19//!
20//! async fn handler(_: u32, worker: WorkerContext) {
21//!     worker.stop().unwrap();
22//! }
23//!
24//! #[tokio::main]
25//! async fn main() {
26//!     let mut store = MemoryStorage::new();
27//!     store.push(42).await.unwrap();
28//!
29//!     let worker = WorkerBuilder::new("int-worker")
30//!         .backend(store)
31//!         .build(handler);
32//!
33//!     worker.run().await.unwrap();
34//! }
35//! ```
36//!
37//! ## Note
38//! This backend is not persistent and is intended for use cases where durability is not required.
39//! For production workloads, consider using a persistent backend such as PostgreSQL or Redis.
40//!
41//! ## See Also
42//! - [`Backend`]
43//! - [`WorkerContext`]
44use crate::backend::finalize::Ephemeral;
45use crate::backend::{Backend, BackendConfig, TryNewBackend};
46use crate::error::BoxDynError;
47use crate::{
48    task::{
49        Task,
50        task_id::{RandomId, TaskId},
51    },
52    worker::context::WorkerContext,
53};
54use futures_channel::mpsc::{SendError, unbounded};
55use futures_core::ready;
56use futures_sink::Sink;
57use futures_util::lock::Mutex;
58use futures_util::{FutureExt, SinkExt, Stream, StreamExt};
59use std::collections::HashSet;
60use std::{
61    pin::Pin,
62    sync::Arc,
63    task::{Context, Poll},
64};
65use tower_layer::Identity;
66
67/// A boxed in-memory task receiver stream
68pub type BoxedReceiver<Args> = Pin<Box<dyn Stream<Item = Task<Args>> + Send>>;
69
70/// In-memory queue that is based on channels
71///
72///
73/// ## Example
74/// ```rust
75/// # use apalis_core::backend::memory::MemoryStorage;
76/// # fn setup() -> MemoryStorage<u32> {
77/// let mut backend = MemoryStorage::new();
78/// # backend
79/// # }
80/// ```
81///
82#[doc = features_table! {
83    setup = r#"
84        # {
85        #   use apalis_core::backend::memory::MemoryStorage;
86        #   MemoryStorage::new()
87        # };
88    "#,
89    Backend => supported("Basic Backend functionality", true),
90    TaskSink => supported("Ability to push new tasks", true),
91    Serialization => not_supported("Serialization support for arguments"),
92
93    PipeExt => not_implemented("Allow other backends to pipe to this backend"),
94    BackendFactory => not_supported("Share the same storage across multiple workers"),
95
96    Update => not_supported("Allow updating a task"),
97    FetchById => not_supported("Allow fetching a task by its ID"),
98    Reschedule => not_supported("Reschedule a task"),
99
100    ResumeById => not_supported("Resume a task by its ID"),
101    ResumeAbandoned => not_supported("Resume abandoned tasks"),
102    Vacuum => not_supported("Vacuum the task storage"),
103
104    Workflow => not_implemented("Flexible enough to support workflows"),
105    WaitForCompletion => not_implemented("Wait for tasks to complete without blocking"), // Requires Clone
106
107    RegisterWorker => not_supported("Allow registering a worker with the backend"),
108    ListWorkers => not_supported("List all workers registered with the backend"),
109    ListTasks => not_supported("List all tasks in the backend"),
110}]
111pub struct MemoryStorage<Args> {
112    pub(super) sender: MemorySink<Args>,
113    pub(super) receiver: std::sync::Mutex<BoxedReceiver<Args>>,
114}
115
116impl<Args: Send + 'static> Default for MemoryStorage<Args> {
117    fn default() -> Self {
118        Self::new()
119    }
120}
121
122/// Error type for MemoryStorage operations
123#[derive(Debug, thiserror::Error)]
124#[non_exhaustive]
125pub enum MemoryStorageError {
126    /// Error occurred while sending a task to the in-memory channel
127    #[error("Failed to send task: {0}")]
128    SendError(#[from] SendError),
129    /// Error occurred while flushing the in-memory channel
130    #[error("Failed to add task to storage: {0}")]
131    Other(BoxDynError),
132}
133
134impl<Args: Send + 'static> MemoryStorage<Args> {
135    /// Create a new in-memory storage
136    #[must_use]
137    pub fn new() -> Self {
138        let (sender, receiver) = unbounded();
139        let sender = Box::new(sender.sink_map_err(|e| e.into()))
140            as Box<dyn Sink<Task<Args>, Error = MemoryStorageError> + Send + Sync + Unpin>;
141        Self {
142            sender: MemorySink {
143                inner: Arc::new(futures_util::lock::Mutex::new(sender)),
144                idempotency_keys: Default::default(),
145            },
146            receiver: receiver.boxed().into(),
147        }
148    }
149    /// Create a storage given a sender and receiver
150    #[must_use]
151    pub fn new_with(sender: MemorySink<Args>, receiver: BoxedReceiver<Args>) -> Self {
152        Self {
153            sender,
154            receiver: receiver.into(),
155        }
156    }
157}
158
159impl<Args> Sink<Task<Args>> for MemoryStorage<Args> {
160    type Error = MemoryStorageError;
161
162    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
163        self.as_mut().sender.poll_ready_unpin(cx)
164    }
165
166    fn start_send(mut self: Pin<&mut Self>, item: Task<Args>) -> Result<(), Self::Error> {
167        self.as_mut().sender.start_send_unpin(item)
168    }
169
170    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
171        self.as_mut().sender.poll_flush_unpin(cx)
172    }
173
174    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
175        self.as_mut().sender.poll_close_unpin(cx)
176    }
177}
178
179type ArcMemorySink<Args> = Arc<
180    Mutex<Box<dyn Sink<Task<Args>, Error = MemoryStorageError> + Send + Sync + Unpin + 'static>>,
181>;
182
183type ArcIdempotencySet = Arc<Mutex<HashSet<String>>>;
184
185/// Memory sink for sending tasks to the in-memory backend
186pub struct MemorySink<Args> {
187    pub(super) inner: ArcMemorySink<Args>,
188    pub(super) idempotency_keys: ArcIdempotencySet,
189}
190
191impl<Args> MemorySink<Args> {
192    /// Build a new memory sink given a sink
193    pub fn new(sink: ArcMemorySink<Args>) -> Self {
194        Self {
195            inner: sink,
196            idempotency_keys: Arc::new(Mutex::new(HashSet::new())),
197        }
198    }
199}
200
201impl<Args> std::fmt::Debug for MemorySink<Args> {
202    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
203        f.debug_struct("MemorySink")
204            .field("inner", &"<Sink>")
205            .field("idempotency_keys", &self.idempotency_keys.lock())
206            .finish()
207    }
208}
209
210impl<Args> Clone for MemorySink<Args> {
211    fn clone(&self) -> Self {
212        Self {
213            inner: Arc::clone(&self.inner),
214            idempotency_keys: Arc::clone(&self.idempotency_keys),
215        }
216    }
217}
218
219impl<Args> Sink<Task<Args>> for MemorySink<Args> {
220    type Error = MemoryStorageError;
221
222    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
223        let mut lock = ready!(self.inner.lock().poll_unpin(cx));
224        Pin::new(&mut *lock).poll_ready_unpin(cx)
225    }
226
227    fn start_send(self: Pin<&mut Self>, mut item: Task<Args>) -> Result<(), Self::Error> {
228        let this = self.get_mut();
229        if let Some(key) = item.idempotency_key() {
230            let mut keys = this.idempotency_keys.try_lock().unwrap();
231
232            if keys.contains(key) {
233                return Ok(());
234            }
235
236            keys.insert(key.to_owned());
237        }
238
239        if item.task_id().is_none() {
240            let task = item
241                .into_builder()
242                .task_id(TaskId::from_string(RandomId::default()));
243            item = task.build();
244        }
245
246        let mut sink = this.inner.try_lock().unwrap();
247        Pin::new(&mut *sink).start_send_unpin(item)
248    }
249
250    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
251        let mut lock = ready!(self.inner.lock().poll_unpin(cx));
252        Pin::new(&mut *lock).poll_flush_unpin(cx)
253    }
254
255    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
256        let mut lock = ready!(self.inner.lock().poll_unpin(cx));
257        Pin::new(&mut *lock).poll_close_unpin(cx)
258    }
259}
260
261impl<Args> std::fmt::Debug for MemoryStorage<Args> {
262    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
263        f.debug_struct("MemoryStorage")
264            .field("sender", &self.sender)
265            .field("receiver", &"<Stream>")
266            .finish()
267    }
268}
269
270impl<Args> Stream for MemoryStorage<Args> {
271    type Item = Task<Args>;
272
273    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
274        self.receiver.lock().unwrap().poll_next_unpin(cx)
275    }
276}
277
278// MemoryStorage as a Backend
279impl<Args> Backend for MemoryStorage<Args> {
280    type Task = Task<Args>;
281
282    type Error = MemoryStorageError;
283
284    fn poll_ready(
285        &mut self,
286        cx: &mut Context<'_>,
287        _: &WorkerContext,
288    ) -> Poll<Result<(), Self::Error>> {
289        self.sender.poll_ready_unpin(cx)
290    }
291
292    fn poll_next(
293        &mut self,
294        cx: &mut Context<'_>,
295        _: &WorkerContext,
296    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
297        self.receiver
298            .lock()
299            .unwrap()
300            .poll_next_unpin(cx)
301            .map(|item| item.map(Ok))
302    }
303
304    fn poll_close(
305        &mut self,
306        cx: &mut Context<'_>,
307        _: &WorkerContext,
308    ) -> Poll<Result<(), Self::Error>> {
309        self.sender.poll_close_unpin(cx)
310    }
311}
312
313impl<Args> BackendConfig for MemoryStorage<Args> {
314    type Id = RandomId;
315
316    type Args = Args;
317
318    type Kind = Ephemeral;
319
320    type Config = ();
321
322    type Layer = Identity;
323
324    fn config(&self) -> &Self::Config {
325        &()
326    }
327
328    fn middleware(&mut self, _: &mut WorkerContext) -> Self::Layer {
329        Identity::new()
330    }
331}
332
333impl<T: Send + 'static> TryNewBackend for MemoryStorage<T> {
334    type Backend = Self;
335    fn try_new(_: Self::Config) -> Result<Self::Backend, Self::Error> {
336        Ok(Self::new())
337    }
338}
apalis_core/backend/impls/memory.rs

apalis_core/backend/impls/
memory.rs
