From 9b1e8ad575613ce4afcf07d5d382c99fc73277c7 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 30 May 2026 00:45:58 +0530 Subject: [PATCH] perf(redis): compute job stats with SINTERCARD, not O(N) scans count_running_by_task / stats_by_queue / stats_all_queues loaded and deserialized every job in the relevant set. Intersect the already-maintained membership and per-status sets server-side with SINTERCARD instead (Redis starts from the smaller set, so the hot running-count stays cheap), and enumerate queues via a bounded SCAN of by_queue keys. No new keys, no extra write-path instrumentation, no backfill. Adds a backend-agnostic contract test covering all three. --- .../src/storage/redis_backend/jobs/query.rs | 119 +++++++++++------- .../taskito-core/tests/rust/storage_tests.rs | 38 ++++++ 2 files changed, 113 insertions(+), 44 deletions(-) diff --git a/crates/taskito-core/src/storage/redis_backend/jobs/query.rs b/crates/taskito-core/src/storage/redis_backend/jobs/query.rs index 26fa3f53..6d1cd0f2 100644 --- a/crates/taskito-core/src/storage/redis_backend/jobs/query.rs +++ b/crates/taskito-core/src/storage/redis_backend/jobs/query.rs @@ -119,68 +119,99 @@ impl RedisStorage { Ok(stats) } + /// `SINTERCARD` of a membership set with a per-status set — the count of + /// jobs that are both in `membership_key` and in the given status, computed + /// server-side without transferring or deserializing any job. Redis starts + /// from the smaller set, so e.g. intersecting against the (small) running + /// set stays cheap regardless of how large the membership set has grown. + fn count_in_status( + &self, + conn: &mut redis::Connection, + membership_key: &str, + status: JobStatus, + ) -> Result { + let status_key = self.key(&["jobs", "status", &(status as i32).to_string()]); + redis::cmd("SINTERCARD") + .arg(2) + .arg(membership_key) + .arg(&status_key) + .query::(conn) + .map_err(map_err) + } + + /// Per-status breakdown for the jobs in `membership_key` (a per-queue or + /// per-task set), one `SINTERCARD` per status. + fn status_breakdown( + &self, + conn: &mut redis::Connection, + membership_key: &str, + ) -> Result { + let mut stats = QueueStats::default(); + for (status, field) in [ + (JobStatus::Pending, &mut stats.pending), + (JobStatus::Running, &mut stats.running), + (JobStatus::Complete, &mut stats.completed), + (JobStatus::Failed, &mut stats.failed), + (JobStatus::Dead, &mut stats.dead), + (JobStatus::Cancelled, &mut stats.cancelled), + ] { + *field = self.count_in_status(conn, membership_key, status)?; + } + Ok(stats) + } + /// Count running jobs for a specific task name (for per-task concurrency limiting). pub fn count_running_by_task(&self, task_name: &str) -> Result { let mut conn = self.conn()?; let by_task_key = self.key(&["jobs", "by_task", task_name]); - let job_ids: Vec = conn.smembers(&by_task_key).map_err(map_err)?; - - let mut count: i64 = 0; - for id in &job_ids { - if let Some(job) = self.load_job(&mut conn, id)? { - if job.status == JobStatus::Running { - count += 1; - } - } - } - - Ok(count) + self.count_in_status(&mut conn, &by_task_key, JobStatus::Running) } pub fn stats_by_queue(&self, queue_name: &str) -> Result { let mut conn = self.conn()?; let by_queue_key = self.key(&["jobs", "by_queue", queue_name]); - let job_ids: Vec = conn.smembers(&by_queue_key).map_err(map_err)?; - - let mut stats = QueueStats::default(); - for id in &job_ids { - if let Some(job) = self.load_job(&mut conn, id)? { - match job.status as i32 { - 0 => stats.pending += 1, - 1 => stats.running += 1, - 2 => stats.completed += 1, - 3 => stats.failed += 1, - 4 => stats.dead += 1, - 5 => stats.cancelled += 1, - _ => {} - } - } - } - - Ok(stats) + self.status_breakdown(&mut conn, &by_queue_key) } pub fn stats_all_queues(&self) -> Result> { let mut conn = self.conn()?; - let all_key = self.key(&["jobs", "all"]); - let job_ids: Vec = conn.zrange(&all_key, 0, -1).map_err(map_err)?; - let mut map = std::collections::HashMap::::new(); - for id in &job_ids { - if let Some(job) = self.load_job(&mut conn, id)? { - let stats = map.entry(job.queue.clone()).or_default(); - match job.status as i32 { - 0 => stats.pending += 1, - 1 => stats.running += 1, - 2 => stats.completed += 1, - 3 => stats.failed += 1, - 4 => stats.dead += 1, - 5 => stats.cancelled += 1, - _ => {} + // Enumerate queues by scanning the per-queue membership keys (one per + // queue with at least one job — empty sets are auto-removed by Redis). + // This is bounded by the number of queues, not the job population. + let pattern = self.key(&["jobs", "by_queue", "*"]); + let strip = self.key(&["jobs", "by_queue", ""]); + let mut queues: Vec = Vec::new(); + let mut cursor = 0u64; + loop { + let (next, batch): (u64, Vec) = redis::cmd("SCAN") + .arg(cursor) + .arg("MATCH") + .arg(&pattern) + .arg("COUNT") + .arg(100) + .query(&mut conn) + .map_err(map_err)?; + for key in batch { + if let Some(name) = key.strip_prefix(&strip) { + queues.push(name.to_string()); } } + cursor = next; + if cursor == 0 { + break; + } } + // SCAN may return duplicates across iterations. + queues.sort_unstable(); + queues.dedup(); + let mut map = std::collections::HashMap::::new(); + for queue in queues { + let by_queue_key = self.key(&["jobs", "by_queue", &queue]); + let stats = self.status_breakdown(&mut conn, &by_queue_key)?; + map.insert(queue, stats); + } Ok(map) } diff --git a/crates/taskito-core/tests/rust/storage_tests.rs b/crates/taskito-core/tests/rust/storage_tests.rs index b689c14d..4be533c0 100644 --- a/crates/taskito-core/tests/rust/storage_tests.rs +++ b/crates/taskito-core/tests/rust/storage_tests.rs @@ -104,6 +104,43 @@ fn test_stats(s: &impl Storage) { assert!(stats.pending >= 2); } +fn test_stats_by_queue_and_task(s: &impl Storage) { + let q = "q-stats-breakdown"; + let task = "stats_breakdown_task"; + s.enqueue(make_job(q, task)).unwrap(); + s.enqueue(make_job(q, task)).unwrap(); + s.enqueue(make_job(q, task)).unwrap(); + + // 3 pending, none running yet. + let st = s.stats_by_queue(q).unwrap(); + assert_eq!(st.pending, 3); + assert_eq!(st.running, 0); + assert_eq!(s.count_running_by_task(task).unwrap(), 0); + + // Run two of them. + let d1 = s.dequeue(q, now_millis() + 1000, None).unwrap().unwrap(); + s.dequeue(q, now_millis() + 1000, None).unwrap().unwrap(); + assert_eq!(s.count_running_by_task(task).unwrap(), 2); + let st = s.stats_by_queue(q).unwrap(); + assert_eq!(st.running, 2); + assert_eq!(st.pending, 1); + + // Complete one — running drops, completed rises. + s.complete(&d1.id, None).unwrap(); + assert_eq!(s.count_running_by_task(task).unwrap(), 1); + let st = s.stats_by_queue(q).unwrap(); + assert_eq!(st.pending, 1); + assert_eq!(st.running, 1); + assert_eq!(st.completed, 1); + + // stats_all_queues reports the same breakdown for this queue. + let all = s.stats_all_queues().unwrap(); + let qs = all.get(q).expect("queue should appear in stats_all_queues"); + assert_eq!(qs.pending, 1); + assert_eq!(qs.running, 1); + assert_eq!(qs.completed, 1); +} + fn test_unique_key_dedup(s: &impl Storage) { let mut job1 = make_job("q-unique", "unique_task"); job1.unique_key = Some("dedup-key".to_string()); @@ -308,6 +345,7 @@ fn run_storage_tests(s: &impl Storage) { test_retry(s); test_cancel_job(s); test_stats(s); + test_stats_by_queue_and_task(s); test_unique_key_dedup(s); test_enqueue_batch(s); test_dead_letter_queue(s);