sova_tasks_store/
lib.rs

1//! Task queue store for Sova.
2//!
3//! Trait is stable (memory + file + sql + redis backends).
4//! Queue claim/lease is **not** plain KvStore.
5
6#[cfg(feature = "file")]
7mod file;
8#[cfg(feature = "redis")]
9mod redis;
10#[cfg(feature = "sql")]
11mod sql;
12
13use bytes::Bytes;
14use std::collections::HashMap;
15use std::future::Future;
16use std::pin::Pin;
17use std::sync::Arc;
18use std::time::{Duration, SystemTime};
19use tokio::sync::Mutex;
20
21pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
22
23#[cfg(feature = "file")]
24pub use file::FileTaskStore;
25
26#[cfg(feature = "redis")]
27pub use redis::RedisTaskStore;
28
29#[cfg(feature = "sql")]
30pub use sql::SqlTaskStore;
31
32#[derive(Debug, thiserror::Error)]
33pub enum TaskError {
34    #[error("{0}")]
35    Msg(String),
36    #[error("not found")]
37    NotFound,
38    #[error("duplicate dedup key")]
39    Duplicate,
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum TaskStatus {
44    Pending,
45    Running,
46    Done,
47    Failed,
48}
49
50#[derive(Debug, Clone)]
51pub struct Task {
52    pub id: String,
53    pub queue: String,
54    pub payload: Bytes,
55    pub run_at: SystemTime,
56    pub attempts: u32,
57    pub lease_until: Option<SystemTime>,
58    pub dedup_key: Option<String>,
59    pub status: TaskStatus,
60    pub worker: Option<String>,
61    /// Higher values are claimed first within a queue (default `0`).
62    pub priority: i32,
63}
64
65pub struct EnqueueOpts {
66    pub queue: String,
67    pub payload: Bytes,
68    pub run_at: Option<SystemTime>,
69    pub dedup_key: Option<String>,
70    /// Higher values are claimed first (default `0`).
71    pub priority: i32,
72}
73
74pub trait TaskStore: Send + Sync + 'static {
75    fn enqueue<'a>(&'a self, opts: EnqueueOpts) -> BoxFuture<'a, Result<String, TaskError>>;
76    fn claim<'a>(
77        &'a self,
78        queue: &'a str,
79        worker: &'a str,
80        lease: Duration,
81        limit: usize,
82    ) -> BoxFuture<'a, Result<Vec<Task>, TaskError>>;
83    fn heartbeat<'a>(
84        &'a self,
85        id: &'a str,
86        lease: Duration,
87    ) -> BoxFuture<'a, Result<(), TaskError>>;
88    fn complete<'a>(&'a self, id: &'a str) -> BoxFuture<'a, Result<(), TaskError>>;
89    fn fail<'a>(
90        &'a self,
91        id: &'a str,
92        retry_at: Option<SystemTime>,
93    ) -> BoxFuture<'a, Result<(), TaskError>>;
94    fn reap<'a>(&'a self, now: SystemTime) -> BoxFuture<'a, Result<u64, TaskError>>;
95    fn list<'a>(
96        &'a self,
97        queue: &'a str,
98        limit: usize,
99    ) -> BoxFuture<'a, Result<Vec<Task>, TaskError>>;
100}
101
102#[derive(Default)]
103struct Inner {
104    tasks: HashMap<String, Task>,
105    dedup: HashMap<String, String>,
106    seq: u64,
107}
108
109#[derive(Clone, Default)]
110pub struct MemoryStore {
111    inner: Arc<Mutex<Inner>>,
112}
113
114impl MemoryStore {
115    pub fn new() -> Self {
116        Self::default()
117    }
118}
119
120impl TaskStore for MemoryStore {
121    fn enqueue<'a>(&'a self, opts: EnqueueOpts) -> BoxFuture<'a, Result<String, TaskError>> {
122        Box::pin(async move {
123            let mut g = self.inner.lock().await;
124            if let Some(ref dk) = opts.dedup_key {
125                if let Some(existing) = g.dedup.get(dk) {
126                    return Ok(existing.clone());
127                }
128            }
129            g.seq += 1;
130            let id = format!("t{}", g.seq);
131            let task = Task {
132                id: id.clone(),
133                queue: opts.queue,
134                payload: opts.payload,
135                run_at: opts.run_at.unwrap_or_else(SystemTime::now),
136                attempts: 0,
137                lease_until: None,
138                dedup_key: opts.dedup_key.clone(),
139                status: TaskStatus::Pending,
140                worker: None,
141                priority: opts.priority,
142            };
143            if let Some(dk) = opts.dedup_key {
144                g.dedup.insert(dk, id.clone());
145            }
146            g.tasks.insert(id.clone(), task);
147            Ok(id)
148        })
149    }
150
151    fn claim<'a>(
152        &'a self,
153        queue: &'a str,
154        worker: &'a str,
155        lease: Duration,
156        limit: usize,
157    ) -> BoxFuture<'a, Result<Vec<Task>, TaskError>> {
158        Box::pin(async move {
159            let mut g = self.inner.lock().await;
160            let now = SystemTime::now();
161            let mut ids: Vec<_> = g
162                .tasks
163                .values()
164                .filter(|t| {
165                    t.queue == queue
166                        && t.status == TaskStatus::Pending
167                        && t.run_at <= now
168                })
169                .map(|t| (std::cmp::Reverse(t.priority), t.run_at, t.id.clone()))
170                .collect();
171            ids.sort();
172            let mut out = Vec::new();
173            for (_, _, id) in ids.into_iter().take(limit) {
174                if let Some(t) = g.tasks.get_mut(&id) {
175                    t.status = TaskStatus::Running;
176                    t.attempts += 1;
177                    t.worker = Some(worker.to_string());
178                    t.lease_until = Some(now + lease);
179                    out.push(t.clone());
180                }
181            }
182            Ok(out)
183        })
184    }
185
186    fn heartbeat<'a>(
187        &'a self,
188        id: &'a str,
189        lease: Duration,
190    ) -> BoxFuture<'a, Result<(), TaskError>> {
191        Box::pin(async move {
192            let mut g = self.inner.lock().await;
193            let t = g.tasks.get_mut(id).ok_or(TaskError::NotFound)?;
194            if t.status != TaskStatus::Running {
195                return Err(TaskError::Msg("not running".into()));
196            }
197            t.lease_until = Some(SystemTime::now() + lease);
198            Ok(())
199        })
200    }
201
202    fn complete<'a>(&'a self, id: &'a str) -> BoxFuture<'a, Result<(), TaskError>> {
203        Box::pin(async move {
204            let mut g = self.inner.lock().await;
205            let t = g.tasks.get_mut(id).ok_or(TaskError::NotFound)?;
206            t.status = TaskStatus::Done;
207            t.lease_until = None;
208            t.worker = None;
209            if let Some(dk) = t.dedup_key.clone() {
210                g.dedup.remove(&dk);
211            }
212            Ok(())
213        })
214    }
215
216    fn fail<'a>(
217        &'a self,
218        id: &'a str,
219        retry_at: Option<SystemTime>,
220    ) -> BoxFuture<'a, Result<(), TaskError>> {
221        Box::pin(async move {
222            let mut g = self.inner.lock().await;
223            let t = g.tasks.get_mut(id).ok_or(TaskError::NotFound)?;
224            if let Some(at) = retry_at {
225                t.status = TaskStatus::Pending;
226                t.run_at = at;
227                t.lease_until = None;
228                t.worker = None;
229            } else {
230                t.status = TaskStatus::Failed;
231                t.lease_until = None;
232                t.worker = None;
233                if let Some(dk) = t.dedup_key.clone() {
234                    g.dedup.remove(&dk);
235                }
236            }
237            Ok(())
238        })
239    }
240
241    fn reap<'a>(&'a self, now: SystemTime) -> BoxFuture<'a, Result<u64, TaskError>> {
242        Box::pin(async move {
243            let mut g = self.inner.lock().await;
244            let mut n = 0u64;
245            for t in g.tasks.values_mut() {
246                if t.status == TaskStatus::Running {
247                    if let Some(until) = t.lease_until {
248                        if until <= now {
249                            t.status = TaskStatus::Pending;
250                            t.lease_until = None;
251                            t.worker = None;
252                            n += 1;
253                        }
254                    }
255                }
256            }
257            Ok(n)
258        })
259    }
260
261    fn list<'a>(
262        &'a self,
263        queue: &'a str,
264        limit: usize,
265    ) -> BoxFuture<'a, Result<Vec<Task>, TaskError>> {
266        Box::pin(async move {
267            let g = self.inner.lock().await;
268            let mut v: Vec<_> = g
269                .tasks
270                .values()
271                .filter(|t| t.queue == queue)
272                .cloned()
273                .collect();
274            v.sort_by(|a, b| a.id.cmp(&b.id));
275            v.truncate(limit);
276            Ok(v)
277        })
278    }
279}
280
281pub mod conformance {
282    use super::*;
283
284    pub async fn run(store: Arc<dyn TaskStore>) {
285        let id = store
286            .enqueue(EnqueueOpts {
287                queue: "default".into(),
288                payload: Bytes::from_static(b"hi"),
289                run_at: None,
290                dedup_key: Some("once".into()),
291                priority: 0,
292            })
293            .await
294            .unwrap();
295        let id2 = store
296            .enqueue(EnqueueOpts {
297                queue: "default".into(),
298                payload: Bytes::from_static(b"hi2"),
299                run_at: None,
300                dedup_key: Some("once".into()),
301                priority: 0,
302            })
303            .await
304            .unwrap();
305        assert_eq!(id, id2);
306
307        let claimed = store
308            .claim("default", "w1", Duration::from_secs(30), 10)
309            .await
310            .unwrap();
311        assert_eq!(claimed.len(), 1);
312        assert_eq!(claimed[0].id, id);
313
314        store.heartbeat(&id, Duration::from_secs(30)).await.unwrap();
315        store.complete(&id).await.unwrap();
316
317        let id3 = store
318            .enqueue(EnqueueOpts {
319                queue: "q".into(),
320                payload: Bytes::from_static(b"x"),
321                run_at: None,
322                dedup_key: None,
323                priority: 0,
324            })
325            .await
326            .unwrap();
327        let c = store
328            .claim("q", "w", Duration::from_millis(10), 1)
329            .await
330            .unwrap();
331        assert_eq!(c[0].id, id3);
332        tokio::time::sleep(Duration::from_millis(30)).await;
333        let n = store.reap(SystemTime::now()).await.unwrap();
334        assert!(n >= 1);
335        let again = store
336            .claim("q", "w2", Duration::from_secs(5), 1)
337            .await
338            .unwrap();
339        assert_eq!(again.len(), 1);
340        store
341            .fail(&again[0].id, Some(SystemTime::now()))
342            .await
343            .unwrap();
344        let listed = store.list("q", 10).await.unwrap();
345        assert!(!listed.is_empty());
346
347        // Higher priority claimed first.
348        let low = store
349            .enqueue(EnqueueOpts {
350                queue: "prio".into(),
351                payload: Bytes::from_static(b"low"),
352                run_at: None,
353                dedup_key: None,
354                priority: -100,
355            })
356            .await
357            .unwrap();
358        let high = store
359            .enqueue(EnqueueOpts {
360                queue: "prio".into(),
361                payload: Bytes::from_static(b"high"),
362                run_at: None,
363                dedup_key: None,
364                priority: 100,
365            })
366            .await
367            .unwrap();
368        let first = store
369            .claim("prio", "w", Duration::from_secs(5), 1)
370            .await
371            .unwrap();
372        assert_eq!(first[0].id, high);
373        assert_eq!(first[0].priority, 100);
374        let second = store
375            .claim("prio", "w", Duration::from_secs(5), 1)
376            .await
377            .unwrap();
378        assert_eq!(second[0].id, low);
379    }
380}
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385
386    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
387    async fn memory_conformance() {
388        conformance::run(Arc::new(MemoryStore::new())).await;
389    }
390}
sova_tasks_store/lib.rs

sova_tasks_store/
lib.rs
