apalis_core/backend/ext/
interleave.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use futures_core::Stream;
7use futures_sink::Sink;
8use futures_util::SinkExt;
9
10use crate::{backend::*, task::Task, worker::context::WorkerContext};
11
12/// A `Backend` wrapper that interleaves tasks from an external/shared `Stream`
13/// with tasks produced by the wrapped backend.
14///
15/// The external stream is polled first on every `poll_next` call, giving it
16/// priority; if it's pending or exhausted, control falls through to the
17/// underlying backend. This is useful for injecting tasks from a source that
18/// isn't itself a `Backend` — e.g. a broadcast channel, webhook listener, or
19/// manually-fed queue — without needing to implement the full `Backend` trait
20/// for that source.
21#[derive(Debug, Clone)]
22pub struct Interleave<B, S> {
23    pub(super) backend: B,
24    pub(super) stream: S,
25    /// Set once the external stream yields `None`, so we stop polling it.
26    pub(super) stream_done: bool,
27}
28
29impl<B, S> Interleave<B, S> {
30    /// Wraps a backend with an external stream of tasks.
31    ///
32    /// The external stream is polled first on every `poll_next` call, giving it
33    /// priority; if it's pending or exhausted, control falls through to the
34    /// underlying backend.
35    pub fn new(backend: B, stream: S) -> Self {
36        Self {
37            backend,
38            stream,
39            stream_done: false,
40        }
41    }
42}
43
44impl<B, S> Backend for Interleave<B, S>
45where
46    B: Backend,
47    S: Stream<Item = Result<B::Task, B::Error>> + Unpin,
48{
49    type Task = B::Task;
50    type Error = B::Error;
51
52    fn poll_ready(
53        &mut self,
54        cx: &mut Context<'_>,
55        worker: &WorkerContext,
56    ) -> Poll<Result<(), Self::Error>> {
57        self.backend.poll_ready(cx, worker)
58    }
59
60    fn poll_next(
61        &mut self,
62        cx: &mut Context<'_>,
63        worker: &WorkerContext,
64    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
65        if !self.stream_done {
66            match Pin::new(&mut self.stream).poll_next(cx) {
67                Poll::Ready(Some(item)) => return Poll::Ready(Some(item)),
68                Poll::Ready(None) => {
69                    self.stream_done = true;
70                }
71                Poll::Pending => {}
72            }
73        }
74
75        self.backend.poll_next(cx, worker)
76    }
77
78    fn poll_close(
79        &mut self,
80        cx: &mut Context<'_>,
81        worker: &WorkerContext,
82    ) -> Poll<Result<(), Self::Error>> {
83        self.backend.poll_close(cx, worker)
84    }
85}
86
87delegate_sink!(Interleave<B, S>, backend);
88
89delegate_deref!(Interleave<B, S>, backend);
90
91delegate_config!(Interleave<B, S>, backend);
92
93delegate_codec!(Interleave<B, S>, backend);
94
95delegate_expose!(
96    impl<B, S> for Interleave<B, S>
97    where {
98        B: Send + Sync,
99        S: Stream<Item = Result<B::Task, B::Error>> + Send + Sync + Unpin,
100    }
101    => backend
102);
apalis_core/backend/ext/interleave.rs

apalis_core/backend/ext/
interleave.rs
