From fb68f297cb368aa98476bcac82056648fa976f0b Mon Sep 17 00:00:00 2001 From: DJ Majumdar Date: Wed, 18 Mar 2026 21:48:52 -0700 Subject: [PATCH 1/3] refactor: decompose spawn_task into focused submodules Break the 380-line spawn_task god function into testable units: - spawn/context.rs: SpawnContext + TaskContext construction - spawn/completion.rs: success path (children, recurring, dependencies) - spawn/failure.rs: failure path (retry, dead-letter, fail-fast cascade) - spawn/parent.rs: parent-child resolution - spawn.rs: slim ~85-line orchestrator Also DRY up ActiveTaskMap bulk operations (preempt_below, pause_module, pause_all) with a shared drain_where helper and cancel_pause_emit. --- src/scheduler/dispatch.rs | 540 ++---------------------------- src/scheduler/mod.rs | 1 + src/scheduler/run_loop.rs | 10 +- src/scheduler/spawn.rs | 146 ++++++++ src/scheduler/spawn/completion.rs | 130 +++++++ src/scheduler/spawn/context.rs | 92 +++++ src/scheduler/spawn/failure.rs | 199 +++++++++++ src/scheduler/spawn/parent.rs | 59 ++++ 8 files changed, 668 insertions(+), 509 deletions(-) create mode 100644 src/scheduler/spawn.rs create mode 100644 src/scheduler/spawn/completion.rs create mode 100644 src/scheduler/spawn/context.rs create mode 100644 src/scheduler/spawn/failure.rs create mode 100644 src/scheduler/spawn/parent.rs diff --git a/src/scheduler/dispatch.rs b/src/scheduler/dispatch.rs index 1922686..5f0dcdb 100644 --- a/src/scheduler/dispatch.rs +++ b/src/scheduler/dispatch.rs @@ -1,17 +1,15 @@ -//! Task spawning, active-task tracking, preemption, and parent-child resolution. +//! Active-task tracking and preemption. use std::collections::{HashMap, HashSet}; -use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::sync::{Arc, Mutex}; use tokio_util::sync::CancellationToken; use crate::priority::Priority; -use crate::registry::{ChildSpawner, IoTracker, ParentContext, StateSnapshot, TaskContext}; +use crate::registry::IoTracker; use crate::store::TaskStore; -use crate::task::{IoBudget, ParentResolution, TaskRecord}; +use crate::task::TaskRecord; -use super::progress::ProgressReporter; use super::SchedulerEvent; // ── Active Task ──────────────────────────────────────────────────── @@ -160,34 +158,16 @@ impl ActiveTaskMap { store: &TaskStore, event_tx: &tokio::sync::broadcast::Sender, ) -> Vec { - // Phase 1: collect + remove under sync lock. - let to_preempt: Vec<(i64, ActiveTask)> = { - let mut active = self.inner.lock().unwrap(); - let ids: Vec = active - .iter() - .filter(|(_, at)| at.record.priority.value() > incoming_priority.value()) - .map(|(id, _)| *id) - .collect(); - ids.into_iter() - .filter_map(|id| active.remove(&id).map(|at| (id, at))) - .collect() - }; - - // Phase 2: async work without the lock held. - let mut preempted = Vec::new(); - for (id, at) in to_preempt { + let drained = self.drain_where(|at| at.record.priority.value() > incoming_priority.value()); + for (id, at) in &drained { tracing::info!( task_id = id, task_type = at.record.task_type, "preempting task for higher-priority work" ); - at.token.cancel(); - let _ = store.pause(id).await; - let _ = event_tx.send(SchedulerEvent::Preempted(at.record.event_header())); - preempted.push(id); } - - preempted + cancel_pause_emit(&drained, store, event_tx).await; + drained.into_iter().map(|(id, _)| id).collect() } /// Check whether any active task would preempt work at the given priority. @@ -240,24 +220,9 @@ impl ActiveTaskMap { store: &TaskStore, event_tx: &tokio::sync::broadcast::Sender, ) -> usize { - let to_pause: Vec<(i64, ActiveTask)> = { - let mut map = self.inner.lock().unwrap(); - let ids: Vec = map - .iter() - .filter(|(_, at)| at.record.task_type.starts_with(prefix)) - .map(|(id, _)| *id) - .collect(); - ids.into_iter() - .filter_map(|id| map.remove(&id).map(|at| (id, at))) - .collect() - }; - let count = to_pause.len(); - for (id, at) in to_pause { - at.token.cancel(); - let _ = store.pause(id).await; - let _ = event_tx.send(SchedulerEvent::Preempted(at.record.event_header())); - } - count + let drained = self.drain_where(|at| at.record.task_type.starts_with(prefix)); + cancel_pause_emit(&drained, store, event_tx).await; + drained.len() } /// Pause all active tasks: cancel their tokens and move them to paused @@ -270,472 +235,39 @@ impl ActiveTaskMap { store: &TaskStore, event_tx: &tokio::sync::broadcast::Sender, ) -> usize { - // Drain under sync lock. - let drained: Vec<(i64, ActiveTask)> = { self.inner.lock().unwrap().drain().collect() }; - let count = drained.len(); - // Async work without the lock held. - for (id, at) in drained { - at.token.cancel(); - let _ = store.pause(id).await; - let _ = event_tx.send(SchedulerEvent::Preempted(at.record.event_header())); - } - count + let drained = self.drain_where(|_| true); + cancel_pause_emit(&drained, store, event_tx).await; + drained.len() } -} - -// ── Spawn ────────────────────────────────────────────────────────── - -/// Whether to call `execute` or `finalize` on the executor. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum ExecutionPhase { - Execute, - Finalize, -} - -/// Shared scheduler resources passed to each spawned task. -pub(crate) struct SpawnContext { - pub store: TaskStore, - pub active: ActiveTaskMap, - pub event_tx: tokio::sync::broadcast::Sender, - pub max_retries: i32, - pub registry: Arc, - pub app_state: crate::registry::StateSnapshot, - pub work_notify: Arc, - pub scheduler: super::WeakScheduler, - #[allow(dead_code)] - pub cancel_hook_timeout: tokio::time::Duration, - /// Per-module live running counts. Incremented on dispatch; decremented on terminal. - pub module_running: Arc>, - /// Pre-snapshotted per-module state (module name → snapshot). Cloned at dispatch time. - pub module_state: Arc>, - /// Registry of all registered modules — shared with spawned tasks so they can - /// construct [`ModuleHandle`](crate::module::ModuleHandle) instances. - pub module_registry: Arc, -} -/// Spawn a task executor and wire up completion/failure handling. -/// -/// Inserts the task into the active map, starts a progress listener, -/// and spawns the executor on a new tokio task. -pub(crate) async fn spawn_task( - task: TaskRecord, - executor: Arc, - ctx: SpawnContext, - phase: ExecutionPhase, -) { - let SpawnContext { - store, - active, - event_tx, - max_retries, - registry, - app_state, - work_notify, - scheduler, - cancel_hook_timeout: _, - module_running, - module_state, - module_registry, - } = ctx; - - // Extract the owning module name from the task type prefix (e.g. "media" from "media::thumb"). - let owning_module: String = task.module_name().unwrap_or_default().to_string(); - - // Clone the pre-snapshotted module state — no lock needed, already lock-free. - let module_state_snapshot: StateSnapshot = task - .module_name() - .and_then(|name| module_state.get(name).cloned()) - .unwrap_or_default(); - let child_token = CancellationToken::new(); - - // Build execution context. - let child_spawner = ChildSpawner::new( - store.clone(), - task.id, - work_notify.clone(), - ParentContext { - created_at: task.created_at, - ttl_seconds: task.ttl_seconds, - ttl_from: task.ttl_from, - started_at: task.started_at, - tags: task.tags.clone(), - }, - ); - let io = Arc::new(IoTracker::new()); - - // Insert into active map before spawning to avoid races. - active.insert( - task.id, - ActiveTask { - record: task.clone(), - token: child_token.clone(), - reported_progress: None, - reported_at: None, - handle: None, - io: io.clone(), - started_at: std::time::Instant::now(), - }, - ); - - // Increment the module running counter for this task. - if let Some(module_name) = task.module_name() { - if let Some(counter) = module_running.get(module_name) { - counter.fetch_add(1, AtomicOrdering::Relaxed); - } + /// Drain tasks matching `predicate` from the active map. + /// + /// Collects matching tasks under the sync lock and removes them + /// atomically. Returns the removed tasks for async follow-up work. + fn drain_where(&self, predicate: impl Fn(&ActiveTask) -> bool) -> Vec<(i64, ActiveTask)> { + let mut map = self.inner.lock().unwrap(); + let ids: Vec = map + .iter() + .filter(|(_, at)| predicate(at)) + .map(|(id, _)| *id) + .collect(); + ids.into_iter() + .filter_map(|id| map.remove(&id).map(|at| (id, at))) + .collect() } - - let ctx = TaskContext { - record: task.clone(), - token: child_token.clone(), - progress: ProgressReporter::new( - task.event_header(), - event_tx.clone(), - active.clone(), - io.clone(), - ), - scheduler, - app_state, - module_state: module_state_snapshot, - child_spawner: Some(child_spawner), - io: io.clone(), - module_registry, - owning_module, - }; - - // Emit dispatched event. - let _ = event_tx.send(SchedulerEvent::Dispatched(task.event_header())); - - // Spawn executor. - let task_id_for_handle = task.id; - let active_for_handle = active.clone(); - let token_for_spawn = child_token.clone(); - let module_running_for_task = module_running; - let handle = tokio::spawn(async move { - let task_id = task.id; - // Helper: decrement the module running counter when this task leaves "running". - let decrement_module = || { - if let Some(name) = task.module_name() { - if let Some(counter) = module_running_for_task.get(name) { - counter.fetch_sub(1, AtomicOrdering::Relaxed); - } - } - }; - let result = match phase { - ExecutionPhase::Execute => executor.execute_erased(&ctx).await, - ExecutionPhase::Finalize => executor.finalize_erased(&ctx).await, - }; - - // Read IO bytes from the context tracker. - let metrics = io.snapshot(); - - // Drop the context (and its progress reporter) — executor is done. - drop(ctx); - - match result { - Ok(()) => { - // For the execute phase, check if the task spawned children. - // If so, transition to waiting instead of completing. - if phase == ExecutionPhase::Execute { - match store.active_children_count(task_id).await { - Ok(count) if count > 0 => { - if let Err(e) = store.set_waiting(task_id).await { - tracing::error!(task_id, error = %e, "failed to set task to waiting"); - } - decrement_module(); - active.remove(task_id); - let _ = event_tx.send(SchedulerEvent::Waiting { - task_id, - children_count: count, - }); - // Children may have completed before we set waiting. - // Re-check to avoid a missed finalization. - handle_parent_resolution( - task_id, - &store, - &active, - &event_tx, - max_retries, - &work_notify, - ) - .await; - // Wake the scheduler to dispatch children (or finalizer). - work_notify.notify_one(); - return; - } - Err(e) => { - tracing::error!(task_id, error = %e, "failed to check children count"); - // Fall through to normal completion. - } - _ => { - // No children — complete normally. - } - } - } - - match store.complete_with_record(&task, &metrics).await { - Ok(recurring_info) => { - // Emit recurring event if this was a recurring task. - if task.recurring_interval_secs.is_some() { - let (next_run, exec_count) = match recurring_info { - Some((next, count)) => (Some(next), count), - None => (None, task.recurring_execution_count + 1), - }; - let _ = event_tx.send(SchedulerEvent::RecurringCompleted { - header: task.event_header(), - execution_count: exec_count, - 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(); - active.remove(task_id); - let _ = event_tx.send(SchedulerEvent::Completed(task.event_header())); - - // Resolve dependency edges: unblock tasks waiting on this one. - match store.resolve_dependents(task_id).await { - Ok(unblocked) => { - for uid in &unblocked { - let _ = event_tx.send(SchedulerEvent::TaskUnblocked { task_id: *uid }); - } - } - Err(e) => { - tracing::error!(task_id, error = %e, "failed to resolve dependents"); - } - } - - work_notify.notify_one(); - - // If this was a child task, check if parent is ready. - if let Some(parent_id) = task.parent_id { - handle_parent_resolution( - parent_id, - &store, - &active, - &event_tx, - max_retries, - &work_notify, - ) - .await; - } - } - Err(te) => { - // If cancelled (preempted), the scheduler already paused it. - if token_for_spawn.is_cancelled() { - decrement_module(); - active.remove(task_id); - return; - } - - // Resolve effective retry policy for this task type. - let policy = registry.type_retry_policy(&task.task_type); - let effective_max_retries = task - .max_retries - .unwrap_or(policy.map(|p| p.max_retries).unwrap_or(max_retries)); - let backoff_strategy = policy.map(|p| &p.strategy); - - let will_retry = te.retryable && task.retry_count < effective_max_retries; - - // Compute retry delay for event reporting. - let retry_delay = if will_retry { - if let Some(ms) = te.retry_after_ms { - Some(std::time::Duration::from_millis(ms)) - } else if let Some(strategy) = backoff_strategy { - let d = strategy.delay_for(task.retry_count); - if d.is_zero() { - None - } else { - Some(d) - } - } else { - None - } - } else { - None - }; - - tracing::warn!( - task_id, - task_type = task.task_type, - error = %te.message, - retryable = te.retryable, - will_retry, - "task failed" - ); - let fail_backoff = crate::store::FailBackoff { - strategy: backoff_strategy, - executor_retry_after_ms: te.retry_after_ms, - }; - if let Err(e) = store - .fail_with_record( - &task, - &te.message, - te.retryable, - effective_max_retries, - &metrics, - &fail_backoff, - ) - .await - { - tracing::error!(task_id, error = %e, "failed to record task failure"); - } - // Remove from active tracking AFTER the store write completes. - decrement_module(); - active.remove(task_id); - let dead_lettered = te.retryable && !will_retry; - if dead_lettered { - let _ = event_tx.send(SchedulerEvent::DeadLettered { - header: task.event_header(), - error: te.message.clone(), - retry_count: task.retry_count + 1, - }); - } else { - let _ = event_tx.send(SchedulerEvent::Failed { - header: task.event_header(), - error: te.message.clone(), - will_retry, - retry_after: retry_delay, - }); - } - work_notify.notify_one(); - - // If permanent failure, propagate to dependency chain. - if !will_retry { - match store.fail_dependents(task_id).await { - Ok((failed_ids, unblocked_ids)) => { - for fid in &failed_ids { - let _ = event_tx.send(SchedulerEvent::DependencyFailed { - task_id: *fid, - failed_dependency: task_id, - }); - } - for uid in &unblocked_ids { - let _ = - event_tx.send(SchedulerEvent::TaskUnblocked { task_id: *uid }); - } - if !unblocked_ids.is_empty() { - work_notify.notify_one(); - } - } - Err(e) => { - tracing::error!(task_id, error = %e, "failed to propagate failure to dependents"); - } - } - - if let Some(parent_id) = task.parent_id { - // Check if parent uses fail_fast. - if let Ok(Some(parent)) = store.task_by_id(parent_id).await { - if parent.fail_fast { - // Cancel remaining siblings. - if let Ok(running_ids) = store.cancel_children(parent_id).await { - for rid in &running_ids { - if let Some(at) = active.remove(*rid) { - at.token.cancel(); - let _ = store.delete(*rid).await; - let _ = event_tx.send(SchedulerEvent::Cancelled( - at.record.event_header(), - )); - } - } - } - // Fail the parent. - let msg = format!("child task {task_id} failed: {}", te.message); - if let Err(e) = store - .fail_with_record( - &parent, - &msg, - false, - 0, - &IoBudget::default(), - &Default::default(), - ) - .await - { - tracing::error!( - parent_id, - error = %e, - "failed to record parent failure" - ); - } - let _ = event_tx.send(SchedulerEvent::Failed { - header: parent.event_header(), - error: msg, - will_retry: false, - retry_after: None, - }); - } else { - // Not fail_fast — check if all children done. - handle_parent_resolution( - parent_id, - &store, - &active, - &event_tx, - max_retries, - &work_notify, - ) - .await; - } - } - } - } - } - } - }); - - // Store the handle so shutdown can join it. - active_for_handle.set_handle(task_id_for_handle, handle); } -/// Check if a waiting parent is ready for finalization or has failed, -/// and dispatch the finalize phase if ready. -async fn handle_parent_resolution( - parent_id: i64, +// ── Helpers ──────────────────────────────────────────────────────── + +/// Cancel tokens, pause in store, and emit `Preempted` events for drained tasks. +async fn cancel_pause_emit( + drained: &[(i64, ActiveTask)], store: &TaskStore, - active: &ActiveTaskMap, event_tx: &tokio::sync::broadcast::Sender, - _max_retries: i32, - work_notify: &Arc, ) { - match store.try_resolve_parent(parent_id).await { - Ok(Some(ParentResolution::ReadyToFinalize)) => { - // Enqueue parent for finalize dispatch. - active.pending_finalizers.lock().unwrap().insert(parent_id); - // Wake the scheduler to dispatch the finalize phase. - work_notify.notify_one(); - } - Ok(Some(ParentResolution::Failed(reason))) => { - // All children done but some failed — fail the parent. - if let Ok(Some(parent)) = store.task_by_id(parent_id).await { - if let Err(e) = store - .fail_with_record( - &parent, - &reason, - false, - 0, - &IoBudget::default(), - &Default::default(), - ) - .await - { - tracing::error!(parent_id, error = %e, "failed to record parent failure"); - } - let _ = event_tx.send(SchedulerEvent::Failed { - header: parent.event_header(), - error: reason, - will_retry: false, - retry_after: None, - }); - } - } - Ok(Some(ParentResolution::StillWaiting)) | Ok(None) => { - // Children still active or parent not found — nothing to do. - } - Err(e) => { - tracing::error!(parent_id, error = %e, "failed to resolve parent"); - } + for (id, at) in drained { + at.token.cancel(); + let _ = store.pause(*id).await; + let _ = event_tx.send(SchedulerEvent::Preempted(at.record.event_header())); } } diff --git a/src/scheduler/mod.rs b/src/scheduler/mod.rs index 1b34998..ddecabb 100644 --- a/src/scheduler/mod.rs +++ b/src/scheduler/mod.rs @@ -28,6 +28,7 @@ pub(crate) mod gate; pub mod progress; mod queries; mod run_loop; +pub(crate) mod spawn; mod submit; #[cfg(test)] mod tests; diff --git a/src/scheduler/run_loop.rs b/src/scheduler/run_loop.rs index baaa6a1..5ecc97f 100644 --- a/src/scheduler/run_loop.rs +++ b/src/scheduler/run_loop.rs @@ -10,8 +10,8 @@ use crate::task::IoBudget; use super::SchedulerEvent; -use super::dispatch::{self, SpawnContext}; use super::gate::GateContext; +use super::spawn::{self, SpawnContext}; use super::{Scheduler, ShutdownMode}; impl Scheduler { @@ -114,11 +114,11 @@ impl Scheduler { // Spawn the task — this inserts into the active map, builds the // context, emits Dispatched, and wires up completion handling. - dispatch::spawn_task( + spawn::spawn_task( task, executor, self.build_spawn_context().await, - dispatch::ExecutionPhase::Execute, + spawn::ExecutionPhase::Execute, ) .await; @@ -168,11 +168,11 @@ impl Scheduler { }; let executor = Arc::clone(executor); - dispatch::spawn_task( + spawn::spawn_task( task, executor, self.build_spawn_context().await, - dispatch::ExecutionPhase::Finalize, + spawn::ExecutionPhase::Finalize, ) .await; diff --git a/src/scheduler/spawn.rs b/src/scheduler/spawn.rs new file mode 100644 index 0000000..1e51b97 --- /dev/null +++ b/src/scheduler/spawn.rs @@ -0,0 +1,146 @@ +//! Task spawning — orchestrator and focused submodules. +//! +//! This module decomposes the former monolithic `spawn_task` function into +//! focused, testable units: +//! +//! - [`context`] — `SpawnContext` and `TaskContext` construction +//! - [`completion`] — success path (children check, completion, recurring) +//! - [`failure`] — failure path (retry, dead-letter, fail-fast cascade) +//! - [`parent`] — parent-child resolution after task completion + +mod completion; +mod context; +mod failure; +mod parent; + +use std::sync::atomic::Ordering as AtomicOrdering; +use std::sync::Arc; + +use crate::registry::ErasedExecutor; +use crate::task::TaskRecord; + +use super::dispatch::ActiveTask; +use super::SchedulerEvent; + +pub(crate) use context::SpawnContext; + +/// Whether to call `execute` or `finalize` on the executor. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ExecutionPhase { + Execute, + Finalize, +} + +/// Spawn a task executor and wire up completion/failure handling. +/// +/// Inserts the task into the active map, starts a progress listener, +/// and spawns the executor on a new tokio task. The actual success and +/// failure handling is delegated to [`completion::handle_success`] and +/// [`failure::handle_failure`]. +pub(crate) async fn spawn_task( + task: TaskRecord, + executor: Arc, + ctx: SpawnContext, + phase: ExecutionPhase, +) { + let prepared = context::build_task_context(&task, &ctx); + + // Insert into active map before spawning to avoid races. + ctx.active.insert( + task.id, + ActiveTask { + record: task.clone(), + token: prepared.token.clone(), + reported_progress: None, + reported_at: None, + handle: None, + io: prepared.io.clone(), + started_at: std::time::Instant::now(), + }, + ); + + // Increment the module running counter for this task. + if let Some(module_name) = task.module_name() { + if let Some(counter) = ctx.module_running.get(module_name) { + counter.fetch_add(1, AtomicOrdering::Relaxed); + } + } + + // Emit dispatched event. + let _ = ctx + .event_tx + .send(SchedulerEvent::Dispatched(task.event_header())); + + // Build deps for handlers (cloned from SpawnContext since they move into the spawned future). + let completion_deps = completion::CompletionDeps { + store: ctx.store.clone(), + active: ctx.active.clone(), + event_tx: ctx.event_tx.clone(), + work_notify: ctx.work_notify.clone(), + max_retries: ctx.max_retries, + }; + let failure_deps = failure::FailureDeps { + store: ctx.store, + active: ctx.active.clone(), + event_tx: ctx.event_tx, + work_notify: ctx.work_notify, + max_retries: ctx.max_retries, + registry: ctx.registry, + }; + + let task_id_for_handle = task.id; + let active_for_handle = ctx.active; + let token_for_spawn = prepared.token.clone(); + let module_running = ctx.module_running; + let io = prepared.io; + + let handle = tokio::spawn(async move { + let task_id = task.id; + + // Helper: decrement the module running counter when this task leaves "running". + let decrement_module = || { + if let Some(name) = task.module_name() { + if let Some(counter) = module_running.get(name) { + counter.fetch_sub(1, AtomicOrdering::Relaxed); + } + } + }; + + let result = match phase { + ExecutionPhase::Execute => executor.execute_erased(&prepared.ctx).await, + ExecutionPhase::Finalize => executor.finalize_erased(&prepared.ctx).await, + }; + + // Read IO bytes from the context tracker. + let metrics = io.snapshot(); + + // Drop the context (and its progress reporter) — executor is done. + drop(prepared.ctx); + + match result { + Ok(()) => { + completion::handle_success( + &task, + phase, + &metrics, + &completion_deps, + decrement_module, + ) + .await; + } + Err(te) => { + // If cancelled (preempted), the scheduler already paused it. + if token_for_spawn.is_cancelled() { + decrement_module(); + failure_deps.active.remove(task_id); + return; + } + + failure::handle_failure(&task, te, &metrics, &failure_deps, decrement_module).await; + } + } + }); + + // Store the handle so shutdown can join it. + active_for_handle.set_handle(task_id_for_handle, handle); +} diff --git a/src/scheduler/spawn/completion.rs b/src/scheduler/spawn/completion.rs new file mode 100644 index 0000000..3459996 --- /dev/null +++ b/src/scheduler/spawn/completion.rs @@ -0,0 +1,130 @@ +//! Success path: children check, completion, recurring re-enqueue, dependency resolution. + +use std::sync::Arc; + +use crate::store::TaskStore; +use crate::task::{IoBudget, TaskRecord}; + +use super::super::dispatch::ActiveTaskMap; +use super::super::SchedulerEvent; +use super::parent::handle_parent_resolution; +use super::ExecutionPhase; + +/// Shared dependencies for the completion handler. +pub(crate) struct CompletionDeps { + pub store: TaskStore, + pub active: ActiveTaskMap, + pub event_tx: tokio::sync::broadcast::Sender, + pub work_notify: Arc, + pub max_retries: i32, +} + +/// Handle a successful task execution. +/// +/// For the execute phase, checks if the task spawned children (transition to +/// waiting). Otherwise records completion, resolves dependents, and handles +/// recurring re-enqueue. +pub(crate) async fn handle_success( + task: &TaskRecord, + phase: ExecutionPhase, + metrics: &IoBudget, + deps: &CompletionDeps, + decrement_module: impl FnOnce(), +) { + let task_id = task.id; + + // For the execute phase, check if the task spawned children. + // If so, transition to waiting instead of completing. + if phase == ExecutionPhase::Execute { + match deps.store.active_children_count(task_id).await { + Ok(count) if count > 0 => { + if let Err(e) = deps.store.set_waiting(task_id).await { + tracing::error!(task_id, error = %e, "failed to set task to waiting"); + } + decrement_module(); + deps.active.remove(task_id); + let _ = deps.event_tx.send(SchedulerEvent::Waiting { + task_id, + children_count: count, + }); + // Children may have completed before we set waiting. + // Re-check to avoid a missed finalization. + handle_parent_resolution( + task_id, + &deps.store, + &deps.active, + &deps.event_tx, + deps.max_retries, + &deps.work_notify, + ) + .await; + // Wake the scheduler to dispatch children (or finalizer). + deps.work_notify.notify_one(); + return; + } + Err(e) => { + tracing::error!(task_id, error = %e, "failed to check children count"); + // Fall through to normal completion. + } + _ => { + // No children — complete normally. + } + } + } + + match deps.store.complete_with_record(task, metrics).await { + Ok(recurring_info) => { + // Emit recurring event if this was a recurring task. + if task.recurring_interval_secs.is_some() { + let (next_run, exec_count) = match recurring_info { + Some((next, count)) => (Some(next), count), + None => (None, task.recurring_execution_count + 1), + }; + let _ = deps.event_tx.send(SchedulerEvent::RecurringCompleted { + header: task.event_header(), + execution_count: exec_count, + 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())); + + // Resolve dependency edges: unblock tasks waiting on this one. + match deps.store.resolve_dependents(task_id).await { + Ok(unblocked) => { + for uid in &unblocked { + let _ = deps + .event_tx + .send(SchedulerEvent::TaskUnblocked { task_id: *uid }); + } + } + Err(e) => { + tracing::error!(task_id, error = %e, "failed to resolve dependents"); + } + } + + deps.work_notify.notify_one(); + + // If this was a child task, check if parent is ready. + if let Some(parent_id) = task.parent_id { + handle_parent_resolution( + parent_id, + &deps.store, + &deps.active, + &deps.event_tx, + deps.max_retries, + &deps.work_notify, + ) + .await; + } +} diff --git a/src/scheduler/spawn/context.rs b/src/scheduler/spawn/context.rs new file mode 100644 index 0000000..40d4464 --- /dev/null +++ b/src/scheduler/spawn/context.rs @@ -0,0 +1,92 @@ +//! Task context construction for spawned tasks. + +use std::collections::HashMap; +use std::sync::atomic::AtomicUsize; +use std::sync::Arc; + +use tokio_util::sync::CancellationToken; + +use crate::registry::{ChildSpawner, IoTracker, ParentContext, StateSnapshot, TaskContext}; +use crate::store::TaskStore; +use crate::task::TaskRecord; + +use super::super::dispatch::ActiveTaskMap; +use super::super::progress::ProgressReporter; +use super::super::SchedulerEvent; + +/// Shared scheduler resources passed to each spawned task. +pub(crate) struct SpawnContext { + pub store: TaskStore, + pub active: ActiveTaskMap, + pub event_tx: tokio::sync::broadcast::Sender, + pub max_retries: i32, + pub registry: Arc, + pub app_state: StateSnapshot, + pub work_notify: Arc, + pub scheduler: super::super::WeakScheduler, + #[allow(dead_code)] + pub cancel_hook_timeout: tokio::time::Duration, + /// Per-module live running counts. Incremented on dispatch; decremented on terminal. + pub module_running: Arc>, + /// Pre-snapshotted per-module state (module name → snapshot). Cloned at dispatch time. + pub module_state: Arc>, + /// Registry of all registered modules — shared with spawned tasks so they can + /// construct [`ModuleHandle`](crate::module::ModuleHandle) instances. + pub module_registry: Arc, +} + +/// Output of task context construction — everything needed to insert into the +/// active map and spawn the executor. +pub(crate) struct PreparedTask { + pub ctx: TaskContext, + pub io: Arc, + pub token: CancellationToken, +} + +/// Build the [`TaskContext`], [`IoTracker`], and [`CancellationToken`] for a task. +pub(crate) fn build_task_context(task: &TaskRecord, spawn_ctx: &SpawnContext) -> PreparedTask { + let owning_module = task.module_name().unwrap_or_default().to_string(); + + // Clone the pre-snapshotted module state — no lock needed, already lock-free. + let module_state_snapshot: StateSnapshot = task + .module_name() + .and_then(|name| spawn_ctx.module_state.get(name).cloned()) + .unwrap_or_default(); + + let token = CancellationToken::new(); + + let child_spawner = ChildSpawner::new( + spawn_ctx.store.clone(), + task.id, + spawn_ctx.work_notify.clone(), + ParentContext { + created_at: task.created_at, + ttl_seconds: task.ttl_seconds, + ttl_from: task.ttl_from, + started_at: task.started_at, + tags: task.tags.clone(), + }, + ); + + let io = Arc::new(IoTracker::new()); + + let ctx = TaskContext { + record: task.clone(), + token: token.clone(), + progress: ProgressReporter::new( + task.event_header(), + spawn_ctx.event_tx.clone(), + spawn_ctx.active.clone(), + io.clone(), + ), + scheduler: spawn_ctx.scheduler.clone(), + app_state: spawn_ctx.app_state.clone(), + module_state: module_state_snapshot, + child_spawner: Some(child_spawner), + io: io.clone(), + module_registry: spawn_ctx.module_registry.clone(), + owning_module: owning_module.clone(), + }; + + PreparedTask { ctx, io, token } +} diff --git a/src/scheduler/spawn/failure.rs b/src/scheduler/spawn/failure.rs new file mode 100644 index 0000000..4a54b88 --- /dev/null +++ b/src/scheduler/spawn/failure.rs @@ -0,0 +1,199 @@ +//! Failure path: retry policy, dead-letter, fail-fast cascade, dependency propagation. + +use std::sync::Arc; + +use crate::store::TaskStore; +use crate::task::{IoBudget, TaskError, TaskRecord}; + +use super::super::dispatch::ActiveTaskMap; +use super::super::SchedulerEvent; +use super::parent::handle_parent_resolution; + +/// Shared dependencies for the failure handler. +pub(crate) struct FailureDeps { + pub store: TaskStore, + pub active: ActiveTaskMap, + pub event_tx: tokio::sync::broadcast::Sender, + pub work_notify: Arc, + pub max_retries: i32, + pub registry: Arc, +} + +/// Handle a failed task execution. +/// +/// Resolves retry policy, records the failure, propagates to dependents, and +/// handles fail-fast parent cascading. +pub(crate) async fn handle_failure( + task: &TaskRecord, + error: TaskError, + metrics: &IoBudget, + deps: &FailureDeps, + decrement_module: impl FnOnce(), +) { + let task_id = task.id; + + // Resolve effective retry policy for this task type. + let policy = deps.registry.type_retry_policy(&task.task_type); + let effective_max_retries = task + .max_retries + .unwrap_or(policy.map(|p| p.max_retries).unwrap_or(deps.max_retries)); + let backoff_strategy = policy.map(|p| &p.strategy); + + let will_retry = error.retryable && task.retry_count < effective_max_retries; + + // Compute retry delay for event reporting. + let retry_delay = if will_retry { + if let Some(ms) = error.retry_after_ms { + Some(std::time::Duration::from_millis(ms)) + } else if let Some(strategy) = backoff_strategy { + let d = strategy.delay_for(task.retry_count); + if d.is_zero() { + None + } else { + Some(d) + } + } else { + None + } + } else { + None + }; + + tracing::warn!( + task_id, + task_type = task.task_type, + error = %error.message, + retryable = error.retryable, + will_retry, + "task failed" + ); + + let fail_backoff = crate::store::FailBackoff { + strategy: backoff_strategy, + executor_retry_after_ms: error.retry_after_ms, + }; + if let Err(e) = deps + .store + .fail_with_record( + task, + &error.message, + error.retryable, + effective_max_retries, + metrics, + &fail_backoff, + ) + .await + { + tracing::error!(task_id, error = %e, "failed to record task failure"); + } + + // Remove from active tracking AFTER the store write completes. + decrement_module(); + deps.active.remove(task_id); + + let dead_lettered = error.retryable && !will_retry; + if dead_lettered { + let _ = deps.event_tx.send(SchedulerEvent::DeadLettered { + header: task.event_header(), + error: error.message.clone(), + retry_count: task.retry_count + 1, + }); + } else { + let _ = deps.event_tx.send(SchedulerEvent::Failed { + header: task.event_header(), + error: error.message.clone(), + will_retry, + retry_after: retry_delay, + }); + } + deps.work_notify.notify_one(); + + // If permanent failure, propagate to dependency chain. + if !will_retry { + propagate_failure(task, &error, deps).await; + } +} + +/// Propagate a permanent failure to dependents and handle fail-fast parent logic. +async fn propagate_failure(task: &TaskRecord, error: &TaskError, deps: &FailureDeps) { + let task_id = task.id; + + match deps.store.fail_dependents(task_id).await { + Ok((failed_ids, unblocked_ids)) => { + for fid in &failed_ids { + let _ = deps.event_tx.send(SchedulerEvent::DependencyFailed { + task_id: *fid, + failed_dependency: task_id, + }); + } + for uid in &unblocked_ids { + let _ = deps + .event_tx + .send(SchedulerEvent::TaskUnblocked { task_id: *uid }); + } + if !unblocked_ids.is_empty() { + deps.work_notify.notify_one(); + } + } + Err(e) => { + tracing::error!(task_id, error = %e, "failed to propagate failure to dependents"); + } + } + + if let Some(parent_id) = task.parent_id { + // Check if parent uses fail_fast. + if let Ok(Some(parent)) = deps.store.task_by_id(parent_id).await { + if parent.fail_fast { + // Cancel remaining siblings. + if let Ok(running_ids) = deps.store.cancel_children(parent_id).await { + for rid in &running_ids { + if let Some(at) = deps.active.remove(*rid) { + at.token.cancel(); + let _ = deps.store.delete(*rid).await; + let _ = deps + .event_tx + .send(SchedulerEvent::Cancelled(at.record.event_header())); + } + } + } + // Fail the parent. + let msg = format!("child task {task_id} failed: {}", error.message); + if let Err(e) = deps + .store + .fail_with_record( + &parent, + &msg, + false, + 0, + &IoBudget::default(), + &Default::default(), + ) + .await + { + tracing::error!( + parent_id, + error = %e, + "failed to record parent failure" + ); + } + let _ = deps.event_tx.send(SchedulerEvent::Failed { + header: parent.event_header(), + error: msg, + will_retry: false, + retry_after: None, + }); + } else { + // Not fail_fast — check if all children done. + handle_parent_resolution( + parent_id, + &deps.store, + &deps.active, + &deps.event_tx, + deps.max_retries, + &deps.work_notify, + ) + .await; + } + } + } +} diff --git a/src/scheduler/spawn/parent.rs b/src/scheduler/spawn/parent.rs new file mode 100644 index 0000000..ee3e941 --- /dev/null +++ b/src/scheduler/spawn/parent.rs @@ -0,0 +1,59 @@ +//! Parent-child resolution after task completion or failure. + +use std::sync::Arc; + +use crate::store::TaskStore; +use crate::task::{IoBudget, ParentResolution}; + +use super::super::dispatch::ActiveTaskMap; +use super::super::SchedulerEvent; + +/// Check if a waiting parent is ready for finalization or has failed, +/// and dispatch the finalize phase if ready. +pub(crate) async fn handle_parent_resolution( + parent_id: i64, + store: &TaskStore, + active: &ActiveTaskMap, + event_tx: &tokio::sync::broadcast::Sender, + _max_retries: i32, + work_notify: &Arc, +) { + match store.try_resolve_parent(parent_id).await { + Ok(Some(ParentResolution::ReadyToFinalize)) => { + // Enqueue parent for finalize dispatch. + active.pending_finalizers.lock().unwrap().insert(parent_id); + // Wake the scheduler to dispatch the finalize phase. + work_notify.notify_one(); + } + Ok(Some(ParentResolution::Failed(reason))) => { + // All children done but some failed — fail the parent. + if let Ok(Some(parent)) = store.task_by_id(parent_id).await { + if let Err(e) = store + .fail_with_record( + &parent, + &reason, + false, + 0, + &IoBudget::default(), + &Default::default(), + ) + .await + { + tracing::error!(parent_id, error = %e, "failed to record parent failure"); + } + let _ = event_tx.send(SchedulerEvent::Failed { + header: parent.event_header(), + error: reason, + will_retry: false, + retry_after: None, + }); + } + } + Ok(Some(ParentResolution::StillWaiting)) | Ok(None) => { + // Children still active or parent not found — nothing to do. + } + Err(e) => { + tracing::error!(parent_id, error = %e, "failed to resolve parent"); + } + } +} From 4c1590835bb4d81072f7547291185a01bab8df97 Mon Sep 17 00:00:00 2001 From: DJ Majumdar Date: Wed, 18 Mar 2026 21:56:36 -0700 Subject: [PATCH 2/3] refactor: decompose TaskStore into focused service modules Move dependency graph operations (resolve_dependents, fail_dependents) out of lifecycle/ into store/dependencies.rs since they operate on the dep graph, not task lifecycle states. Consolidate pop, complete, and fail into lifecycle/transitions.rs to make the state machine visible in one place. Add clear section headers to cancel_expire.rs. --- src/store/{lifecycle => }/dependencies.rs | 2 +- src/store/lifecycle/cancel_expire.rs | 12 + src/store/lifecycle/complete.rs | 210 --------- src/store/lifecycle/fail.rs | 181 -------- src/store/lifecycle/mod.rs | 7 +- src/store/lifecycle/pop.rs | 113 ----- src/store/lifecycle/transitions.rs | 512 ++++++++++++++++++++++ src/store/mod.rs | 1 + 8 files changed, 528 insertions(+), 510 deletions(-) rename src/store/{lifecycle => }/dependencies.rs (99%) delete mode 100644 src/store/lifecycle/complete.rs delete mode 100644 src/store/lifecycle/fail.rs delete mode 100644 src/store/lifecycle/pop.rs create mode 100644 src/store/lifecycle/transitions.rs diff --git a/src/store/lifecycle/dependencies.rs b/src/store/dependencies.rs similarity index 99% rename from src/store/lifecycle/dependencies.rs rename to src/store/dependencies.rs index edbff94..15f452f 100644 --- a/src/store/lifecycle/dependencies.rs +++ b/src/store/dependencies.rs @@ -4,7 +4,7 @@ use crate::store::row_mapping::row_to_task_record; use crate::store::{StoreError, TaskStore}; use crate::task::{DependencyFailurePolicy, IoBudget}; -use super::{insert_history, HistoryStatus}; +use super::lifecycle::{insert_history, HistoryStatus}; impl TaskStore { /// After a task completes, check if any blocked tasks are now unblocked. diff --git a/src/store/lifecycle/cancel_expire.rs b/src/store/lifecycle/cancel_expire.rs index 7ee487d..2b55970 100644 --- a/src/store/lifecycle/cancel_expire.rs +++ b/src/store/lifecycle/cancel_expire.rs @@ -1,4 +1,8 @@ //! Pause, resume, cancellation, and TTL expiry. +//! +//! Both cancellation (user-initiated) and expiry (time-driven) share the same +//! terminal transition pattern: record in history → clean up edges/tags → delete. +//! They are kept together because of this shared structure. use crate::store::row_mapping::row_to_task_record; use crate::store::{StoreError, TaskStore}; @@ -6,6 +10,8 @@ use crate::task::{IoBudget, TaskRecord}; use super::{compute_duration_ms, insert_history, HistoryStatus}; +// ── Pause / Resume ────────────────────────────────────────────────── + impl TaskStore { /// Pause a running task (for preemption). Sets status to paused. pub async fn pause(&self, id: i64) -> Result<(), StoreError> { @@ -25,6 +31,8 @@ impl TaskStore { Ok(()) } + // ── Cancellation (user-initiated) ────────────────────────────── + /// Move a task to history as cancelled and delete it from the active queue. /// Also cleans up dependency edges and cascades failure to dependents. /// @@ -115,6 +123,8 @@ impl TaskStore { Ok(()) } + // ── Bulk pause / resume by type prefix ───────────────────────── + /// Pause all pending tasks whose `task_type` starts with `prefix`. /// /// Updates their status from `pending` to `paused` in a single SQL statement. @@ -145,6 +155,8 @@ impl TaskStore { Ok(result.rows_affected()) } + // ── Expiry (time-driven) ─────────────────────────────────────── + /// Sweep for expired tasks and move them to history. /// /// Finds tasks whose `expires_at` has passed and that are still pending diff --git a/src/store/lifecycle/complete.rs b/src/store/lifecycle/complete.rs deleted file mode 100644 index 8508647..0000000 --- a/src/store/lifecycle/complete.rs +++ /dev/null @@ -1,210 +0,0 @@ -//! Task completion: move to history, handle recurring re-creation, requeue. - -use crate::store::row_mapping::row_to_task_record; -use crate::store::{StoreError, TaskStore}; -use crate::task::IoBudget; - -use super::{compute_duration_ms, insert_history, HistoryStatus}; - -impl TaskStore { - /// Mark a task as completed and move it to history. - pub async fn complete(&self, id: i64, metrics: &IoBudget) -> Result<(), StoreError> { - tracing::debug!(task_id = id, "store.complete: BEGIN tx"); - let mut conn = self.begin_write().await?; - - // Fetch the task to move. - let row = sqlx::query("SELECT * FROM tasks WHERE id = ?") - .bind(id) - .fetch_optional(&mut *conn) - .await?; - - let Some(row) = row else { return Ok(()) }; - let task = row_to_task_record(&row); - - let _recurring = Self::complete_inner(&mut conn, &task, metrics).await?; - - sqlx::query("COMMIT").execute(&mut *conn).await?; - drop(conn); - tracing::debug!(task_id = id, "store.complete: COMMIT ok"); - - self.maybe_prune().await; - - Ok(()) - } - - /// Mark a task as completed using an in-memory record, avoiding the - /// redundant `SELECT *` round-trip. The `requeue` flag is still checked - /// from the database row since it may have been set by a concurrent - /// `submit()` while the task was running. - /// - /// Returns `Some((next_run, execution_count))` if a recurring next - /// instance was created, `None` otherwise. - pub async fn complete_with_record( - &self, - task: &crate::task::TaskRecord, - metrics: &IoBudget, - ) -> Result, i64)>, StoreError> { - tracing::debug!(task_id = task.id, "store.complete_with_record: BEGIN tx"); - let mut conn = self.begin_write().await?; - - let recurring_info = Self::complete_inner(&mut conn, task, metrics).await?; - - sqlx::query("COMMIT").execute(&mut *conn).await?; - drop(conn); - tracing::debug!(task_id = task.id, "store.complete_with_record: COMMIT ok"); - - self.maybe_prune().await; - - Ok(recurring_info) - } - - /// Shared completion logic: insert history, handle recurring next instance, - /// then handle requeue or delete. - /// - /// Returns `Some((next_run, exec_count))` if a recurring next instance was - /// created, `None` otherwise. - async fn complete_inner( - conn: &mut sqlx::pool::PoolConnection, - task: &crate::task::TaskRecord, - metrics: &IoBudget, - ) -> Result, i64)>, StoreError> { - let duration_ms = compute_duration_ms(task); - - // Insert into history. - insert_history( - conn, - task, - HistoryStatus::Completed, - metrics, - duration_ms, - task.last_error.as_deref(), - ) - .await?; - - // Read tags into memory before potential deletion (needed for recurring re-creation). - let saved_tags: Vec<(String, String)> = if task.recurring_interval_secs.is_some() { - sqlx::query_as("SELECT key, value FROM task_tags WHERE task_id = ?") - .bind(task.id) - .fetch_all(&mut **conn) - .await? - } else { - Vec::new() - }; - - // Try to delete (normal completion, requeue = 0). - let del = sqlx::query("DELETE FROM tasks WHERE id = ? AND requeue = 0") - .bind(task.id) - .execute(&mut **conn) - .await?; - - if del.rows_affected() == 0 { - // Requeue flag was set by a concurrent submit — reset to pending. - // No-op if the task was already deleted (cancelled). - sqlx::query( - "UPDATE tasks SET status = 'pending', - priority = COALESCE(requeue_priority, priority), - started_at = NULL, retry_count = 0, last_error = NULL, - requeue = 0, requeue_priority = NULL - WHERE id = ?", - ) - .bind(task.id) - .execute(&mut **conn) - .await?; - // Don't create recurring next instance if requeued. - return Ok(None); - } - - // Task was deleted — clean up orphaned tags. - crate::store::delete_task_tags(conn, task.id).await?; - - // Handle recurring tasks: create the next instance after deleting - // the completed one (to avoid UNIQUE constraint on key). - let mut recurring_info = None; - if let Some(interval) = task.recurring_interval_secs { - if !task.recurring_paused { - let execution_count = task.recurring_execution_count + 1; - let should_create = task - .recurring_max_executions - .map_or(true, |max| execution_count < max); - - if should_create { - // Pile-up prevention: check if a pending instance already exists - // (e.g. from a concurrent submit with the same key). - let existing: Option<(i64,)> = - sqlx::query_as("SELECT id FROM tasks WHERE key = ? AND status = 'pending'") - .bind(&task.key) - .fetch_optional(&mut **conn) - .await?; - - if existing.is_none() { - let next_run = chrono::Utc::now() + chrono::Duration::seconds(interval); - let next_run_str = next_run.format("%Y-%m-%d %H:%M:%S").to_string(); - let fail_fast_val: i32 = if task.fail_fast { 1 } else { 0 }; - - // Compute TTL columns for the next instance. - let expires_at_str: Option = match (task.ttl_seconds, task.ttl_from) - { - (Some(ttl_secs), crate::task::TtlFrom::Submission) => { - let exp = chrono::Utc::now() + chrono::Duration::seconds(ttl_secs); - Some(exp.format("%Y-%m-%d %H:%M:%S").to_string()) - } - _ => None, - }; - - let recurring_result = sqlx::query( - "INSERT INTO tasks (task_type, key, label, priority, status, payload, - expected_read_bytes, expected_write_bytes, - expected_net_rx_bytes, expected_net_tx_bytes, - parent_id, fail_fast, group_key, - ttl_seconds, ttl_from, expires_at, - run_after, recurring_interval_secs, - recurring_max_executions, recurring_execution_count, - recurring_paused, max_retries) - VALUES (?, ?, ?, ?, 'pending', ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?)", - ) - .bind(&task.task_type) - .bind(&task.key) - .bind(&task.label) - .bind(task.priority.value() as i32) - .bind(&task.payload) - .bind(task.expected_io.disk_read) - .bind(task.expected_io.disk_write) - .bind(task.expected_io.net_rx) - .bind(task.expected_io.net_tx) - .bind(task.parent_id) - .bind(fail_fast_val) - .bind(&task.group_key) - .bind(task.ttl_seconds) - .bind(task.ttl_from.as_str()) - .bind(&expires_at_str) - .bind(&next_run_str) - .bind(task.recurring_interval_secs) - .bind(task.recurring_max_executions) - .bind(execution_count) - .bind(task.max_retries) - .execute(&mut **conn) - .await?; - - // Copy tags to the new recurring instance. - let next_id = recurring_result.last_insert_rowid(); - for (key, value) in &saved_tags { - sqlx::query( - "INSERT INTO task_tags (task_id, key, value) VALUES (?, ?, ?)", - ) - .bind(next_id) - .bind(key) - .bind(value) - .execute(&mut **conn) - .await?; - } - - recurring_info = Some((next_run, execution_count)); - } - // If existing.is_some(), skip (pile-up prevention). - } - } - } - - Ok(recurring_info) - } -} diff --git a/src/store/lifecycle/fail.rs b/src/store/lifecycle/fail.rs deleted file mode 100644 index 172ab80..0000000 --- a/src/store/lifecycle/fail.rs +++ /dev/null @@ -1,181 +0,0 @@ -//! Task failure: retry with backoff or move to history. - -use crate::store::row_mapping::row_to_task_record; -use crate::store::{StoreError, TaskStore}; -use crate::task::{BackoffStrategy, IoBudget}; - -use super::{compute_duration_ms, insert_history, HistoryStatus}; - -/// Backoff parameters for retry delay computation. -/// -/// Bundles the optional backoff strategy and executor-signaled override into a -/// single argument to keep `fail()` / `fail_with_record()` under the clippy -/// argument-count lint. -#[derive(Debug, Default, Clone)] -pub struct FailBackoff<'a> { - /// Per-type backoff strategy. `None` means immediate retry. - pub strategy: Option<&'a BackoffStrategy>, - /// Executor-requested retry delay in milliseconds. Overrides the strategy - /// when set. - pub executor_retry_after_ms: Option, -} - -impl TaskStore { - /// Mark a task as failed. If `retryable` and under max retries, requeue - /// it as pending with the same priority. Otherwise move to history as failed. - /// - /// `backoff` controls the delay before the next retry attempt. See - /// `fail_inner` for details. - pub async fn fail( - &self, - id: i64, - error: &str, - retryable: bool, - max_retries: i32, - metrics: &IoBudget, - backoff: &FailBackoff<'_>, - ) -> Result<(), StoreError> { - tracing::debug!(task_id = id, "store.fail: BEGIN tx"); - let mut conn = self.begin_write().await?; - tracing::debug!(task_id = id, "store.fail: BEGIN acquired"); - - let row = sqlx::query("SELECT * FROM tasks WHERE id = ?") - .bind(id) - .fetch_optional(&mut *conn) - .await?; - - let Some(row) = row else { return Ok(()) }; - let task = row_to_task_record(&row); - - Self::fail_inner( - &mut conn, - &task, - error, - retryable, - max_retries, - metrics, - backoff, - ) - .await?; - - sqlx::query("COMMIT").execute(&mut *conn).await?; - drop(conn); - tracing::debug!(task_id = id, "store.fail: COMMIT ok"); - - self.maybe_prune().await; - - Ok(()) - } - - /// Mark a task as failed using an in-memory record, avoiding the - /// redundant `SELECT *` round-trip. - pub async fn fail_with_record( - &self, - task: &crate::task::TaskRecord, - error: &str, - retryable: bool, - max_retries: i32, - metrics: &IoBudget, - backoff: &FailBackoff<'_>, - ) -> Result<(), StoreError> { - tracing::debug!(task_id = task.id, "store.fail_with_record: BEGIN tx"); - let mut conn = self.begin_write().await?; - tracing::debug!(task_id = task.id, "store.fail_with_record: BEGIN acquired"); - - Self::fail_inner( - &mut conn, - task, - error, - retryable, - max_retries, - metrics, - backoff, - ) - .await?; - - sqlx::query("COMMIT").execute(&mut *conn).await?; - drop(conn); - tracing::debug!(task_id = task.id, "store.fail_with_record: COMMIT ok"); - - self.maybe_prune().await; - - Ok(()) - } - - /// Shared failure logic: retry or move to history. - /// - /// When retrying, computes the backoff delay from (in priority order): - /// 1. `executor_retry_after_ms` — executor-signaled override - /// 2. `backoff` strategy — per-type backoff computation - /// 3. Immediate retry (no delay) — backward-compatible default - /// - /// The delay is applied by setting `run_after` on the requeued task. - async fn fail_inner( - conn: &mut sqlx::pool::PoolConnection, - task: &crate::task::TaskRecord, - error: &str, - retryable: bool, - max_retries: i32, - metrics: &IoBudget, - backoff: &FailBackoff<'_>, - ) -> Result<(), StoreError> { - if retryable && task.retry_count < max_retries { - // Compute delay: executor override > backoff strategy > immediate. - let delay = if let Some(ms) = backoff.executor_retry_after_ms { - std::time::Duration::from_millis(ms) - } else if let Some(strategy) = backoff.strategy { - strategy.delay_for(task.retry_count) - } else { - std::time::Duration::ZERO - }; - - if delay.is_zero() { - // Immediate retry — current behavior. - sqlx::query( - "UPDATE tasks SET status = 'pending', started_at = NULL, - retry_count = retry_count + 1, last_error = ? - WHERE id = ?", - ) - .bind(error) - .bind(task.id) - .execute(&mut **conn) - .await?; - } else { - // Delayed retry — set run_after. - let run_after = - chrono::Utc::now() + chrono::Duration::milliseconds(delay.as_millis() as i64); - let run_after_str = run_after.format("%Y-%m-%d %H:%M:%S%.3f").to_string(); - sqlx::query( - "UPDATE tasks SET status = 'pending', started_at = NULL, - retry_count = retry_count + 1, last_error = ?, - run_after = ? - WHERE id = ?", - ) - .bind(error) - .bind(&run_after_str) - .bind(task.id) - .execute(&mut **conn) - .await?; - } - } else { - // Terminal failure — move to history. - // Distinguish: retryable + exhausted → dead_letter; non-retryable → failed. - let status = if retryable { - HistoryStatus::DeadLetter - } else { - HistoryStatus::Failed - }; - let duration_ms = compute_duration_ms(task); - - insert_history(conn, task, status, metrics, duration_ms, Some(error)).await?; - - crate::store::delete_task_tags(conn, task.id).await?; - sqlx::query("DELETE FROM tasks WHERE id = ?") - .bind(task.id) - .execute(&mut **conn) - .await?; - } - - Ok(()) - } -} diff --git a/src/store/lifecycle/mod.rs b/src/store/lifecycle/mod.rs index aae4244..a9294d9 100644 --- a/src/store/lifecycle/mod.rs +++ b/src/store/lifecycle/mod.rs @@ -2,15 +2,12 @@ //! dependency resolution. mod cancel_expire; -mod complete; -mod dependencies; -mod fail; -mod pop; +mod transitions; #[cfg(test)] mod tests; -pub use fail::FailBackoff; +pub use transitions::FailBackoff; use crate::task::{IoBudget, TaskRecord}; diff --git a/src/store/lifecycle/pop.rs b/src/store/lifecycle/pop.rs deleted file mode 100644 index 7f8c060..0000000 --- a/src/store/lifecycle/pop.rs +++ /dev/null @@ -1,113 +0,0 @@ -//! Pop / peek / requeue operations for task dispatch. - -use crate::store::row_mapping::row_to_task_record; -use crate::store::{StoreError, TaskStore}; -use crate::task::TaskRecord; - -impl TaskStore { - /// Peek at the highest-priority pending task without modifying it. - /// Returns `None` if the queue is empty. Tasks with a future `run_after` - /// timestamp are excluded (not yet eligible for dispatch). - pub async fn peek_next(&self) -> Result, StoreError> { - let row = sqlx::query( - "SELECT * FROM tasks - WHERE id = ( - SELECT id FROM tasks - WHERE status = 'pending' - AND (run_after IS NULL OR run_after <= strftime('%Y-%m-%d %H:%M:%f', 'now')) - ORDER BY priority ASC, id ASC - LIMIT 1 - )", - ) - .fetch_optional(&self.pool) - .await?; - - let mut record = row.as_ref().map(row_to_task_record); - if let Some(ref mut r) = record { - self.populate_tags(std::slice::from_mut(r)).await?; - } - Ok(record) - } - - /// Atomically claim a specific pending task by id, setting it to running. - /// Returns `None` if the task is no longer pending (e.g. claimed by another - /// dispatcher or cancelled). - /// - /// For tasks with `ttl_from = 'first_attempt'`, sets `expires_at` on the - /// first pop (when `expires_at IS NULL` and `ttl_seconds IS NOT NULL`). - pub async fn pop_by_id(&self, id: i64) -> Result, StoreError> { - tracing::debug!(task_id = id, "store.pop_by_id: 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: UPDATE end"); - - let mut record = row.as_ref().map(row_to_task_record); - if let Some(ref mut r) = record { - self.populate_tags(std::slice::from_mut(r)).await?; - } - Ok(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. - /// - /// For tasks with `ttl_from = 'first_attempt'`, sets `expires_at` on - /// the first pop. - pub async fn pop_next(&self) -> Result, StoreError> { - 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 = ( - SELECT id FROM tasks - WHERE status = 'pending' - AND (run_after IS NULL OR run_after <= strftime('%Y-%m-%d %H:%M:%f', 'now')) - ORDER BY priority ASC, id ASC - LIMIT 1 - ) - RETURNING *", - ) - .fetch_optional(&self.pool) - .await?; - - let mut record = row.map(|r| row_to_task_record(&r)); - if let Some(ref mut r) = record { - self.populate_tags(std::slice::from_mut(r)).await?; - } - Ok(record) - } - - /// Atomically requeue a running task back to pending. - /// - /// Used when a task is popped but then rejected by backpressure or IO - /// budget checks. Unlike pause+resume, this is a single atomic operation - /// that never puts the task in an intermediate state visible to queries. - pub async fn requeue(&self, id: i64) -> Result<(), StoreError> { - sqlx::query( - "UPDATE tasks SET status = 'pending', started_at = NULL WHERE id = ? AND status = 'running'", - ) - .bind(id) - .execute(&self.pool) - .await?; - Ok(()) - } -} diff --git a/src/store/lifecycle/transitions.rs b/src/store/lifecycle/transitions.rs new file mode 100644 index 0000000..5077fce --- /dev/null +++ b/src/store/lifecycle/transitions.rs @@ -0,0 +1,512 @@ +//! Core state machine transitions: pop, complete, fail. +//! +//! Grouping these in one file makes the state machine visible at a glance: +//! +//! ```text +//! pending → running (pop_next / pop_by_id) +//! running → pending (requeue — backpressure rejection) +//! running → completed (complete — moved to history) +//! running → failed (fail, non-retryable — moved to history) +//! running → dead_letter (fail, retries exhausted — moved to history) +//! running → pending (fail, retryable — requeued with backoff) +//! ``` +//! +//! The `waiting` transition lives in `hierarchy.rs`, and pause/resume/cancel/expire +//! live in `cancel_expire.rs`. + +use crate::store::row_mapping::row_to_task_record; +use crate::store::{StoreError, TaskStore}; +use crate::task::{BackoffStrategy, IoBudget, TaskRecord}; + +use super::{compute_duration_ms, insert_history, HistoryStatus}; + +// ── Pop / Peek / Requeue ──────────────────────────────────────────── + +impl TaskStore { + /// Peek at the highest-priority pending task without modifying it. + /// Returns `None` if the queue is empty. Tasks with a future `run_after` + /// timestamp are excluded (not yet eligible for dispatch). + pub async fn peek_next(&self) -> Result, StoreError> { + let row = sqlx::query( + "SELECT * FROM tasks + WHERE id = ( + SELECT id FROM tasks + WHERE status = 'pending' + AND (run_after IS NULL OR run_after <= strftime('%Y-%m-%d %H:%M:%f', 'now')) + ORDER BY priority ASC, id ASC + LIMIT 1 + )", + ) + .fetch_optional(&self.pool) + .await?; + + let mut record = row.as_ref().map(row_to_task_record); + if let Some(ref mut r) = record { + self.populate_tags(std::slice::from_mut(r)).await?; + } + Ok(record) + } + + /// Atomically claim a specific pending task by id, setting it to running. + /// Returns `None` if the task is no longer pending (e.g. claimed by another + /// dispatcher or cancelled). + /// + /// For tasks with `ttl_from = 'first_attempt'`, sets `expires_at` on the + /// first pop (when `expires_at IS NULL` and `ttl_seconds IS NOT NULL`). + pub async fn pop_by_id(&self, id: i64) -> Result, StoreError> { + tracing::debug!(task_id = id, "store.pop_by_id: 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: UPDATE end"); + + let mut record = row.as_ref().map(row_to_task_record); + if let Some(ref mut r) = record { + self.populate_tags(std::slice::from_mut(r)).await?; + } + Ok(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. + /// + /// For tasks with `ttl_from = 'first_attempt'`, sets `expires_at` on + /// the first pop. + pub async fn pop_next(&self) -> Result, StoreError> { + 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 = ( + SELECT id FROM tasks + WHERE status = 'pending' + AND (run_after IS NULL OR run_after <= strftime('%Y-%m-%d %H:%M:%f', 'now')) + ORDER BY priority ASC, id ASC + LIMIT 1 + ) + RETURNING *", + ) + .fetch_optional(&self.pool) + .await?; + + let mut record = row.map(|r| row_to_task_record(&r)); + if let Some(ref mut r) = record { + self.populate_tags(std::slice::from_mut(r)).await?; + } + Ok(record) + } + + /// Atomically requeue a running task back to pending. + /// + /// Used when a task is popped but then rejected by backpressure or IO + /// budget checks. Unlike pause+resume, this is a single atomic operation + /// that never puts the task in an intermediate state visible to queries. + pub async fn requeue(&self, id: i64) -> Result<(), StoreError> { + sqlx::query( + "UPDATE tasks SET status = 'pending', started_at = NULL WHERE id = ? AND status = 'running'", + ) + .bind(id) + .execute(&self.pool) + .await?; + Ok(()) + } +} + +// ── Complete ──────────────────────────────────────────────────────── + +impl TaskStore { + /// Mark a task as completed and move it to history. + pub async fn complete(&self, id: i64, metrics: &IoBudget) -> Result<(), StoreError> { + tracing::debug!(task_id = id, "store.complete: BEGIN tx"); + let mut conn = self.begin_write().await?; + + // Fetch the task to move. + let row = sqlx::query("SELECT * FROM tasks WHERE id = ?") + .bind(id) + .fetch_optional(&mut *conn) + .await?; + + let Some(row) = row else { return Ok(()) }; + let task = row_to_task_record(&row); + + let _recurring = Self::complete_inner(&mut conn, &task, metrics).await?; + + sqlx::query("COMMIT").execute(&mut *conn).await?; + drop(conn); + tracing::debug!(task_id = id, "store.complete: COMMIT ok"); + + self.maybe_prune().await; + + Ok(()) + } + + /// Mark a task as completed using an in-memory record, avoiding the + /// redundant `SELECT *` round-trip. The `requeue` flag is still checked + /// from the database row since it may have been set by a concurrent + /// `submit()` while the task was running. + /// + /// Returns `Some((next_run, execution_count))` if a recurring next + /// instance was created, `None` otherwise. + pub async fn complete_with_record( + &self, + task: &crate::task::TaskRecord, + metrics: &IoBudget, + ) -> Result, i64)>, StoreError> { + tracing::debug!(task_id = task.id, "store.complete_with_record: BEGIN tx"); + let mut conn = self.begin_write().await?; + + let recurring_info = Self::complete_inner(&mut conn, task, metrics).await?; + + sqlx::query("COMMIT").execute(&mut *conn).await?; + drop(conn); + tracing::debug!(task_id = task.id, "store.complete_with_record: COMMIT ok"); + + self.maybe_prune().await; + + Ok(recurring_info) + } + + /// Shared completion logic: insert history, handle recurring next instance, + /// then handle requeue or delete. + /// + /// Returns `Some((next_run, exec_count))` if a recurring next instance was + /// created, `None` otherwise. + async fn complete_inner( + conn: &mut sqlx::pool::PoolConnection, + task: &crate::task::TaskRecord, + metrics: &IoBudget, + ) -> Result, i64)>, StoreError> { + let duration_ms = compute_duration_ms(task); + + // Insert into history. + insert_history( + conn, + task, + HistoryStatus::Completed, + metrics, + duration_ms, + task.last_error.as_deref(), + ) + .await?; + + // Read tags into memory before potential deletion (needed for recurring re-creation). + let saved_tags: Vec<(String, String)> = if task.recurring_interval_secs.is_some() { + sqlx::query_as("SELECT key, value FROM task_tags WHERE task_id = ?") + .bind(task.id) + .fetch_all(&mut **conn) + .await? + } else { + Vec::new() + }; + + // Try to delete (normal completion, requeue = 0). + let del = sqlx::query("DELETE FROM tasks WHERE id = ? AND requeue = 0") + .bind(task.id) + .execute(&mut **conn) + .await?; + + if del.rows_affected() == 0 { + // Requeue flag was set by a concurrent submit — reset to pending. + // No-op if the task was already deleted (cancelled). + sqlx::query( + "UPDATE tasks SET status = 'pending', + priority = COALESCE(requeue_priority, priority), + started_at = NULL, retry_count = 0, last_error = NULL, + requeue = 0, requeue_priority = NULL + WHERE id = ?", + ) + .bind(task.id) + .execute(&mut **conn) + .await?; + // Don't create recurring next instance if requeued. + return Ok(None); + } + + // Task was deleted — clean up orphaned tags. + crate::store::delete_task_tags(conn, task.id).await?; + + // Handle recurring tasks: create the next instance after deleting + // the completed one (to avoid UNIQUE constraint on key). + let mut recurring_info = None; + if let Some(interval) = task.recurring_interval_secs { + if !task.recurring_paused { + let execution_count = task.recurring_execution_count + 1; + let should_create = task + .recurring_max_executions + .map_or(true, |max| execution_count < max); + + if should_create { + // Pile-up prevention: check if a pending instance already exists + // (e.g. from a concurrent submit with the same key). + let existing: Option<(i64,)> = + sqlx::query_as("SELECT id FROM tasks WHERE key = ? AND status = 'pending'") + .bind(&task.key) + .fetch_optional(&mut **conn) + .await?; + + if existing.is_none() { + let next_run = chrono::Utc::now() + chrono::Duration::seconds(interval); + let next_run_str = next_run.format("%Y-%m-%d %H:%M:%S").to_string(); + let fail_fast_val: i32 = if task.fail_fast { 1 } else { 0 }; + + // Compute TTL columns for the next instance. + let expires_at_str: Option = match (task.ttl_seconds, task.ttl_from) + { + (Some(ttl_secs), crate::task::TtlFrom::Submission) => { + let exp = chrono::Utc::now() + chrono::Duration::seconds(ttl_secs); + Some(exp.format("%Y-%m-%d %H:%M:%S").to_string()) + } + _ => None, + }; + + let recurring_result = sqlx::query( + "INSERT INTO tasks (task_type, key, label, priority, status, payload, + expected_read_bytes, expected_write_bytes, + expected_net_rx_bytes, expected_net_tx_bytes, + parent_id, fail_fast, group_key, + ttl_seconds, ttl_from, expires_at, + run_after, recurring_interval_secs, + recurring_max_executions, recurring_execution_count, + recurring_paused, max_retries) + VALUES (?, ?, ?, ?, 'pending', ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?)", + ) + .bind(&task.task_type) + .bind(&task.key) + .bind(&task.label) + .bind(task.priority.value() as i32) + .bind(&task.payload) + .bind(task.expected_io.disk_read) + .bind(task.expected_io.disk_write) + .bind(task.expected_io.net_rx) + .bind(task.expected_io.net_tx) + .bind(task.parent_id) + .bind(fail_fast_val) + .bind(&task.group_key) + .bind(task.ttl_seconds) + .bind(task.ttl_from.as_str()) + .bind(&expires_at_str) + .bind(&next_run_str) + .bind(task.recurring_interval_secs) + .bind(task.recurring_max_executions) + .bind(execution_count) + .bind(task.max_retries) + .execute(&mut **conn) + .await?; + + // Copy tags to the new recurring instance. + let next_id = recurring_result.last_insert_rowid(); + for (key, value) in &saved_tags { + sqlx::query( + "INSERT INTO task_tags (task_id, key, value) VALUES (?, ?, ?)", + ) + .bind(next_id) + .bind(key) + .bind(value) + .execute(&mut **conn) + .await?; + } + + recurring_info = Some((next_run, execution_count)); + } + // If existing.is_some(), skip (pile-up prevention). + } + } + } + + Ok(recurring_info) + } +} + +// ── Fail ──────────────────────────────────────────────────────────── + +/// Backoff parameters for retry delay computation. +/// +/// Bundles the optional backoff strategy and executor-signaled override into a +/// single argument to keep `fail()` / `fail_with_record()` under the clippy +/// argument-count lint. +#[derive(Debug, Default, Clone)] +pub struct FailBackoff<'a> { + /// Per-type backoff strategy. `None` means immediate retry. + pub strategy: Option<&'a BackoffStrategy>, + /// Executor-requested retry delay in milliseconds. Overrides the strategy + /// when set. + pub executor_retry_after_ms: Option, +} + +impl TaskStore { + /// Mark a task as failed. If `retryable` and under max retries, requeue + /// it as pending with the same priority. Otherwise move to history as failed. + /// + /// `backoff` controls the delay before the next retry attempt. See + /// `fail_inner` for details. + pub async fn fail( + &self, + id: i64, + error: &str, + retryable: bool, + max_retries: i32, + metrics: &IoBudget, + backoff: &FailBackoff<'_>, + ) -> Result<(), StoreError> { + tracing::debug!(task_id = id, "store.fail: BEGIN tx"); + let mut conn = self.begin_write().await?; + tracing::debug!(task_id = id, "store.fail: BEGIN acquired"); + + let row = sqlx::query("SELECT * FROM tasks WHERE id = ?") + .bind(id) + .fetch_optional(&mut *conn) + .await?; + + let Some(row) = row else { return Ok(()) }; + let task = row_to_task_record(&row); + + Self::fail_inner( + &mut conn, + &task, + error, + retryable, + max_retries, + metrics, + backoff, + ) + .await?; + + sqlx::query("COMMIT").execute(&mut *conn).await?; + drop(conn); + tracing::debug!(task_id = id, "store.fail: COMMIT ok"); + + self.maybe_prune().await; + + Ok(()) + } + + /// Mark a task as failed using an in-memory record, avoiding the + /// redundant `SELECT *` round-trip. + pub async fn fail_with_record( + &self, + task: &crate::task::TaskRecord, + error: &str, + retryable: bool, + max_retries: i32, + metrics: &IoBudget, + backoff: &FailBackoff<'_>, + ) -> Result<(), StoreError> { + tracing::debug!(task_id = task.id, "store.fail_with_record: BEGIN tx"); + let mut conn = self.begin_write().await?; + tracing::debug!(task_id = task.id, "store.fail_with_record: BEGIN acquired"); + + Self::fail_inner( + &mut conn, + task, + error, + retryable, + max_retries, + metrics, + backoff, + ) + .await?; + + sqlx::query("COMMIT").execute(&mut *conn).await?; + drop(conn); + tracing::debug!(task_id = task.id, "store.fail_with_record: COMMIT ok"); + + self.maybe_prune().await; + + Ok(()) + } + + /// Shared failure logic: retry or move to history. + /// + /// When retrying, computes the backoff delay from (in priority order): + /// 1. `executor_retry_after_ms` — executor-signaled override + /// 2. `backoff` strategy — per-type backoff computation + /// 3. Immediate retry (no delay) — backward-compatible default + /// + /// The delay is applied by setting `run_after` on the requeued task. + async fn fail_inner( + conn: &mut sqlx::pool::PoolConnection, + task: &crate::task::TaskRecord, + error: &str, + retryable: bool, + max_retries: i32, + metrics: &IoBudget, + backoff: &FailBackoff<'_>, + ) -> Result<(), StoreError> { + if retryable && task.retry_count < max_retries { + // Compute delay: executor override > backoff strategy > immediate. + let delay = if let Some(ms) = backoff.executor_retry_after_ms { + std::time::Duration::from_millis(ms) + } else if let Some(strategy) = backoff.strategy { + strategy.delay_for(task.retry_count) + } else { + std::time::Duration::ZERO + }; + + if delay.is_zero() { + // Immediate retry — current behavior. + sqlx::query( + "UPDATE tasks SET status = 'pending', started_at = NULL, + retry_count = retry_count + 1, last_error = ? + WHERE id = ?", + ) + .bind(error) + .bind(task.id) + .execute(&mut **conn) + .await?; + } else { + // Delayed retry — set run_after. + let run_after = + chrono::Utc::now() + chrono::Duration::milliseconds(delay.as_millis() as i64); + let run_after_str = run_after.format("%Y-%m-%d %H:%M:%S%.3f").to_string(); + sqlx::query( + "UPDATE tasks SET status = 'pending', started_at = NULL, + retry_count = retry_count + 1, last_error = ?, + run_after = ? + WHERE id = ?", + ) + .bind(error) + .bind(&run_after_str) + .bind(task.id) + .execute(&mut **conn) + .await?; + } + } else { + // Terminal failure — move to history. + // Distinguish: retryable + exhausted → dead_letter; non-retryable → failed. + let status = if retryable { + HistoryStatus::DeadLetter + } else { + HistoryStatus::Failed + }; + let duration_ms = compute_duration_ms(task); + + insert_history(conn, task, status, metrics, duration_ms, Some(error)).await?; + + crate::store::delete_task_tags(conn, task.id).await?; + sqlx::query("DELETE FROM tasks WHERE id = ?") + .bind(task.id) + .execute(&mut **conn) + .await?; + } + + Ok(()) + } +} diff --git a/src/store/mod.rs b/src/store/mod.rs index cb5f1f8..533e75e 100644 --- a/src/store/mod.rs +++ b/src/store/mod.rs @@ -22,6 +22,7 @@ //! via [`Scheduler::store()`](crate::Scheduler::store) for queries and //! diagnostics. +mod dependencies; mod hierarchy; mod lifecycle; mod query; From 350cc7f09fd2ba7263aed5e9804ae347fc237bee Mon Sep 17 00:00:00 2001 From: DJ Majumdar Date: Wed, 18 Mar 2026 22:03:44 -0700 Subject: [PATCH 3/3] refactor: decompose SubmitBuilder::resolve into focused precedence methods MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split the monolithic resolve() into four named methods that each own one layer of the 5-layer precedence chain: - apply_prefix: module name prefixing - apply_defaults: layers 2–4 (TypedTask / submission / module defaults) - apply_module_scalar_defaults: shared priority/group/TTL logic for both typed and untyped paths (eliminates duplicated conditional blocks) - apply_overrides: layer 1 per-call builder overrides No external API or behavioral change — all existing tests pass unchanged. --- src/task/submit_builder.rs | 174 ++++++++++++++++++++----------------- 1 file changed, 93 insertions(+), 81 deletions(-) diff --git a/src/task/submit_builder.rs b/src/task/submit_builder.rs index 55b58be..6d9e910 100644 --- a/src/task/submit_builder.rs +++ b/src/task/submit_builder.rs @@ -220,111 +220,123 @@ impl SubmitBuilder { /// Apply all default layers and per-call overrides, returning the /// scheduler and the fully resolved [`TaskSubmission`]. /// - /// Two modes determined by whether [`with_typed_defaults`](Self::with_typed_defaults) - /// was called: + /// Delegates to three focused methods that each handle one layer of the + /// precedence chain: /// - /// **`submit_typed()` path** (typed_defaults present): - /// 1. TypedTask values are set as the base (layer 4). - /// 2. Module defaults override them where the module has an explicit - /// setting (layer 3). - /// 3. SubmitBuilder per-call overrides trump everything (layer 1). - /// - /// **`submit()` path** (typed_defaults absent): - /// 1. Module defaults fill in only where the submission is at its - /// zero/`None` value (layer 3 fills in layer 2 gaps). - /// 2. SubmitBuilder per-call overrides trump everything (layer 1). + /// 1. [`apply_prefix`](Self::apply_prefix) — prefix `task_type` with the + /// module name. + /// 2. [`apply_defaults`](Self::apply_defaults) — layers 2–4 (TypedTask + /// baseline, module defaults, or submission-explicit values). + /// 3. [`apply_overrides`](Self::apply_overrides) — layer 1 (per-call + /// chained overrides, always highest priority). /// /// Layer 5 (Scheduler global defaults, e.g. global TTL) is applied later /// inside `Scheduler::submit()`. - fn resolve(self) -> (Scheduler, TaskSubmission) { - let scheduler = self.scheduler; - let mut sub = self.submission; + fn resolve(mut self) -> (Scheduler, TaskSubmission) { + self.apply_prefix(); + self.apply_defaults(); + self.apply_overrides(); + (self.scheduler, self.submission) + } - // ── 1. Prefix task_type with the module name ───────────────────────── + /// Prefix `task_type` with the module name (e.g. `"thumbnail"` → + /// `"media::thumbnail"`). Updates `label` when it matches the old + /// unprefixed type. + fn apply_prefix(&mut self) { if !self.module_name.is_empty() { - let old_type = sub.task_type.clone(); - sub.task_type = format!("{}::{}", self.module_name, old_type); - // Update label if it was the default (equal to the old task_type). - if sub.label == old_type { - sub.label = sub.task_type.clone(); + let old_type = self.submission.task_type.clone(); + self.submission.task_type = format!("{}::{}", self.module_name, old_type); + if self.submission.label == old_type { + self.submission.label = self.submission.task_type.clone(); } } + } - // ── 2+3. Apply TypedTask defaults then module overrides ─────────────── - if let Some(td) = self.typed_defaults { - // ── submit_typed() path ─────────────────────────────────────────── - // TypedTask values are the baseline (layer 4). Module defaults - // unconditionally override them when the module has an explicit - // setting (layer 3). - sub.priority = td.priority; - sub.group_key = td.group; - sub.ttl = td.ttl; - // TypedTask tags are the base; module tags add new keys only. - sub.tags = td.tags; - Self::merge_module_tags(&mut sub, &self.module_defaults.tags); - if let Some(p) = self.module_defaults.priority { - sub.priority = p; - } - if let Some(g) = self.module_defaults.group { - sub.group_key = Some(g); - } - if let Some(t) = self.module_defaults.ttl { - sub.ttl = Some(t); - } + /// Layers 2–4: apply typed/module/submission defaults. + /// + /// **`submit_typed()` path** (`typed_defaults` present): + /// TypedTask values are the baseline (layer 4). Module defaults + /// unconditionally override them (layer 3). + /// + /// **`submit()` path** (`typed_defaults` absent): + /// Submission values are the baseline (layer 2). Module defaults fill in + /// only where the submission is at its zero/`None` value (layer 3). + fn apply_defaults(&mut self) { + if let Some(td) = self.typed_defaults.take() { + // ── submit_typed() path ─────────────────────────────────────── + // TypedTask values are the baseline (layer 4). + self.submission.priority = td.priority; + self.submission.group_key = td.group; + self.submission.ttl = td.ttl; + self.submission.tags = td.tags; + // Module defaults: tags merge first, then scalars override + // unconditionally (layer 3 beats layer 4). + Self::merge_module_tags(&mut self.submission, &self.module_defaults.tags); + self.apply_module_scalar_defaults(true); } else { - // ── submit() path ───────────────────────────────────────────────── - // Module defaults fill in only where the submission is at its - // zero/None value (submission explicit values beat module defaults). - // - // Priority: treat `NORMAL` as "not explicitly set" — the same - // convention used by `BatchSubmission::build`. - if sub.priority == Priority::NORMAL { - if let Some(p) = self.module_defaults.priority { - sub.priority = p; - } + // ── submit() path ───────────────────────────────────────────── + // Module defaults fill gaps only (layer 3 fills layer 2 gaps). + self.apply_module_scalar_defaults(false); + Self::merge_module_tags(&mut self.submission, &self.module_defaults.tags); + } + } + + /// Apply module-level scalar defaults (priority, group, TTL). + /// + /// When `unconditional` is `true` (typed path), module values always + /// override the current submission value. When `false` (untyped path), + /// they fill in only where the submission is at its zero/`None` value + /// (`Priority::NORMAL` is treated as "not explicitly set"). + fn apply_module_scalar_defaults(&mut self, unconditional: bool) { + if let Some(p) = self.module_defaults.priority { + if unconditional || self.submission.priority == Priority::NORMAL { + self.submission.priority = p; } - if sub.group_key.is_none() { - if let Some(g) = self.module_defaults.group { - sub.group_key = Some(g); - } + } + if let Some(ref g) = self.module_defaults.group { + if unconditional || self.submission.group_key.is_none() { + self.submission.group_key = Some(g.clone()); } - if sub.ttl.is_none() { - if let Some(t) = self.module_defaults.ttl { - sub.ttl = Some(t); - } + } + if let Some(t) = self.module_defaults.ttl { + if unconditional || self.submission.ttl.is_none() { + self.submission.ttl = Some(t); } - // Module tags: add keys not already on the submission (submission wins). - Self::merge_module_tags(&mut sub, &self.module_defaults.tags); } + } - // ── 4. Apply per-call overrides (layer 1 — always highest priority) ── - if let Some(p) = self.override_priority { - sub.priority = p; + /// Layer 1: per-call overrides — always highest priority. + /// + /// Values set via chained builder methods (`.priority()`, `.group()`, + /// etc.) unconditionally overwrite whatever layers 2–4 produced. + fn apply_overrides(&mut self) { + if let Some(p) = self.override_priority.take() { + self.submission.priority = p; } - if let Some(g) = self.override_group { - sub.group_key = Some(g); + if let Some(g) = self.override_group.take() { + self.submission.group_key = Some(g); } - if let Some(k) = self.override_key { - sub.label = k.clone(); - sub.dedup_key = Some(k); + if let Some(k) = self.override_key.take() { + self.submission.label = k.clone(); + self.submission.dedup_key = Some(k); } - if let Some(ra) = self.override_run_after { - sub.run_after = Some(ra); + if let Some(ra) = self.override_run_after.take() { + self.submission.run_after = Some(ra); } - if let Some(t) = self.override_ttl { - sub.ttl = Some(t); + if let Some(t) = self.override_ttl.take() { + self.submission.ttl = Some(t); } if !self.override_depends_on.is_empty() { - sub.dependencies.extend(self.override_depends_on); + self.submission + .dependencies + .append(&mut self.override_depends_on); } - for (k, v) in self.override_tags { - sub.tags.insert(k, v); + for (k, v) in std::mem::take(&mut self.override_tags) { + self.submission.tags.insert(k, v); } - if let Some(pid) = self.override_parent_id { - sub.parent_id = Some(pid); + if let Some(pid) = self.override_parent_id.take() { + self.submission.parent_id = Some(pid); } - - (scheduler, sub) } /// Submit the task, returning the outcome.