apalis_core/backend/impls/json/
shared.rs

1/// Sharable JSON based backend.
2///
3/// The [`SharedJsonStore`] allows multiple task types to be stored
4/// and processed concurrently using a single JSON-based in-memory backend.
5/// It is useful for testing, prototyping,
6/// or sharing state between workers in a single process.
7///
8/// # Example
9///
10/// ```rust
11/// # use apalis_core::backend::shared::MakeShared;
12/// # use apalis_core::task::Task;
13/// # use apalis_core::worker::context::WorkerContext;
14/// # use apalis_core::worker::builder::WorkerBuilder;
15/// # use apalis_core::backend::json::SharedJsonStore;
16/// # use apalis_core::error::BoxDynError;
17/// # use std::time::Duration;
18/// # use apalis_core::backend::TaskSink;
19///
20/// #[tokio::main]
21/// async fn main() {
22///     let mut store = SharedJsonStore::new();
23///     let mut int_store = store.make_shared().unwrap();
24///     int_store.push(42).await.unwrap();
25///
26///     async fn task(
27///         task: u32,
28///         ctx: WorkerContext,
29///     ) -> Result<(), BoxDynError> {
30///         tokio::time::sleep(Duration::from_millis(2)).await;
31///         ctx.stop()?;
32///         Ok(())
33///     }
34///
35///     let int_worker = WorkerBuilder::new("int-worker")
36///         .backend(int_store)
37///         .build(task)
38///         .run();
39///
40///     int_worker.await.unwrap();
41/// }
42/// ```
43///
44/// See the tests for more advanced usage with multiple types and event listeners.
45use std::{
46    pin::Pin,
47    sync::Arc,
48    task::{Context, Poll},
49};
50
51use futures_channel::mpsc::SendError;
52use futures_core::Stream;
53use futures_sink::Sink;
54use futures_util::{SinkExt, StreamExt};
55use serde::{Serialize, de::DeserializeOwned};
56use serde_json::Value;
57
58use crate::{
59    backend::impls::{
60        json::{
61            JsonStorage,
62            meta::JsonMapMetadata,
63            util::{FindFirstWith, TaskKey, TaskWithMeta},
64        },
65        memory::{MemorySink, MemoryStorage},
66    },
67    task::{Task, status::Status, task_id::TaskId},
68};
69
70#[derive(Debug)]
71struct SharedJsonStream<T, Ctx> {
72    inner: JsonStorage<Value>,
73    req_type: std::marker::PhantomData<(T, Ctx)>,
74}
75
76impl<Args: DeserializeOwned + Unpin> Stream for SharedJsonStream<Args, JsonMapMetadata> {
77    type Item = Task<Args, JsonMapMetadata>;
78    fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
79        use crate::task::builder::TaskBuilder;
80        let map = self.inner.tasks.try_read().expect("Failed to read tasks");
81        if let Some((key, _)) = map.find_first_with(|k, _| {
82            &k.queue == std::any::type_name::<Args>() && k.status == Status::Pending
83        }) {
84            let task = map.get(&key).unwrap();
85            let args = match Args::deserialize(&task.args) {
86                Ok(value) => value,
87                Err(_) => return Poll::Pending,
88            };
89            let task = TaskBuilder::new(args)
90                .with_task_id(key.task_id.clone())
91                .with_ctx(task.ctx.clone())
92                .build();
93            let key = key.clone();
94            drop(map);
95            let this = &mut self.get_mut().inner;
96            this.update_status(&key, Status::Running)
97                .expect("Failed to update status");
98            this.persist_to_disk().expect("Failed to persist to disk");
99            Poll::Ready(Some(task))
100        } else {
101            Poll::Pending
102        }
103    }
104}
105/// Sharable JSON based backend.
106///
107/// # Features
108///
109/// - Concurrent processing of multiple task types
110/// - In-memory storage with optional disk persistence
111/// - Metadata support for tasks
112#[derive(Debug)]
113pub struct SharedJsonStore {
114    inner: JsonStorage<serde_json::Value>,
115}
116
117impl SharedJsonStore {
118    /// Create a new instance of the shared JSON store.
119    pub fn new() -> Self {
120        Self {
121            inner: JsonStorage::new_temp().unwrap(),
122        }
123    }
124}
125
126impl<Args: Send + Serialize + DeserializeOwned + Unpin + 'static>
127    crate::backend::shared::MakeShared<Args> for SharedJsonStore
128{
129    type Backend = MemoryStorage<Args, JsonMapMetadata>;
130
131    type Config = ();
132
133    type MakeError = String;
134
135    fn make_shared_with_config(
136        &mut self,
137        _: Self::Config,
138    ) -> Result<Self::Backend, Self::MakeError> {
139        let (sender, receiver) = self.inner.create_channel::<Args>();
140        Ok(MemoryStorage {
141            sender: MemorySink {
142                inner: Arc::new(futures_util::lock::Mutex::new(sender)),
143            },
144            receiver,
145        })
146    }
147}
148
149impl JsonStorage<Value> {
150    fn create_channel<Args: 'static + DeserializeOwned + Serialize + Send + Unpin>(
151        &self,
152    ) -> (
153        Box<
154            dyn Sink<Task<Args, JsonMapMetadata>, Error = SendError>
155                + Send
156                + Sync
157                + Unpin
158                + 'static,
159        >,
160        Pin<Box<dyn Stream<Item = Task<Args, JsonMapMetadata>> + Send>>,
161    ) {
162        // Create a channel for communication
163        let sender = self.clone();
164
165        // Create a wrapped sender that will insert into the in-memory store
166        let wrapped_sender = {
167            let mut store = self.clone();
168
169            sender.with_flat_map(move |task: Task<Args, JsonMapMetadata>| {
170                use crate::task::task_id::RandomId;
171                let task_id = task
172                    .parts
173                    .task_id
174                    .clone()
175                    .unwrap_or(TaskId::new(RandomId::default()));
176                let task = task.map(|args| serde_json::to_value(args).unwrap());
177                store
178                    .insert(
179                        &TaskKey {
180                            task_id,
181                            queue: std::any::type_name::<Args>().to_owned(),
182                            status: Status::Pending,
183                        },
184                        TaskWithMeta {
185                            args: task.args.clone(),
186                            ctx: task.parts.ctx.clone(),
187                            result: None,
188                        },
189                    )
190                    .unwrap();
191                futures_util::stream::iter(vec![Ok(task)])
192            })
193        };
194
195        // Create a stream that filters by type T
196        let filtered_stream = {
197            let inner = self.clone();
198            SharedJsonStream {
199                inner,
200                req_type: std::marker::PhantomData,
201            }
202        };
203
204        // Combine the sender and receiver
205        let sender = Box::new(wrapped_sender)
206            as Box<dyn Sink<Task<Args, JsonMapMetadata>, Error = SendError> + Send + Sync + Unpin>;
207        let receiver = filtered_stream.boxed();
208
209        (sender, receiver)
210    }
211}
212#[cfg(test)]
213mod tests {
214    use std::time::Duration;
215
216    use crate::error::BoxDynError;
217
218    use crate::worker::context::WorkerContext;
219    use crate::{
220        backend::{TaskSink, shared::MakeShared},
221        worker::{builder::WorkerBuilder, ext::event_listener::EventListenerExt},
222    };
223
224    use super::*;
225
226    const ITEMS: u32 = 10;
227
228    #[tokio::test]
229    async fn basic_shared() {
230        let mut store = SharedJsonStore::new();
231        let mut string_store = store.make_shared().unwrap();
232        let mut int_store = store.make_shared().unwrap();
233        for i in 0..ITEMS {
234            string_store.push(format!("ITEM: {i}")).await.unwrap();
235            int_store.push(i).await.unwrap();
236        }
237
238        async fn task(task: u32, ctx: WorkerContext) -> Result<(), BoxDynError> {
239            tokio::time::sleep(Duration::from_millis(2)).await;
240            if task == ITEMS - 1 {
241                ctx.stop()?;
242                return Err("Worker stopped!")?;
243            }
244            Ok(())
245        }
246
247        let string_worker = WorkerBuilder::new("rango-tango-string")
248            .backend(string_store)
249            .on_event(|ctx, ev| {
250                println!("CTX {:?}, On Event = {:?}", ctx.name(), ev);
251            })
252            .build(|req: String, ctx: WorkerContext| async move {
253                tokio::time::sleep(Duration::from_millis(2)).await;
254                println!("{req}");
255                if req.ends_with(&(ITEMS - 1).to_string()) {
256                    ctx.stop().unwrap();
257                }
258            })
259            .run();
260
261        let int_worker = WorkerBuilder::new("rango-tango-int")
262            .backend(int_store)
263            .on_event(|ctx, ev| {
264                println!("CTX {:?}, On Event = {:?}", ctx.name(), ev);
265            })
266            .build(task)
267            .run();
268
269        let _ = futures_util::future::join(int_worker, string_worker).await;
270    }
271}
apalis_core/backend/impls/json/shared.rs

apalis_core/backend/impls/json/
shared.rs
