From 49bbb59b5b7880e89640438e384be4d60574cd69 Mon Sep 17 00:00:00 2001 From: DJ Majumdar Date: Wed, 18 Mar 2026 22:18:06 -0700 Subject: [PATCH 1/5] perf: replace iterative BFS cycle detection with recursive CTE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The iterative BFS issued one SQL round-trip per graph node, producing O(n²) total queries for linear dependency chains (~19,900 for depth 200). A single recursive CTE collapses each cycle check to one query, yielding an 82% speedup at depth 200. --- src/store/submit/dependencies.rs | 39 +++++++++++++++++--------------- 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/src/store/submit/dependencies.rs b/src/store/submit/dependencies.rs index fca2200..b7ea8c1 100644 --- a/src/store/submit/dependencies.rs +++ b/src/store/submit/dependencies.rs @@ -1,7 +1,5 @@ //! Dependency resolution and cycle detection for task submission. -use std::collections::{HashSet, VecDeque}; - use crate::store::StoreError; use crate::task::DependencyFailurePolicy; @@ -81,30 +79,35 @@ pub(super) async fn resolve_dependency_edges( Ok((active_deps, status)) } -/// Cycle detection: iterative BFS from each dep upward through -/// the dependency graph. If we encounter `new_task_id`, there's a cycle. +/// Cycle detection via recursive CTE. +/// +/// Walks the entire upstream ancestry of `deps` in a single SQL query +/// instead of issuing one SELECT per BFS level. This reduces the number +/// of Rust↔SQLite round-trips from O(chain_depth) to O(1). pub(super) async fn detect_cycle( conn: &mut sqlx::pool::PoolConnection, new_task_id: i64, deps: &[i64], ) -> Result<(), StoreError> { - let mut visited = HashSet::new(); - let mut queue: VecDeque = deps.iter().copied().collect(); + for &dep_id in deps { + let found: Option<(i64,)> = sqlx::query_as( + "WITH RECURSIVE ancestors(id) AS ( + SELECT ? AS id + UNION + SELECT td.depends_on_id + FROM task_deps td + JOIN ancestors a ON td.task_id = a.id + ) + SELECT id FROM ancestors WHERE id = ? LIMIT 1", + ) + .bind(dep_id) + .bind(new_task_id) + .fetch_optional(&mut **conn) + .await?; - while let Some(current) = queue.pop_front() { - if current == new_task_id { + if found.is_some() { return Err(StoreError::CyclicDependency); } - if !visited.insert(current) { - continue; - } - // Find what `current` depends on. - let upstream: Vec<(i64,)> = - sqlx::query_as("SELECT depends_on_id FROM task_deps WHERE task_id = ?") - .bind(current) - .fetch_all(&mut **conn) - .await?; - queue.extend(upstream.into_iter().map(|(id,)| id)); } Ok(()) } From 221e80003eef72d00b6df2eed95a6a9c70e288cd Mon Sep 17 00:00:00 2001 From: DJ Majumdar Date: Wed, 18 Mar 2026 22:43:33 -0700 Subject: [PATCH 2/5] perf: reduce SQL round-trips in dispatch and completion hot paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Combine completion + dependency resolution into a single transaction, cache next_run_after to skip the query when no scheduled tasks exist, carry tags from peek to pop to avoid a redundant populate_tags query, skip tag queries entirely when the store has never had tags inserted, and skip the paused_tasks query when no tasks have been preempted. Benchmarked at ~40% improvement on dep_chain_dispatch/50 (42ms → 26ms) and ~20-34% on fan-in dispatch benchmarks. --- src/scheduler/mod.rs | 6 +++ src/scheduler/run_loop.rs | 61 ++++++++++++++++++++---------- src/scheduler/spawn/completion.rs | 33 ++++++++-------- src/scheduler/submit.rs | 16 +++++++- src/store/dependencies.rs | 18 ++++++--- src/store/lifecycle/transitions.rs | 56 +++++++++++++++++++++++++++ src/store/mod.rs | 22 ++++++++++- src/store/query/mod.rs | 2 +- src/store/submit/mod.rs | 11 ++++-- 9 files changed, 177 insertions(+), 48 deletions(-) diff --git a/src/scheduler/mod.rs b/src/scheduler/mod.rs index ddecabb..89bb959 100644 --- a/src/scheduler/mod.rs +++ b/src/scheduler/mod.rs @@ -117,6 +117,10 @@ pub(crate) struct SchedulerInner { /// Incremented when a task is dispatched; decremented on every terminal transition. /// Shared with spawned tasks via `Arc` so they can decrement on completion. pub(crate) module_running: Arc>, + /// Fast-path flag: set to `true` when a task is paused (preempted). + /// Cleared when `paused_tasks()` returns empty. Avoids a SQL round-trip + /// per dispatch cycle when no tasks are paused. + pub(crate) has_paused_tasks: AtomicBool, } /// IO-aware priority scheduler. @@ -243,6 +247,8 @@ impl Scheduler { module_paused, module_caps: RwLock::new(module_caps), module_running, + // Conservative: true on startup so the first cycle checks. + has_paused_tasks: AtomicBool::new(true), }), } } diff --git a/src/scheduler/run_loop.rs b/src/scheduler/run_loop.rs index 5ecc97f..6cf442c 100644 --- a/src/scheduler/run_loop.rs +++ b/src/scheduler/run_loop.rs @@ -85,11 +85,12 @@ impl Scheduler { } drop(reader_guard); - // Atomically claim the task. Returns None if another dispatcher - // claimed it (or it was cancelled) between peek and now. - let Some(task) = self.inner.store.pop_by_id(candidate.id).await? else { + // Atomically claim the task. Skip tag population since we already + // have the tags from the peek above. + let Some(mut task) = self.inner.store.pop_by_id_no_tags(candidate.id).await? else { return Ok(false); }; + task.tags = candidate.tags; // Look up executor. let Some(executor) = self.inner.registry.get(&task.task_type) else { @@ -205,16 +206,28 @@ impl Scheduler { )); } + // Track whether we need to query next_run_after. Starts true + // (conservative). Cleared when the query returns None (no scheduled + // tasks). Re-set when a notification arrives (new submit may have + // set run_after) or when we previously found scheduled tasks. + let mut check_scheduled = true; + loop { - // Compute sleep duration: min(poll_interval, time until next scheduled task). - let sleep_dur = match self.inner.store.next_run_after().await { - Ok(Some(next)) => { - let until_next = (next - chrono::Utc::now()) - .to_std() - .unwrap_or(std::time::Duration::ZERO); - std::cmp::min(self.inner.poll_interval, until_next) + let sleep_dur = if check_scheduled { + match self.inner.store.next_run_after().await { + Ok(Some(next)) => { + let until_next = (next - chrono::Utc::now()) + .to_std() + .unwrap_or(std::time::Duration::ZERO); + std::cmp::min(self.inner.poll_interval, until_next) + } + _ => { + check_scheduled = false; + self.inner.poll_interval + } } - _ => self.inner.poll_interval, + } else { + self.inner.poll_interval }; tokio::select! { @@ -224,6 +237,8 @@ impl Scheduler { break; } _ = self.inner.work_notify.notified() => { + // New work submitted — may include run_after tasks. + check_scheduled = true; self.poll_and_dispatch().await; } _ = tokio::time::sleep(sleep_dur) => { @@ -278,14 +293,22 @@ impl Scheduler { self.maybe_expire_tasks().await; // Resume paused tasks only if no active preemptors exist. - if let Ok(paused) = self.inner.store.paused_tasks().await { - for task in paused { - if !self - .inner - .active - .has_preemptors_for(task.priority, self.inner.preempt_priority) - { - let _ = self.inner.store.resume(task.id).await; + // Skip the query entirely when no tasks have been paused. + if self.inner.has_paused_tasks.load(AtomicOrdering::Relaxed) { + if let Ok(paused) = self.inner.store.paused_tasks().await { + if paused.is_empty() { + self.inner + .has_paused_tasks + .store(false, AtomicOrdering::Relaxed); + } + for task in paused { + if !self + .inner + .active + .has_preemptors_for(task.priority, self.inner.preempt_priority) + { + let _ = self.inner.store.resume(task.id).await; + } } } } diff --git a/src/scheduler/spawn/completion.rs b/src/scheduler/spawn/completion.rs index 3459996..63061eb 100644 --- a/src/scheduler/spawn/completion.rs +++ b/src/scheduler/spawn/completion.rs @@ -72,8 +72,12 @@ pub(crate) async fn handle_success( } } - match deps.store.complete_with_record(task, metrics).await { - Ok(recurring_info) => { + match deps + .store + .complete_with_record_and_resolve(task, metrics) + .await + { + Ok((recurring_info, unblocked)) => { // Emit recurring event if this was a recurring task. if task.recurring_interval_secs.is_some() { let (next_run, exec_count) = match recurring_info { @@ -86,22 +90,15 @@ pub(crate) async fn handle_success( next_run, }); } - } - Err(e) => { - tracing::error!(task_id, error = %e, "failed to record task completion"); - } - } - // Remove from active tracking AFTER the store write completes. - decrement_module(); - deps.active.remove(task_id); - let _ = deps - .event_tx - .send(SchedulerEvent::Completed(task.event_header())); + // Remove from active tracking AFTER the store write completes. + decrement_module(); + deps.active.remove(task_id); + let _ = deps + .event_tx + .send(SchedulerEvent::Completed(task.event_header())); - // Resolve dependency edges: unblock tasks waiting on this one. - match deps.store.resolve_dependents(task_id).await { - Ok(unblocked) => { + // Emit unblocked events for resolved dependents. for uid in &unblocked { let _ = deps .event_tx @@ -109,7 +106,9 @@ pub(crate) async fn handle_success( } } Err(e) => { - tracing::error!(task_id, error = %e, "failed to resolve dependents"); + tracing::error!(task_id, error = %e, "failed to complete task and resolve dependents"); + decrement_module(); + deps.active.remove(task_id); } } diff --git a/src/scheduler/submit.rs b/src/scheduler/submit.rs index 19fe65f..e1da13f 100644 --- a/src/scheduler/submit.rs +++ b/src/scheduler/submit.rs @@ -73,10 +73,16 @@ impl Scheduler { if !matches!(outcome, SubmitOutcome::Duplicate | SubmitOutcome::Rejected) { // Preempt if this is a high-priority task. if sub.priority.value() <= self.inner.preempt_priority.value() { - self.inner + let preempted = self + .inner .active .preempt_below(sub.priority, &self.inner.store, &self.inner.event_tx) .await; + if !preempted.is_empty() { + self.inner + .has_paused_tasks + .store(true, std::sync::atomic::Ordering::Relaxed); + } } // Wake the scheduler loop so it picks up the new/upgraded task. @@ -147,10 +153,16 @@ impl Scheduler { if let Some(priority) = best_priority { if priority.value() <= self.inner.preempt_priority.value() { - self.inner + let preempted = self + .inner .active .preempt_below(priority, &self.inner.store, &self.inner.event_tx) .await; + if !preempted.is_empty() { + self.inner + .has_paused_tasks + .store(true, std::sync::atomic::Ordering::Relaxed); + } } } diff --git a/src/store/dependencies.rs b/src/store/dependencies.rs index 15f452f..35c189f 100644 --- a/src/store/dependencies.rs +++ b/src/store/dependencies.rs @@ -14,18 +14,27 @@ impl TaskStore { /// Returns IDs of newly-unblocked tasks (for event emission). pub async fn resolve_dependents(&self, completed_task_id: i64) -> Result, StoreError> { let mut conn = self.begin_write().await?; + let unblocked = Self::resolve_dependents_inner(&mut conn, completed_task_id).await?; + sqlx::query("COMMIT").execute(&mut *conn).await?; + Ok(unblocked) + } + /// Inner dependency resolution that runs within an existing transaction. + pub(crate) async fn resolve_dependents_inner( + conn: &mut sqlx::pool::PoolConnection, + completed_task_id: i64, + ) -> Result, StoreError> { // Find tasks that depend on the completed task. let dependent_ids: Vec<(i64,)> = sqlx::query_as("SELECT task_id FROM task_deps WHERE depends_on_id = ?") .bind(completed_task_id) - .fetch_all(&mut *conn) + .fetch_all(&mut **conn) .await?; // Remove the satisfied edges. sqlx::query("DELETE FROM task_deps WHERE depends_on_id = ?") .bind(completed_task_id) - .execute(&mut *conn) + .execute(&mut **conn) .await?; let mut unblocked = Vec::new(); @@ -35,7 +44,7 @@ impl TaskStore { let (remaining,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM task_deps WHERE task_id = ?") .bind(dep_id) - .fetch_one(&mut *conn) + .fetch_one(&mut **conn) .await?; if remaining == 0 { @@ -44,7 +53,7 @@ impl TaskStore { "UPDATE tasks SET status = 'pending' WHERE id = ? AND status = 'blocked'", ) .bind(dep_id) - .execute(&mut *conn) + .execute(&mut **conn) .await?; if result.rows_affected() > 0 { unblocked.push(dep_id); @@ -52,7 +61,6 @@ impl TaskStore { } } - sqlx::query("COMMIT").execute(&mut *conn).await?; Ok(unblocked) } diff --git a/src/store/lifecycle/transitions.rs b/src/store/lifecycle/transitions.rs index 5077fce..bd17f3e 100644 --- a/src/store/lifecycle/transitions.rs +++ b/src/store/lifecycle/transitions.rs @@ -79,6 +79,36 @@ impl TaskStore { Ok(record) } + /// Atomically claim a specific pending task by id without populating tags. + /// + /// Callers must provide tags separately (e.g. carried from a prior + /// `peek_next`). This avoids a redundant `populate_tags` round-trip + /// when the caller already has the tags. + pub(crate) async fn pop_by_id_no_tags( + &self, + id: i64, + ) -> Result, StoreError> { + tracing::debug!(task_id = id, "store.pop_by_id_no_tags: UPDATE start"); + let row = sqlx::query( + "UPDATE tasks SET + status = 'running', + started_at = datetime('now'), + expires_at = CASE + WHEN ttl_from = 'first_attempt' AND ttl_seconds IS NOT NULL AND expires_at IS NULL + THEN datetime('now', '+' || ttl_seconds || ' seconds') + ELSE expires_at + END + WHERE id = ? AND status = 'pending' + RETURNING *", + ) + .bind(id) + .fetch_optional(&self.pool) + .await?; + tracing::debug!(task_id = id, "store.pop_by_id_no_tags: UPDATE end"); + + Ok(row.as_ref().map(row_to_task_record)) + } + /// Pop the highest-priority pending task and mark it as running. /// Returns `None` if the queue is empty. Tasks with a future `run_after` /// timestamp are excluded. @@ -184,6 +214,32 @@ impl TaskStore { Ok(recurring_info) } + /// Mark a task as completed and resolve its dependents in a single transaction. + /// + /// Combines `complete_with_record` and `resolve_dependents` to avoid + /// two separate `BEGIN IMMEDIATE` / `COMMIT` cycles. + /// + /// Returns `(recurring_info, unblocked_ids)`. + pub async fn complete_with_record_and_resolve( + &self, + task: &crate::task::TaskRecord, + metrics: &IoBudget, + ) -> Result<(Option<(chrono::DateTime, i64)>, Vec), StoreError> { + tracing::debug!(task_id = task.id, "store.complete_and_resolve: BEGIN tx"); + let mut conn = self.begin_write().await?; + + let recurring_info = Self::complete_inner(&mut conn, task, metrics).await?; + let unblocked = Self::resolve_dependents_inner(&mut conn, task.id).await?; + + sqlx::query("COMMIT").execute(&mut *conn).await?; + drop(conn); + tracing::debug!(task_id = task.id, "store.complete_and_resolve: COMMIT ok"); + + self.maybe_prune().await; + + Ok((recurring_info, unblocked)) + } + /// Shared completion logic: insert history, handle recurring next instance, /// then handle requeue or delete. /// diff --git a/src/store/mod.rs b/src/store/mod.rs index 533e75e..f6f5bc6 100644 --- a/src/store/mod.rs +++ b/src/store/mod.rs @@ -31,7 +31,7 @@ mod submit; pub use lifecycle::FailBackoff; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use serde::{Deserialize, Serialize}; use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous}; @@ -165,6 +165,9 @@ pub struct TaskStore { pub(crate) retention_policy: Option, pub(crate) prune_interval: u64, pub(crate) completion_count: std::sync::Arc, + /// Fast-path flag: `false` means no tags have been inserted into + /// `task_tags`, so `populate_tags` can skip the query entirely. + pub(crate) has_tags: std::sync::Arc, } impl TaskStore { @@ -192,6 +195,8 @@ impl TaskStore { retention_policy: config.retention_policy, prune_interval: config.prune_interval, completion_count: std::sync::Arc::new(AtomicU64::new(0)), + // Conservative for file-backed stores that may have existing tags. + has_tags: std::sync::Arc::new(AtomicBool::new(true)), }; store.migrate().await?; store.recover_running().await?; @@ -216,6 +221,8 @@ impl TaskStore { retention_policy: Some(RetentionPolicy::MaxCount(10_000)), prune_interval: 100, completion_count: std::sync::Arc::new(AtomicU64::new(0)), + // In-memory stores start empty — no tags to query. + has_tags: std::sync::Arc::new(AtomicBool::new(false)), }; store.migrate().await?; Ok(store) @@ -504,3 +511,16 @@ pub(crate) async fn insert_tags( } Ok(()) } + +/// Insert tags and mark the store's `has_tags` flag if non-empty. +pub(crate) async fn insert_tags_flagged( + conn: &mut sqlx::pool::PoolConnection, + task_id: i64, + tags: &std::collections::HashMap, + has_tags_flag: &AtomicBool, +) -> Result<(), StoreError> { + if !tags.is_empty() { + has_tags_flag.store(true, Ordering::Relaxed); + } + insert_tags(conn, task_id, tags).await +} diff --git a/src/store/query/mod.rs b/src/store/query/mod.rs index f75e3a3..e413af1 100644 --- a/src/store/query/mod.rs +++ b/src/store/query/mod.rs @@ -19,7 +19,7 @@ impl TaskStore { /// Populate tags for a slice of task records from the task_tags table. pub(crate) async fn populate_tags(&self, records: &mut [TaskRecord]) -> Result<(), StoreError> { - if records.is_empty() { + if records.is_empty() || !self.has_tags.load(std::sync::atomic::Ordering::Relaxed) { return Ok(()); } let ids: Vec = records.iter().map(|r| r.id).collect(); diff --git a/src/store/submit/mod.rs b/src/store/submit/mod.rs index 8b9fd10..f73c06e 100644 --- a/src/store/submit/mod.rs +++ b/src/store/submit/mod.rs @@ -54,6 +54,7 @@ fn validate_tags(tags: &HashMap) -> Result<(), StoreError> { pub(crate) async fn submit_one( conn: &mut sqlx::pool::PoolConnection, sub: &TaskSubmission, + has_tags_flag: Option<&std::sync::atomic::AtomicBool>, ) -> Result { if let Some(ref err) = sub.payload_error { return Err(StoreError::Serialization(err.clone())); @@ -127,7 +128,11 @@ pub(crate) async fn submit_one( let task_id = result.last_insert_rowid(); // Insert tags. - super::insert_tags(conn, task_id, &sub.tags).await?; + if let Some(flag) = has_tags_flag { + super::insert_tags_flagged(conn, task_id, &sub.tags, flag).await?; + } else { + super::insert_tags(conn, task_id, &sub.tags).await?; + } // Handle dependencies if any. if !sub.dependencies.is_empty() { @@ -183,7 +188,7 @@ impl TaskStore { let mut conn = self.begin_write().await?; tracing::debug!(task_type = %sub.task_type, "store.submit: INSERT start"); - let outcome = submit_one(&mut conn, sub).await?; + let outcome = submit_one(&mut conn, sub, Some(&self.has_tags)).await?; tracing::debug!(task_type = %sub.task_type, "store.submit: INSERT end"); sqlx::query("COMMIT").execute(&mut *conn).await?; Ok(outcome) @@ -237,7 +242,7 @@ impl TaskStore { if last_occurrence[&sub.effective_key()] != global_i { results.push(SubmitOutcome::Duplicate); } else { - results.push(submit_one(&mut conn, sub).await?); + results.push(submit_one(&mut conn, sub, Some(&self.has_tags)).await?); } } From 7fe59f2b0780a2556907a862e098abf3b01edbcb Mon Sep 17 00:00:00 2001 From: DJ Majumdar Date: Wed, 18 Mar 2026 22:44:22 -0700 Subject: [PATCH 3/5] docs: add migration guide for 0.4.x to 0.5.0 --- docs/migrating-to-0.5.md | 313 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100644 docs/migrating-to-0.5.md diff --git a/docs/migrating-to-0.5.md b/docs/migrating-to-0.5.md new file mode 100644 index 0000000..c8cc121 --- /dev/null +++ b/docs/migrating-to-0.5.md @@ -0,0 +1,313 @@ +# Migrating from 0.4.x to 0.5.0 + +0.5.0 replaces the stringly-typed `Module` / `ModuleHandle` API with a **domain-centric API** that enforces module identity and task ownership at compile time. Every `Module` becomes a `Domain` keyed on a Rust type; every `ModuleHandle` becomes a `DomainHandle`; and executors receive a typed, deserialized payload rather than a raw `TaskContext`. + +The changes are mechanical — there are no new database migrations and no changes to stored task type strings. + +--- + +## 1. Declare a `DomainKey` + +Every module needs a zero-sized marker type that implements `DomainKey`. This is the compile-time identity for the domain. + +**Before (0.4):** +```rust +// Module name lived only in a string literal +Module::new("media") +``` + +**After (0.5):** +```rust +pub struct Media; +impl DomainKey for Media { + const NAME: &'static str = "media"; +} +``` + +`NAME` must match the string you previously passed to `Module::new(...)`. The scheduler still uses this string to prefix task types in the database (`"media::thumbnail"`), so existing records are unaffected. + +--- + +## 2. Replace `Module` with `Domain` + +`Module` is no longer part of the public API. Replace every `Module::new(...)` call with `Domain::::new()`. + +**Before:** +```rust +let module = Module::new("media") + .executor("thumbnail", Arc::new(ThumbnailExec::new(cdn))) + .typed_executor::(Arc::new(TranscodeExec::new())) + .executor_with_ttl("upload", Arc::new(UploadExec), Duration::from_secs(600)) + .default_priority(Priority::NORMAL) + .default_retry_policy(RetryPolicy { ... }) + .default_group("pipeline") + .default_ttl(Duration::from_secs(3600)) + .default_tag("team", "media") + .max_concurrency(4) + .app_state(MediaConfig { cdn_url: "...".into() }); +``` + +**After:** +```rust +let domain = Domain::::new() + .task::(ThumbnailExec::new(cdn)) // no Arc needed + .task::(TranscodeExec::new()) + .task_with::(UploadExec, TaskTypeOptions { ttl: Some(Duration::from_secs(600)), ..Default::default() }) + .default_priority(Priority::NORMAL) + .default_retry(RetryPolicy { ... }) // renamed: default_retry_policy → default_retry + .default_group("pipeline") + .default_ttl(Duration::from_secs(3600)) + .default_tag("team", "media") + .max_concurrency(4) + .state(MediaConfig { cdn_url: "...".into() }); // renamed: app_state → state +``` + +**Method renames on the domain builder:** + +| 0.4 (`Module`) | 0.5 (`Domain`) | +|-----------------------------------------------|-------------------------------------------| +| `.executor("name", Arc::new(e))` | `.task::(e)` (see §3) | +| `.typed_executor::(Arc::new(e))` | `.task::(e)` | +| `.executor_with_ttl("name", Arc::new(e), d)` | `.task_with::(e, TaskTypeOptions { ttl: Some(d), .. })` | +| `.default_retry_policy(p)` | `.default_retry(p)` | +| `.app_state(v)` | `.state(v)` (or `.state_arc(arc)`) | + +### Register with the scheduler + +`SchedulerBuilder::module()` is now private. Use `SchedulerBuilder::domain()`: + +**Before:** +```rust +Scheduler::builder().module(module) +``` + +**After:** +```rust +Scheduler::builder().domain(domain) +``` + +--- + +## 3. Implement `TypedTask` with `type Domain` + +`TypedTask` gains a required associated type `Domain` that binds each task to its domain. Static defaults (priority, IO, TTL, retry, etc.) move from per-method overrides into a single `config()` method that returns `TaskTypeConfig`. + +**Before:** +```rust +impl TypedTask for Thumbnail { + const TASK_TYPE: &'static str = "thumbnail"; + + fn priority(&self) -> Priority { Priority::NORMAL } + fn expected_io(&self) -> IoBudget { IoBudget::disk(4096, 1024) } + fn ttl(&self) -> Option { Some(Duration::from_secs(3600)) } + fn key(&self) -> Option { Some(format!("thumb:{}:{}", self.path, self.size)) } +} +``` + +**After:** +```rust +impl TypedTask for Thumbnail { + type Domain = Media; // ← required + const TASK_TYPE: &'static str = "thumbnail"; + + fn config() -> TaskTypeConfig { // ← static, not &self + TaskTypeConfig::new() + .priority(Priority::NORMAL) + .expected_io(IoBudget::disk(4096, 1024)) + .ttl(Duration::from_secs(3600)) + .retry(RetryPolicy::exponential(3, Duration::from_secs(1), Duration::from_secs(60))) + .on_duplicate(DuplicateStrategy::Supersede) + } + + fn key(&self) -> Option { // instance methods unchanged + Some(format!("thumb:{}:{}", self.path, self.size)) + } +} +``` + +`key()`, `label()`, and `tags()` remain as `&self` instance methods because their values depend on the payload. Everything else moves to `config()`. + +The compiler enforces that you only register `T` with the domain whose `DomainKey` matches `T::Domain`: + +```rust +Domain::::new() + .task::(...) // ok — Thumbnail::Domain = Media + .task::(...) // ok — Upload::Domain = Media + // .task::(...) // compile error — SendEmail::Domain ≠ Media +``` + +--- + +## 4. Implement `TypedExecutor` instead of `TaskExecutor` + +Executors no longer need to deserialize the payload themselves. Implement `TypedExecutor` and receive the typed payload directly. + +**Before:** +```rust +impl TaskExecutor for ThumbnailExec { + async fn execute<'a>(&'a self, ctx: &'a TaskContext) -> Result<(), TaskError> { + let thumb: Thumbnail = ctx.payload()?; // manual deserialization + process(&thumb, ctx).await + } +} +``` + +**After:** +```rust +impl TypedExecutor for ThumbnailExec { + async fn execute(&self, thumb: Thumbnail, ctx: &TaskContext) -> Result<(), TaskError> { + process(&thumb, ctx).await + } +} +``` + +The `finalize` and `on_cancel` hooks follow the same pattern: + +```rust +impl TypedExecutor for ThumbnailExec { + async fn finalize(&self, thumb: Thumbnail, ctx: &TaskContext) -> Result<(), TaskError> { + // called after all children settle + Ok(()) + } + + async fn on_cancel(&self, thumb: Thumbnail, ctx: &TaskContext) -> Result<(), TaskError> { + // cleanup on preemption/cancellation + Ok(()) + } +} +``` + +The `TaskExecutor` trait remains available as an escape hatch via `Domain::raw_executor("name", exec)`, but prefer `TypedExecutor` for all new code. + +--- + +## 5. Replace `ModuleHandle` with `DomainHandle` + +`scheduler.module("media")` is now private. Use `scheduler.domain::()`. + +**Before:** +```rust +let media = scheduler.module("media"); // ModuleHandle +``` + +**After:** +```rust +let media = scheduler.domain::(); // DomainHandle +``` + +### Submission + +| 0.4 (`ModuleHandle`) | 0.5 (`DomainHandle`) | +|------------------------------------------------|-------------------------------------------| +| `handle.submit_typed(&task).await?` | `domain.submit(task).await?` | +| `handle.submit_typed(&task).priority(p).await?`| `domain.submit_with(task).priority(p).await?` | +| `handle.submit(sub).await?` | `domain.submit_raw(sub).await?` | + +Note that `submit` takes the task **by value** in 0.5, not by reference. Use `.clone()` if you need the value afterward. + +```rust +// 0.4 +media.submit_typed(&thumb).priority(Priority::HIGH).await?; + +// 0.5 +media.submit_with(thumb).priority(Priority::HIGH).await?; +// or zero-ceremony: +media.submit(thumb).await?; +``` + +All other handle methods (`cancel`, `pause`, `resume`, `snapshot`, `active_tasks`, `dead_letter_tasks`, `retry_dead_letter`, `cancel_all`, `cancel_where`, `tasks_by_tags`, `set_max_concurrency`, `pause_recurring`, `resume_recurring`, `cancel_recurring`) are unchanged on `DomainHandle`. + +--- + +## 6. Update cross-domain access in executors + +`ctx.current_module()` and `ctx.module("name")` are now internal (`pub(crate)`). Use `ctx.domain::()` for cross-domain submission from within an executor. + +**Before:** +```rust +// same-module follow-up +ctx.current_module().submit_typed(&NextStep { ... }).await?; + +// cross-module +ctx.module("notifications").submit_typed(&Notify { ... }).await?; +if let Some(h) = ctx.try_module("analytics") { + h.submit_typed(&Track { ... }).await?; +} +``` + +**After:** +```rust +// same-domain follow-up +ctx.domain::().submit(NextStep { ... }).await?; + +// cross-domain +ctx.domain::().submit(Notify { ... }).await?; +if let Some(h) = ctx.try_domain::() { + h.submit(Track { ... }).await?; +} +``` + +`spawn_child` is unchanged — it still auto-prefixes the task type and inherits TTL/tags from the parent: + +```rust +ctx.spawn_child(TaskSubmission::new("postprocess").payload_json(&p)).await?; +``` + +--- + +## 7. Update event subscriptions + +**Before:** +```rust +let mut rx = scheduler.module("media").subscribe(); // ModuleReceiver +while let Ok(event) = rx.recv().await { + match event { + SchedulerEvent::Completed(header) if header.task_type.ends_with("thumbnail") => { ... } + _ => {} + } +} +``` + +**After:** +```rust +// All events for the domain: +let mut rx = scheduler.domain::().events(); // ModuleReceiver + +// Or per-type typed stream (preferred): +let mut stream = media.task_events::(); // TypedEventStream +while let Ok(event) = stream.recv().await { + match event { + TaskEvent::Completed { id, record } => { + let thumb: Thumbnail = serde_json::from_slice( + record.payload.as_deref().unwrap() + ).unwrap(); + println!("done: {}", thumb.path); + } + TaskEvent::Failed { id, error, will_retry, .. } => { ... } + TaskEvent::DeadLettered { id, record, .. } => { ... } + TaskEvent::Progress { id, percent, message } => { ... } + _ => {} + } +} +``` + +`TypedEventStream` filters by both domain and task type, so no manual `task_type` string matching is needed. Terminal variants include an `Arc` for zero-cost access to the history entry. + +--- + +## Summary checklist + +- [ ] For each `Module::new("name")`: declare a `struct Name; impl DomainKey for Name { const NAME = "name"; }`. +- [ ] Replace `Module::new(...)` with `Domain::::new()`. +- [ ] Replace `.executor(...)` / `.typed_executor::(Arc::new(e))` with `.task::(e)`. +- [ ] Add `type Domain = Name;` to every `TypedTask` impl. +- [ ] Replace per-method `TypedTask` defaults with `fn config() -> TaskTypeConfig { ... }`. +- [ ] Replace `impl TaskExecutor for E { async fn execute(&self, ctx) }` with `impl TypedExecutor for E { async fn execute(&self, payload: T, ctx) }`. +- [ ] Replace `SchedulerBuilder::module(m)` with `SchedulerBuilder::domain(d)`. +- [ ] Replace `scheduler.module("name")` / `scheduler.try_module("name")` with `scheduler.domain::()` / `scheduler.try_domain::()`. +- [ ] Replace `handle.submit_typed(&task)` with `domain.submit(task)` (note: by value). +- [ ] Replace `handle.submit_typed(&task).priority(p)` with `domain.submit_with(task).priority(p)`. +- [ ] Replace `handle.subscribe()` with `domain.events()` or `domain.task_events::()`. +- [ ] Replace `ctx.current_module()` and `ctx.module("name")` with `ctx.domain::()`. +- [ ] Rename `.default_retry_policy(p)` → `.default_retry(p)` on the domain builder. +- [ ] Rename `.app_state(v)` → `.state(v)` on the domain builder. From 4b987b0a320482a91a67b795f6de53cde17efcf2 Mon Sep 17 00:00:00 2001 From: DJ Majumdar Date: Wed, 18 Mar 2026 22:58:26 -0700 Subject: [PATCH 4/5] perf: batch dependency queries and eliminate redundant dispatch round-trips MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Batch resolve_dependency_edges: replace 3N per-dep queries (history check, active check, edge insert) with 3 batch queries using IN clauses - Single-pass cycle detection: seed one recursive CTE with all dep IDs instead of running N separate CTEs - Replace pop_by_id_no_tags (UPDATE RETURNING *) with claim_task (UPDATE only), reusing the TaskRecord already fetched by peek_next Benchmarked improvement on dep_fan_in_dispatch: width=10: -15% (5.2ms → 4.4ms) width=50: -19% (21.7ms → 17.5ms) width=100: -14% (40.5ms → 34.8ms) --- src/scheduler/run_loop.rs | 20 +++- src/store/lifecycle/transitions.rs | 25 ++--- src/store/submit/dependencies.rs | 164 +++++++++++++++++++---------- 3 files changed, 130 insertions(+), 79 deletions(-) diff --git a/src/scheduler/run_loop.rs b/src/scheduler/run_loop.rs index 6cf442c..6781a18 100644 --- a/src/scheduler/run_loop.rs +++ b/src/scheduler/run_loop.rs @@ -85,12 +85,22 @@ impl Scheduler { } drop(reader_guard); - // Atomically claim the task. Skip tag population since we already - // have the tags from the peek above. - let Some(mut task) = self.inner.store.pop_by_id_no_tags(candidate.id).await? else { + // Atomically claim the task. We already have the full record from + // peek_next, so use claim_task (no RETURNING *) and patch in-memory. + if !self.inner.store.claim_task(candidate.id).await? { return Ok(false); - }; - task.tags = candidate.tags; + } + let mut task = candidate; + task.status = crate::task::TaskStatus::Running; + task.started_at = Some(chrono::Utc::now()); + // Mirror the SQL TTL logic for first-attempt tasks. + if task.ttl_from == crate::task::TtlFrom::FirstAttempt + && task.ttl_seconds.is_some() + && task.expires_at.is_none() + { + task.expires_at = + Some(chrono::Utc::now() + chrono::Duration::seconds(task.ttl_seconds.unwrap())); + } // Look up executor. let Some(executor) = self.inner.registry.get(&task.task_type) else { diff --git a/src/store/lifecycle/transitions.rs b/src/store/lifecycle/transitions.rs index bd17f3e..91c1364 100644 --- a/src/store/lifecycle/transitions.rs +++ b/src/store/lifecycle/transitions.rs @@ -79,17 +79,13 @@ impl TaskStore { Ok(record) } - /// Atomically claim a specific pending task by id without populating tags. + /// Atomically claim a pending task by id, returning `true` if claimed. /// - /// Callers must provide tags separately (e.g. carried from a prior - /// `peek_next`). This avoids a redundant `populate_tags` round-trip - /// when the caller already has the tags. - pub(crate) async fn pop_by_id_no_tags( - &self, - id: i64, - ) -> Result, StoreError> { - tracing::debug!(task_id = id, "store.pop_by_id_no_tags: UPDATE start"); - let row = sqlx::query( + /// The caller is expected to already hold the full [`TaskRecord`] from a + /// prior `peek_next` and will update its in-memory fields directly. This + /// avoids the `RETURNING *` round-trip of [`pop_by_id`](Self::pop_by_id). + pub(crate) async fn claim_task(&self, id: i64) -> Result { + let result = sqlx::query( "UPDATE tasks SET status = 'running', started_at = datetime('now'), @@ -98,15 +94,12 @@ impl TaskStore { THEN datetime('now', '+' || ttl_seconds || ' seconds') ELSE expires_at END - WHERE id = ? AND status = 'pending' - RETURNING *", + WHERE id = ? AND status = 'pending'", ) .bind(id) - .fetch_optional(&self.pool) + .execute(&self.pool) .await?; - tracing::debug!(task_id = id, "store.pop_by_id_no_tags: UPDATE end"); - - Ok(row.as_ref().map(row_to_task_record)) + Ok(result.rows_affected() > 0) } /// Pop the highest-priority pending task and mark it as running. diff --git a/src/store/submit/dependencies.rs b/src/store/submit/dependencies.rs index b7ea8c1..6c2a79e 100644 --- a/src/store/submit/dependencies.rs +++ b/src/store/submit/dependencies.rs @@ -1,5 +1,7 @@ //! Dependency resolution and cycle detection for task submission. +use std::collections::{HashMap, HashSet}; + use crate::store::StoreError; use crate::task::DependencyFailurePolicy; @@ -20,54 +22,92 @@ pub(super) async fn resolve_dependency_edges( deps: &[i64], policy: DependencyFailurePolicy, ) -> Result<(Vec, crate::task::TaskStatus), StoreError> { - let mut active_deps = Vec::new(); + if deps.is_empty() { + return Ok((Vec::new(), crate::task::TaskStatus::Pending)); + } + // --- Step 1: Batch-check history for all deps (one query). --- + // History is authoritative: SQLite may reuse row IDs of deleted tasks, + // so a completed dep's ID could now belong to a different active task. + let placeholders = deps.iter().map(|_| "?").collect::>().join(","); + let history_query = format!( + "SELECT h.id, h.status FROM task_history h + WHERE h.id IN ({placeholders}) + AND h.completed_at = ( + SELECT MAX(h2.completed_at) FROM task_history h2 WHERE h2.id = h.id + )" + ); + let mut q = sqlx::query_as::<_, (i64, String)>(&history_query); for &dep_id in deps { - // Check history FIRST. SQLite may reuse row IDs of deleted tasks, - // so a completed dep's ID could now belong to a different active task. - // History is authoritative for previously-completed/failed tasks. - let history_status: Option<(String,)> = sqlx::query_as( - "SELECT status FROM task_history WHERE id = ? ORDER BY completed_at DESC LIMIT 1", - ) - .bind(dep_id) - .fetch_optional(&mut **conn) - .await?; - - if let Some((ref status,)) = history_status { + q = q.bind(dep_id); + } + let history_rows = q.fetch_all(&mut **conn).await?; + + let mut history_map: HashMap = HashMap::with_capacity(history_rows.len()); + for (id, status) in history_rows { + history_map.insert(id, status); + } + + // Process history results; collect deps that need an active-queue check. + let mut need_active_check = Vec::new(); + for &dep_id in deps { + if let Some(status) = history_map.get(&dep_id) { match status.as_str() { "completed" => { /* already done, no edge needed */ } - _ => { - // Dep failed/cancelled/expired — apply failure policy. - match policy { - DependencyFailurePolicy::Cancel | DependencyFailurePolicy::Fail => { - return Err(StoreError::DependencyFailed(dep_id)); - } - DependencyFailurePolicy::Ignore => { /* skip */ } + _ => match policy { + DependencyFailurePolicy::Cancel | DependencyFailurePolicy::Fail => { + return Err(StoreError::DependencyFailed(dep_id)); } - } + DependencyFailurePolicy::Ignore => { /* skip */ } + }, } - continue; + } else { + need_active_check.push(dep_id); } + } - // Not in history — check if dep exists in active queue. - let active: Option<(i64,)> = sqlx::query_as("SELECT id FROM tasks WHERE id = ?") - .bind(dep_id) - .fetch_optional(&mut **conn) - .await?; - - if active.is_some() { - // Dep is still active — insert edge. - sqlx::query("INSERT INTO task_deps (task_id, depends_on_id) VALUES (?, ?)") - .bind(task_id) - .bind(dep_id) - .execute(&mut **conn) - .await?; + if need_active_check.is_empty() { + return Ok((Vec::new(), crate::task::TaskStatus::Pending)); + } + + // --- Step 2: Batch-check active tasks (one query). --- + let placeholders2 = need_active_check + .iter() + .map(|_| "?") + .collect::>() + .join(","); + let active_query = format!("SELECT id FROM tasks WHERE id IN ({placeholders2})"); + let mut q2 = sqlx::query_as::<_, (i64,)>(&active_query); + for &dep_id in &need_active_check { + q2 = q2.bind(dep_id); + } + let active_rows = q2.fetch_all(&mut **conn).await?; + let active_set: HashSet = active_rows.into_iter().map(|(id,)| id).collect(); + + // Validate: any dep not in history AND not active is invalid. + let mut active_deps = Vec::with_capacity(active_set.len()); + for &dep_id in &need_active_check { + if active_set.contains(&dep_id) { active_deps.push(dep_id); - continue; + } else { + return Err(StoreError::InvalidDependency(dep_id)); } + } - // Not in history and not in active queue. - return Err(StoreError::InvalidDependency(dep_id)); + // --- Step 3: Batch-insert edges (one query). --- + if !active_deps.is_empty() { + let values = active_deps + .iter() + .map(|_| "(?, ?)") + .collect::>() + .join(", "); + let insert_query = + format!("INSERT INTO task_deps (task_id, depends_on_id) VALUES {values}"); + let mut q3 = sqlx::query(&insert_query); + for &dep_id in &active_deps { + q3 = q3.bind(task_id).bind(dep_id); + } + q3.execute(&mut **conn).await?; } let status = if active_deps.is_empty() { @@ -81,33 +121,41 @@ pub(super) async fn resolve_dependency_edges( /// Cycle detection via recursive CTE. /// -/// Walks the entire upstream ancestry of `deps` in a single SQL query -/// instead of issuing one SELECT per BFS level. This reduces the number -/// of Rust↔SQLite round-trips from O(chain_depth) to O(1). +/// Seeds one CTE with **all** dependency IDs and walks the entire upstream +/// ancestry in a single SQL round-trip. If `new_task_id` appears anywhere +/// in the ancestor set, a cycle exists (because we just added +/// `new_task_id → deps`, so any path from deps back to new_task_id is a +/// cycle). pub(super) async fn detect_cycle( conn: &mut sqlx::pool::PoolConnection, new_task_id: i64, deps: &[i64], ) -> Result<(), StoreError> { + if deps.is_empty() { + return Ok(()); + } + + let seeds = deps.iter().map(|_| "(?)").collect::>().join(", "); + let query = format!( + "WITH RECURSIVE ancestors(id) AS ( + VALUES {seeds} + UNION + SELECT td.depends_on_id + FROM task_deps td + JOIN ancestors a ON td.task_id = a.id + ) + SELECT 1 FROM ancestors WHERE id = ? LIMIT 1" + ); + + let mut q = sqlx::query_as::<_, (i32,)>(&query); for &dep_id in deps { - let found: Option<(i64,)> = sqlx::query_as( - "WITH RECURSIVE ancestors(id) AS ( - SELECT ? AS id - UNION - SELECT td.depends_on_id - FROM task_deps td - JOIN ancestors a ON td.task_id = a.id - ) - SELECT id FROM ancestors WHERE id = ? LIMIT 1", - ) - .bind(dep_id) - .bind(new_task_id) - .fetch_optional(&mut **conn) - .await?; - - if found.is_some() { - return Err(StoreError::CyclicDependency); - } + q = q.bind(dep_id); } + q = q.bind(new_task_id); + + if q.fetch_optional(&mut **conn).await?.is_some() { + return Err(StoreError::CyclicDependency); + } + Ok(()) } From eed3c2f89c2535a50198bd230e524c67129f140c Mon Sep 17 00:00:00 2001 From: DJ Majumdar Date: Wed, 18 Mar 2026 23:02:18 -0700 Subject: [PATCH 5/5] perf: gate pprof profiler behind optional `profile` feature for CI compatibility Move pprof to an optional dependency gated by a `profile` feature so benchmarks can run in CI without requiring perf_event_open. Local profiling is still available via `cargo bench --features profile`. Also simplify TTL expiry check in run_loop to use if-let instead of unwrap. --- Cargo.toml | 3 ++- benches/dependencies.rs | 4 ++++ src/scheduler/run_loop.rs | 10 ++++------ 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 9a194e0..c5f496d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ bench = false [features] default = ["sysinfo-monitor"] sysinfo-monitor = ["dep:sysinfo"] +profile = ["dep:pprof"] [dependencies] tokio = { version = "1", features = ["sync", "time", "rt", "macros"] } @@ -28,11 +29,11 @@ serde_json = "1" sha2 = "0.10" fastrand = "2" sysinfo = { version = "0.33", optional = true } +pprof = { version = "0.13", features = ["flamegraph", "criterion"], optional = true } [dev-dependencies] tokio = { version = "1", features = ["full", "test-util"] } criterion = { version = "0.5", features = ["async_tokio"] } -pprof = { version = "0.13", features = ["flamegraph", "criterion"] } [[bench]] name = "scheduler" diff --git a/benches/dependencies.rs b/benches/dependencies.rs index e0f8bc2..46fda7f 100644 --- a/benches/dependencies.rs +++ b/benches/dependencies.rs @@ -5,6 +5,7 @@ use std::time::{Duration, Instant}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +#[cfg(feature = "profile")] use pprof::criterion::{Output, PProfProfiler}; use taskmill::{ Domain, DomainKey, Scheduler, SchedulerEvent, TaskContext, TaskError, TaskExecutor, TaskStore, @@ -211,12 +212,15 @@ fn bench_dep_fan_in_dispatch(c: &mut Criterion) { group.finish(); } +#[cfg(feature = "profile")] criterion_group! { name = submit_benches; config = Criterion::default() .with_profiler(PProfProfiler::new(1000, Output::Flamegraph(None))); targets = bench_dep_chain_submit } +#[cfg(not(feature = "profile"))] +criterion_group!(submit_benches, bench_dep_chain_submit); criterion_group!( dispatch_benches, bench_dep_chain_dispatch, diff --git a/src/scheduler/run_loop.rs b/src/scheduler/run_loop.rs index 6781a18..0810756 100644 --- a/src/scheduler/run_loop.rs +++ b/src/scheduler/run_loop.rs @@ -94,12 +94,10 @@ impl Scheduler { task.status = crate::task::TaskStatus::Running; task.started_at = Some(chrono::Utc::now()); // Mirror the SQL TTL logic for first-attempt tasks. - if task.ttl_from == crate::task::TtlFrom::FirstAttempt - && task.ttl_seconds.is_some() - && task.expires_at.is_none() - { - task.expires_at = - Some(chrono::Utc::now() + chrono::Duration::seconds(task.ttl_seconds.unwrap())); + if task.ttl_from == crate::task::TtlFrom::FirstAttempt && task.expires_at.is_none() { + if let Some(ttl) = task.ttl_seconds { + task.expires_at = Some(chrono::Utc::now() + chrono::Duration::seconds(ttl)); + } } // Look up executor.