tako_rs_core/queue/
backend.rs

1//! Pluggable queue backend abstraction.
2//!
3//! `QueueBackend` is the v2 trait that lets a `Queue` swap its in-process
4//! storage for an external broker (Redis, Postgres, NATS, …). The bundled
5//! [`MemoryBackend`](crate::queue::backend::MemoryBackend) keeps the existing in-process `Queue` semantics behind
6//! the same trait so consumer code can move to the trait at its own pace.
7
8use std::sync::Arc;
9use std::time::Instant;
10
11use async_trait::async_trait;
12use parking_lot::Mutex;
13
14/// Job identifier returned by [`QueueBackend::push`].
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
16pub struct JobId(pub u64);
17
18/// Per-push options.
19#[derive(Debug, Clone, Default)]
20pub struct PushOptions {
21  /// Idempotency / dedup key — duplicate pushes with the same key are coalesced.
22  pub dedup_key: Option<String>,
23  /// Schedule the job to execute no earlier than this instant.
24  pub run_after: Option<Instant>,
25  /// Optional named attempt count (for retries that re-push instead of mutating in place).
26  pub attempt: u32,
27}
28
29/// Reserved job returned by [`QueueBackend::reserve`].
30#[derive(Debug, Clone)]
31pub struct ReservedJob {
32  pub id: JobId,
33  pub queue: String,
34  pub payload: Vec<u8>,
35  pub attempt: u32,
36}
37
38/// Queue backend trait — async by design, since real backends are remote.
39#[async_trait]
40pub trait QueueBackend: Send + Sync + 'static {
41  /// Push a payload onto a named queue, returning the job id.
42  async fn push(
43    &self,
44    queue: &str,
45    payload: &[u8],
46    opts: PushOptions,
47  ) -> Result<JobId, BackendError>;
48
49  /// Reserve the next ready job from a queue (`FIFO`, `ready_at` <= now).
50  async fn reserve(&self, queue: &str) -> Result<Option<ReservedJob>, BackendError>;
51
52  /// Mark a reserved job complete — no retries, no DLQ entry.
53  async fn complete(&self, id: JobId) -> Result<(), BackendError>;
54
55  /// Mark a reserved job failed; if `retry_at` is set, requeue for that time.
56  async fn fail(&self, id: JobId, retry_at: Option<Instant>) -> Result<(), BackendError>;
57
58  /// Move a job to the dead-letter queue (terminal failure).
59  async fn dead_letter(&self, id: JobId) -> Result<(), BackendError>;
60}
61
62/// Backend-level error.
63#[derive(Debug, Clone)]
64pub enum BackendError {
65  /// The transport (Redis, Postgres, …) returned an error.
66  Transport(String),
67  /// The job referenced by `JobId` does not exist.
68  NotFound,
69}
70
71impl std::fmt::Display for BackendError {
72  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73    match self {
74      Self::Transport(e) => write!(f, "queue transport error: {e}"),
75      Self::NotFound => write!(f, "job not found"),
76    }
77  }
78}
79
80impl std::error::Error for BackendError {}
81
82/// In-process backend keeping job state in a single `Mutex<Vec<…>>`. Drop-in
83/// replacement for the bundled `Queue` storage; suitable for tests and
84/// single-node deployments. Replace with a remote backend for multi-pod use.
85///
86/// **Memory bound**: the pending queue is capped at [`MemoryBackend::max_pending`]
87/// (`None` = unlimited, default). Configure via [`MemoryBackend::with_max_pending`]
88/// to protect against unbounded growth when producers outpace consumers.
89#[derive(Default)]
90pub struct MemoryBackend {
91  inner: Arc<Mutex<MemoryInner>>,
92  next_id: std::sync::atomic::AtomicU64,
93  max_pending: Option<usize>,
94}
95
96#[derive(Default)]
97struct MemoryInner {
98  pending: Vec<MemoryJob>,
99  reserved: Vec<MemoryJob>,
100  dlq: Vec<MemoryJob>,
101  /// Active dedup keys.
102  ///
103  /// **Window semantics**: a key is inserted by [`MemoryBackend::push`] when
104  /// the caller supplied `PushOptions::dedup_key`, and is removed only by
105  /// [`MemoryBackend::complete`] (success) or [`MemoryBackend::dead_letter`]
106  /// (terminal failure). [`MemoryBackend::fail`] (transient retry) deliberately
107  /// keeps the key alive so a duplicate push during the retry window is
108  /// idempotent — the original job is still going to run.
109  ///
110  /// Callers that need a shorter window can call
111  /// [`MemoryBackend::purge_dedup_key`] explicitly.
112  dedup: std::collections::HashSet<String>,
113}
114
115#[derive(Clone)]
116struct MemoryJob {
117  id: JobId,
118  queue: String,
119  payload: Vec<u8>,
120  attempt: u32,
121  run_after: Option<Instant>,
122  dedup_key: Option<String>,
123}
124
125impl MemoryBackend {
126  pub fn new() -> Self {
127    Self::default()
128  }
129
130  /// Cap the pending vector at `max` jobs. Subsequent
131  /// [`MemoryBackend::push`] calls that would exceed the cap return
132  /// [`BackendError::Transport`] with the message `queue full`.
133  ///
134  /// Use this to bound memory growth when a misconfigured retry loop or
135  /// runaway producer would otherwise fill the heap.
136  #[must_use]
137  pub fn with_max_pending(mut self, max: usize) -> Self {
138    self.max_pending = Some(max);
139    self
140  }
141}
142
143#[async_trait]
144impl QueueBackend for MemoryBackend {
145  async fn push(
146    &self,
147    queue: &str,
148    payload: &[u8],
149    opts: PushOptions,
150  ) -> Result<JobId, BackendError> {
151    let mut inner = self.inner.lock();
152    if let Some(cap) = self.max_pending
153      && inner.pending.len() >= cap
154    {
155      return Err(BackendError::Transport("queue full".into()));
156    }
157    if let Some(key) = opts.dedup_key.as_ref()
158      && !inner.dedup.insert(key.clone())
159    {
160      // Treat as a successful idempotent no-op; report a synthetic id.
161      return Ok(JobId(0));
162    }
163    let id = JobId(
164      self
165        .next_id
166        .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
167        + 1,
168    );
169    inner.pending.push(MemoryJob {
170      id,
171      queue: queue.to_string(),
172      payload: payload.to_vec(),
173      attempt: opts.attempt,
174      run_after: opts.run_after,
175      dedup_key: opts.dedup_key,
176    });
177    Ok(id)
178  }
179
180  async fn reserve(&self, queue: &str) -> Result<Option<ReservedJob>, BackendError> {
181    let mut inner = self.inner.lock();
182    let now = Instant::now();
183    let pos = inner
184      .pending
185      .iter()
186      .position(|j| j.queue == queue && j.run_after.is_none_or(|t| now >= t));
187    let Some(idx) = pos else {
188      return Ok(None);
189    };
190    let job = inner.pending.remove(idx);
191    let reserved = ReservedJob {
192      id: job.id,
193      queue: job.queue.clone(),
194      payload: job.payload.clone(),
195      attempt: job.attempt,
196    };
197    inner.reserved.push(job);
198    Ok(Some(reserved))
199  }
200
201  async fn complete(&self, id: JobId) -> Result<(), BackendError> {
202    let mut inner = self.inner.lock();
203    let pos = inner.reserved.iter().position(|j| j.id == id);
204    let Some(idx) = pos else {
205      return Err(BackendError::NotFound);
206    };
207    let job = inner.reserved.remove(idx);
208    if let Some(key) = job.dedup_key {
209      inner.dedup.remove(&key);
210    }
211    Ok(())
212  }
213
214  async fn fail(&self, id: JobId, retry_at: Option<Instant>) -> Result<(), BackendError> {
215    let mut inner = self.inner.lock();
216    let pos = inner.reserved.iter().position(|j| j.id == id);
217    let Some(idx) = pos else {
218      return Err(BackendError::NotFound);
219    };
220    let mut job = inner.reserved.remove(idx);
221    job.attempt = job.attempt.saturating_add(1);
222    job.run_after = retry_at;
223    inner.pending.push(job);
224    Ok(())
225  }
226
227  async fn dead_letter(&self, id: JobId) -> Result<(), BackendError> {
228    let mut inner = self.inner.lock();
229    let pos = inner.reserved.iter().position(|j| j.id == id);
230    let Some(idx) = pos else {
231      return Err(BackendError::NotFound);
232    };
233    let job = inner.reserved.remove(idx);
234    if let Some(key) = job.dedup_key.as_ref() {
235      inner.dedup.remove(key);
236    }
237    inner.dlq.push(job);
238    Ok(())
239  }
240}
241
242impl MemoryBackend {
243  /// Snapshot of dead-letter jobs.
244  pub fn dead_letters(&self) -> Vec<(JobId, String, Vec<u8>, u32)> {
245    self
246      .inner
247      .lock()
248      .dlq
249      .iter()
250      .map(|j| (j.id, j.queue.clone(), j.payload.clone(), j.attempt))
251      .collect()
252  }
253
254  /// Number of pending jobs across every queue.
255  pub fn pending_count(&self) -> usize {
256    self.inner.lock().pending.len()
257  }
258
259  /// Number of currently reserved jobs.
260  pub fn reserved_count(&self) -> usize {
261    self.inner.lock().reserved.len()
262  }
263
264  /// Number of active dedup keys. Useful for diagnostics — if this grows
265  /// unboundedly your jobs are not reaching `complete`/`dead_letter`.
266  pub fn dedup_size(&self) -> usize {
267    self.inner.lock().dedup.len()
268  }
269
270  /// Drop a single dedup key out of band. Returns `true` if the key was
271  /// present. Use this when a long retry tail needs to be cut short and a
272  /// fresh push (same key) should be allowed to enqueue immediately.
273  pub fn purge_dedup_key(&self, key: &str) -> bool {
274    self.inner.lock().dedup.remove(key)
275  }
276}
tako_rs_core/queue/backend.rs

tako_rs_core/queue/
backend.rs
