1use crate::kv::MemoryKv;
25use crate::pubsub::MemoryPubsub;
26use async_trait::async_trait;
27use bytes::Bytes;
28use klieo_core::bus::{ClaimHandle, ClaimHandleImpl, ClaimedJob, Job, JobQueue, Lease, LeaseImpl};
29use klieo_core::error::BusError;
30use klieo_core::ids::{JobId, RunId};
31use std::collections::{HashMap, HashSet, VecDeque};
32use std::sync::atomic::{AtomicU64, Ordering};
33use std::sync::Arc;
34use std::time::{Duration, Instant};
35use tokio::sync::Mutex;
36
37#[derive(Debug)]
38struct JobEntry {
39 payload: Bytes,
40 attempts: u32,
41 max_attempts: u32,
42 lease_until: Option<Instant>,
45 not_before: Option<Instant>,
48 causation_run_id: Option<RunId>,
51}
52
53#[derive(Default)]
54struct PerQueue {
55 entries: HashMap<JobId, JobEntry>,
57 ready: VecDeque<JobId>,
60 leased: HashSet<JobId>,
63}
64
65#[derive(Default)]
66pub(crate) struct State {
67 queues: HashMap<String, PerQueue>,
68}
69
70pub struct MemoryJobQueue {
72 state: Arc<Mutex<State>>,
73 pub(crate) pubsub: Arc<MemoryPubsub>,
74 kv: Arc<MemoryKv>,
75 next_id: AtomicU64,
76}
77
78impl MemoryJobQueue {
79 pub fn new(pubsub: Arc<MemoryPubsub>, kv: Arc<MemoryKv>) -> Self {
82 Self {
83 state: Arc::new(Mutex::new(State::default())),
84 pubsub,
85 kv,
86 next_id: AtomicU64::new(1),
87 }
88 }
89
90 fn fresh_id(&self) -> JobId {
91 let n = self.next_id.fetch_add(1, Ordering::Relaxed);
92 JobId(format!("mem-{n}"))
93 }
94}
95
96fn sweep_expired_leases(pq: &mut PerQueue, now: Instant) {
99 let expired: Vec<JobId> = pq
100 .leased
101 .iter()
102 .filter(|id| {
103 pq.entries
104 .get(*id)
105 .and_then(|e| e.lease_until)
106 .is_none_or(|t| t <= now)
107 })
108 .cloned()
109 .collect();
110 for id in expired {
111 pq.leased.remove(&id);
112 if let Some(entry) = pq.entries.get_mut(&id) {
113 entry.lease_until = None;
114 }
115 pq.ready.push_back(id);
116 }
117}
118
119#[async_trait]
120impl JobQueue for MemoryJobQueue {
121 async fn enqueue(&self, queue: &str, job: Job) -> Result<JobId, BusError> {
122 let id = self.fresh_id();
123
124 if let Some(dedup_key) = &job.dedup_key {
126 use klieo_core::KvStore;
127 let bucket = format!("dedup.{queue}");
128 let id_bytes: Bytes = id.0.clone().into_bytes().into();
129 match self.kv.cas(&bucket, dedup_key, id_bytes, None).await {
130 Ok(_) => {}
131 Err(BusError::CasConflict { .. }) => {
132 let existing = self.kv.get(&bucket, dedup_key).await?.ok_or_else(|| {
134 BusError::Permanent("dedup record missing after CAS conflict".into())
135 })?;
136 let id_str = String::from_utf8(existing.value.to_vec())
137 .map_err(|e| BusError::Permanent(format!("dedup id bytes: {e}")))?;
138 return Ok(JobId(id_str));
139 }
140 Err(e) => return Err(e),
141 }
142 }
143
144 let max_attempts = job.max_attempts.unwrap_or(5);
145 let mut g = self.state.lock().await;
146 let pq = g.queues.entry(queue.into()).or_default();
147 pq.entries.insert(
148 id.clone(),
149 JobEntry {
150 payload: job.payload,
151 attempts: 0,
152 max_attempts,
153 lease_until: None,
154 not_before: None,
155 causation_run_id: job.causation_run_id,
156 },
157 );
158 pq.ready.push_back(id.clone());
159 Ok(id)
160 }
161
162 async fn claim(
163 &self,
164 queue: &str,
165 _worker_id: &str,
166 lease_ttl: Duration,
167 ) -> Result<Option<ClaimedJob>, BusError> {
168 let now = Instant::now();
169 let mut g = self.state.lock().await;
170 let pq = g.queues.entry(queue.into()).or_default();
171
172 if let Some(claim) = try_claim_from_ready(pq, now, lease_ttl) {
177 drop(g);
178 return Ok(Some(build_claimed_job(
179 self.state.clone(),
180 self.pubsub.clone(),
181 queue.to_string(),
182 claim,
183 lease_ttl,
184 )));
185 }
186
187 sweep_expired_leases(pq, now);
190 if let Some(claim) = try_claim_from_ready(pq, now, lease_ttl) {
191 drop(g);
192 return Ok(Some(build_claimed_job(
193 self.state.clone(),
194 self.pubsub.clone(),
195 queue.to_string(),
196 claim,
197 lease_ttl,
198 )));
199 }
200
201 Ok(None)
202 }
203}
204
205fn try_claim_from_ready(
211 pq: &mut PerQueue,
212 now: Instant,
213 lease_ttl: Duration,
214) -> Option<ClaimSnapshot> {
215 let mut rotated = 0;
216 let initial_len = pq.ready.len();
217 while rotated < initial_len {
218 let id = pq.ready.pop_front()?;
219 let claimable = pq
220 .entries
221 .get(&id)
222 .map(|e| e.not_before.is_none_or(|t| t <= now))
223 .unwrap_or(false);
224
225 if !claimable {
226 pq.ready.push_back(id);
227 rotated += 1;
228 continue;
229 }
230
231 let entry = pq.entries.get_mut(&id).expect("ready id without entry");
232 entry.attempts += 1;
233 entry.lease_until = Some(now + lease_ttl);
234 entry.not_before = None;
235 let snapshot = ClaimSnapshot {
236 id: id.clone(),
237 payload: entry.payload.clone(),
238 max_attempts: entry.max_attempts,
239 causation_run_id: entry.causation_run_id,
240 };
241 pq.leased.insert(id);
242 return Some(snapshot);
243 }
244 None
245}
246
247struct ClaimSnapshot {
249 id: JobId,
250 payload: Bytes,
251 max_attempts: u32,
252 causation_run_id: Option<RunId>,
253}
254
255fn build_claimed_job(
256 state: Arc<Mutex<State>>,
257 pubsub: Arc<MemoryPubsub>,
258 queue: String,
259 snapshot: ClaimSnapshot,
260 lease_ttl: Duration,
261) -> ClaimedJob {
262 let id = snapshot.id.clone();
263 let lease = Lease::new(Box::new(JobLease {
264 state: state.clone(),
265 queue: queue.clone(),
266 job_id: id.clone(),
267 ttl: lease_ttl,
268 }));
269 let claim = ClaimHandle::new(Box::new(MemoryClaimHandle {
270 state,
271 pubsub,
272 queue,
273 job_id: id.clone(),
274 max_attempts: snapshot.max_attempts,
275 }));
276 ClaimedJob::new(id, snapshot.payload, lease, claim).with_causation(snapshot.causation_run_id)
277}
278
279struct JobLease {
281 state: Arc<Mutex<State>>,
282 queue: String,
283 job_id: JobId,
284 ttl: Duration,
285}
286
287#[async_trait]
288impl LeaseImpl for JobLease {
289 async fn heartbeat(&self) -> Result<(), BusError> {
290 let mut g = self.state.lock().await;
291 let pq = g
292 .queues
293 .get_mut(&self.queue)
294 .ok_or_else(|| BusError::NotFound(self.queue.clone()))?;
295 let entry = pq
296 .entries
297 .get_mut(&self.job_id)
298 .ok_or_else(|| BusError::NotFound(format!("job {}", self.job_id.0)))?;
299 entry.lease_until = Some(Instant::now() + self.ttl);
300 Ok(())
301 }
302}
303
304pub(crate) struct MemoryClaimHandle {
308 pub(crate) state: Arc<Mutex<State>>,
309 pub(crate) pubsub: Arc<MemoryPubsub>,
310 pub(crate) queue: String,
311 pub(crate) job_id: JobId,
312 pub(crate) max_attempts: u32,
313}
314
315#[async_trait]
316impl ClaimHandleImpl for MemoryClaimHandle {
317 async fn ack(self: Box<Self>) -> Result<(), BusError> {
318 let mut g = self.state.lock().await;
319 if let Some(pq) = g.queues.get_mut(&self.queue) {
320 pq.entries.remove(&self.job_id);
321 pq.leased.remove(&self.job_id);
322 pq.ready.retain(|id| id != &self.job_id);
325 }
326 Ok(())
327 }
328
329 async fn nak(self: Box<Self>, delay: Duration) -> Result<(), BusError> {
330 let payload_to_dlq;
331 {
332 let mut g = self.state.lock().await;
333 let pq = match g.queues.get_mut(&self.queue) {
334 Some(q) => q,
335 None => return Ok(()),
336 };
337 let entry = match pq.entries.get_mut(&self.job_id) {
338 Some(e) => e,
339 None => return Ok(()),
340 };
341 if entry.attempts >= self.max_attempts {
342 let removed = pq
343 .entries
344 .remove(&self.job_id)
345 .expect("entry exists, just checked");
346 pq.leased.remove(&self.job_id);
347 pq.ready.retain(|id| id != &self.job_id);
348 payload_to_dlq = Some(removed.payload);
349 } else {
350 entry.lease_until = None;
351 entry.not_before = Some(Instant::now() + delay);
352 pq.leased.remove(&self.job_id);
353 pq.ready.push_back(self.job_id.clone());
354 payload_to_dlq = None;
355 }
356 }
357 if let Some(payload) = payload_to_dlq {
358 use klieo_core::bus::Headers;
359 use klieo_core::Pubsub;
360 let mut headers = Headers::new();
361 headers.insert("dlq_reason".into(), "max_attempts_exceeded".into());
362 let subject = format!("{}.dlq", self.queue);
363 self.pubsub.publish(&subject, payload, headers).await?;
364 }
365 Ok(())
366 }
367
368 async fn dead_letter(self: Box<Self>, reason: &str) -> Result<(), BusError> {
369 let payload = {
370 let mut g = self.state.lock().await;
371 let pq = match g.queues.get_mut(&self.queue) {
372 Some(q) => q,
373 None => return Ok(()),
374 };
375 let removed = match pq.entries.remove(&self.job_id) {
376 Some(e) => e,
377 None => return Ok(()),
378 };
379 pq.leased.remove(&self.job_id);
380 pq.ready.retain(|id| id != &self.job_id);
381 removed.payload
382 };
383 use klieo_core::bus::Headers;
384 use klieo_core::Pubsub;
385 let mut headers = Headers::new();
386 headers.insert("dlq_reason".into(), reason.into());
387 let subject = format!("{}.dlq", self.queue);
388 self.pubsub.publish(&subject, payload, headers).await
389 }
390}
391
392#[cfg(test)]
393mod tests {
394 use super::*;
395 use crate::kv::MemoryKv;
396 use crate::pubsub::MemoryPubsub;
397 use klieo_core::bus::Job;
398 use klieo_core::ids::{DurableName, RunId};
399 use klieo_core::{JobQueue, Pubsub};
400 use std::sync::Arc;
401 use std::time::Duration;
402
403 fn build_jq() -> MemoryJobQueue {
404 MemoryJobQueue::new(Arc::new(MemoryPubsub::new()), Arc::new(MemoryKv::new()))
405 }
406
407 #[tokio::test]
408 async fn claim_returns_none_when_empty() {
409 let jq = build_jq();
410 let claimed = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
411 assert!(claimed.is_none());
412 }
413
414 #[tokio::test]
415 async fn job_causation_round_trips_through_memory_queue() {
416 let jq = build_jq();
417 let run = RunId::new();
418 let job = Job::new(Bytes::from_static(b"x")).caused_by(run);
419 jq.enqueue("q", job).await.unwrap();
420 let claimed = jq
421 .claim("q", "w1", Duration::from_secs(5))
422 .await
423 .unwrap()
424 .unwrap();
425 assert_eq!(
426 claimed.causation_run_id,
427 Some(run),
428 "causation must survive enqueue->claim"
429 );
430 }
431
432 #[tokio::test]
433 async fn job_without_causation_claims_with_none() {
434 let jq = build_jq();
435 jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
436 .await
437 .unwrap();
438 let claimed = jq
439 .claim("q", "w1", Duration::from_secs(5))
440 .await
441 .unwrap()
442 .unwrap();
443 assert_eq!(claimed.causation_run_id, None);
444 }
445
446 #[tokio::test]
447 async fn enqueue_then_claim_returns_job() {
448 let jq = build_jq();
449 let id = jq
450 .enqueue("q", Job::new(Bytes::from_static(b"payload")))
451 .await
452 .unwrap();
453 let claimed = jq
454 .claim("q", "w1", Duration::from_secs(1))
455 .await
456 .unwrap()
457 .unwrap();
458 assert_eq!(claimed.id, id);
459 assert_eq!(claimed.payload, Bytes::from_static(b"payload"));
460 }
461
462 #[tokio::test]
463 async fn ack_removes_job() {
464 let jq = build_jq();
465 jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
466 .await
467 .unwrap();
468 let claimed = jq
469 .claim("q", "w1", Duration::from_secs(1))
470 .await
471 .unwrap()
472 .unwrap();
473 claimed.ack().await.unwrap();
474 let none = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
475 assert!(none.is_none());
476 }
477
478 #[tokio::test]
479 async fn lease_expiry_redelivers_job() {
480 let jq = build_jq();
481 jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
482 .await
483 .unwrap();
484 let claim1 = jq
485 .claim("q", "w1", Duration::from_millis(20))
486 .await
487 .unwrap()
488 .unwrap();
489 drop(claim1);
491 tokio::time::sleep(Duration::from_millis(40)).await;
492 let claim2 = jq
493 .claim("q", "w2", Duration::from_secs(1))
494 .await
495 .unwrap()
496 .unwrap();
497 assert_eq!(claim2.payload, Bytes::from_static(b"x"));
498 }
499
500 #[tokio::test]
501 async fn nak_with_delay_suppresses_claim_until_deadline() {
502 let jq = build_jq();
503 jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
504 .await
505 .unwrap();
506 let claim1 = jq
507 .claim("q", "w1", Duration::from_secs(1))
508 .await
509 .unwrap()
510 .unwrap();
511 claim1.nak(Duration::from_millis(40)).await.unwrap();
512 let none = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
514 assert!(
515 none.is_none(),
516 "should be suppressed during not_before window"
517 );
518 tokio::time::sleep(Duration::from_millis(60)).await;
519 let some = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
520 assert!(
521 some.is_some(),
522 "should be reclaimable after not_before elapses"
523 );
524 }
525
526 #[tokio::test]
527 async fn nak_after_max_attempts_routes_to_dlq() {
528 let jq = build_jq();
529 let pubsub_handle = jq.pubsub.clone();
530 let mut dlq_sub = pubsub_handle
531 .subscribe("q.dlq", DurableName::new("dlq"))
532 .await
533 .unwrap();
534
535 let mut job = Job::new(Bytes::from_static(b"x"));
536 job.max_attempts = Some(2);
537 jq.enqueue("q", job).await.unwrap();
538
539 for _ in 0..2 {
540 let claim = jq
541 .claim("q", "w1", Duration::from_secs(1))
542 .await
543 .unwrap()
544 .unwrap();
545 claim.nak(Duration::from_millis(0)).await.unwrap();
546 }
547
548 let none = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
551 assert!(none.is_none(), "queue should be empty after DLQ routing");
552
553 use tokio_stream::StreamExt;
554 let dlq_msg = tokio::time::timeout(Duration::from_millis(200), dlq_sub.next())
555 .await
556 .expect("DLQ delivery should arrive within 200ms")
557 .expect("stream did not close")
558 .expect("no error");
559 assert_eq!(dlq_msg.payload, Bytes::from_static(b"x"));
560 assert_eq!(
561 dlq_msg.headers.get("dlq_reason").map(|s| s.as_str()),
562 Some("max_attempts_exceeded")
563 );
564 }
565
566 #[tokio::test]
567 async fn dead_letter_removes_job_and_publishes_with_reason() {
568 let jq = build_jq();
569 use klieo_core::Pubsub;
570 let mut dlq_sub = jq
571 .pubsub
572 .subscribe("q.dlq", klieo_core::ids::DurableName::new("dlq"))
573 .await
574 .unwrap();
575 jq.enqueue("q", Job::new(Bytes::from_static(b"x")))
576 .await
577 .unwrap();
578 let claim = jq
579 .claim("q", "w1", Duration::from_secs(1))
580 .await
581 .unwrap()
582 .unwrap();
583 claim.dead_letter("explicit").await.unwrap();
584 let none = jq.claim("q", "w1", Duration::from_secs(1)).await.unwrap();
585 assert!(none.is_none(), "queue should be empty after DLQ");
586 use tokio_stream::StreamExt;
587 let msg = tokio::time::timeout(Duration::from_millis(200), dlq_sub.next())
588 .await
589 .unwrap()
590 .unwrap()
591 .unwrap();
592 assert_eq!(msg.payload, Bytes::from_static(b"x"));
593 assert_eq!(
594 msg.headers.get("dlq_reason").map(|s| s.as_str()),
595 Some("explicit")
596 );
597 }
598
599 #[tokio::test]
600 async fn dedup_key_returns_existing_id_on_duplicate_enqueue() {
601 let jq = build_jq();
602 let mut j1 = Job::new(Bytes::from_static(b"x"));
603 j1.dedup_key = Some("d1".into());
604 let id1 = jq.enqueue("q", j1).await.unwrap();
605 let mut j2 = Job::new(Bytes::from_static(b"y"));
606 j2.dedup_key = Some("d1".into());
607 let id2 = jq.enqueue("q", j2).await.unwrap();
608 assert_eq!(id1, id2, "duplicate dedup_key returns existing id");
609 let _claimed = jq
611 .claim("q", "w", Duration::from_secs(1))
612 .await
613 .unwrap()
614 .unwrap();
615 let none = jq.claim("q", "w", Duration::from_secs(1)).await.unwrap();
616 assert!(none.is_none(), "second enqueue must not have queued a job");
617 }
618
619 #[tokio::test]
620 async fn dedup_key_distinct_keys_enqueue_independently() {
621 let jq = build_jq();
622 let mut j1 = Job::new(Bytes::from_static(b"x"));
623 j1.dedup_key = Some("a".into());
624 let id1 = jq.enqueue("q", j1).await.unwrap();
625 let mut j2 = Job::new(Bytes::from_static(b"y"));
626 j2.dedup_key = Some("b".into());
627 let id2 = jq.enqueue("q", j2).await.unwrap();
628 assert_ne!(id1, id2);
629
630 let claim1 = jq
632 .claim("q", "w", Duration::from_secs(1))
633 .await
634 .unwrap()
635 .expect("first job should be claimable");
636 let claim2 = jq
637 .claim("q", "w", Duration::from_secs(1))
638 .await
639 .unwrap()
640 .expect("second job should be claimable");
641 assert_eq!(claim1.id, id1);
643 assert_eq!(claim2.id, id2);
644 assert_eq!(claim1.payload, Bytes::from_static(b"x"));
645 assert_eq!(claim2.payload, Bytes::from_static(b"y"));
646 }
647
648 #[tokio::test]
653 async fn claim_skips_leased_entries_fast_path() {
654 let jq = build_jq();
655 for i in 0..100u8 {
657 jq.enqueue("q", Job::new(Bytes::from(vec![i])))
658 .await
659 .unwrap();
660 }
661 let mut held = Vec::new();
663 for _ in 0..50 {
664 held.push(
665 jq.claim("q", "w", Duration::from_secs(10))
666 .await
667 .unwrap()
668 .unwrap(),
669 );
670 }
671 let claim51 = jq
673 .claim("q", "w", Duration::from_secs(10))
674 .await
675 .unwrap()
676 .expect("51st claim returns next ready entry");
677 assert_eq!(claim51.payload, Bytes::from(vec![50]));
678 drop(held);
680 }
681}