apalis_core/backend/impls/
dequeue.rs

1use std::{
2    collections::VecDeque,
3    pin::Pin,
4    sync::{Arc, Mutex},
5    task::{Context, Poll, Waker},
6};
7
8use futures_sink::Sink;
9use tower_layer::Identity;
10
11use crate::{
12    backend::{Backend, BackendConfig, finalize::Ephemeral},
13    error::BoxDynError,
14    task::{
15        Task,
16        task_id::{RandomId, TaskId},
17    },
18    worker::context::WorkerContext,
19};
20
21/// A simple in-memory backend that uses a `VecDeque` to store tasks.
22///
23/// This backend is primarily intended for testing and demonstration purposes. It does not persist tasks and is not suitable for production use.
24#[derive(Debug, Clone)]
25pub struct VecDequeBackend<T> {
26    queue: Arc<Mutex<VecDeque<Task<T>>>>,
27    waker: Arc<Mutex<Option<Waker>>>,
28}
29
30impl<T> Default for VecDequeBackend<T> {
31    fn default() -> Self {
32        Self {
33            queue: Arc::new(Mutex::new(VecDeque::new())),
34            waker: Arc::new(Mutex::new(None)),
35        }
36    }
37}
38
39impl<T> VecDequeBackend<T> {
40    /// Create a new `VecDequeBackend` with default capacity.
41    #[must_use]
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    /// Create a new `VecDequeBackend` with a specified capacity for the internal queue.
47    #[must_use]
48    pub fn with_capacity(capacity: usize) -> Self {
49        Self {
50            queue: Arc::new(Mutex::new(VecDeque::with_capacity(capacity))),
51            waker: Arc::new(Mutex::new(None)),
52        }
53    }
54}
55
56/// Error type for VecDequeBackend
57#[derive(Debug, thiserror::Error, Clone)]
58#[non_exhaustive]
59pub enum VecDequeError {
60    /// Error occurred during polling
61    #[error("Polling error: {0}")]
62    PollError(Arc<BoxDynError>),
63    /// Error occurred during sending
64    #[error("Sending error: {0}")]
65    SendError(Arc<BoxDynError>),
66}
67
68impl<T> Backend for VecDequeBackend<T> {
69    type Task = Task<T>;
70    type Error = VecDequeError;
71
72    fn poll_ready(
73        &mut self,
74        cx: &mut Context<'_>,
75        _worker: &WorkerContext,
76    ) -> Poll<Result<(), Self::Error>> {
77        if self
78            .queue
79            .lock()
80            .map_err(|e| VecDequeError::PollError(Arc::new(e.to_string().into())))?
81            .is_empty()
82        {
83            *self.waker.lock().unwrap() = Some(cx.waker().clone());
84            Poll::Pending
85        } else {
86            Poll::Ready(Ok(()))
87        }
88    }
89
90    fn poll_next(
91        &mut self,
92        _cx: &mut Context<'_>,
93        _worker: &WorkerContext,
94    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
95        match self
96            .queue
97            .lock()
98            .map_err(|e| VecDequeError::PollError(Arc::new(e.to_string().into())))?
99            .pop_front()
100        {
101            Some(task) => Poll::Ready(Some(Ok(task))),
102            None => Poll::Ready(None),
103        }
104    }
105
106    fn poll_close(
107        &mut self,
108        _: &mut Context<'_>,
109        _: &WorkerContext,
110    ) -> Poll<Result<(), Self::Error>> {
111        Poll::Ready(Ok(()))
112    }
113}
114
115impl<T> Sink<Task<T>> for VecDequeBackend<T>
116where
117    T: Send + Unpin + 'static,
118{
119    type Error = VecDequeError;
120
121    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
122        Poll::Ready(Ok(()))
123    }
124
125    fn start_send(self: Pin<&mut Self>, mut item: Task<T>) -> Result<(), Self::Error> {
126        let this = self.get_mut();
127
128        let mut tasks = this
129            .queue
130            .lock()
131            .map_err(|e| VecDequeError::SendError(Arc::new(e.to_string().into())))?;
132
133        if let Some(ref key) = item.idempotency_key() {
134            let exists = tasks.iter().any(|task| {
135                task.idempotency_key()
136                    .as_ref()
137                    .map(|existing| existing == key)
138                    .unwrap_or(false)
139            });
140
141            if exists {
142                return Ok(());
143            }
144        }
145
146        if item.task_id().is_none() {
147            item = item
148                .into_builder()
149                .task_id(TaskId::from_string(RandomId::default()))
150                .build();
151        }
152
153        tasks.push_back(item);
154
155        if let Some(waker) = this.waker.lock().unwrap().as_ref() {
156            waker.wake_by_ref();
157        }
158
159        Ok(())
160    }
161
162    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
163        Poll::Ready(Ok(()))
164    }
165
166    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
167        Poll::Ready(Ok(()))
168    }
169}
170
171impl<Args> BackendConfig for VecDequeBackend<Args> {
172    type Args = Args;
173    type Id = RandomId;
174
175    type Kind = Ephemeral;
176
177    type Config = ();
178
179    type Layer = Identity;
180
181    fn config(&self) -> &Self::Config {
182        &()
183    }
184
185    fn middleware(&mut self, _worker: &mut WorkerContext) -> Self::Layer {
186        Identity::new()
187    }
188}
apalis_core/backend/impls/dequeue.rs

apalis_core/backend/impls/
dequeue.rs
