apalis_core/backend/persistence/
sink.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use futures_sink::Sink;
7use futures_util::FutureExt;
8
9use crate::{
10    backend::persistence::{Persisted, Persistence},
11    task::Task,
12};
13
14impl<P> Sink<Task<P::Compact>> for Persisted<P>
15where
16    P: Persistence + Unpin,
17    P::Compact: Send + Unpin + 'static,
18{
19    type Error = P::Error;
20
21    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
22        Poll::Ready(Ok(()))
23    }
24
25    fn start_send(self: Pin<&mut Self>, item: Task<P::Compact>) -> Result<(), Self::Error> {
26        self.get_mut().sink_buffer.push(item);
27        Ok(())
28    }
29
30    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
31        let this = self.get_mut();
32
33        // If there's no existing future and buffer is empty, we're done
34        if this.sink_future.is_none() && this.sink_buffer.is_empty() {
35            return Poll::Ready(Ok(()));
36        }
37
38        // Create the future only if we don't have one and there's work to do
39        if this.sink_future.is_none() && !this.sink_buffer.is_empty() {
40            let mut provider = this.inner.clone();
41            let buffer = std::mem::take(&mut this.sink_buffer);
42            let sink_fut = async move { provider.push_tasks(buffer).await };
43            this.sink_future = Some(sink_fut.boxed().into());
44        }
45
46        if let Some(mut fut) = this.sink_future.take() {
47            match fut.poll_unpin(cx) {
48                Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
49                Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
50                Poll::Pending => {
51                    // Future is still pending, put it back and return Pending
52                    this.sink_future = Some(fut);
53                    Poll::Pending
54                }
55            }
56        } else {
57            Poll::Ready(Ok(()))
58        }
59    }
60
61    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
62        self.poll_flush(cx)
63    }
64}
apalis_core/backend/persistence/sink.rs

apalis_core/backend/persistence/
sink.rs
