apalis_core/backend/ext/
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::{ext::pipe::PipeExt, memory::MemoryStorage};
13//! # use apalis_core::worker::{builder::WorkerBuilder, context::WorkerContext};
14//! # use apalis_core::error::BoxDynError;
15//! # use std::time::Duration;
16//! # use futures_util::StreamExt;
17//! # use crate::apalis_core::worker::ext::event_listener::EventListenerExt;
18//! #[tokio::main]
19//! async fn main() {
20//!     let stm = stream::iter(0..10).map(|s| Ok::<_, std::io::Error>(s));
21//!
22//!     let in_memory = MemoryStorage::new();
23//!     let backend = stm.pipe_to(in_memory);
24//!
25//!     async fn task(task: u32, worker: WorkerContext) -> Result<(), BoxDynError> {
26//!         tokio::time::sleep(Duration::from_secs(1)).await;
27//! #        if task == 9 {
28//! #            worker.stop().unwrap();
29//! #        }
30//!         Ok(())
31//!     }
32//!
33//!     let worker = WorkerBuilder::new("rango-tango")
34//!         .backend(backend)
35//!         .on_event(|_worker, ev| {
36//!             println!("On Event = {:?}", ev);
37//!         })
38//!         .build(task);
39//!     worker.run().await.unwrap();
40//! }
41//! ```
42//!
43//! This example pipes a stream of numbers into an in-memory backend and processes them with a worker.
44//!
45//! See also:
46//! - [`apalis-cron`](https://docs.rs/apalis-cron)
47use std::fmt::Debug;
48use std::fmt::{self};
49use std::{
50    pin::Pin,
51    task::{Context, Poll},
52};
53
54use crate::backend::*;
55use crate::error::BoxDynError;
56use crate::task::Task;
57use crate::worker::context::WorkerContext;
58use futures_core::stream::BoxStream;
59use futures_core::{Stream, ready};
60use futures_sink::Sink;
61use futures_util::SinkExt;
62use futures_util::TryStreamExt;
63
64/// A generic pipe that wraps a [`Stream`] and passes it to a backend
65#[doc = features_table! {
66    setup = "{ unreachable!() }",
67    TaskSink => supported("Ability to push new tasks", false),
68    InheritsFeatures => limited("Inherits features from the underlying backend", false),
69}]
70pub struct Pipe<Dst, S> {
71    pub(crate) from: S,
72    pub(crate) into: Dst,
73}
74
75impl<S, Dst> Pipe<Dst, S> {
76    /// Create a new `Pipe` from a raw `from` source and an `into` sink.
77    /// Prefer [`PipeExt::pipe_to`] or [`BackendExt::pipe_to`] over calling
78    /// this directly.
79    ///
80    /// [`BackendExt::pipe_to`]: crate::backend::ext::BackendExt
81    pub fn new(from: S, into: Dst) -> Self {
82        Self { from, into }
83    }
84}
85
86impl<S: fmt::Debug, Dst: fmt::Debug> fmt::Debug for Pipe<Dst, S> {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        f.debug_struct("Pipe")
89            .field("from", &self.from)
90            .field("into", &self.into)
91            .finish()
92    }
93}
94
95impl<Dst: Clone, S: Clone> Clone for Pipe<Dst, S> {
96    fn clone(&self) -> Self {
97        Self {
98            from: self.from.clone(),
99            into: self.into.clone(),
100        }
101    }
102}
103
104impl<S, TSink, Args, Kind, Err> Backend for Pipe<TSink, S>
105where
106    S: Backend<Task = Task<Args>, Error = Err> + Send + 'static,
107    Err: std::error::Error + Send + Sync + 'static,
108    TSink: BackendConfig<Kind = Kind> + Backend + TaskSink<Args, Kind> + Unpin + Send + 'static,
109    <TSink as Backend>::Error: std::error::Error + Send + Sync + 'static,
110    Args: Send + 'static,
111{
112    type Task = TSink::Task;
113
114    type Error = PipeError;
115
116    fn poll_ready(
117        &mut self,
118        cx: &mut Context<'_>,
119        wkr: &WorkerContext,
120    ) -> Poll<Result<(), Self::Error>> {
121        trace!("poll_ready: polling source");
122
123        ready!(S::poll_ready(&mut self.from, cx, wkr)).map_err(|e| {
124            trace!(error = ?e, "poll_ready: source returned error");
125            PipeError::Inner(e.into())
126        })?;
127
128        trace!("poll_ready: source ready, polling destination");
129
130        ready!(Backend::poll_ready(&mut self.into, cx, wkr)).map_err(|e| {
131            trace!(error = ?e, "poll_ready: destination returned error");
132            PipeError::Inner(e.into())
133        })?;
134
135        trace!("poll_ready: ready");
136
137        Poll::Ready(Ok(()))
138    }
139
140    fn poll_next(
141        &mut self,
142        cx: &mut Context<'_>,
143        worker: &WorkerContext,
144    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
145        loop {
146            match TaskSink::poll_ready(Pin::new(&mut self.into), cx) {
147                Poll::Ready(Ok(())) => {
148                    trace!("poll_next: destination ready");
149                }
150                Poll::Ready(Err(e)) => {
151                    trace!(error = ?e, "poll_next: destination poll_ready failed");
152                    return Poll::Ready(Some(Err(PipeError::Inner(e.into()))));
153                }
154                Poll::Pending => break,
155            }
156
157            match self.from.poll_next(cx, worker) {
158                Poll::Ready(Some(Ok(task))) => {
159                    trace!("poll_next: received task from source");
160
161                    if let Err(e) = TaskSink::start_send(Pin::new(&mut self.into), task) {
162                        trace!(error = ?e, "poll_next: destination start_send failed");
163                        return Poll::Ready(Some(Err(PipeError::Inner(e.into()))));
164                    }
165
166                    trace!("poll_next: task sent to destination");
167                }
168                Poll::Ready(Some(Err(e))) => {
169                    trace!(error = ?e, "poll_next: source returned error");
170                    return Poll::Ready(Some(Err(PipeError::Inner(e.into()))));
171                }
172                Poll::Ready(None) => {
173                    trace!("poll_next: source closed");
174                    break;
175                }
176                Poll::Pending => break,
177            }
178        }
179
180        match TSink::poll_flush(Pin::new(&mut self.into), cx) {
181            Poll::Ready(Err(e)) => {
182                trace!(error = ?e, "poll_next: destination flush failed");
183                return Poll::Ready(Some(Err(PipeError::Inner(e.into()))));
184            }
185            Poll::Ready(Ok(())) => {
186                trace!("poll_next: destination flushed");
187            }
188            Poll::Pending => {}
189        }
190
191        if let Poll::Ready(Err(e)) = self.into.poll_ready(cx, worker) {
192            trace!(error = ?e, "poll_next: destination backend poll_ready failed");
193            return Poll::Ready(Some(Err(PipeError::Inner(e.into()))));
194        }
195
196        self.into.poll_next(cx, worker).map_err(|e| {
197            trace!(error = ?e, "poll_next: destination backend returned error");
198            PipeError::Inner(e.into())
199        })
200    }
201
202    fn poll_close(
203        &mut self,
204        cx: &mut Context<'_>,
205        worker: &WorkerContext,
206    ) -> Poll<Result<(), Self::Error>> {
207        trace!("poll_close: closing destination");
208
209        self.into.poll_close(cx, worker).map_err(|e| {
210            trace!(error = ?e, "poll_close: destination returned error");
211            PipeError::Inner(e.into())
212        })
213    }
214}
215
216/// Utility for piping a plain stream of `Result<Args, Err>` into a backend.
217pub trait PipeExt<B, Args>
218where
219    B: Backend,
220{
221    /// Pipe the current stream into the provided sink backend.
222    fn pipe_to(self, backend: B) -> Pipe<B, BoxStream<'static, Result<Args, PipeError>>>;
223}
224
225impl<B, Args, Err, S> PipeExt<B, Args> for S
226where
227    B: Backend + Unpin + Send + 'static,
228    S: Stream<Item = Result<Args, Err>> + Send + Unpin + 'static,
229    Err: Into<BoxDynError> + Send + Sync + 'static,
230    Args: 'static,
231{
232    fn pipe_to(self, backend: B) -> Pipe<B, BoxStream<'static, Result<Args, PipeError>>> {
233        Pipe::new(
234            Box::pin(self.map_err(|e| PipeError::Inner(e.into()))),
235            backend,
236        )
237    }
238}
239
240/// Error encountered while piping streams
241#[derive(Debug, thiserror::Error)]
242pub enum PipeError {
243    /// The cron stream provided a None
244    #[error("The inner stream provided a None")]
245    EmptyStream,
246    /// An inner stream error occurred
247    #[error("The inner stream error: {0}")]
248    Inner(BoxDynError),
249}
250
251impl<Dst, S, T, Err> Sink<T> for Pipe<Dst, S>
252where
253    Dst: Sink<T, Error = Err> + Unpin,
254    S: Unpin,
255    Err: Into<BoxDynError> + Send + Sync,
256{
257    type Error = PipeError;
258    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
259        self.get_mut()
260            .into
261            .start_send_unpin(item)
262            .map_err(|e| PipeError::Inner(e.into()))
263    }
264    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
265        self.get_mut()
266            .into
267            .poll_ready_unpin(cx)
268            .map_err(|e| PipeError::Inner(e.into()))
269    }
270    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
271        self.get_mut()
272            .into
273            .poll_flush_unpin(cx)
274            .map_err(|e| PipeError::Inner(e.into()))
275    }
276    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
277        self.get_mut()
278            .into
279            .poll_close_unpin(cx)
280            .map_err(|e| PipeError::Inner(e.into()))
281    }
282}
283
284delegate_config!(Pipe<Dst, S>, into);
285
286delegate_deref!(Pipe<Dst, S>, into);
287
288delegate_codec!(Pipe<Dst, S>, into);
289
290delegate_expose!(
291    impl<B, S, Args, Kind> for Pipe<B, S>
292    where {
293        S: Backend<Task = Task<Args>> + Send + Sync + 'static,
294        S::Error: std::error::Error + Send + Sync + 'static,
295        B: BackendConfig<Kind = Kind>
296        + Backend
297        + TaskSink<Args, Kind>
298        + Unpin
299        + Send
300        + 'static,
301        <B as Backend>::Error: std::error::Error + Send + Sync + 'static,
302        Args: Send + 'static,
303    }
304    => into,
305    wrap = |this, result| result.map_err(|e| PipeError::Inner(e.into()))
306);
307
308#[cfg(test)]
309mod tests {
310    use std::{io, time::Duration};
311
312    use futures_util::{StreamExt, stream};
313
314    use crate::{
315        backend::{dequeue::VecDequeBackend, ext::BackendExt, memory::MemoryStorage},
316        error::BoxDynError,
317        worker::{
318            builder::WorkerBuilder, context::WorkerContext, ext::event_listener::EventListenerExt,
319        },
320    };
321
322    use super::*;
323
324    const ITEMS: u32 = 10;
325
326    #[tokio::test]
327    async fn basic_worker() {
328        let stm = stream::iter(0..ITEMS).map(Ok::<_, io::Error>);
329        let in_memory = MemoryStorage::new();
330
331        let backend = stm.pipe_to(in_memory);
332
333        async fn task(task: u32, worker: WorkerContext) -> Result<(), BoxDynError> {
334            tokio::time::sleep(Duration::from_secs(1)).await;
335            if task == ITEMS - 1 {
336                worker.stop().unwrap();
337                return Err("Graceful Exit".into());
338            }
339            Ok(())
340        }
341
342        let worker = WorkerBuilder::new("rango-tango")
343            .backend(backend)
344            .on_event(|_worker, ev| {
345                println!("On Event = {ev:?}");
346            })
347            .build(task);
348        worker.run().await.unwrap();
349    }
350
351    #[tokio::test]
352    async fn dequeue_to_memory_worker() {
353        let dequeue = VecDequeBackend::new();
354
355        let mut in_memory = MemoryStorage::new();
356
357        in_memory.push(42).await.unwrap();
358
359        let mut backend = in_memory.pipe_to(dequeue);
360
361        backend.push(43).await.unwrap();
362
363        async fn task(task: u32, worker: WorkerContext) -> Result<(), BoxDynError> {
364            tokio::time::sleep(Duration::from_secs(1)).await;
365            if task == 42 {
366                worker.stop().unwrap();
367                return Err("Graceful Exit".into());
368            }
369            Ok(())
370        }
371
372        let worker = WorkerBuilder::new("rango-tango")
373            .backend(backend)
374            .on_event(|_worker, ev| {
375                println!("On Event = {ev:?}");
376            })
377            .build(task);
378        worker.run().await.unwrap();
379    }
380}
apalis_core/backend/ext/pipe.rs

apalis_core/backend/ext/
pipe.rs
