apalis_core/backend/ext/
wake_on_push.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll, Waker},
4};
5
6use futures_sink::Sink;
7use futures_util::SinkExt;
8
9use crate::{backend::*, task::Task, worker::context::WorkerContext};
10
11/// A backend wrapper that wakes the worker when a task is pushed.
12///
13/// `WakeOnPush` stores the most recent worker waker received through
14/// [`Backend::poll_ready`] and wakes it whenever a task is successfully sent
15/// through the backend's [`Sink`] implementation.
16///
17/// This is useful for backends where pushing a task should immediately cause
18/// a worker that is waiting for work to poll the backend again.
19///
20/// # Examples
21///
22/// ```ignore
23/// let backend = backend.wake_on_push();
24/// ```
25#[derive(Debug, Clone)]
26pub struct WakeOnPush<B> {
27    pub(super) backend: B,
28    pub(super) waker: Option<Waker>,
29}
30
31impl<B> WakeOnPush<B> {
32    /// Creates a new `WakeOnPush` wrapper around the given backend.
33    ///
34    /// The worker waker is captured when [`Backend::poll_ready`] is first
35    /// called.
36    pub fn new(backend: B) -> Self {
37        Self {
38            backend,
39            waker: None,
40        }
41    }
42}
43
44impl<B> Backend for WakeOnPush<B>
45where
46    B: Backend,
47{
48    type Task = B::Task;
49    type Error = B::Error;
50
51    fn poll_ready(
52        &mut self,
53        cx: &mut Context<'_>,
54        worker: &WorkerContext,
55    ) -> Poll<Result<(), Self::Error>> {
56        self.backend.poll_ready(cx, worker)
57    }
58
59    fn poll_next(
60        &mut self,
61        cx: &mut Context<'_>,
62        worker: &WorkerContext,
63    ) -> Poll<Option<Result<Self::Task, Self::Error>>> {
64        if self.waker.as_ref().is_none_or(|w| !w.will_wake(cx.waker())) {
65            self.waker = Some(cx.waker().clone());
66        }
67        self.backend.poll_next(cx, worker)
68    }
69
70    fn poll_close(
71        &mut self,
72        cx: &mut Context<'_>,
73        worker: &WorkerContext,
74    ) -> Poll<Result<(), Self::Error>> {
75        self.backend.poll_close(cx, worker)
76    }
77}
78
79impl<B, T, Err> Sink<T> for WakeOnPush<B>
80where
81    B: Sink<T, Error = Err> + Unpin,
82{
83    type Error = Err;
84    /// Sends a task to the backend and wakes the worker if the task was
85    /// pushed successfully.
86    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
87        let this = self.get_mut();
88        this.backend.start_send_unpin(item).map(|_| {
89            if let Some(s) = this.waker.as_ref() {
90                s.wake_by_ref()
91            }
92        })
93    }
94    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
95        self.get_mut().backend.poll_ready_unpin(cx)
96    }
97    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
98        self.get_mut().backend.poll_flush_unpin(cx)
99    }
100    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
101        self.get_mut().backend.poll_close_unpin(cx)
102    }
103}
104
105delegate_config!(WakeOnPush<B>, backend);
106
107delegate_codec!(WakeOnPush<B>, backend);
108
109delegate_deref!(WakeOnPush<B>, backend);
110
111delegate_expose!(
112    impl<B> for WakeOnPush<B>
113    where {
114        B: Send + Sync,
115    }
116    => backend
117);
apalis_core/backend/ext/wake_on_push.rs

apalis_core/backend/ext/
wake_on_push.rs
