1mod config;
81mod enqueue_rate;
82mod fleet;
83pub mod keys;
84
85use std::sync::Arc;
86
87use async_trait::async_trait;
88use boson_core::{
89 BosonError, IdempotencyMode, Job, JobEnqueueDisposition, JobStatus, QueueBackend, Result, Run,
90 RunStatus, TaskConfig, TaskRunStats,
91};
92use chrono::{DateTime, Utc};
93use enqueue_rate::EnqueueRateLimiter;
94use redis::aio::ConnectionManager;
95use redis::{AsyncCommands, Script};
96use serde::{Deserialize, Serialize};
97use uuid::Uuid;
98
99pub use config::RedisQueueConfig;
100pub use fleet::connect_fleet_from_env;
101
102#[derive(Debug, Clone, Serialize, Deserialize)]
104struct LeaseRow {
105 lease_id: String,
106 job_id: String,
107 worker_id: String,
108 expires_at: DateTime<Utc>,
109}
110
111pub struct RedisQueueBackend {
116 conn: ConnectionManager,
117 keys: keys::Keyspace,
118 enqueue_rate: EnqueueRateLimiter,
119}
120
121impl std::fmt::Debug for RedisQueueBackend {
122 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123 f.debug_struct("RedisQueueBackend").finish_non_exhaustive()
124 }
125}
126
127impl RedisQueueBackend {
128 pub async fn connect(config: RedisQueueConfig) -> Result<Self> {
146 Self::connect_with_keyspace(
147 &config.url,
148 keys::Keyspace::new(config.key_prefix),
149 )
150 .await
151 }
152
153 pub async fn connect_url(url: &str) -> Result<Self> {
159 Self::connect(RedisQueueConfig {
160 url: url.into(),
161 ..Default::default()
162 })
163 .await
164 }
165
166 pub async fn connect_with_keyspace(url: &str, keyspace: keys::Keyspace) -> Result<Self> {
172 let client = redis::Client::open(url).map_err(map_err)?;
173 let conn = ConnectionManager::new(client).await.map_err(map_err)?;
174 Ok(Self {
175 conn,
176 keys: keyspace,
177 enqueue_rate: EnqueueRateLimiter::new(),
178 })
179 }
180
181 #[must_use]
183 pub fn test_url() -> String {
184 std::env::var("BOSON_TEST_REDIS_URL")
185 .unwrap_or_else(|_| "redis://127.0.0.1:6379".into())
186 }
187
188 pub async fn flush_boson_keys(&self) -> Result<()> {
194 let pattern = self.keys.scan_pattern();
195 let mut conn = self.conn.clone();
196 let mut cursor = 0_u64;
197 loop {
198 let (next, batch): (u64, Vec<String>) = redis::cmd("SCAN")
199 .arg(cursor)
200 .arg("MATCH")
201 .arg(&pattern)
202 .arg("COUNT")
203 .arg(500)
204 .query_async(&mut conn)
205 .await
206 .map_err(map_err)?;
207 if !batch.is_empty() {
208 let _: () = conn.del(batch).await.map_err(map_err)?;
209 }
210 cursor = next;
211 if cursor == 0 {
212 break;
213 }
214 }
215 Ok(())
216 }
217
218 async fn load_job(&self, job_id: &str) -> Result<Option<Job>> {
219 let mut conn = self.conn.clone();
220 let raw: Option<String> = conn.get(self.keys.job(job_id)).await.map_err(map_err)?;
221 raw.map_or(Ok(None), |s| serde_json::from_str(&s).map_err(map_err).map(Some))
222 }
223
224 async fn save_job(&self, job: &Job) -> Result<()> {
225 let mut conn = self.conn.clone();
226 let json = serde_json::to_string(job).map_err(map_err)?;
227 let _: () = conn.set(self.keys.job(&job.job_id), json).await.map_err(map_err)?;
228 Ok(())
229 }
230
231 async fn add_ready(&self, job: &Job) -> Result<()> {
232 if job.status != JobStatus::Queued {
233 return Ok(());
234 }
235 let mut conn = self.conn.clone();
236 let score = keys::ready_score(job.priority, job.created_at.timestamp_millis());
237 let ready_key = self.keys.ready(&job.pool);
238 let _: () = conn
239 .zadd(&ready_key, job.job_id.as_str(), score)
240 .await
241 .map_err(map_err)?;
242 let _: () = conn
243 .sadd(self.keys.pools_set(), job.pool.as_str())
244 .await
245 .map_err(map_err)?;
246 Ok(())
247 }
248
249 async fn remove_ready(&self, pool: &str, job_id: &str) -> Result<()> {
250 let mut conn = self.conn.clone();
251 let ready_key = self.keys.ready(pool);
252 let _: () = conn.zrem(&ready_key, job_id).await.map_err(map_err)?;
253 let len: i64 = conn.zcard(&ready_key).await.map_err(map_err)?;
254 if len == 0 {
255 let _: () = conn
256 .srem(self.keys.pools_set(), pool)
257 .await
258 .map_err(map_err)?;
259 }
260 Ok(())
261 }
262
263 async fn load_run(&self, run_id: &str) -> Result<Option<Run>> {
264 let mut conn = self.conn.clone();
265 let raw: Option<String> = conn.get(self.keys.run(run_id)).await.map_err(map_err)?;
266 raw.map_or(Ok(None), |s| serde_json::from_str(&s).map_err(map_err).map(Some))
267 }
268
269 async fn save_run(&self, run: &Run) -> Result<()> {
270 let mut conn = self.conn.clone();
271 let json = serde_json::to_string(run).map_err(map_err)?;
272 let _: () = conn.set(self.keys.run(&run.run_id), json).await.map_err(map_err)?;
273 Ok(())
274 }
275}
276
277const CLAIM_SCRIPT: &str = r#"
278local raw = redis.call('GET', KEYS[1])
279if not raw then return nil end
280if not string.find(raw, '"status":"queued"', 1, true) then return nil end
281local updated = string.gsub(raw, '"status":"queued"', '"status":"running"', 1)
282redis.call('SET', KEYS[1], updated)
283redis.call('ZREM', KEYS[2], ARGV[1])
284return updated
285"#;
286
287const POP_CLAIM_SCRIPT: &str = r#"
288local ids = redis.call('ZRANGE', KEYS[1], 0, 0)
289if #ids == 0 then return nil end
290local job_id = ids[1]
291local job_key = KEYS[2] .. job_id
292local raw = redis.call('GET', job_key)
293if not raw then
294 redis.call('ZREM', KEYS[1], job_id)
295 return nil
296end
297if not string.find(raw, '"status":"queued"', 1, true) then
298 redis.call('ZREM', KEYS[1], job_id)
299 return nil
300end
301local updated = string.gsub(raw, '"status":"queued"', '"status":"running"', 1)
302redis.call('SET', job_key, updated)
303redis.call('ZREM', KEYS[1], job_id)
304return updated
305"#;
306
307fn map_err(e: impl std::fmt::Display) -> BosonError {
308 BosonError::Backend(e.to_string())
309}
310
311#[async_trait]
312impl QueueBackend for RedisQueueBackend {
313 async fn upsert_job(&self, job: &Job) -> Result<()> {
314 let existing = self.load_job(&job.job_id).await?;
315 if let Some(ref old) = existing {
316 if old.status == JobStatus::Queued && job.status != JobStatus::Queued {
317 self.remove_ready(&old.pool, &job.job_id).await?;
318 } else if job.status == JobStatus::Queued {
319 self.remove_ready(&old.pool, &job.job_id).await?;
320 self.add_ready(job).await?;
321 }
322 } else if job.status == JobStatus::Queued {
323 self.add_ready(job).await?;
324 }
325 self.save_job(job).await
326 }
327
328 async fn enqueue_with_policies(
329 &self,
330 job: Job,
331 task_config: &TaskConfig,
332 ) -> Result<(String, JobEnqueueDisposition)> {
333 let idempotency = task_config.resolved_idempotency_mode(IdempotencyMode::Lwt);
334 let mut job = job;
335 if idempotency == IdempotencyMode::Lwt {
336 if let Some(ref key) = job.idempotency_key {
337 if !key.is_empty() {
338 let mut conn = self.conn.clone();
339 let idem_key = self.keys.idempotency(key);
340 let inserted: bool = conn
341 .set_nx(&idem_key, job.job_id.as_str())
342 .await
343 .map_err(map_err)?;
344 if !inserted {
345 let existing_id: Option<String> =
346 conn.get(&idem_key).await.map_err(map_err)?;
347 if let Some(ref prior_id) = existing_id {
348 if let Some(prior) = self.load_job(prior_id).await? {
349 if matches!(prior.status, JobStatus::Queued | JobStatus::Running)
350 {
351 return Ok((
352 prior_id.clone(),
353 JobEnqueueDisposition::ReusedIdempotent,
354 ));
355 }
356 }
357 let _: () = conn
358 .set(&idem_key, job.job_id.as_str())
359 .await
360 .map_err(map_err)?;
361 }
362 }
363 }
364 }
365 } else {
366 job.idempotency_key = None;
367 }
368
369 let policy = &task_config.rate_limit_policy;
370 if policy.max_in_flight > 0 {
371 let count = self.count_active_jobs_for_task(&job.task_name).await?;
372 if count >= policy.max_in_flight {
373 return Err(BosonError::RateLimited(job.task_name.clone()));
374 }
375 }
376 if policy.max_enqueue_per_second > 0
377 && !self
378 .enqueue_rate
379 .try_record(&job.task_name, policy.max_enqueue_per_second)
380 {
381 return Err(BosonError::RateLimited(job.task_name.clone()));
382 }
383
384 let job_id = job.job_id.clone();
385 self.save_job(&job).await?;
386 self.add_ready(&job).await?;
387 Ok((job_id, JobEnqueueDisposition::InsertedNew))
388 }
389
390 async fn get_job(&self, job_id: &str) -> Result<Option<Job>> {
391 self.load_job(job_id).await
392 }
393
394 async fn list_jobs(
395 &self,
396 status_filter: Option<JobStatus>,
397 offset: usize,
398 limit: usize,
399 ) -> Result<Vec<Job>> {
400 let pattern = self.keys.job_pattern();
401 let mut conn = self.conn.clone();
402 let mut jobs = Vec::new();
403 let mut cursor = 0_u64;
404 loop {
405 let (next, batch): (u64, Vec<String>) = redis::cmd("SCAN")
406 .arg(cursor)
407 .arg("MATCH")
408 .arg(&pattern)
409 .arg("COUNT")
410 .arg(200)
411 .query_async(&mut conn)
412 .await
413 .map_err(map_err)?;
414 for key in batch {
415 let raw: Option<String> = conn.get(&key).await.map_err(map_err)?;
416 if let Some(s) = raw {
417 if let Ok(job) = serde_json::from_str::<Job>(&s) {
418 if status_filter.is_none_or(|st| job.status == st) {
419 jobs.push(job);
420 }
421 }
422 }
423 }
424 cursor = next;
425 if cursor == 0 {
426 break;
427 }
428 }
429 jobs.sort_by_key(|j| j.created_at);
430 Ok(jobs.into_iter().skip(offset).take(limit).collect())
431 }
432
433 async fn cancel_job_if_active(&self, job_id: &str) -> Result<()> {
434 let Some(mut job) = self.load_job(job_id).await? else {
435 return Err(BosonError::JobNotFound(job_id.to_string()));
436 };
437 if !matches!(job.status, JobStatus::Queued | JobStatus::Running) {
438 return Ok(());
439 }
440 if job.status == JobStatus::Queued {
441 self.remove_ready(&job.pool, job_id).await?;
442 }
443 job.status = JobStatus::Canceled;
444 self.save_job(&job).await
445 }
446
447 async fn try_claim_job(&self, job_id: &str) -> Result<Option<Job>> {
448 let Some(job) = self.load_job(job_id).await? else {
449 return Ok(None);
450 };
451 if job.status != JobStatus::Queued {
452 return Ok(None);
453 }
454 let script = Script::new(CLAIM_SCRIPT);
455 let result: Option<String> = script
456 .key(self.keys.job(job_id))
457 .key(self.keys.ready(&job.pool))
458 .arg(job_id)
459 .invoke_async(&mut self.conn.clone())
460 .await
461 .map_err(map_err)?;
462 result
463 .and_then(|s| serde_json::from_str(&s).ok())
464 .map_or(Ok(None), |j| Ok(Some(j)))
465 }
466
467 async fn revert_job_to_queued(&self, job_id: &str) -> Result<()> {
468 let Some(mut job) = self.load_job(job_id).await? else {
469 return Ok(());
470 };
471 if job.status != JobStatus::Running {
472 return Ok(());
473 }
474 job.status = JobStatus::Queued;
475 self.save_job(&job).await?;
476 self.add_ready(&job).await
477 }
478
479 async fn distinct_pools_queued(&self) -> Result<Vec<String>> {
480 let mut conn = self.conn.clone();
481 let pools: Vec<String> = conn.smembers(self.keys.pools_set()).await.map_err(map_err)?;
482 let mut out = pools;
483 out.sort();
484 Ok(out)
485 }
486
487 async fn list_queued_for_pool_sorted(&self, pool: &str, limit: usize) -> Result<Vec<Job>> {
488 let limit = limit.max(1);
489 let mut conn = self.conn.clone();
490 let ids: Vec<String> = conn
491 .zrange(self.keys.ready(pool), 0, isize::try_from(limit.saturating_sub(1)).unwrap_or(0))
492 .await
493 .map_err(map_err)?;
494 if ids.is_empty() {
495 return Ok(Vec::new());
496 }
497 let job_keys: Vec<String> = ids.iter().map(|id| self.keys.job(id)).collect();
498 let mut pipe = redis::pipe();
499 for key in &job_keys {
500 pipe.get(key);
501 }
502 let raws: Vec<Option<String>> = pipe.query_async(&mut conn).await.map_err(map_err)?;
503 let mut jobs = Vec::new();
504 for (id, raw) in ids.iter().zip(raws) {
505 if let Some(s) = raw {
506 if let Ok(job) = serde_json::from_str::<Job>(&s) {
507 if job.status == JobStatus::Queued && job.pool == pool {
508 jobs.push(job);
509 }
510 }
511 } else {
512 let _: () = conn.zrem(self.keys.ready(pool), id).await.map_err(map_err)?;
513 }
514 }
515 Ok(jobs)
516 }
517
518 async fn pop_claim_from_pool(&self, pool: &str) -> Result<Option<Job>> {
519 let script = Script::new(POP_CLAIM_SCRIPT);
520 let result: Option<String> = script
521 .key(self.keys.ready(pool))
522 .key(self.keys.job_key_prefix())
523 .invoke_async(&mut self.conn.clone())
524 .await
525 .map_err(map_err)?;
526 result
527 .and_then(|s| serde_json::from_str(&s).ok())
528 .map_or(Ok(None), |j| Ok(Some(j)))
529 }
530
531 async fn count_jobs(&self, status_filter: Option<JobStatus>) -> Result<u64> {
532 let jobs = self.list_jobs(status_filter, 0, usize::MAX).await?;
533 Ok(u64::try_from(jobs.len()).unwrap_or(u64::MAX))
534 }
535
536 async fn count_jobs_for_task(
537 &self,
538 task_name: &str,
539 status: Option<JobStatus>,
540 ) -> Result<u64> {
541 let jobs = self.list_jobs(status, 0, usize::MAX).await?;
542 let count = jobs.iter().filter(|j| j.task_name == task_name).count();
543 Ok(u64::try_from(count).unwrap_or(u64::MAX))
544 }
545
546 async fn count_active_jobs_for_task(&self, task_name: &str) -> Result<u32> {
547 let jobs = self.list_jobs(None, 0, usize::MAX).await?;
548 let count = jobs
549 .iter()
550 .filter(|j| {
551 j.task_name == task_name
552 && matches!(j.status, JobStatus::Queued | JobStatus::Running)
553 })
554 .count();
555 Ok(u32::try_from(count).unwrap_or(u32::MAX))
556 }
557
558 async fn find_nonterminal_by_idempotency_key(&self, key: &str) -> Result<Option<String>> {
559 if key.is_empty() {
560 return Ok(None);
561 }
562 let mut conn = self.conn.clone();
563 let job_id: Option<String> = conn.get(self.keys.idempotency(key)).await.map_err(map_err)?;
564 let Some(job_id) = job_id else {
565 return Ok(None);
566 };
567 if let Some(job) = self.load_job(&job_id).await? {
568 if matches!(job.status, JobStatus::Queued | JobStatus::Running) {
569 return Ok(Some(job_id));
570 }
571 }
572 Ok(None)
573 }
574
575 async fn upsert_run(&self, run: &Run) -> Result<()> {
576 self.save_run(run).await
577 }
578
579 async fn get_run(&self, run_id: &str) -> Result<Option<Run>> {
580 self.load_run(run_id).await
581 }
582
583 async fn list_runs(
584 &self,
585 job_id_filter: Option<&str>,
586 offset: usize,
587 limit: usize,
588 ) -> Result<Vec<Run>> {
589 let pattern = self.keys.run_pattern();
590 let mut conn = self.conn.clone();
591 let mut runs = Vec::new();
592 let mut cursor = 0_u64;
593 loop {
594 let (next, batch): (u64, Vec<String>) = redis::cmd("SCAN")
595 .arg(cursor)
596 .arg("MATCH")
597 .arg(&pattern)
598 .arg("COUNT")
599 .arg(200)
600 .query_async(&mut conn)
601 .await
602 .map_err(map_err)?;
603 for key in batch {
604 let raw: Option<String> = conn.get(&key).await.map_err(map_err)?;
605 if let Some(s) = raw {
606 if let Ok(run) = serde_json::from_str::<Run>(&s) {
607 if job_id_filter.is_none_or(|id| run.job_id == id) {
608 runs.push(run);
609 }
610 }
611 }
612 }
613 cursor = next;
614 if cursor == 0 {
615 break;
616 }
617 }
618 runs.sort_by_key(|r| r.started_at);
619 Ok(runs.into_iter().skip(offset).take(limit).collect())
620 }
621
622 async fn finish_run(
623 &self,
624 run_id: &str,
625 status: RunStatus,
626 duration_ms: Option<i64>,
627 error_message: Option<String>,
628 ) -> Result<()> {
629 let Some(mut run) = self.load_run(run_id).await? else {
630 return Ok(());
631 };
632 run.status = status;
633 run.finished_at = Some(Utc::now());
634 run.duration_ms = duration_ms;
635 run.error_message = error_message;
636 self.save_run(&run).await
637 }
638
639 async fn count_runs(&self, job_id_filter: Option<&str>) -> Result<u64> {
640 let runs = self.list_runs(job_id_filter, 0, usize::MAX).await?;
641 Ok(u64::try_from(runs.len()).unwrap_or(u64::MAX))
642 }
643
644 async fn count_runs_since(&self, since: DateTime<Utc>) -> Result<u64> {
645 let runs = self.list_runs(None, 0, usize::MAX).await?;
646 let count = runs.iter().filter(|r| r.started_at >= since).count();
647 Ok(u64::try_from(count).unwrap_or(u64::MAX))
648 }
649
650 async fn task_run_stats(&self, task_name: &str) -> Result<TaskRunStats> {
651 let runs = self.list_runs(None, 0, usize::MAX).await?;
652 let filtered: Vec<_> = runs.iter().filter(|r| r.task_name == task_name).collect();
653 let runs_total = u32::try_from(filtered.len()).unwrap_or(u32::MAX);
654 let success_count = u32::try_from(
655 filtered
656 .iter()
657 .filter(|r| r.status == RunStatus::Success)
658 .count(),
659 )
660 .unwrap_or(u32::MAX);
661 Ok(TaskRunStats {
662 runs_total,
663 success_count,
664 })
665 }
666
667 async fn get_task_config(&self, task_name: &str) -> Result<Option<TaskConfig>> {
668 let mut conn = self.conn.clone();
669 let raw: Option<String> = conn.get(self.keys.task_config(task_name)).await.map_err(map_err)?;
670 raw.map_or(Ok(None), |s| serde_json::from_str(&s).map_err(map_err).map(Some))
671 }
672
673 async fn upsert_task_config(&self, config: &TaskConfig) -> Result<()> {
674 let mut conn = self.conn.clone();
675 let json = serde_json::to_string(config).map_err(map_err)?;
676 let _: () = conn
677 .set(self.keys.task_config(&config.task_name), json)
678 .await
679 .map_err(map_err)?;
680 Ok(())
681 }
682
683 async fn try_claim_run_lease(
684 &self,
685 job_id: &str,
686 worker_id: &str,
687 ttl_secs: i64,
688 ) -> Result<Option<String>> {
689 let mut conn = self.conn.clone();
690 let existing: Option<String> = conn.get(self.keys.lease_by_job(job_id)).await.map_err(map_err)?;
691 if let Some(ref lid) = existing {
692 let raw: Option<String> = conn.get(self.keys.lease(lid)).await.map_err(map_err)?;
693 if let Some(s) = raw {
694 if let Ok(row) = serde_json::from_str::<LeaseRow>(&s) {
695 if row.expires_at > Utc::now() {
696 return Ok(None);
697 }
698 }
699 }
700 }
701 let lease_id = Uuid::new_v4().to_string();
702 let row = LeaseRow {
703 lease_id: lease_id.clone(),
704 job_id: job_id.to_string(),
705 worker_id: worker_id.to_string(),
706 expires_at: Utc::now() + chrono::Duration::seconds(ttl_secs),
707 };
708 let json = serde_json::to_string(&row).map_err(map_err)?;
709 let inserted: bool = conn
710 .set_nx(self.keys.lease_by_job(job_id), lease_id.as_str())
711 .await
712 .map_err(map_err)?;
713 if !inserted {
714 return Ok(None);
715 }
716 let _: () = conn.set(self.keys.lease(&lease_id), json).await.map_err(map_err)?;
717 Ok(Some(lease_id))
718 }
719
720 async fn extend_lease(&self, lease_id: &str, ttl_secs: i64) -> Result<()> {
721 let Some(mut row) = self
722 .load_lease_row(lease_id)
723 .await?
724 else {
725 return Ok(());
726 };
727 row.expires_at = Utc::now() + chrono::Duration::seconds(ttl_secs);
728 let mut conn = self.conn.clone();
729 let json = serde_json::to_string(&row).map_err(map_err)?;
730 let _: () = conn.set(self.keys.lease(lease_id), json).await.map_err(map_err)?;
731 Ok(())
732 }
733
734 async fn release_lease(&self, lease_id: &str) -> Result<()> {
735 let Some(row) = self.load_lease_row(lease_id).await? else {
736 return Ok(());
737 };
738 let mut conn = self.conn.clone();
739 let _: () = conn.del(self.keys.lease(lease_id)).await.map_err(map_err)?;
740 let _: () = conn
741 .del(self.keys.lease_by_job(&row.job_id))
742 .await
743 .map_err(map_err)?;
744 Ok(())
745 }
746
747 async fn expired_lease_job_pairs(&self) -> Result<Vec<(String, String)>> {
748 let pattern = self.keys.lease_pattern();
749 let mut conn = self.conn.clone();
750 let now = Utc::now();
751 let mut out = Vec::new();
752 let mut cursor = 0_u64;
753 loop {
754 let (next, batch): (u64, Vec<String>) = redis::cmd("SCAN")
755 .arg(cursor)
756 .arg("MATCH")
757 .arg(&pattern)
758 .arg("COUNT")
759 .arg(200)
760 .query_async(&mut conn)
761 .await
762 .map_err(map_err)?;
763 for key in batch {
764 if key.contains(":lease-by-job:") {
765 continue;
766 }
767 let raw: Option<String> = conn.get(&key).await.map_err(map_err)?;
768 if let Some(s) = raw {
769 if let Ok(row) = serde_json::from_str::<LeaseRow>(&s) {
770 if row.expires_at <= now {
771 out.push((row.lease_id, row.job_id));
772 }
773 }
774 }
775 }
776 cursor = next;
777 if cursor == 0 {
778 break;
779 }
780 }
781 Ok(out)
782 }
783}
784
785impl RedisQueueBackend {
786 async fn load_lease_row(&self, lease_id: &str) -> Result<Option<LeaseRow>> {
787 let mut conn = self.conn.clone();
788 let raw: Option<String> = conn.get(self.keys.lease(lease_id)).await.map_err(map_err)?;
789 raw.map_or(Ok(None), |s| serde_json::from_str(&s).map_err(map_err).map(Some))
790 }
791}
792
793pub async fn install_default_redis_backend(url: &str) -> Result<Arc<RedisQueueBackend>> {
799 let backend = Arc::new(RedisQueueBackend::connect_url(url).await?);
800 boson_core::QueueRouter::set_global(boson_core::QueueRouter::with_default(
801 Arc::clone(&backend) as Arc<dyn QueueBackend>,
802 ));
803 Ok(backend)
804}