apalis_core/backend/persistence/
service.rs

1use std::task::{Context, Poll};
2
3use futures_channel::mpsc::Sender;
4use futures_core::future::BoxFuture;
5use futures_util::{FutureExt, SinkExt};
6use tower_service::Service;
7
8use crate::{
9    backend::{TaskResult, codec::Codec, persistence::event::TaskEvent},
10    error::BoxDynError,
11    task::Task,
12};
13
14/// A middleware layer that persists task lifecycle events.
15///
16/// `TaskPersistLayer` sends task events to a persistence channel and can be
17/// configured to persist task results and lock tasks while they are being
18/// processed.
19///
20/// # Type Parameters
21///
22/// * `C` - The codec used to encode task data.
23/// * `Compact` - The compact representation used for persisted task events.
24#[derive(Clone, Debug)]
25pub struct TaskPersistLayer<C, Compact> {
26    sender: Sender<TaskEvent<Compact>>,
27    codec: C,
28    persist_results: bool,
29    lock_tasks: bool,
30}
31
32impl<Compact, C> TaskPersistLayer<C, Compact> {
33    /// Creates a new task persistence layer.
34    ///
35    /// Task results and task locking are disabled by default.
36    #[must_use]
37    pub(super) fn new(sender: Sender<TaskEvent<Compact>>, codec: C) -> Self {
38        Self {
39            sender,
40            codec,
41            persist_results: false,
42            lock_tasks: false,
43        }
44    }
45
46    /// Configures whether task results should be persisted.
47    ///
48    /// By default, task results are not persisted.
49    #[must_use]
50    pub fn persist_results(self, persist_results: bool) -> Self {
51        Self {
52            sender: self.sender,
53            codec: self.codec,
54            persist_results,
55            lock_tasks: self.lock_tasks,
56        }
57    }
58
59    /// Configures whether tasks should be locked while being processed.
60    ///
61    /// By default, tasks are not locked.
62    #[must_use]
63    pub fn lock_tasks(self, lock_tasks: bool) -> Self {
64        Self {
65            sender: self.sender,
66            codec: self.codec,
67            persist_results: self.persist_results,
68            lock_tasks,
69        }
70    }
71}
72
73impl<S, Compact, C> tower_layer::Layer<S> for TaskPersistLayer<C, Compact>
74where
75    C: Clone,
76{
77    type Service = TaskPersistService<S, Compact, C>;
78
79    fn layer(&self, inner: S) -> Self::Service {
80        TaskPersistService {
81            inner,
82            sender: self.sender.clone(),
83            codec: self.codec.clone(),
84            should_ack: self.persist_results,
85            should_lock: self.lock_tasks,
86        }
87    }
88}
89
90#[derive(Clone, Debug)]
91pub struct TaskPersistService<S, Compact, C> {
92    inner: S,
93    sender: Sender<TaskEvent<Compact>>,
94    codec: C,
95    should_ack: bool,
96    should_lock: bool,
97}
98
99impl<S, Args, C> Service<Task<Args>> for TaskPersistService<S, C::Compact, C>
100where
101    S: Service<Task<Args>> + Send + 'static,
102    S::Future: Send + 'static,
103    S::Error: Into<BoxDynError>,
104    Args: Send + 'static,
105    C: Codec<S::Response> + Clone + Send + 'static,
106    C::Error: std::error::Error,
107    S::Response: Send,
108    C::Compact: Send,
109{
110    type Response = S::Response;
111    type Error = BoxDynError;
112    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
113
114    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
115        match self.sender.poll_ready(cx) {
116            Poll::Ready(Ok(_)) => self.inner.poll_ready(cx).map_err(|e| e.into()),
117            Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
118            Poll::Pending => Poll::Pending,
119        }
120    }
121
122    fn call(&mut self, req: Task<Args>) -> Self::Future {
123        let mut sender = self.sender.clone();
124        let ctx = req.ctx().clone();
125        let fut = self.inner.call(req);
126        let codec = self.codec.clone();
127        let should_ack = self.should_ack;
128        let should_lock = self.should_lock;
129        async move {
130            let task_id = ctx.task_id().unwrap().clone();
131            if should_lock {
132                sender
133                    .send(TaskEvent::Lock {
134                        task_id: task_id.clone(),
135                    })
136                    .await?;
137            }
138
139            let res = fut.await.map_err(|e| e.into());
140            if should_ack {
141                let status = ctx.status();
142                let result = res
143                    .as_ref()
144                    .map_err(|e| e.to_string())
145                    .and_then(|res| codec.encode(res).map_err(|e| e.to_string()));
146                let attempt = ctx.attempt();
147                sender
148                    .send(TaskEvent::Complete(TaskResult {
149                        task_id,
150                        attempt,
151                        result,
152                        status,
153                    }))
154                    .await?;
155            }
156            res
157        }
158        .boxed()
159    }
160}
apalis_core/backend/persistence/service.rs

apalis_core/backend/persistence/
service.rs
