From 9644988d99ec03663a9911252a592b1ec6c8b3d3 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sun, 12 Jul 2026 14:02:28 +0530 Subject: [PATCH] feat(pubsub): per-subscription backlog and lag metrics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Attribute each delivery job (and its dead-letter row) to its subscription via indexed topic/subscription_name columns, populated from the delivery notes at insert. topic_backlog_stats() aggregates pending/running/dead counts and oldest-pending age per subscription off a partial index — live, drift-free, never a full-table scan. Redis keeps per-subscription pending/running/dead sets. --- crates/taskito-core/src/pubsub.rs | 12 ++ .../src/storage/diesel_common/dead_letter.rs | 14 ++ .../src/storage/diesel_common/jobs.rs | 27 +++- .../src/storage/diesel_common/migrations.rs | 15 ++ .../src/storage/diesel_common/pubsub.rs | 58 +++++++ crates/taskito-core/src/storage/mod.rs | 90 +++++++++++ crates/taskito-core/src/storage/models.rs | 6 + .../src/storage/redis_backend/dead_letter.rs | 60 +++++-- .../src/storage/redis_backend/jobs/dequeue.rs | 7 + .../src/storage/redis_backend/jobs/enqueue.rs | 28 +++- .../src/storage/redis_backend/jobs/helpers.rs | 13 ++ .../src/storage/redis_backend/jobs/state.rs | 6 + .../src/storage/redis_backend/pubsub.rs | 153 ++++++++++++++++++ crates/taskito-core/src/storage/schema.rs | 4 + crates/taskito-core/src/storage/traits.rs | 7 +- .../taskito-core/tests/rust/storage_tests.rs | 76 +++++++++ crates/taskito-python/src/py_queue/pubsub.rs | 41 +++++ sdks/python/taskito/_taskito.pyi | 3 + sdks/python/taskito/mixins/pubsub.py | 25 +++ sdks/python/tests/core/test_pubsub.py | 54 +++++++ 20 files changed, 681 insertions(+), 18 deletions(-) diff --git a/crates/taskito-core/src/pubsub.rs b/crates/taskito-core/src/pubsub.rs index 5cf382ef..f1cb0aad 100644 --- a/crates/taskito-core/src/pubsub.rs +++ b/crates/taskito-core/src/pubsub.rs @@ -116,6 +116,18 @@ fn delivery_job(request: &PublishRequest, sub: &SubscriptionRow) -> NewJob { } } +/// Inverse of [`delivery_notes`]: pull `topic`/`subscription` back out of a +/// job's notes JSON when both are present. Feeds the indexed +/// `jobs.topic`/`jobs.subscription_name` columns (and their `dead_letter` +/// mirrors) at insert time, so backlog/lag aggregation runs off an index +/// instead of a JSON scan — without adding pub/sub-only fields to `NewJob`. +pub fn extract_topic_subscription(notes: Option<&str>) -> Option<(String, String)> { + let obj: serde_json::Map = serde_json::from_str(notes?).ok()?; + let topic = obj.get("topic")?.as_str()?.to_string(); + let subscription = obj.get("subscription")?.as_str()?.to_string(); + Some((topic, subscription)) +} + /// Stamp `topic` and `subscription` into the caller's notes object so every /// delivery is filterable per subscriber without a schema change. fn delivery_notes(request: &PublishRequest, sub: &SubscriptionRow) -> String { diff --git a/crates/taskito-core/src/storage/diesel_common/dead_letter.rs b/crates/taskito-core/src/storage/diesel_common/dead_letter.rs index 4e3d5027..876639f9 100644 --- a/crates/taskito-core/src/storage/diesel_common/dead_letter.rs +++ b/crates/taskito-core/src/storage/diesel_common/dead_letter.rs @@ -25,6 +25,11 @@ macro_rules! impl_diesel_dead_letter_ops { .and_then(|m| serde_json::from_str::(m).ok()) .and_then(|v| v.get("__dlq_retry_count")?.as_i64()) .unwrap_or(0) as i32; + // Carry the delivery's subscription attribution into the DLQ so + // per-subscription dead-letter depth is countable. + let (topic, subscription_name) = + $crate::pubsub::extract_topic_subscription(job.notes.as_deref()) + .map_or((None, None), |(t, s)| (Some(t), Some(s))); self.write_transaction(|conn| { let dlq_row = NewDeadLetterRow { @@ -46,6 +51,8 @@ macro_rules! impl_diesel_dead_letter_ops { result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), dlq_retry_count, + topic: topic.as_deref(), + subscription_name: subscription_name.as_deref(), }; diesel::insert_into(dead_letter::table) @@ -190,6 +197,11 @@ macro_rules! impl_diesel_dead_letter_ops { }; let job = new_job.into_job(); + // Re-attribute the resurrected delivery so it counts against the + // subscription's backlog again, not just its DLQ depth. + let (topic, subscription_name) = + $crate::pubsub::extract_topic_subscription(job.notes.as_deref()) + .map_or((None, None), |(t, s)| (Some(t), Some(s))); let result = conn.transaction(|conn| { let row = super::super::models::NewJobRow { @@ -212,6 +224,8 @@ macro_rules! impl_diesel_dead_letter_ops { result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), has_deps: job.has_deps, + topic: topic.as_deref(), + subscription_name: subscription_name.as_deref(), }; diesel::insert_into(jobs::table) diff --git a/crates/taskito-core/src/storage/diesel_common/jobs.rs b/crates/taskito-core/src/storage/diesel_common/jobs.rs index 28dd4a53..f62c340c 100644 --- a/crates/taskito-core/src/storage/diesel_common/jobs.rs +++ b/crates/taskito-core/src/storage/diesel_common/jobs.rs @@ -162,6 +162,11 @@ macro_rules! impl_diesel_job_ops { Self::validate_dependency(conn, dep_id)?; } + // Attribute pub/sub deliveries to their subscription so + // backlog stats index by it; `None` for ordinary jobs. + let (topic, subscription_name) = + $crate::pubsub::extract_topic_subscription(job.notes.as_deref()) + .map_or((None, None), |(t, s)| (Some(t), Some(s))); let row = NewJobRow { id: &job.id, queue: &job.queue, @@ -182,6 +187,8 @@ macro_rules! impl_diesel_job_ops { result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), has_deps: job.has_deps, + topic: topic.as_deref(), + subscription_name: subscription_name.as_deref(), }; diesel::insert_into(jobs::table) @@ -224,10 +231,21 @@ macro_rules! impl_diesel_job_ops { let jobs: Vec = new_jobs.into_iter().map(|nj| nj.into_job()).collect(); + // Pre-compute subscription attribution so the owned strings + // outlive the borrowing rows built below. + let attribution: Vec<(Option, Option)> = jobs + .iter() + .map(|job| { + $crate::pubsub::extract_topic_subscription(job.notes.as_deref()) + .map_or((None, None), |(t, s)| (Some(t), Some(s))) + }) + .collect(); + self.write_transaction(|conn| { let rows: Vec = jobs .iter() - .map(|job| NewJobRow { + .zip(&attribution) + .map(|(job, (topic, subscription_name))| NewJobRow { id: &job.id, queue: &job.queue, task_name: &job.task_name, @@ -247,6 +265,8 @@ macro_rules! impl_diesel_job_ops { result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), has_deps: job.has_deps, + topic: topic.as_deref(), + subscription_name: subscription_name.as_deref(), }) .collect(); @@ -299,6 +319,9 @@ macro_rules! impl_diesel_job_ops { Self::validate_dependency(conn, dep_id)?; } + let (topic, subscription_name) = + $crate::pubsub::extract_topic_subscription(job.notes.as_deref()) + .map_or((None, None), |(t, s)| (Some(t), Some(s))); let row = NewJobRow { id: &job.id, queue: &job.queue, @@ -319,6 +342,8 @@ macro_rules! impl_diesel_job_ops { result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), has_deps: job.has_deps, + topic: topic.as_deref(), + subscription_name: subscription_name.as_deref(), }; diesel::insert_into(jobs::table) diff --git a/crates/taskito-core/src/storage/diesel_common/migrations.rs b/crates/taskito-core/src/storage/diesel_common/migrations.rs index b2f34f6b..06caf73d 100644 --- a/crates/taskito-core/src/storage/diesel_common/migrations.rs +++ b/crates/taskito-core/src/storage/diesel_common/migrations.rs @@ -325,6 +325,13 @@ pub fn create_indexes() -> &'static [&'static str] { // the partial index keeps durable rows (owner NULL) out of the reaper's way. "CREATE INDEX IF NOT EXISTS idx_topic_subs_owner ON topic_subscriptions(owner_worker_id) WHERE owner_worker_id IS NOT NULL", + // Backlog/lag aggregation groups pub/sub deliveries by (topic, subscription, + // status). Partial so only the small set of pub/sub jobs is indexed — the + // dashboard query never touches the mass of ordinary jobs. + "CREATE INDEX IF NOT EXISTS idx_jobs_subscription_status + ON jobs(topic, subscription_name, status) WHERE subscription_name IS NOT NULL", + "CREATE INDEX IF NOT EXISTS idx_dead_letter_subscription + ON dead_letter(topic, subscription_name) WHERE subscription_name IS NOT NULL", ] } @@ -396,6 +403,14 @@ pub fn alter_statements(d: &Dialect) -> Vec { format!("ALTER TABLE topic_subscriptions ADD COLUMN {ife}priority INTEGER"), format!("ALTER TABLE topic_subscriptions ADD COLUMN {ife}max_retries INTEGER"), format!("ALTER TABLE topic_subscriptions ADD COLUMN {ife}timeout_ms {bi}"), + // Pub/sub delivery attribution: publish stamps the topic + subscription + // onto each delivery job (and its dead_letter row) so backlog/lag stats + // aggregate off an index instead of scanning the notes JSON. NULL for + // every non-pub/sub job, so the partial indexes stay small. + format!("ALTER TABLE jobs ADD COLUMN {ife}topic TEXT"), + format!("ALTER TABLE jobs ADD COLUMN {ife}subscription_name TEXT"), + format!("ALTER TABLE dead_letter ADD COLUMN {ife}topic TEXT"), + format!("ALTER TABLE dead_letter ADD COLUMN {ife}subscription_name TEXT"), ] } diff --git a/crates/taskito-core/src/storage/diesel_common/pubsub.rs b/crates/taskito-core/src/storage/diesel_common/pubsub.rs index 3b135820..f94a332b 100644 --- a/crates/taskito-core/src/storage/diesel_common/pubsub.rs +++ b/crates/taskito-core/src/storage/diesel_common/pubsub.rs @@ -92,6 +92,64 @@ macro_rules! impl_diesel_pubsub_ops { Ok(affected as u64) } + + /// Backlog/lag snapshot per registered subscription. Four bounded, + /// index-backed queries (subscription skeleton + pending/running + /// counts + oldest-pending age + dead counts) merged in Rust — the + /// same live+aggregate shape as `stats_all_queues`, never a full + /// table scan. Counts are always a direct read of real state, so + /// they cannot drift the way a maintained counter would. + pub fn topic_backlog_stats( + &self, + ) -> Result> { + use $crate::job::JobStatus; + use $crate::storage::schema::{dead_letter, jobs}; + let mut conn = self.conn()?; + + let subs = topic_subscriptions::table + .select($crate::storage::models::SubscriptionRow::as_select()) + .load::<$crate::storage::models::SubscriptionRow>(&mut conn)?; + + // Pending + running counts, scoped to pub/sub-tagged rows only. + let counts: Vec<(String, String, i32, i64)> = jobs::table + .filter(jobs::subscription_name.is_not_null()) + .filter( + jobs::status.eq_any([JobStatus::Pending as i32, JobStatus::Running as i32]), + ) + .group_by((jobs::topic, jobs::subscription_name, jobs::status)) + .select(( + jobs::topic.assume_not_null(), + jobs::subscription_name.assume_not_null(), + jobs::status, + diesel::dsl::count(jobs::id), + )) + .load(&mut conn)?; + + let oldest: Vec<(String, String, Option)> = jobs::table + .filter(jobs::subscription_name.is_not_null()) + .filter(jobs::status.eq(JobStatus::Pending as i32)) + .group_by((jobs::topic, jobs::subscription_name)) + .select(( + jobs::topic.assume_not_null(), + jobs::subscription_name.assume_not_null(), + diesel::dsl::min(jobs::created_at), + )) + .load(&mut conn)?; + + let dead: Vec<(String, String, i64)> = dead_letter::table + .filter(dead_letter::subscription_name.is_not_null()) + .group_by((dead_letter::topic, dead_letter::subscription_name)) + .select(( + dead_letter::topic.assume_not_null(), + dead_letter::subscription_name.assume_not_null(), + diesel::dsl::count(dead_letter::id), + )) + .load(&mut conn)?; + + Ok($crate::storage::merge_backlog_stats( + subs, counts, oldest, dead, + )) + } } }; } diff --git a/crates/taskito-core/src/storage/mod.rs b/crates/taskito-core/src/storage/mod.rs index b983a4d9..cbea07c8 100644 --- a/crates/taskito-core/src/storage/mod.rs +++ b/crates/taskito-core/src/storage/mod.rs @@ -42,6 +42,88 @@ pub struct QueueStats { pub cancelled: i64, } +/// Per-subscription backlog/lag snapshot for the pub/sub dashboard. One entry +/// per *registered* subscription — durable or ephemeral, active or paused — +/// even at zero backlog, so the dashboard renders the full subscriber list +/// from a single call. Counts are computed live off the delivery-attribution +/// indexes; they can never drift the way a maintained counter would. +#[derive(Debug, Clone)] +pub struct SubscriptionBacklogStats { + pub topic: String, + pub subscription_name: String, + pub task_name: String, + pub queue: String, + pub active: bool, + pub durable: bool, + pub pending: i64, + pub running: i64, + pub dead: i64, + /// Milliseconds since the oldest still-pending delivery was created. + /// `None` when the subscription currently has no pending backlog. + pub oldest_pending_age_ms: Option, +} + +/// Merge the four `topic_backlog_stats` aggregate queries into one row per +/// registered subscription. Seeds a zeroed entry per subscription so idle +/// subscriptions still appear, then folds in the pending/running counts, +/// oldest-pending age (converted to a millisecond age), and dead counts. +/// Shared by the Diesel and Redis backends. +pub(crate) fn merge_backlog_stats( + subs: Vec, + counts: Vec<(String, String, i32, i64)>, + oldest: Vec<(String, String, Option)>, + dead: Vec<(String, String, i64)>, +) -> Vec { + use crate::job::{now_millis, JobStatus}; + use std::collections::HashMap; + + let mut by_key: HashMap<(String, String), SubscriptionBacklogStats> = subs + .into_iter() + .map(|s| { + ( + (s.topic.clone(), s.subscription_name.clone()), + SubscriptionBacklogStats { + topic: s.topic, + subscription_name: s.subscription_name, + task_name: s.task_name, + queue: s.queue, + active: s.active, + durable: s.durable, + pending: 0, + running: 0, + dead: 0, + oldest_pending_age_ms: None, + }, + ) + }) + .collect(); + + for (topic, name, status, count) in counts { + if let Some(stats) = by_key.get_mut(&(topic, name)) { + if status == JobStatus::Pending as i32 { + stats.pending = count; + } else if status == JobStatus::Running as i32 { + stats.running = count; + } + } + } + + let now = now_millis(); + for (topic, name, oldest_created) in oldest { + if let (Some(stats), Some(created)) = (by_key.get_mut(&(topic, name)), oldest_created) { + stats.oldest_pending_age_ms = Some((now - created).max(0)); + } + } + + for (topic, name, count) in dead { + if let Some(stats) = by_key.get_mut(&(topic, name)) { + stats.dead = count; + } + } + + by_key.into_values().collect() +} + #[derive(Debug, Clone)] pub struct DeadJob { pub id: String, @@ -387,6 +469,11 @@ macro_rules! impl_storage { ) -> $crate::error::Result { self.reap_ephemeral_subscriptions(live_worker_ids) } + fn topic_backlog_stats( + &self, + ) -> $crate::error::Result> { + self.topic_backlog_stats() + } fn record_metric( &self, task_name: &str, @@ -950,6 +1037,9 @@ impl Storage for StorageBackend { fn reap_ephemeral_subscriptions(&self, live_worker_ids: &[String]) -> Result { delegate!(self, reap_ephemeral_subscriptions, live_worker_ids) } + fn topic_backlog_stats(&self) -> Result> { + delegate!(self, topic_backlog_stats) + } fn record_metric( &self, task_name: &str, diff --git a/crates/taskito-core/src/storage/models.rs b/crates/taskito-core/src/storage/models.rs index 40796536..b5368b39 100644 --- a/crates/taskito-core/src/storage/models.rs +++ b/crates/taskito-core/src/storage/models.rs @@ -93,6 +93,10 @@ pub struct NewJobRow<'a> { pub result_ttl_ms: Option, pub namespace: Option<&'a str>, pub has_deps: bool, + /// Set on pub/sub deliveries (derived from `notes`) so backlog/lag stats + /// index by subscription; `None` for ordinary jobs. + pub topic: Option<&'a str>, + pub subscription_name: Option<&'a str>, } /// A row in the `dead_letter` table. @@ -137,6 +141,8 @@ pub struct NewDeadLetterRow<'a> { pub result_ttl_ms: Option, pub namespace: Option<&'a str>, pub dlq_retry_count: i32, + pub topic: Option<&'a str>, + pub subscription_name: Option<&'a str>, } /// A row in the `rate_limits` table. diff --git a/crates/taskito-core/src/storage/redis_backend/dead_letter.rs b/crates/taskito-core/src/storage/redis_backend/dead_letter.rs index 9c519b30..0c634cfc 100644 --- a/crates/taskito-core/src/storage/redis_backend/dead_letter.rs +++ b/crates/taskito-core/src/storage/redis_backend/dead_letter.rs @@ -207,7 +207,9 @@ impl RedisStorage { // ids whose entry matches `task_name`. let ids: Vec = conn.zrevrange(&dlq_all, 0, -1).map_err(map_err)?; - let mut to_delete = Vec::new(); + // (dlq_id, notes, original_job_id) — notes + original id let the + // sub:dead index shrink with the purge (no-op for non-pub/sub rows). + let mut to_delete: Vec<(String, Option, String)> = Vec::new(); for id in ids { let dlq_key = self.key(&["dlq", &id]); let data: Option = conn.get(&dlq_key).map_err(map_err)?; @@ -217,7 +219,7 @@ impl RedisStorage { let entry: DeadJobEntry = serde_json::from_str(&d).map_err(|e| QueueError::Other(e.to_string()))?; if entry.task_name == task_name { - to_delete.push(id); + to_delete.push((id, entry.notes, entry.original_job_id)); } } } @@ -227,9 +229,10 @@ impl RedisStorage { } let pipe = &mut redis::pipe(); - for id in &to_delete { + for (id, notes, member) in &to_delete { pipe.del(self.key(&["dlq", id])); pipe.zrem(&dlq_all, id.as_str()); + self.push_pubsub_dead_remove(pipe, notes.as_deref(), member); } pipe.query::<()>(&mut conn).map_err(map_err)?; Ok(to_delete.len() as u64) @@ -249,6 +252,12 @@ impl RedisStorage { let entry: DeadJobEntry = serde_json::from_str(&data).map_err(|e| QueueError::Other(e.to_string()))?; + // Attribution + original job id for the sub:dead removal below, captured + // before `entry`'s fields are moved into `new_job`. The re-enqueue below + // re-attributes the fresh job into sub:pending via its carried notes. + let dead_notes = entry.notes.clone(); + let dead_member = entry.original_job_id.clone(); + let retry_metadata = { let next_count = entry.dlq_retry_count + 1; let mut obj = entry @@ -284,11 +293,13 @@ impl RedisStorage { let job = self.enqueue(new_job)?; - // Remove from DLQ + // Remove from DLQ, and drop it from its subscription's sub:dead index + // so the dead count doesn't keep the retried delivery (no-op otherwise). let dlq_all = self.key(&["dlq", "all"]); let pipe = &mut redis::pipe(); pipe.del(&dlq_key); pipe.zrem(&dlq_all, dead_id); + self.push_pubsub_dead_remove(pipe, dead_notes.as_deref(), &dead_member); pipe.query::<()>(&mut conn).map_err(map_err)?; Ok(job.id) @@ -307,11 +318,24 @@ impl RedisStorage { return Ok(0); } + // Load blobs to attribute each dead row to its subscription so the + // sub:dead index shrinks with the purge (no-op for non-pub/sub rows). + let blob_keys: Vec = ids.iter().map(|id| self.key(&["dlq", id])).collect(); + let blobs: Vec> = conn.mget(&blob_keys).map_err(map_err)?; + let pipe = &mut redis::pipe(); - for id in &ids { - let dlq_key = self.key(&["dlq", id]); - pipe.del(&dlq_key); - pipe.zrem(&dlq_all, id); + for (id, blob) in ids.iter().zip(&blobs) { + pipe.del(self.key(&["dlq", id])); + pipe.zrem(&dlq_all, id.as_str()); + if let Some(d) = blob { + if let Ok(entry) = serde_json::from_str::(d) { + self.push_pubsub_dead_remove( + pipe, + entry.notes.as_deref(), + &entry.original_job_id, + ); + } + } } pipe.query::<()>(&mut conn).map_err(map_err)?; @@ -323,14 +347,19 @@ impl RedisStorage { let dlq_key = self.key(&["dlq", dead_id]); let dlq_all = self.key(&["dlq", "all"]); - let existed: bool = conn.exists(&dlq_key).map_err(map_err)?; - if !existed { + // Load the row (rather than a bare EXISTS) so its subscription + // attribution is known and the sub:dead index shrinks with the delete. + let data: Option = conn.get(&dlq_key).map_err(map_err)?; + let Some(d) = data else { return Ok(false); - } + }; + let entry: DeadJobEntry = + serde_json::from_str(&d).map_err(|e| QueueError::Other(e.to_string()))?; let pipe = &mut redis::pipe(); pipe.del(&dlq_key); pipe.zrem(&dlq_all, dead_id); + self.push_pubsub_dead_remove(pipe, entry.notes.as_deref(), &entry.original_job_id); pipe.query::<()>(&mut conn).map_err(map_err)?; Ok(true) } @@ -344,7 +373,9 @@ impl RedisStorage { .zrangebyscore(&dlq_all, "-inf", "+inf") .map_err(map_err)?; - let mut to_delete = Vec::new(); + // (dlq_id, notes, original_job_id) — keeps the sub:dead index in step + // with the TTL purge (no-op for non-pub/sub rows). + let mut to_delete: Vec<(String, Option, String)> = Vec::new(); for id in &ids { let dlq_key = self.key(&["dlq", id]); let data: Option = conn.get(&dlq_key).map_err(map_err)?; @@ -355,7 +386,7 @@ impl RedisStorage { None => entry.failed_at < global_cutoff_ms, }; if expired { - to_delete.push(id.clone()); + to_delete.push((id.clone(), entry.notes, entry.original_job_id)); } } } @@ -366,9 +397,10 @@ impl RedisStorage { } let pipe = &mut redis::pipe(); - for id in &to_delete { + for (id, notes, member) in &to_delete { pipe.del(self.key(&["dlq", id])); pipe.zrem(&dlq_all, id.as_str()); + self.push_pubsub_dead_remove(pipe, notes.as_deref(), member); } pipe.query::<()>(&mut conn).map_err(map_err)?; Ok(to_delete.len() as u64) diff --git a/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs b/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs index 511ed93e..6c64a0a8 100644 --- a/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs +++ b/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs @@ -153,6 +153,10 @@ impl RedisStorage { job.status = JobStatus::Running; job.started_at = Some(now); if self.claim_pending(&mut conn, &job, &queue_key)? { + // Best-effort pub/sub backlog reindex Pending→Running; a + // follow-up rather than folded into the correctness-critical + // claim script (see `reindex_pubsub_best_effort`). + self.reindex_pubsub_best_effort(&mut conn, &job, JobStatus::Running); return Ok(Some(job)); } } @@ -291,6 +295,9 @@ impl RedisStorage { job.status = JobStatus::Running; job.started_at = Some(now); if self.claim_pending(&mut conn, &job, &queue_key)? { + // Best-effort pub/sub backlog reindex Pending→Running (see the + // single-claim `dequeue` for the rationale). + self.reindex_pubsub_best_effort(&mut conn, &job, JobStatus::Running); claimed.push(job); } } diff --git a/crates/taskito-core/src/storage/redis_backend/jobs/enqueue.rs b/crates/taskito-core/src/storage/redis_backend/jobs/enqueue.rs index 8a422752..e6bd953c 100644 --- a/crates/taskito-core/src/storage/redis_backend/jobs/enqueue.rs +++ b/crates/taskito-core/src/storage/redis_backend/jobs/enqueue.rs @@ -77,6 +77,10 @@ impl RedisStorage { pipe.sadd(&dependents_key, &job.id); } + // A pub/sub delivery enters its subscription's pending backlog index + // (no-op for ordinary jobs). Same pipe as the job's own indices. + self.push_pubsub_transition(pipe, &job, JobStatus::Pending); + pipe.query::<()>(&mut conn).map_err(map_err)?; Ok(job) @@ -122,6 +126,10 @@ impl RedisStorage { pipe.sadd(&depends_on_key, dep_id); pipe.sadd(&dependents_key, &job.id); } + + // Pub/sub deliveries enter their subscription's pending backlog + // index (no-op for ordinary jobs), atomically with the batch insert. + self.push_pubsub_transition(pipe, job, JobStatus::Pending); } pipe.query::<()>(&mut conn).map_err(map_err)?; @@ -234,6 +242,14 @@ impl RedisStorage { redis.call('ZADD', all_key, -created_at, job_id) redis.call('SET', unique_key, job_id) + -- Pub/sub delivery: mirror into its subscription's pending + -- backlog index (KEYS[8], present only for pub/sub deliveries), + -- scored by created_at. Folded into this atomic store so the + -- backlog index cannot desync from the job on a crash. + if KEYS[8] then + redis.call('ZADD', KEYS[8], created_at, job_id) + end + -- Store dependencies (3 ARGVs per dep: dep_on_key, dep_id, dependents_key). for i = 1, num_deps do local offset = dep_args_base + (i - 1) * 3 @@ -257,8 +273,11 @@ impl RedisStorage { let score = dequeue_score(job.priority, job.scheduled_at); let job_key_prefix = self.key(&["job", ""]); - // Build keys and args vectors to avoid temporary lifetime issues - let keys = vec![ + // Build keys and args vectors to avoid temporary lifetime issues. + // KEYS[8] (the subscription's pending backlog index) is appended + // only for pub/sub deliveries, so ordinary jobs pass 7 keys and the + // Lua's `if KEYS[8]` guard is false. + let mut keys = vec![ unique_key.clone(), job_key, status_key, @@ -267,6 +286,11 @@ impl RedisStorage { by_task_key, all_key, ]; + if let Some((topic, name)) = + crate::pubsub::extract_topic_subscription(job.notes.as_deref()) + { + keys.push(self.key(&["sub", "pending", &topic, &name])); + } let mut args: Vec = vec![ job.id.clone(), job_json.clone(), diff --git a/crates/taskito-core/src/storage/redis_backend/jobs/helpers.rs b/crates/taskito-core/src/storage/redis_backend/jobs/helpers.rs index aa862a87..a129f2e9 100644 --- a/crates/taskito-core/src/storage/redis_backend/jobs/helpers.rs +++ b/crates/taskito-core/src/storage/redis_backend/jobs/helpers.rs @@ -82,6 +82,9 @@ impl RedisStorage { pipe.srem(&old_status_key, &job.id); pipe.sadd(&new_status_key, &job.id); } + // Mirror the status move on the pub/sub backlog indices, keyed on the + // job's new status (no-op for ordinary jobs). Same pipe as the move. + self.push_pubsub_transition(pipe, job, job.status); pipe.query::<()>(conn).map_err(map_err)?; Ok(()) @@ -117,6 +120,9 @@ impl RedisStorage { pipe.sadd(&pending_status_key, &job.id); } pipe.zadd(&queue_key, &job.id, score); + // Back to Pending: leave the running index and (re)enter the pending + // backlog index. No-op for ordinary jobs; same atomic pipe. + self.push_pubsub_transition(pipe, job, JobStatus::Pending); pipe.query::<()>(conn).map_err(map_err)?; Ok(()) @@ -183,6 +189,13 @@ impl RedisStorage { pipe.sadd(&archived_by_queue, &job.id).ignore(); pipe.zadd(&archived_all, &job.id, completed_at as f64) .ignore(); + + // Mirror the terminal move on the pub/sub backlog indices (no-op for + // ordinary jobs). `job.status` is the terminal status the caller set: + // Complete/Failed/Cancelled leave the backlog; Dead (from `move_to_dlq`) + // also enters the sub:dead index. Same atomic pipe as the archive move, + // so the backlog index can never be left desynced from the archive. + self.push_pubsub_transition(pipe, job, job.status); } /// Move a terminal job out of the live indices into the archive in one diff --git a/crates/taskito-core/src/storage/redis_backend/jobs/state.rs b/crates/taskito-core/src/storage/redis_backend/jobs/state.rs index 1d8a9c01..aec00283 100644 --- a/crates/taskito-core/src/storage/redis_backend/jobs/state.rs +++ b/crates/taskito-core/src/storage/redis_backend/jobs/state.rs @@ -172,6 +172,12 @@ impl RedisStorage { .invoke(&mut conn) .map_err(map_err)?; + if applied == 1 { + // Best-effort backlog reindex Running→Pending; a follow-up rather + // than folded into the correctness-critical requeue-stuck script + // (see `reindex_pubsub_best_effort`). + self.reindex_pubsub_best_effort(&mut conn, &job, JobStatus::Pending); + } Ok(applied == 1) } diff --git a/crates/taskito-core/src/storage/redis_backend/pubsub.rs b/crates/taskito-core/src/storage/redis_backend/pubsub.rs index 83205b42..fe937652 100644 --- a/crates/taskito-core/src/storage/redis_backend/pubsub.rs +++ b/crates/taskito-core/src/storage/redis_backend/pubsub.rs @@ -5,7 +5,9 @@ use serde::{Deserialize, Serialize}; use super::{map_err, RedisStorage}; use crate::error::{QueueError, Result}; +use crate::job::{Job, JobStatus}; use crate::storage::models::{NewSubscriptionRow, SubscriptionRow}; +use crate::storage::SubscriptionBacklogStats; /// JSON blob stored at `sub::`. Mirrors /// [`SubscriptionRow`] so the CRUD path needs no Lua — index sets alone keep @@ -282,4 +284,155 @@ impl RedisStorage { Ok(removed) } + + // ── Per-subscription backlog/lag indices ─────────────────────────────── + // + // Three index keys per subscription mirror the live delivery backlog so + // `topic_backlog_stats` costs O(subscriptions), never a job scan: + // sub:pending:: — ZSET of job ids scored by `created_at` (ms) + // sub:running:: — SET of job ids + // sub:dead:: — SET of (original) job ids + // Every write is gated on `extract_topic_subscription`, so an ordinary job + // never touches them: the whole feature stays zero-cost off the pub/sub path. + + /// Key for a subscription's per-status backlog index. `kind` is + /// `"pending"`, `"running"`, or `"dead"`. + fn sub_index_key(&self, kind: &str, topic: &str, subscription_name: &str) -> String { + self.key(&["sub", kind, topic, subscription_name]) + } + + /// Append the backlog-index ops for `job` landing in status `to` onto + /// `pipe`, if (and only if) `job` is a pub/sub delivery. Keyed on the + /// *destination* status alone: every transition additionally scrubs the job + /// from the indices it must not be in, so the indices self-heal any drift a + /// best-effort claim reindex (see `reindex_pubsub_best_effort`) may leave, + /// on the very next state change. All ops are `.ignore()`d so the helper is + /// safe to fold into either a plain or a MULTI/EXEC pipeline without + /// disturbing the caller's result type. + pub(in crate::storage::redis_backend) fn push_pubsub_transition( + &self, + pipe: &mut redis::Pipeline, + job: &Job, + to: JobStatus, + ) { + let Some((topic, name)) = crate::pubsub::extract_topic_subscription(job.notes.as_deref()) + else { + return; + }; + let pending = self.sub_index_key("pending", &topic, &name); + let running = self.sub_index_key("running", &topic, &name); + let dead = self.sub_index_key("dead", &topic, &name); + match to { + JobStatus::Pending => { + pipe.srem(&running, &job.id).ignore(); + pipe.zadd(&pending, &job.id, job.created_at as f64).ignore(); + } + JobStatus::Running => { + pipe.zrem(&pending, &job.id).ignore(); + pipe.sadd(&running, &job.id).ignore(); + } + JobStatus::Dead => { + pipe.zrem(&pending, &job.id).ignore(); + pipe.srem(&running, &job.id).ignore(); + pipe.sadd(&dead, &job.id).ignore(); + } + // Complete / Failed / Cancelled: the delivery leaves the backlog. + _ => { + pipe.zrem(&pending, &job.id).ignore(); + pipe.srem(&running, &job.id).ignore(); + } + } + } + + /// Append a `sub:dead` removal onto `pipe` for a dead-letter row whose + /// attribution comes from its persisted `notes` (retry/purge/delete paths + /// hold a `DeadJobEntry`, not a live `Job`). `member` is the original job id + /// — the value `move_to_dlq` added via [`Self::push_pubsub_transition`]. + /// No-op for non-pub/sub rows. Keeps the `dead` count from drifting upward + /// after a dead-letter entry is retried or purged. + pub(in crate::storage::redis_backend) fn push_pubsub_dead_remove( + &self, + pipe: &mut redis::Pipeline, + notes: Option<&str>, + member: &str, + ) { + if let Some((topic, name)) = crate::pubsub::extract_topic_subscription(notes) { + pipe.srem(self.sub_index_key("dead", &topic, &name), member) + .ignore(); + } + } + + /// Best-effort backlog reindex for a status change committed by a bespoke, + /// correctness-critical Lua script we deliberately keep minimal (the + /// Pending→Running claim and the requeue-stuck rescue). The index move runs + /// in a follow-up pipe: a crash between the script and this pipe can leave a + /// delivery in the wrong backlog set, but that only skews a dashboard metric + /// — never job state — and the next transition scrubs it. A Redis error is + /// logged and swallowed so it cannot fail an operation that already + /// succeeded. No-op (and no round trip) for ordinary jobs. + pub(in crate::storage::redis_backend) fn reindex_pubsub_best_effort( + &self, + conn: &mut redis::Connection, + job: &Job, + to: JobStatus, + ) { + if crate::pubsub::extract_topic_subscription(job.notes.as_deref()).is_none() { + return; + } + let mut pipe = redis::pipe(); + self.push_pubsub_transition(&mut pipe, job, to); + if let Err(e) = pipe.query::<()>(conn) { + log::warn!("pub/sub backlog reindex failed (metrics only): {e}"); + } + } + + /// Backlog/lag snapshot per registered subscription. One bounded round trip + /// per subscription (never a job scan): the three index cardinalities plus + /// the oldest pending score. `oldest_pending_age_ms` is `now - score` + /// floored at 0, or `None` when the subscription has no pending backlog. + /// Mirrors the Diesel backends' shape; counts are a direct read of the live + /// indices, so they cannot drift the way a maintained counter would. + pub fn topic_backlog_stats(&self) -> Result> { + let subs = self.list_subscriptions()?; + if subs.is_empty() { + return Ok(Vec::new()); + } + let mut conn = self.conn()?; + let now = crate::job::now_millis(); + + let mut out = Vec::with_capacity(subs.len()); + for sub in subs { + let pending_key = self.sub_index_key("pending", &sub.topic, &sub.subscription_name); + let running_key = self.sub_index_key("running", &sub.topic, &sub.subscription_name); + let dead_key = self.sub_index_key("dead", &sub.topic, &sub.subscription_name); + + // ZRANGE 0 0 WITHSCORES yields at most the single oldest pending job. + let (pending, running, dead, oldest): (i64, i64, i64, Vec<(String, f64)>) = + redis::pipe() + .zcard(&pending_key) + .scard(&running_key) + .scard(&dead_key) + .zrange_withscores(&pending_key, 0, 0) + .query(&mut conn) + .map_err(map_err)?; + + let oldest_pending_age_ms = oldest + .first() + .map(|(_, score)| (now - *score as i64).max(0)); + + out.push(SubscriptionBacklogStats { + topic: sub.topic, + subscription_name: sub.subscription_name, + task_name: sub.task_name, + queue: sub.queue, + active: sub.active, + durable: sub.durable, + pending, + running, + dead, + oldest_pending_age_ms, + }); + } + Ok(out) + } } diff --git a/crates/taskito-core/src/storage/schema.rs b/crates/taskito-core/src/storage/schema.rs index db7b6197..3f573b13 100644 --- a/crates/taskito-core/src/storage/schema.rs +++ b/crates/taskito-core/src/storage/schema.rs @@ -24,6 +24,8 @@ diesel::table! { result_ttl_ms -> Nullable, namespace -> Nullable, has_deps -> Bool, + topic -> Nullable, + subscription_name -> Nullable, } } @@ -45,6 +47,8 @@ diesel::table! { result_ttl_ms -> Nullable, namespace -> Nullable, dlq_retry_count -> Integer, + topic -> Nullable, + subscription_name -> Nullable, } } diff --git a/crates/taskito-core/src/storage/traits.rs b/crates/taskito-core/src/storage/traits.rs index 66aa3878..a947f18c 100644 --- a/crates/taskito-core/src/storage/traits.rs +++ b/crates/taskito-core/src/storage/traits.rs @@ -1,7 +1,7 @@ use crate::error::Result; use crate::job::{Job, NewJob}; use crate::storage::models::*; -use crate::storage::{DeadJob, QueueStats}; +use crate::storage::{DeadJob, QueueStats, SubscriptionBacklogStats}; /// Trait abstracting the storage backend for the task queue. /// @@ -148,6 +148,11 @@ pub trait Storage: Send + Sync + Clone { /// count removed. fn reap_ephemeral_subscriptions(&self, live_worker_ids: &[String]) -> Result; + /// Backlog/lag snapshot for every registered subscription across all topics. + /// Bounded by the pub/sub-tagged rows (partial index), never a full `jobs` + /// table scan — safe to poll on a dashboard cadence. + fn topic_backlog_stats(&self) -> Result>; + // ── Metrics operations ────────────────────────────────────────── fn record_metric( diff --git a/crates/taskito-core/tests/rust/storage_tests.rs b/crates/taskito-core/tests/rust/storage_tests.rs index e1d2ed10..bb26de1c 100644 --- a/crates/taskito-core/tests/rust/storage_tests.rs +++ b/crates/taskito-core/tests/rust/storage_tests.rs @@ -1030,6 +1030,81 @@ fn test_concurrent_dequeue_no_double_claim(s: &impl Storage) { ); } +fn test_topic_backlog_stats(s: &impl Storage) { + use taskito_core::pubsub::{publish_to_topic, DeliveryDefaults, PublishRequest}; + use taskito_core::storage::models::NewSubscriptionRow; + + let sub = |name: &'static str, task: &'static str| NewSubscriptionRow { + topic: "tbs-orders", + subscription_name: name, + task_name: task, + queue: "default", + active: true, + durable: true, + owner_worker_id: None, + created_at: now_millis(), + priority: None, + max_retries: None, + timeout_ms: None, + }; + s.register_subscription(&sub("tbs-email", "tbs_send")) + .unwrap(); + s.register_subscription(&sub("tbs-analytics", "tbs_track")) + .unwrap(); + + let request = |topic: &str| PublishRequest { + topic: topic.to_string(), + payload: vec![0x02, 0xf5], + idempotency_key: None, + metadata: None, + notes: None, + priority: None, + scheduled_at: now_millis(), + max_retries: None, + timeout_ms: None, + expires_at: None, + result_ttl_ms: None, + namespace: None, + queue_defaults: DeliveryDefaults { + priority: 0, + max_retries: 3, + timeout_ms: 300_000, + }, + }; + publish_to_topic(s, &request("tbs-orders")).unwrap(); + publish_to_topic(s, &request("tbs-orders")).unwrap(); + + let stats = s.topic_backlog_stats().unwrap(); + let by_name: std::collections::HashMap<_, _> = stats + .iter() + .filter(|st| st.topic == "tbs-orders") + .map(|st| (st.subscription_name.as_str(), st)) + .collect(); + assert_eq!(by_name.len(), 2, "both subscriptions appear in the stats"); + assert_eq!(by_name["tbs-email"].pending, 2); + assert_eq!(by_name["tbs-analytics"].pending, 2); + assert_eq!(by_name["tbs-email"].running, 0); + assert_eq!(by_name["tbs-email"].dead, 0); + assert!( + by_name["tbs-email"].oldest_pending_age_ms.is_some(), + "a pending backlog yields an oldest-pending age" + ); + + // A dequeued delivery moves from pending to running. + let claimed = s.dequeue("default", now_millis(), None).unwrap().unwrap(); + let stats = s.topic_backlog_stats().unwrap(); + let claimed_sub = stats + .iter() + .find(|st| st.running == 1) + .expect("one delivery is now running"); + assert_eq!( + claimed_sub.pending, 1, + "its backlog dropped by the claimed one" + ); + // The claimed job belongs to one of our subscriptions. + assert!(claimed.task_name == "tbs_send" || claimed.task_name == "tbs_track"); +} + fn run_storage_tests(s: &impl Storage) { test_enqueue_and_get(s); test_dequeue(s); @@ -1054,6 +1129,7 @@ fn run_storage_tests(s: &impl Storage) { test_pause_resume_queue(s); test_periodic_crud(s); test_topic_subscriptions_crud(s); + test_topic_backlog_stats(s); test_circuit_breakers(s); test_execution_claims_purge(s); test_reap_stale_jobs(s); diff --git a/crates/taskito-python/src/py_queue/pubsub.rs b/crates/taskito-python/src/py_queue/pubsub.rs index a5d44041..1d41c7db 100644 --- a/crates/taskito-python/src/py_queue/pubsub.rs +++ b/crates/taskito-python/src/py_queue/pubsub.rs @@ -12,6 +12,22 @@ use crate::py_job::PyJob; /// `(topic, subscription_name, task_name, queue, active, durable)`. type SubscriptionTuple = (String, String, String, String, bool, bool); +/// Per-subscription backlog snapshot surfaced to Python: +/// `(topic, subscription, task_name, queue, active, durable, pending, running, +/// dead, oldest_pending_age_ms)`. +type SubscriptionBacklogTuple = ( + String, + String, + String, + String, + bool, + bool, + i64, + i64, + i64, + Option, +); + #[pymethods] impl PyQueue { /// Insert or update a topic subscription (idempotent on topic + name). @@ -100,6 +116,31 @@ impl PyQueue { .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) } + /// Backlog/lag snapshot per subscription: + /// `(topic, subscription, task_name, queue, active, durable, pending, + /// running, dead, oldest_pending_age_ms)`. + pub fn topic_backlog_stats(&self) -> PyResult> { + self.storage + .topic_backlog_stats() + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))? + .into_iter() + .map(|s| { + Ok(( + s.topic, + s.subscription_name, + s.task_name, + s.queue, + s.active, + s.durable, + s.pending, + s.running, + s.dead, + s.oldest_pending_age_ms, + )) + }) + .collect() + } + /// Drop ephemeral subscriptions whose owning worker is gone. Runs on the /// heartbeat cadence, after `reap_dead_workers` has pruned the registry. pub fn reap_ephemeral_subscriptions(&self) -> PyResult { diff --git a/sdks/python/taskito/_taskito.pyi b/sdks/python/taskito/_taskito.pyi index c9aaf422..660a5364 100644 --- a/sdks/python/taskito/_taskito.pyi +++ b/sdks/python/taskito/_taskito.pyi @@ -206,6 +206,9 @@ class PyQueue: self, topic: str, subscription_name: str, active: bool ) -> bool: ... def reap_ephemeral_subscriptions(self) -> int: ... + def topic_backlog_stats( + self, + ) -> list[tuple[str, str, str, str, bool, bool, int, int, int, int | None]]: ... def publish( self, topic: str, diff --git a/sdks/python/taskito/mixins/pubsub.py b/sdks/python/taskito/mixins/pubsub.py index 4d1dfb6e..b774ad22 100644 --- a/sdks/python/taskito/mixins/pubsub.py +++ b/sdks/python/taskito/mixins/pubsub.py @@ -188,6 +188,31 @@ def list_topics(self) -> list[str]: seen.setdefault(row[0], None) return list(seen) + def topic_stats(self, topic: str | None = None) -> list[dict[str, Any]]: + """Backlog/lag snapshot per subscription, optionally filtered to a topic. + + Each entry: ``topic``, ``subscription``, ``task_name``, ``queue``, + ``active``, ``durable``, ``pending``, ``running``, ``dead``, and + ``oldest_pending_age_ms`` (``None`` when the subscription has no pending + backlog). Computed live off indexed columns — safe to poll. + """ + stats = [ + { + "topic": row[0], + "subscription": row[1], + "task_name": row[2], + "queue": row[3], + "active": row[4], + "durable": row[5], + "pending": row[6], + "running": row[7], + "dead": row[8], + "oldest_pending_age_ms": row[9], + } + for row in self._inner.topic_backlog_stats() + ] + return stats if topic is None else [s for s in stats if s["topic"] == topic] + def _delivery_task_defaults(self) -> dict[str, tuple[int, int, int]]: """Per-task ``(priority, max_retries, timeout_ms)`` from this process's task registry. Persisted on the subscription row at registration so a diff --git a/sdks/python/tests/core/test_pubsub.py b/sdks/python/tests/core/test_pubsub.py index e5c9c9ee..52e5a8ac 100644 --- a/sdks/python/tests/core/test_pubsub.py +++ b/sdks/python/tests/core/test_pubsub.py @@ -211,6 +211,60 @@ def handle(order_id: int) -> None: ... assert job is not None and job._py_job.max_retries == 1 +class TestTopicStats: + def test_backlog_counts_and_zero_rows(self, queue: Queue) -> None: + @queue.subscriber("orders", name="email") + def email(order_id: int) -> None: ... + + @queue.subscriber("orders", name="analytics") + def track(order_id: int) -> None: ... + + queue.declare_subscriptions() + queue.publish("orders", 1) + queue.publish("orders", 2) + + stats = {s["subscription"]: s for s in queue.topic_stats("orders")} + # Both registered subscriptions appear, each with 2 pending deliveries. + assert set(stats) == {"email", "analytics"} + assert stats["email"]["pending"] == 2 + assert stats["analytics"]["pending"] == 2 + assert stats["email"]["dead"] == 0 + assert stats["email"]["oldest_pending_age_ms"] is not None + assert stats["email"]["oldest_pending_age_ms"] >= 0 + + def test_idle_subscription_reports_zeros(self, queue: Queue) -> None: + @queue.subscriber("orders", name="email") + def email(order_id: int) -> None: ... + + queue.declare_subscriptions() + (stat,) = queue.topic_stats("orders") + assert stat["pending"] == stat["running"] == stat["dead"] == 0 + assert stat["oldest_pending_age_ms"] is None + + def test_failed_delivery_counts_as_dead( + self, queue: Queue, run_worker: threading.Thread, poll_until: PollUntil + ) -> None: + @queue.subscriber("orders", name="flaky", max_retries=0) + def flaky(order_id: int) -> None: + raise RuntimeError("boom") + + queue.declare_subscriptions() + queue.publish("orders", 1) + + def dead_counted() -> bool: + stats = queue.topic_stats("orders") + return bool(stats) and stats[0]["dead"] == 1 and stats[0]["pending"] == 0 + + poll_until(dead_counted, message="failed delivery should count as dead") + + def test_non_pubsub_jobs_are_invisible(self, queue: Queue) -> None: + @queue.task() + def plain() -> None: ... + + plain.delay() + assert queue.topic_stats() == [] + + class TestEphemeralSubscriptions: def test_ephemeral_registers_only_with_worker(self, queue: Queue) -> None: @queue.subscriber("orders", name="debug-tail", durable=False)