1use crate::{
2 backend::{NotificationStream, StorageBackend, StreamBackend},
3 model::{ConsumerGroupStatus, Event, Job, JobListItem, JobStatus, NewEvent, NewJob},
4};
5use async_trait::async_trait;
6use chrono::{DateTime, Utc};
7use std::{
8 collections::HashMap,
9 sync::{Arc, RwLock},
10};
11use tokio_stream::{wrappers::BroadcastStream, StreamExt};
12use uuid::Uuid;
13
14#[derive(Debug, Clone)]
16pub struct MemoryAttempt {
17 pub id: Uuid,
18 pub dataset_id: String,
19 pub job_id: Uuid,
20 pub attempt_no: i32,
21 pub status: String,
22 pub worker_id: String,
23 pub started_at: DateTime<Utc>,
24 pub finished_at: Option<DateTime<Utc>>,
25 pub latency_ms: Option<i32>,
26 pub error_code: Option<String>,
27 pub error_message: Option<String>,
28}
29
30#[derive(Debug, Default)]
31struct InnerState {
32 jobs: HashMap<Uuid, Job>,
33 archive: HashMap<Uuid, Job>,
34 attempts: HashMap<Uuid, MemoryAttempt>,
35 streams: HashMap<String, Vec<Event>>,
36 stream_offsets: HashMap<(String, String), ConsumerGroupStatus>,
37}
38
39#[derive(Debug, Clone, Default)]
43pub struct MemoryBackend {
44 state: Arc<RwLock<InnerState>>,
45 notifiers: Arc<RwLock<HashMap<String, tokio::sync::broadcast::Sender<()>>>>,
46 stream_notifiers: Arc<RwLock<HashMap<String, tokio::sync::broadcast::Sender<()>>>>,
47}
48
49impl MemoryBackend {
50 pub fn new() -> Self {
52 Self::default()
53 }
54
55 pub fn clear(&self) {
57 let mut state = self.state.write().unwrap();
58 state.jobs.clear();
59 state.archive.clear();
60 state.attempts.clear();
61 state.streams.clear();
62 state.stream_offsets.clear();
63 }
64
65 fn notify_queue(&self, queue: &str) {
66 let notifiers = self.notifiers.read().unwrap();
67 if let Some(tx) = notifiers.get(queue) {
68 let _ = tx.send(());
69 }
70 }
71
72 fn notify_stream(&self, stream: &str) {
73 let notifiers = self.stream_notifiers.read().unwrap();
74 if let Some(tx) = notifiers.get(stream) {
75 let _ = tx.send(());
76 }
77 }
78}
79
80#[async_trait]
81impl StorageBackend for MemoryBackend {
82 fn as_stream(&self) -> Option<&dyn StreamBackend> {
83 Some(self)
84 }
85 async fn run_migrations(&self) -> anyhow::Result<()> {
86 Ok(())
87 }
88
89 async fn health_check(&self) -> anyhow::Result<()> {
90 Ok(())
91 }
92
93 async fn enqueue(&self, job: NewJob) -> anyhow::Result<Uuid> {
94 let job_id = Uuid::new_v4();
95 let now = Utc::now();
96 let queue_name = job.queue.clone();
97
98 let job_entity = Job {
99 dataset_id: "default".to_string(),
100 replay_of_job_id: None,
101 id: job_id,
102 queue: job.queue,
103 job_type: job.job_type,
104 payload: job.payload_json,
105 run_at: job.run_at,
106 status: JobStatus::Queued.as_str().to_string(),
107 priority: job.priority,
108 max_attempts: job.max_attempts,
109 locked_at: None,
110 locked_by: None,
111 lock_expires_at: None,
112 dlq_reason_code: None,
113 dlq_at: None,
114 created_at: now,
115 updated_at: now,
116 };
117
118 {
119 let mut state = self.state.write().unwrap();
120 state.jobs.insert(job_id, job_entity);
121 }
122
123 self.notify_queue(&queue_name);
124 Ok(job_id)
125 }
126
127 async fn subscribe(&self, queue: &str) -> anyhow::Result<NotificationStream> {
128 let rx = {
129 let mut notifiers = self.notifiers.write().unwrap();
130 let tx = notifiers
131 .entry(queue.to_string())
132 .or_insert_with(|| tokio::sync::broadcast::channel(128).0);
133 tx.subscribe()
134 };
135
136 let stream = BroadcastStream::new(rx).filter_map(|res| res.ok());
137 Ok(Box::pin(stream))
138 }
139
140 async fn lease_jobs_batch(
141 &self,
142 queue: &str,
143 worker_id: &str,
144 lease_seconds: i64,
145 batch_size: i64,
146 ) -> anyhow::Result<Vec<Job>> {
147 self.lease_jobs_batch_with_ordering(
148 queue,
149 worker_id,
150 lease_seconds,
151 batch_size,
152 crate::model::QueueOrdering::Fifo,
153 )
154 .await
155 }
156
157 async fn lease_jobs_batch_with_ordering(
158 &self,
159 queue: &str,
160 worker_id: &str,
161 lease_seconds: i64,
162 batch_size: i64,
163 ordering: crate::model::QueueOrdering,
164 ) -> anyhow::Result<Vec<Job>> {
165 let mut state = self.state.write().unwrap();
166 let now = Utc::now();
167
168 let mut candidates: Vec<Job> = state
169 .jobs
170 .values()
171 .filter(|j| j.queue == queue && j.status == "queued" && j.run_at <= now)
172 .cloned()
173 .collect();
174
175 match ordering {
176 crate::model::QueueOrdering::Fifo => {
177 candidates.sort_by(|a, b| {
178 b.priority
179 .cmp(&a.priority)
180 .then_with(|| a.run_at.cmp(&b.run_at))
181 .then_with(|| a.created_at.cmp(&b.created_at))
182 .then_with(|| a.id.cmp(&b.id))
183 });
184 }
185 crate::model::QueueOrdering::Fastest => {
186 candidates.sort_by_key(|a| std::cmp::Reverse(a.priority));
187 }
188 }
189
190 let candidates: Vec<Job> = candidates.into_iter().take(batch_size as usize).collect();
191 if candidates.is_empty() {
192 return Ok(Vec::new());
193 }
194
195 let lock_expires_at = now + chrono::Duration::seconds(lease_seconds);
196 let mut leased = Vec::with_capacity(candidates.len());
197
198 for mut candidate in candidates {
199 if let Some(j) = state.jobs.get_mut(&candidate.id) {
200 j.status = JobStatus::Running.as_str().to_string();
201 j.locked_at = Some(now);
202 j.locked_by = Some(worker_id.to_string());
203 j.lock_expires_at = Some(lock_expires_at);
204 j.updated_at = now;
205
206 candidate.status = j.status.clone();
207 candidate.locked_at = j.locked_at;
208 candidate.locked_by = j.locked_by.clone();
209 candidate.lock_expires_at = j.lock_expires_at;
210 candidate.updated_at = j.updated_at;
211
212 leased.push(candidate);
213 }
214 }
215
216 Ok(leased)
217 }
218
219 async fn reap_expired_locks(&self) -> anyhow::Result<u64> {
220 let mut state = self.state.write().unwrap();
221 let now = Utc::now();
222 let mut reaped = 0u64;
223
224 for job in state.jobs.values_mut() {
225 if job.status == "running" {
226 if let Some(exp) = job.lock_expires_at {
227 if exp <= now {
228 job.status = JobStatus::Queued.as_str().to_string();
229 job.locked_at = None;
230 job.locked_by = None;
231 job.lock_expires_at = None;
232 job.updated_at = now;
233 reaped += 1;
234 }
235 }
236 }
237 }
238
239 Ok(reaped)
240 }
241
242 async fn start_attempts_batch(
243 &self,
244 _dataset_ids: &[String],
245 job_ids: &[Uuid],
246 worker_id: &str,
247 ) -> anyhow::Result<Vec<(Uuid, Uuid, i32)>> {
248 if job_ids.is_empty() {
249 return Ok(Vec::new());
250 }
251
252 let mut state = self.state.write().unwrap();
253 let now = Utc::now();
254 let mut results = Vec::with_capacity(job_ids.len());
255
256 for &job_id in job_ids {
257 let max_attempt = state
258 .attempts
259 .values()
260 .filter(|a| a.job_id == job_id)
261 .map(|a| a.attempt_no)
262 .max()
263 .unwrap_or(0);
264
265 let next_attempt_no = max_attempt + 1;
266 let attempt_id = Uuid::new_v4();
267
268 let attempt = MemoryAttempt {
269 id: attempt_id,
270 dataset_id: "default".to_string(),
271 job_id,
272 attempt_no: next_attempt_no,
273 status: "running".to_string(),
274 worker_id: worker_id.to_string(),
275 started_at: now,
276 finished_at: None,
277 latency_ms: None,
278 error_code: None,
279 error_message: None,
280 };
281
282 state.attempts.insert(attempt_id, attempt);
283 results.push((job_id, attempt_id, next_attempt_no));
284 }
285
286 Ok(results)
287 }
288
289 async fn mark_succeeded(
290 &self,
291 job_id: Uuid,
292 attempt_id: Uuid,
293 _worker_id: &str,
294 latency_ms: i32,
295 ) -> anyhow::Result<()> {
296 let mut state = self.state.write().unwrap();
297 let now = Utc::now();
298
299 if let Some(att) = state.attempts.get_mut(&attempt_id) {
300 att.status = "succeeded".to_string();
301 att.finished_at = Some(now);
302 att.latency_ms = Some(latency_ms);
303 }
304
305 if let Some(job) = state.jobs.get_mut(&job_id) {
306 job.status = JobStatus::Succeeded.as_str().to_string();
307 job.locked_at = None;
308 job.locked_by = None;
309 job.lock_expires_at = None;
310 job.updated_at = now;
311 }
312
313 Ok(())
314 }
315
316 async fn mark_succeeded_batch(
317 &self,
318 _dataset_id: &str,
319 updates: &[(Uuid, Uuid, i32)],
320 worker_id: &str,
321 ) -> anyhow::Result<()> {
322 for &(job_id, attempt_id, latency_ms) in updates {
323 self.mark_succeeded(job_id, attempt_id, worker_id, latency_ms)
324 .await?;
325 }
326 Ok(())
327 }
328
329 #[allow(clippy::too_many_arguments)]
330 async fn reschedule_for_retry(
331 &self,
332 job_id: Uuid,
333 attempt_id: Uuid,
334 _worker_id: &str,
335 latency_ms: i32,
336 next_run_at: DateTime<Utc>,
337 error_code: &str,
338 error_message: &str,
339 _attempt_no: i32,
340 ) -> anyhow::Result<()> {
341 let mut state = self.state.write().unwrap();
342 let now = Utc::now();
343
344 if let Some(att) = state.attempts.get_mut(&attempt_id) {
345 att.status = "failed".to_string();
346 att.finished_at = Some(now);
347 att.latency_ms = Some(latency_ms);
348 att.error_code = Some(error_code.to_string());
349 att.error_message = Some(error_message.to_string());
350 }
351
352 if let Some(job) = state.jobs.get_mut(&job_id) {
353 job.status = JobStatus::Queued.as_str().to_string();
354 job.run_at = next_run_at;
355 job.locked_at = None;
356 job.locked_by = None;
357 job.lock_expires_at = None;
358 job.updated_at = now;
359 }
360
361 Ok(())
362 }
363
364 #[allow(clippy::too_many_arguments)]
365 async fn mark_dlq(
366 &self,
367 job_id: Uuid,
368 attempt_id: Uuid,
369 _worker_id: &str,
370 latency_ms: i32,
371 reason_code: &str,
372 error_code: &str,
373 error_message: &str,
374 _attempt_no: i32,
375 ) -> anyhow::Result<()> {
376 let mut state = self.state.write().unwrap();
377 let now = Utc::now();
378
379 if let Some(att) = state.attempts.get_mut(&attempt_id) {
380 att.status = "failed".to_string();
381 att.finished_at = Some(now);
382 att.latency_ms = Some(latency_ms);
383 att.error_code = Some(error_code.to_string());
384 att.error_message = Some(error_message.to_string());
385 }
386
387 if let Some(job) = state.jobs.get_mut(&job_id) {
388 job.status = JobStatus::Dlq.as_str().to_string();
389 job.dlq_reason_code = Some(reason_code.to_string());
390 job.dlq_at = Some(now);
391 job.locked_at = None;
392 job.locked_by = None;
393 job.lock_expires_at = None;
394 job.updated_at = now;
395 }
396
397 Ok(())
398 }
399
400 async fn archive_succeeded_older_than(
401 &self,
402 cutoff: DateTime<Utc>,
403 limit: i64,
404 ) -> anyhow::Result<u64> {
405 let mut state = self.state.write().unwrap();
406
407 let to_archive: Vec<Uuid> = state
408 .jobs
409 .values()
410 .filter(|j| j.status == "succeeded" && j.updated_at < cutoff)
411 .take(limit as usize)
412 .map(|j| j.id)
413 .collect();
414
415 let count = to_archive.len() as u64;
416 for id in to_archive {
417 if let Some(job) = state.jobs.remove(&id) {
418 state.archive.insert(id, job);
419 }
420 }
421
422 Ok(count)
423 }
424
425 async fn delete_history_for_succeeded_older_than(
426 &self,
427 cutoff: DateTime<Utc>,
428 limit: i64,
429 ) -> anyhow::Result<(u64, u64)> {
430 let mut state = self.state.write().unwrap();
431
432 let archived_ids: Vec<Uuid> = state
433 .archive
434 .values()
435 .filter(|j| j.updated_at < cutoff)
436 .map(|j| j.id)
437 .collect();
438
439 let to_remove: Vec<Uuid> = state
440 .attempts
441 .values()
442 .filter(|a| a.started_at < cutoff && archived_ids.contains(&a.job_id))
443 .take(limit as usize)
444 .map(|a| a.id)
445 .collect();
446
447 let count = to_remove.len() as u64;
448 for aid in to_remove {
449 state.attempts.remove(&aid);
450 }
451
452 Ok((count, 0))
453 }
454
455 async fn perform_maintenance(&self) -> anyhow::Result<()> {
456 Ok(())
457 }
458
459 async fn extend_lease(
460 &self,
461 job_id: Uuid,
462 worker_id: &str,
463 lease_seconds: i64,
464 ) -> anyhow::Result<bool> {
465 let mut state = self.state.write().unwrap();
466 if let Some(job) = state.jobs.get_mut(&job_id) {
467 if job.status == "running" && job.locked_by.as_deref() == Some(worker_id) {
468 let now = Utc::now();
469 job.lock_expires_at = Some(now + chrono::Duration::seconds(lease_seconds));
470 job.updated_at = now;
471 return Ok(true);
472 }
473 }
474 Ok(false)
475 }
476
477 async fn get_job(&self, job_id: Uuid) -> anyhow::Result<Option<Job>> {
478 let state = self.state.read().unwrap();
479 let job = state
480 .jobs
481 .get(&job_id)
482 .or_else(|| state.archive.get(&job_id))
483 .cloned();
484 Ok(job)
485 }
486
487 async fn list_jobs(
488 &self,
489 queue: Option<&str>,
490 status: Option<&str>,
491 limit: i64,
492 _cursor_created_at: Option<DateTime<Utc>>,
493 _cursor_id: Option<Uuid>,
494 ) -> anyhow::Result<Vec<JobListItem>> {
495 let state = self.state.read().unwrap();
496 let limit = limit.clamp(1, 500) as usize;
497
498 let mut items: Vec<JobListItem> = state
499 .jobs
500 .values()
501 .filter(|j| {
502 if let Some(q) = queue {
503 if j.queue != q {
504 return false;
505 }
506 }
507 if let Some(st) = status {
508 if j.status != st {
509 return false;
510 }
511 }
512 true
513 })
514 .map(|j| JobListItem {
515 id: j.id,
516 queue: j.queue.clone(),
517 job_type: j.job_type.clone(),
518 status: j.status.clone(),
519 run_at: j.run_at,
520 priority: j.priority,
521 max_attempts: j.max_attempts,
522 last_error_code: None,
523 last_error_message: None,
524 dlq_reason_code: j.dlq_reason_code.clone(),
525 created_at: j.created_at,
526 updated_at: j.updated_at,
527 })
528 .collect();
529
530 items.sort_by(|a, b| {
531 b.created_at
532 .cmp(&a.created_at)
533 .then_with(|| b.id.cmp(&a.id))
534 });
535
536 items.truncate(limit);
537 Ok(items)
538 }
539
540 async fn replay_job(
541 &self,
542 job_id: Uuid,
543 override_queue: Option<&str>,
544 override_run_at: Option<DateTime<Utc>>,
545 ) -> anyhow::Result<Uuid> {
546 let mut state = self.state.write().unwrap();
547
548 let src = state
549 .jobs
550 .get(&job_id)
551 .or_else(|| state.archive.get(&job_id))
552 .cloned()
553 .ok_or_else(|| anyhow::anyhow!("Job {job_id} not found"))?;
554
555 let new_id = Uuid::new_v4();
556 let now = Utc::now();
557 let target_queue = override_queue.unwrap_or(&src.queue).to_string();
558 let target_run_at = override_run_at.unwrap_or(now);
559
560 let target_queue_clone = target_queue.clone();
561 let new_job = Job {
562 dataset_id: "default".to_string(),
563 replay_of_job_id: Some(job_id),
564 id: new_id,
565 queue: target_queue,
566 job_type: src.job_type,
567 payload: src.payload,
568 run_at: target_run_at,
569 status: JobStatus::Queued.as_str().to_string(),
570 priority: src.priority,
571 max_attempts: src.max_attempts,
572 locked_at: None,
573 locked_by: None,
574 lock_expires_at: None,
575 dlq_reason_code: None,
576 dlq_at: None,
577 created_at: now,
578 updated_at: now,
579 };
580
581 state.jobs.insert(new_id, new_job);
582 drop(state);
583 self.notify_queue(&target_queue_clone);
584 Ok(new_id)
585 }
586}
587
588#[async_trait]
589impl StreamBackend for MemoryBackend {
590 async fn publish(&self, stream: &str, event: NewEvent) -> anyhow::Result<i64> {
591 let mut state = self.state.write().unwrap();
592 let log = state.streams.entry(stream.to_string()).or_default();
593 let sequence_no = (log.len() + 1) as i64;
594 let now = Utc::now();
595
596 let event_entity = Event {
597 sequence_no,
598 stream_name: stream.to_string(),
599 event_type: event.event_type,
600 payload_json: event.payload_json,
601 created_at: now,
602 };
603
604 log.push(event_entity);
605 drop(state);
606
607 self.notify_stream(stream);
608 Ok(sequence_no)
609 }
610
611 async fn subscribe_stream(
612 &self,
613 stream: &str,
614 _consumer_group: &str,
615 _last_seq: Option<i64>,
616 ) -> anyhow::Result<NotificationStream> {
617 let rx = {
618 let mut notifiers = self.stream_notifiers.write().unwrap();
619 let tx = notifiers
620 .entry(stream.to_string())
621 .or_insert_with(|| tokio::sync::broadcast::channel(128).0);
622 tx.subscribe()
623 };
624
625 let stream = BroadcastStream::new(rx).filter_map(|res| res.ok());
626 Ok(Box::pin(stream))
627 }
628
629 async fn ack(&self, stream: &str, consumer_group: &str, seq: i64) -> anyhow::Result<()> {
630 let mut state = self.state.write().unwrap();
631 let key = (stream.to_string(), consumer_group.to_string());
632 let now = Utc::now();
633
634 let entry = state
635 .stream_offsets
636 .entry(key)
637 .or_insert_with(|| ConsumerGroupStatus {
638 consumer_group: consumer_group.to_string(),
639 stream_name: stream.to_string(),
640 last_acked_seq: 0,
641 updated_at: now,
642 });
643
644 if seq > entry.last_acked_seq {
645 entry.last_acked_seq = seq;
646 entry.updated_at = now;
647 }
648
649 Ok(())
650 }
651
652 async fn read_events(
653 &self,
654 stream: &str,
655 after_seq: i64,
656 limit: i64,
657 ) -> anyhow::Result<Vec<Event>> {
658 let state = self.state.read().unwrap();
659 let limit = limit.clamp(1, 1000) as usize;
660
661 if let Some(log) = state.streams.get(stream) {
662 let events: Vec<Event> = log
663 .iter()
664 .filter(|e| e.sequence_no > after_seq)
665 .take(limit)
666 .cloned()
667 .collect();
668 Ok(events)
669 } else {
670 Ok(Vec::new())
671 }
672 }
673
674 async fn consumer_group_info(&self, stream: &str) -> anyhow::Result<Vec<ConsumerGroupStatus>> {
675 let state = self.state.read().unwrap();
676 let info: Vec<ConsumerGroupStatus> = state
677 .stream_offsets
678 .values()
679 .filter(|cg| cg.stream_name == stream)
680 .cloned()
681 .collect();
682 Ok(info)
683 }
684}