apalis_core/backend/
pipe.rs

1//! # Pipe streams to backends
2//!
3//! This backend allows you to pipe tasks from any stream into another backend.
4//! It is useful for connecting different backends together, such as piping tasks
5//! from a cron stream into a database backend, or transforming and forwarding tasks
6//! between systems.
7//!
8//! ## Example
9//!
10//! ```rust
11//! # use futures_util::stream;
12//! # use apalis_core::backend::pipe::PipeExt;
13//! # use apalis_core::backend::json::JsonStorage;
14//! # use apalis_core::worker::{builder::WorkerBuilder, context::WorkerContext};
15//! # use apalis_core::error::BoxDynError;
16//! # use std::time::Duration;
17//! # use futures_util::StreamExt;
18//! # use crate::apalis_core::worker::ext::event_listener::EventListenerExt;
19//! #[tokio::main]
20//! async fn main() {
21//!     let stm = stream::iter(0..10).map(|s| Ok::<_, std::io::Error>(s));
22//!
23//!     let in_memory = JsonStorage::new_temp().unwrap();
24//!     let backend = stm.pipe_to(in_memory);
25//!
26//!     async fn task(task: u32, ctx: WorkerContext) -> Result<(), BoxDynError> {
27//!         tokio::time::sleep(Duration::from_secs(1)).await;
28//! #        if task == 9 {
29//! #            ctx.stop().unwrap();
30//! #        }
31//!         Ok(())
32//!     }
33//!
34//!     let worker = WorkerBuilder::new("rango-tango")
35//!         .backend(backend)
36//!         .on_event(|_ctx, ev| {
37//!             println!("On Event = {:?}", ev);
38//!         })
39//!         .build(task);
40//!     worker.run().await.unwrap();
41//! }
42//! ```
43//!
44//! This example pipes a stream of numbers into an in-memory backend and processes them with a worker.
45//!
46//! See also:
47//! - [`apalis-cron`](https://docs.rs/apalis-cron)
48
49use crate::backend::{BackendExt, TaskSink};
50use crate::error::BoxDynError;
51use crate::features_table;
52use crate::task::Task;
53use crate::{backend::Backend, backend::codec::Codec, worker::context::WorkerContext};
54use futures_sink::Sink;
55use futures_util::stream::{once, select};
56use futures_util::{SinkExt, Stream, TryStreamExt};
57use futures_util::{StreamExt, stream::BoxStream};
58use std::fmt;
59use std::fmt::Debug;
60use std::marker::PhantomData;
61use std::ops::{Deref, DerefMut};
62
63/// A generic pipe that wraps a [`Stream`] and passes it to a backend
64#[doc = features_table! {
65    setup = "{ unreachable!() }",
66    TaskSink => supported("Ability to push new tasks", false),
67    InheritsFeatures => limited("Inherits features from the underlying backend", false),
68}]
69pub struct Pipe<S, Into, Args, Ctx> {
70    pub(crate) from: S,
71    pub(crate) into: Into,
72    pub(crate) _req: PhantomData<(Args, Ctx)>,
73}
74
75impl<S: Clone, Into: Clone, Args, Ctx> Clone for Pipe<S, Into, Args, Ctx> {
76    fn clone(&self) -> Self {
77        Self {
78            from: self.from.clone(),
79            into: self.into.clone(),
80            _req: PhantomData,
81        }
82    }
83}
84
85impl<S, Into, Args, Ctx> Deref for Pipe<S, Into, Args, Ctx> {
86    type Target = Into;
87
88    fn deref(&self) -> &Self::Target {
89        &self.into
90    }
91}
92
93impl<S, Into, Args, Ctx> DerefMut for Pipe<S, Into, Args, Ctx> {
94    fn deref_mut(&mut self) -> &mut Self::Target {
95        &mut self.into
96    }
97}
98
99impl<S, Into, Args, Ctx> Pipe<S, Into, Args, Ctx> {
100    /// Create a new Pipe instance
101    pub fn new(stream: S, backend: Into) -> Self {
102        Self {
103            from: stream,
104            into: backend,
105            _req: PhantomData,
106        }
107    }
108}
109
110impl<S: fmt::Debug, Into: fmt::Debug, Args, Ctx> fmt::Debug for Pipe<S, Into, Args, Ctx> {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        f.debug_struct("Pipe")
113            .field("inner", &self.from)
114            .field("into", &self.into)
115            .finish()
116    }
117}
118
119impl<Args, Ctx, S, TSink, Err> Backend for Pipe<S, TSink, Args, Ctx>
120where
121    S: Stream<Item = Result<Args, Err>> + Send + 'static,
122    TSink: Backend<Args = Args, Context = Ctx>
123        + BackendExt
124        + TaskSink<Args>
125        + Clone
126        + Unpin
127        + Send
128        + 'static
129        + Sink<Task<TSink::Compact, Ctx, TSink::IdType>>,
130    <TSink as Backend>::Error: std::error::Error + Send + Sync + 'static,
131    TSink::Beat: Send + 'static,
132    TSink::IdType: Send + Clone + 'static,
133    TSink::Stream: Send + 'static,
134    Args: Send + 'static,
135    Ctx: Send + 'static + Default,
136    Err: std::error::Error + Send + Sync + 'static,
137    <TSink as Sink<Task<TSink::Compact, Ctx, TSink::IdType>>>::Error:
138        std::error::Error + Send + Sync + 'static,
139    <<TSink as BackendExt>::Codec as Codec<Args>>::Error: std::error::Error + Send + Sync + 'static,
140    TSink::Compact: Send,
141{
142    type Args = Args;
143
144    type IdType = TSink::IdType;
145
146    type Context = Ctx;
147
148    type Stream = BoxStream<'static, Result<Option<Task<Args, Ctx, Self::IdType>>, PipeError>>;
149
150    type Layer = TSink::Layer;
151
152    type Beat = BoxStream<'static, Result<(), PipeError>>;
153
154    type Error = PipeError;
155
156    fn heartbeat(&self, worker: &WorkerContext) -> Self::Beat {
157        self.into
158            .heartbeat(worker)
159            .map_err(|e| PipeError::Inner(e.into()))
160            .boxed()
161    }
162
163    fn middleware(&self) -> Self::Layer {
164        self.into.middleware()
165    }
166
167    fn poll(self, worker: &WorkerContext) -> Self::Stream {
168        let mut sink = self
169            .into
170            .clone()
171            .sink_map_err(|e| PipeError::Inner(e.into()));
172
173        let mut sink_stream = self
174            .from
175            .map_err(|e| PipeError::Inner(e.into()))
176            .map_ok(|s| {
177                Task::new(s)
178                    .try_map(|s| TSink::Codec::encode(&s).map_err(|e| PipeError::Inner(e.into())))
179            })
180            .map(|t| t.and_then(|t| t))
181            .boxed();
182
183        let sender_stream = self.into.poll(worker);
184        select(
185            once(async move {
186                let fut = sink.send_all(&mut sink_stream);
187                fut.await.map_err(|e| PipeError::Inner(e.into()))?;
188                Ok(None)
189            }),
190            sender_stream.map_err(|e| PipeError::Inner(e.into())),
191        )
192        .boxed()
193    }
194}
195
196/// Represents utility for piping streams into a backend
197pub trait PipeExt<B, Args, Ctx>
198where
199    Self: Sized,
200{
201    /// Pipe the current stream into the provided backend
202    fn pipe_to(self, backend: B) -> Pipe<Self, B, Args, Ctx>;
203}
204
205impl<B, Args, Ctx, Err, S> PipeExt<B, Args, Ctx> for S
206where
207    S: Stream<Item = Result<Args, Err>> + Send + 'static,
208    <B as Backend>::Error: Into<BoxDynError> + Send + Sync + 'static,
209    B: Backend<Args = Args> + TaskSink<Args>,
210{
211    fn pipe_to(self, backend: B) -> Pipe<Self, B, Args, Ctx> {
212        Pipe::new(self, backend)
213    }
214}
215
216/// Error encountered while piping streams
217#[derive(Debug, thiserror::Error)]
218pub enum PipeError {
219    /// The cron stream provided a None
220    #[error("The inner stream provided a None")]
221    EmptyStream,
222    /// An inner stream error occurred
223    #[error("The inner stream error: {0}")]
224    Inner(BoxDynError),
225}
226
227#[cfg(test)]
228#[cfg(feature = "json")]
229mod tests {
230    use std::{io, time::Duration};
231
232    use futures_util::stream;
233
234    use crate::{
235        backend::json::JsonStorage,
236        error::BoxDynError,
237        worker::{
238            builder::WorkerBuilder, context::WorkerContext, ext::event_listener::EventListenerExt,
239        },
240    };
241
242    use super::*;
243
244    const ITEMS: u32 = 10;
245
246    #[tokio::test]
247    async fn basic_worker() {
248        let stm = stream::iter(0..ITEMS).map(Ok::<_, io::Error>);
249        let in_memory = JsonStorage::new_temp().unwrap();
250
251        let backend = Pipe::new(stm, in_memory);
252
253        async fn task(task: u32, ctx: WorkerContext) -> Result<(), BoxDynError> {
254            tokio::time::sleep(Duration::from_secs(1)).await;
255            if task == ITEMS - 1 {
256                ctx.stop().unwrap();
257                return Err("Graceful Exit".into());
258            }
259            Ok(())
260        }
261
262        let worker = WorkerBuilder::new("rango-tango")
263            .backend(backend)
264            .on_event(|_ctx, ev| {
265                println!("On Event = {ev:?}");
266            })
267            .build(task);
268        worker.run().await.unwrap();
269    }
270}