apalis_workflow/in_memory/
mod.rs

1use apalis_codec::json::JsonCodec;
2use apalis_core::backend::ext::shared::Shared;
3use apalis_core::backend::finalize::Durable;
4use apalis_core::backend::memory::{BoxedReceiver, MemorySink, MemoryStorage, MemoryStorageError};
5use apalis_core::backend::{
6    Backend, BackendConfig, TaskResult, WaitForCompletion, WireFormatBackend,
7};
8use apalis_core::features_table;
9use apalis_core::{
10    task::{
11        Task,
12        task_id::{RandomId, TaskId},
13    },
14    worker::context::WorkerContext,
15};
16use futures_sink::Sink;
17use futures_util::SinkExt;
18use serde::de::DeserializeOwned;
19use std::marker::PhantomData;
20use std::{
21    pin::Pin,
22    sync::Arc,
23    task::{Context, Poll},
24};
25
26use crate::in_memory::result_store::ResultStore;
27use crate::in_memory::service::StoreResultsLayer;
28use crate::in_memory::stream::WaitForStream;
29
30mod result_store;
31
32mod service;
33
34mod stream;
35
36/// In-memory queue that is based on channels
37///
38///
39/// ## Example
40/// ```rust
41/// # use apalis_workflow::in_memory::InMemoryWorkflow;
42/// # use apalis_core::backend::ext::shared::Shared;
43/// # fn setup() -> Shared<InMemoryWorkflow<u32>> {
44/// let mut backend = InMemoryWorkflow::create();
45/// # backend
46/// # }
47/// ```
48///
49#[doc = features_table! {
50    setup = r#"
51        # {
52        #   use apalis_workflow::in_memory::InMemoryWorkflow;
53        #   InMemoryWorkflow::create()
54        # };
55    "#,
56    Backend => supported("Basic Backend functionality", true),
57    TaskSink => supported("Ability to push new tasks", true),
58    Serialization => not_supported("Serialization support for arguments"),
59
60    PipeExt => not_implemented("Allow other backends to pipe to this backend"),
61    BackendFactory => not_supported("Share the same storage across multiple workers"),
62
63    Update => not_supported("Allow updating a task"),
64    FetchById => not_supported("Allow fetching a task by its ID"),
65    Reschedule => not_supported("Reschedule a task"),
66
67    ResumeById => not_supported("Resume a task by its ID"),
68    ResumeAbandoned => not_supported("Resume abandoned tasks"),
69    Vacuum => not_supported("Vacuum the task storage"),
70
71    Workflow => not_implemented("Flexible enough to support workflows"),
72    WaitForCompletion => not_implemented("Wait for tasks to complete without blocking"), // Requires Clone
73
74    RegisterWorker => not_supported("Allow registering a worker with the backend"),
75    ListWorkers => not_supported("List all workers registered with the backend"),
76    ListTasks => not_supported("List all tasks in the backend"),
77}]
78pub struct InMemoryWorkflow<Args> {
79    pub(super) inner: MemoryStorage<Vec<u8>>,
80    _marker: PhantomData<Args>,
81    codec: JsonCodec,
82    store: Arc<ResultStore>,
83}
84
85impl<Args> InMemoryWorkflow<Args> {
86    /// Create a new in-memory storage
87    #[must_use]
88    pub fn create() -> Shared<Self> {
89        Shared::new(Self {
90            _marker: PhantomData,
91            codec: JsonCodec::default(),
92            inner: MemoryStorage::new(),
93            store: Arc::default(),
94        })
95    }
96}
97
98impl<Args> InMemoryWorkflow<Args> {
99    /// Create a storage given a sender and receiver
100    #[must_use]
101    pub fn new_with(sender: MemorySink<Vec<u8>>, receiver: BoxedReceiver<Vec<u8>>) -> Shared<Self> {
102        Shared::new(Self {
103            inner: MemoryStorage::new_with(sender, receiver),
104            _marker: PhantomData,
105            codec: JsonCodec::default(),
106            store: Arc::default(),
107        })
108    }
109}
110
111impl<Args> Sink<Task<Vec<u8>>> for InMemoryWorkflow<Args>
112where
113    Args: Unpin,
114{
115    type Error = MemoryStorageError;
116
117    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
118        self.as_mut().inner.poll_ready_unpin(cx)
119    }
120
121    fn start_send(mut self: Pin<&mut Self>, item: Task<Vec<u8>>) -> Result<(), Self::Error> {
122        self.as_mut().inner.start_send_unpin(item)
123    }
124
125    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
126        self.as_mut().inner.poll_flush_unpin(cx)
127    }
128
129    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
130        self.as_mut().inner.poll_close_unpin(cx)
131    }
132}
133
134impl<Args> std::fmt::Debug for InMemoryWorkflow<Args> {
135    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136        f.debug_struct("InMemoryWorkflow")
137            .field("inner", &self.inner)
138            .finish()
139    }
140}
141
142// InMemoryWorkflow as a Backend
143impl<Args> Backend for InMemoryWorkflow<Args> {
144    type Task = Task<Vec<u8>>;
145
146    type Error = MemoryStorageError;
147
148    fn poll_ready(
149        &mut self,
150        cx: &mut Context<'_>,
151        worker: &WorkerContext,
152    ) -> Poll<Result<(), Self::Error>> {
153        self.inner.poll_ready(cx, worker)
154    }
155
156    fn poll_next(
157        &mut self,
158        cx: &mut Context<'_>,
159        worker: &WorkerContext,
160    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
161        self.inner.poll_next(cx, worker)
162    }
163
164    fn poll_close(
165        &mut self,
166        cx: &mut Context<'_>,
167        worker: &WorkerContext,
168    ) -> Poll<Result<(), Self::Error>> {
169        self.inner.poll_close(cx, worker)
170    }
171}
172
173impl<Args> BackendConfig for InMemoryWorkflow<Args> {
174    type Id = RandomId;
175
176    type Args = Args;
177
178    type Kind = Durable;
179
180    type Config = ();
181
182    type Layer = StoreResultsLayer;
183
184    fn config(&self) -> &Self::Config {
185        &()
186    }
187
188    fn middleware(&mut self, _: &mut WorkerContext) -> Self::Layer {
189        StoreResultsLayer::new(Arc::clone(&self.store))
190    }
191}
192
193impl<Args> WireFormatBackend for InMemoryWorkflow<Args> {
194    type Codec = JsonCodec;
195
196    type Compact = Vec<u8>;
197
198    fn codec(&self) -> &Self::Codec {
199        &self.codec
200    }
201}
202
203impl<Args, Output> WaitForCompletion<Output> for InMemoryWorkflow<Args>
204where
205    Output: DeserializeOwned + Unpin + Send + 'static,
206{
207    type ResultStream = WaitForStream<Output>;
208
209    fn wait_for(&self, task_ids: impl IntoIterator<Item = TaskId>) -> Self::ResultStream {
210        WaitForStream::new(task_ids, Arc::clone(&self.store))
211    }
212
213    fn check_status(
214        &self,
215        task_ids: impl IntoIterator<Item = TaskId> + Send,
216    ) -> impl Future<Output = Result<Vec<TaskResult<Output>>, Self::Error>> + Send {
217        let store = Arc::clone(&self.store);
218        let task_ids: Vec<_> = task_ids.into_iter().collect();
219
220        async move {
221            let results = &store.results;
222
223            task_ids
224                .iter()
225                .filter_map(|id| results.get(id))
226                .map(|result| {
227                    let result = result.value();
228                    let decoded = result
229                        .result
230                        .as_ref()
231                        .map(|a| Output::deserialize(a).map_err(|e| e.to_string()))
232                        .map_err(|e| MemoryStorageError::Other(e.as_str().into()))?;
233                    Ok(TaskResult {
234                        task_id: result.task_id.clone(),
235                        attempt: result.attempt,
236                        status: result.status.clone(),
237                        result: decoded,
238                    })
239                })
240                .collect()
241        }
242    }
243}
apalis_workflow/in_memory/mod.rs

apalis_workflow/in_memory/
mod.rs
