diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index efe02bdb..1c35c623 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,6 +6,10 @@ on: pull_request: branches: [master] +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + jobs: lint: runs-on: ubuntu-latest @@ -92,7 +96,15 @@ jobs: - name: Install Rust uses: dtolnay/rust-toolchain@stable + - name: Build wheels (with postgres) + if: runner.os == 'Linux' + uses: PyO3/maturin-action@v1 + with: + args: --release --out dist --features postgres + manylinux: auto + - name: Build wheels + if: runner.os != 'Linux' uses: PyO3/maturin-action@v1 with: args: --release --out dist diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 00000000..5fbf1f74 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,47 @@ +name: Deploy Docs + +on: + push: + branches: [master] + paths: + - "docs/**" + - "zensical.toml" + - ".github/workflows/docs.yml" + workflow_dispatch: + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: pages + cancel-in-progress: false + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: astral-sh/setup-uv@v5 + + - name: Install dependencies + run: uv sync --group docs + + - name: Build docs + run: uv run zensical build + + - uses: actions/upload-pages-artifact@v3 + with: + path: site + + deploy: + needs: build + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - id: deployment + uses: actions/deploy-pages@v4 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..a91998b3 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,37 @@ +repos: + - repo: local + hooks: + - id: cargo-fmt + name: cargo fmt + entry: cargo fmt --all --check + language: system + types: [rust] + pass_filenames: false + + - id: cargo-clippy + name: cargo clippy + entry: cargo clippy --all-targets --all-features -- -D warnings + language: system + types: [rust] + pass_filenames: false + + - id: ruff-check + name: ruff check + entry: uv run ruff check . + language: system + types: [python] + pass_filenames: false + + - id: ruff-format + name: ruff format + entry: uv run ruff format --check . + language: system + types: [python] + pass_filenames: false + + - id: mypy + name: mypy + entry: uv run mypy py_src/ + language: system + types: [python] + pass_filenames: false diff --git a/Cargo.toml b/Cargo.toml index 0cbc17fa..0d3dae79 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ resolver = "2" [workspace.dependencies] diesel = { version = "2.2", features = ["sqlite", "r2d2", "returning_clauses_for_sqlite_3_35"] } libsqlite3-sys = { version = "0.30", features = ["bundled"] } +pq-sys = { version = "0.6", features = ["bundled"] } uuid = { version = "1", features = ["v7"] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } diff --git a/README.md b/README.md index 6308ad71..965e44b3 100644 --- a/README.md +++ b/README.md @@ -127,7 +127,7 @@ chord([download.s(u) for u in urls], merge.s()).apply() ### Periodic Tasks ```python -@queue.periodic(cron="0 */6 * * *") +@queue.periodic(cron="0 0 */6 * * *") def cleanup_temp_files(): ... ``` diff --git a/crates/taskito-core/Cargo.toml b/crates/taskito-core/Cargo.toml index 56a37e5c..a7452fd7 100644 --- a/crates/taskito-core/Cargo.toml +++ b/crates/taskito-core/Cargo.toml @@ -3,9 +3,14 @@ name = "taskito-core" version = "0.2.2" edition = "2021" +[features] +default = [] +postgres = ["diesel/postgres", "pq-sys"] + [dependencies] diesel = { workspace = true } libsqlite3-sys = { workspace = true } +pq-sys = { workspace = true, optional = true } uuid = { workspace = true } tokio = { workspace = true } serde = { workspace = true } diff --git a/crates/taskito-core/src/job.rs b/crates/taskito-core/src/job.rs index b244d418..e2e2c283 100644 --- a/crates/taskito-core/src/job.rs +++ b/crates/taskito-core/src/job.rs @@ -1,5 +1,5 @@ use serde::{Deserialize, Serialize}; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use uuid::Uuid; use crate::storage::models::JobRow; @@ -142,6 +142,6 @@ impl NewJob { pub fn now_millis() -> i64 { SystemTime::now() .duration_since(UNIX_EPOCH) - .expect("time went backwards") + .unwrap_or(Duration::ZERO) .as_millis() as i64 } diff --git a/crates/taskito-core/src/lib.rs b/crates/taskito-core/src/lib.rs index 0bab0676..19f8b362 100644 --- a/crates/taskito-core/src/lib.rs +++ b/crates/taskito-core/src/lib.rs @@ -1,13 +1,16 @@ -pub mod job; pub mod error; -pub mod storage; -pub mod scheduler; +pub mod job; pub mod periodic; pub mod resilience; +pub mod scheduler; +pub mod storage; // Primary public API — the types most consumers need. pub use error::QueueError; pub use job::{Job, JobStatus, NewJob}; pub use scheduler::{Scheduler, SchedulerConfig}; +#[cfg(feature = "postgres")] +pub use storage::postgres::PostgresStorage; pub use storage::sqlite::SqliteStorage; pub use storage::Storage; +pub use storage::StorageBackend; diff --git a/crates/taskito-core/src/periodic.rs b/crates/taskito-core/src/periodic.rs index 011e65b3..d3ac85bf 100644 --- a/crates/taskito-core/src/periodic.rs +++ b/crates/taskito-core/src/periodic.rs @@ -8,12 +8,10 @@ use crate::error::{QueueError, Result}; /// Compute the next run time (in UNIX milliseconds) for a cron expression, /// starting from `after_ms` (also UNIX milliseconds). pub fn next_cron_time(cron_expr: &str, after_ms: i64) -> Result { - let schedule = Schedule::from_str(cron_expr).map_err(|e| { - QueueError::Config(format!("invalid cron expression '{cron_expr}': {e}")) - })?; + let schedule = Schedule::from_str(cron_expr) + .map_err(|e| QueueError::Config(format!("invalid cron expression '{cron_expr}': {e}")))?; - let after_dt = chrono::DateTime::from_timestamp_millis(after_ms) - .unwrap_or_else(Utc::now); + let after_dt = chrono::DateTime::from_timestamp_millis(after_ms).unwrap_or_else(Utc::now); let next = schedule .after(&after_dt) diff --git a/crates/taskito-core/src/resilience/circuit_breaker.rs b/crates/taskito-core/src/resilience/circuit_breaker.rs index 6a2cf19f..ccf743f8 100644 --- a/crates/taskito-core/src/resilience/circuit_breaker.rs +++ b/crates/taskito-core/src/resilience/circuit_breaker.rs @@ -1,7 +1,7 @@ use crate::error::Result; use crate::job::now_millis; use crate::storage::models::CircuitBreakerRow; -use crate::storage::sqlite::SqliteStorage; +use crate::storage::{Storage, StorageBackend}; /// Circuit breaker states. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -40,11 +40,11 @@ pub struct CircuitBreakerConfig { /// Circuit breaker manager backed by SQLite. pub struct CircuitBreaker { - storage: SqliteStorage, + storage: StorageBackend, } impl CircuitBreaker { - pub fn new(storage: SqliteStorage) -> Self { + pub fn new(storage: StorageBackend) -> Self { Self { storage } } diff --git a/crates/taskito-core/src/resilience/dlq.rs b/crates/taskito-core/src/resilience/dlq.rs index e559dfb2..2331498a 100644 --- a/crates/taskito-core/src/resilience/dlq.rs +++ b/crates/taskito-core/src/resilience/dlq.rs @@ -1,14 +1,14 @@ use crate::error::Result; use crate::job::Job; -use crate::storage::sqlite::{SqliteStorage, DeadJob}; +use crate::storage::{DeadJob, Storage, StorageBackend}; /// Dead letter queue manager. pub struct DeadLetterQueue { - storage: SqliteStorage, + storage: StorageBackend, } impl DeadLetterQueue { - pub fn new(storage: SqliteStorage) -> Self { + pub fn new(storage: StorageBackend) -> Self { Self { storage } } diff --git a/crates/taskito-core/src/resilience/rate_limiter.rs b/crates/taskito-core/src/resilience/rate_limiter.rs index 83ccca1a..dd57b8fc 100644 --- a/crates/taskito-core/src/resilience/rate_limiter.rs +++ b/crates/taskito-core/src/resilience/rate_limiter.rs @@ -1,9 +1,9 @@ use crate::error::Result; -use crate::storage::sqlite::SqliteStorage; +use crate::storage::{Storage, StorageBackend}; /// Token bucket rate limiter backed by SQLite for persistence. pub struct RateLimiter { - storage: SqliteStorage, + storage: StorageBackend, } /// Parsed rate limit configuration (e.g., "100/m" → 100 tokens, refill 1.667/s). @@ -37,7 +37,7 @@ impl RateLimitConfig { } impl RateLimiter { - pub fn new(storage: SqliteStorage) -> Self { + pub fn new(storage: StorageBackend) -> Self { Self { storage } } @@ -71,7 +71,8 @@ mod tests { #[test] fn test_token_bucket() { - let storage = SqliteStorage::in_memory().unwrap(); + let storage = + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()); let limiter = RateLimiter::new(storage); let config = RateLimitConfig { max_tokens: 3.0, @@ -89,13 +90,15 @@ mod tests { #[test] fn test_concurrent_token_acquisition() { - use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; use std::sync::Barrier; let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("rate_limit_test.db"); - let storage = SqliteStorage::new(db_path.to_str().unwrap()).unwrap(); + let storage = StorageBackend::Sqlite( + crate::storage::sqlite::SqliteStorage::new(db_path.to_str().unwrap()).unwrap(), + ); let limiter = Arc::new(RateLimiter::new(storage)); let config = RateLimitConfig { max_tokens: 10.0, diff --git a/crates/taskito-core/src/resilience/retry.rs b/crates/taskito-core/src/resilience/retry.rs index 6e69f8cf..17e29726 100644 --- a/crates/taskito-core/src/resilience/retry.rs +++ b/crates/taskito-core/src/resilience/retry.rs @@ -15,8 +15,8 @@ impl Default for RetryPolicy { fn default() -> Self { Self { max_retries: 3, - base_delay_ms: 1_000, // 1 second - max_delay_ms: 300_000, // 5 minutes + base_delay_ms: 1_000, // 1 second + max_delay_ms: 300_000, // 5 minutes } } } diff --git a/crates/taskito-core/src/scheduler.rs b/crates/taskito-core/src/scheduler.rs index e98b59e0..b3a94452 100644 --- a/crates/taskito-core/src/scheduler.rs +++ b/crates/taskito-core/src/scheduler.rs @@ -6,14 +6,14 @@ use crossbeam_channel::Sender; use log::{error, info, warn}; use tokio::sync::Notify; -use crate::resilience::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; -use crate::resilience::dlq::DeadLetterQueue; use crate::error::Result; -use crate::job::{Job, NewJob, now_millis}; +use crate::job::{now_millis, Job, NewJob}; use crate::periodic::next_cron_time; +use crate::resilience::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig}; +use crate::resilience::dlq::DeadLetterQueue; use crate::resilience::rate_limiter::{RateLimitConfig, RateLimiter}; use crate::resilience::retry::RetryPolicy; -use crate::storage::sqlite::SqliteStorage; +use crate::storage::{Storage, StorageBackend}; /// Delay before re-scheduling a circuit-broken job (ms). const CIRCUIT_BREAKER_RETRY_DELAY_MS: i64 = 5000; @@ -88,7 +88,7 @@ pub struct TaskConfig { /// The central scheduler that coordinates job dispatch, retries, rate limiting, and circuit breakers. pub struct Scheduler { - storage: SqliteStorage, + storage: StorageBackend, rate_limiter: RateLimiter, dlq: DeadLetterQueue, circuit_breaker: CircuitBreaker, @@ -107,7 +107,7 @@ struct TickCounters { } impl Scheduler { - pub fn new(storage: SqliteStorage, queues: Vec, config: SchedulerConfig) -> Self { + pub fn new(storage: StorageBackend, queues: Vec, config: SchedulerConfig) -> Self { let rate_limiter = RateLimiter::new(storage.clone()); let dlq = DeadLetterQueue::new(storage.clone()); let circuit_breaker = CircuitBreaker::new(storage.clone()); @@ -124,7 +124,7 @@ impl Scheduler { } } - pub fn storage(&self) -> &SqliteStorage { + pub fn storage(&self) -> &StorageBackend { &self.storage } @@ -176,13 +176,19 @@ impl Scheduler { } } - if counters.periodic.is_multiple_of(self.config.periodic_check_interval) { + if counters + .periodic + .is_multiple_of(self.config.periodic_check_interval) + { if let Err(e) = self.check_periodic() { error!("periodic check error: {e}"); } } - if counters.cleanup.is_multiple_of(self.config.cleanup_interval) { + if counters + .cleanup + .is_multiple_of(self.config.cleanup_interval) + { if let Err(e) = self.auto_cleanup() { error!("auto-cleanup error: {e}"); } @@ -199,16 +205,16 @@ impl Scheduler { // Check circuit breaker for this task if let Some(config) = self.task_configs.get(&job.task_name) { - if config.circuit_breaker.is_some() - && !self.circuit_breaker.allow(&job.task_name)? - { - self.storage.retry(&job.id, now + CIRCUIT_BREAKER_RETRY_DELAY_MS)?; + if config.circuit_breaker.is_some() && !self.circuit_breaker.allow(&job.task_name)? { + self.storage + .retry(&job.id, now + CIRCUIT_BREAKER_RETRY_DELAY_MS)?; return Ok(false); } if let Some(ref rl_config) = config.rate_limit { if !self.rate_limiter.try_acquire(&job.task_name, rl_config)? { - self.storage.retry(&job.id, now + RATE_LIMIT_RETRY_DELAY_MS)?; + self.storage + .retry(&job.id, now + RATE_LIMIT_RETRY_DELAY_MS)?; return Ok(false); } } @@ -225,10 +231,18 @@ impl Scheduler { /// Handle a completed or failed job result from a worker. pub fn handle_result(&self, result: JobResult) -> Result<()> { match result { - JobResult::Success { job_id, result, ref task_name, wall_time_ns } => { + JobResult::Success { + job_id, + result, + ref task_name, + wall_time_ns, + } => { self.storage.complete(&job_id, result)?; - if let Err(e) = self.storage.record_metric(task_name, &job_id, wall_time_ns, 0, true) { + if let Err(e) = + self.storage + .record_metric(task_name, &job_id, wall_time_ns, 0, true) + { error!("failed to record metric for job {job_id}: {e}"); } @@ -249,7 +263,10 @@ impl Scheduler { log::error!("failed to record error for job {job_id}: {e}"); } - if let Err(e) = self.storage.record_metric(&task_name, &job_id, wall_time_ns, 0, false) { + if let Err(e) = + self.storage + .record_metric(&task_name, &job_id, wall_time_ns, 0, false) + { log::error!("failed to record metric for job {job_id}: {e}"); } @@ -259,8 +276,9 @@ impl Scheduler { // If should_retry is false (exception filtering), skip straight to DLQ if !should_retry { - if let Some(job) = self.storage.get_job(&job_id)? { - self.dlq.move_to_dlq(&job, &error, None)?; + match self.storage.get_job(&job_id)? { + Some(job) => self.dlq.move_to_dlq(&job, &error, None)?, + None => warn!("job {job_id} disappeared before DLQ move"), } return Ok(()); } @@ -282,17 +300,25 @@ impl Scheduler { self.storage.retry(&job_id, next_at)?; } else { // Move to DLQ - if let Some(job) = self.storage.get_job(&job_id)? { - self.dlq.move_to_dlq(&job, &error, None)?; + match self.storage.get_job(&job_id)? { + Some(job) => self.dlq.move_to_dlq(&job, &error, None)?, + None => warn!("job {job_id} disappeared before DLQ move"), } } } - JobResult::Cancelled { job_id, task_name, wall_time_ns } => { + JobResult::Cancelled { + job_id, + task_name, + wall_time_ns, + } => { // Mark as cancelled, no retry if let Err(e) = self.storage.mark_cancelled(&job_id) { error!("failed to mark job {job_id} as cancelled: {e}"); } - if let Err(e) = self.storage.record_metric(&task_name, &job_id, wall_time_ns, 0, false) { + if let Err(e) = + self.storage + .record_metric(&task_name, &job_id, wall_time_ns, 0, false) + { error!("failed to record metric for cancelled job {job_id}: {e}"); } } @@ -306,7 +332,6 @@ impl Scheduler { for job in stale_jobs { let error = format!("job timed out after {}ms", job.timeout_ms); - let _ = self.storage.fail(&job.id, &error); self.handle_result(JobResult::Failure { job_id: job.id.clone(), error, @@ -352,22 +377,23 @@ impl Scheduler { let due_tasks = self.storage.get_due_periodic(now)?; for task in due_tasks { + let unique_key = Some(format!("periodic:{}:{}", task.name, now)); let new_job = NewJob { queue: task.queue.clone(), task_name: task.task_name.clone(), - payload: Self::build_periodic_payload(&task.args, &task.kwargs), + payload: Self::build_periodic_payload(&task.args), priority: 0, scheduled_at: now, max_retries: PERIODIC_DEFAULT_MAX_RETRIES, timeout_ms: PERIODIC_DEFAULT_TIMEOUT_MS, - unique_key: None, + unique_key, metadata: None, depends_on: vec![], expires_at: None, result_ttl_ms: None, }; - if let Err(e) = self.storage.enqueue(new_job) { + if let Err(e) = self.storage.enqueue_unique(new_job) { error!("failed to enqueue periodic task '{}': {e}", task.name); continue; } @@ -375,27 +401,25 @@ impl Scheduler { let next_run = match next_cron_time(&task.cron_expr, now) { Ok(t) => t, Err(e) => { - error!( - "failed to compute next run for '{}': {e}", - task.name - ); + error!("failed to compute next run for '{}': {e}", task.name); continue; } }; - if let Err(e) = self.storage.update_periodic_schedule(&task.name, now, next_run) { - error!( - "failed to update schedule for '{}': {e}", - task.name - ); + if let Err(e) = self + .storage + .update_periodic_schedule(&task.name, now, next_run) + { + error!("failed to update schedule for '{}': {e}", task.name); } } Ok(()) } - /// Build a cloudpickle-compatible payload from stored args/kwargs. - fn build_periodic_payload(args: &Option>, _kwargs: &Option>) -> Vec { + /// Build a cloudpickle-compatible payload from stored args. + /// kwargs are embedded in the args blob by the Python side. + fn build_periodic_payload(args: &Option>) -> Vec { match args { Some(blob) => blob.clone(), None => Vec::new(), @@ -410,7 +434,8 @@ mod tests { use crate::storage::models::NewPeriodicTaskRow; fn test_scheduler() -> Scheduler { - let storage = SqliteStorage::in_memory().unwrap(); + let storage = + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()); Scheduler::new( storage, vec!["default".to_string()], @@ -450,12 +475,14 @@ mod tests { let job = enqueue_and_run(&scheduler, "success_task"); assert_eq!(job.status, JobStatus::Running); - scheduler.handle_result(JobResult::Success { - job_id: job.id.clone(), - result: Some(vec![42]), - task_name: "success_task".to_string(), - wall_time_ns: 1_000_000, - }).unwrap(); + scheduler + .handle_result(JobResult::Success { + job_id: job.id.clone(), + result: Some(vec![42]), + task_name: "success_task".to_string(), + wall_time_ns: 1_000_000, + }) + .unwrap(); let completed = scheduler.storage.get_job(&job.id).unwrap().unwrap(); assert_eq!(completed.status, JobStatus::Complete); @@ -465,27 +492,32 @@ mod tests { #[test] fn test_handle_failure_with_retry() { let mut scheduler = test_scheduler(); - scheduler.register_task("retry_task".to_string(), TaskConfig { - retry_policy: RetryPolicy { - max_retries: 3, - base_delay_ms: 100, - max_delay_ms: 1000, + scheduler.register_task( + "retry_task".to_string(), + TaskConfig { + retry_policy: RetryPolicy { + max_retries: 3, + base_delay_ms: 100, + max_delay_ms: 1000, + }, + rate_limit: None, + circuit_breaker: None, }, - rate_limit: None, - circuit_breaker: None, - }); + ); let job = enqueue_and_run(&scheduler, "retry_task"); - scheduler.handle_result(JobResult::Failure { - job_id: job.id.clone(), - error: "transient error".to_string(), - retry_count: 0, - max_retries: 3, - task_name: "retry_task".to_string(), - wall_time_ns: 500_000, - should_retry: true, - }).unwrap(); + scheduler + .handle_result(JobResult::Failure { + job_id: job.id.clone(), + error: "transient error".to_string(), + retry_count: 0, + max_retries: 3, + task_name: "retry_task".to_string(), + wall_time_ns: 500_000, + should_retry: true, + }) + .unwrap(); let retried = scheduler.storage.get_job(&job.id).unwrap().unwrap(); assert_eq!(retried.status, JobStatus::Pending); @@ -497,15 +529,17 @@ mod tests { let scheduler = test_scheduler(); let job = enqueue_and_run(&scheduler, "exhausted_task"); - scheduler.handle_result(JobResult::Failure { - job_id: job.id.clone(), - error: "fatal".to_string(), - retry_count: 3, - max_retries: 3, - task_name: "exhausted_task".to_string(), - wall_time_ns: 100, - should_retry: true, - }).unwrap(); + scheduler + .handle_result(JobResult::Failure { + job_id: job.id.clone(), + error: "fatal".to_string(), + retry_count: 3, + max_retries: 3, + task_name: "exhausted_task".to_string(), + wall_time_ns: 100, + should_retry: true, + }) + .unwrap(); let dead = scheduler.storage.get_job(&job.id).unwrap().unwrap(); assert_eq!(dead.status, JobStatus::Dead); @@ -520,15 +554,17 @@ mod tests { let scheduler = test_scheduler(); let job = enqueue_and_run(&scheduler, "no_retry_task"); - scheduler.handle_result(JobResult::Failure { - job_id: job.id.clone(), - error: "non-retryable".to_string(), - retry_count: 0, - max_retries: 3, - task_name: "no_retry_task".to_string(), - wall_time_ns: 100, - should_retry: false, - }).unwrap(); + scheduler + .handle_result(JobResult::Failure { + job_id: job.id.clone(), + error: "non-retryable".to_string(), + retry_count: 0, + max_retries: 3, + task_name: "no_retry_task".to_string(), + wall_time_ns: 100, + should_retry: false, + }) + .unwrap(); let dead = scheduler.storage.get_job(&job.id).unwrap().unwrap(); assert_eq!(dead.status, JobStatus::Dead); @@ -539,11 +575,13 @@ mod tests { let scheduler = test_scheduler(); let job = enqueue_and_run(&scheduler, "cancel_task"); - scheduler.handle_result(JobResult::Cancelled { - job_id: job.id.clone(), - task_name: "cancel_task".to_string(), - wall_time_ns: 100, - }).unwrap(); + scheduler + .handle_result(JobResult::Cancelled { + job_id: job.id.clone(), + task_name: "cancel_task".to_string(), + wall_time_ns: 100, + }) + .unwrap(); let cancelled = scheduler.storage.get_job(&job.id).unwrap().unwrap(); assert_eq!(cancelled.status, JobStatus::Cancelled); @@ -552,11 +590,17 @@ mod tests { #[test] fn test_try_dispatch_rate_limited() { let mut scheduler = test_scheduler(); - scheduler.register_task("rl_task".to_string(), TaskConfig { - retry_policy: RetryPolicy::default(), - rate_limit: Some(RateLimitConfig { max_tokens: 1.0, refill_rate: 0.0 }), - circuit_breaker: None, - }); + scheduler.register_task( + "rl_task".to_string(), + TaskConfig { + retry_policy: RetryPolicy::default(), + rate_limit: Some(RateLimitConfig { + max_tokens: 1.0, + refill_rate: 0.0, + }), + circuit_breaker: None, + }, + ); // Enqueue two jobs scheduler.storage.enqueue(make_job("rl_task")).unwrap(); @@ -574,9 +618,10 @@ mod tests { assert_eq!(rx.len(), 1); // still just 1 dispatched // The second job should be back in pending with a future scheduled_at - let jobs = scheduler.storage.list_jobs( - Some(JobStatus::Pending as i32), None, None, 10, 0 - ).unwrap(); + let jobs = scheduler + .storage + .list_jobs(Some(JobStatus::Pending as i32), None, None, 10, 0) + .unwrap(); assert_eq!(jobs.len(), 1); assert!(jobs[0].scheduled_at > now_millis()); } @@ -589,11 +634,14 @@ mod tests { window_ms: 60_000, cooldown_ms: 300_000, }; - scheduler.register_task("cb_task".to_string(), TaskConfig { - retry_policy: RetryPolicy::default(), - rate_limit: None, - circuit_breaker: Some(cb_config), - }); + scheduler.register_task( + "cb_task".to_string(), + TaskConfig { + retry_policy: RetryPolicy::default(), + rate_limit: None, + circuit_breaker: Some(cb_config), + }, + ); // Trip the circuit breaker scheduler.circuit_breaker.record_failure("cb_task").unwrap(); @@ -608,9 +656,10 @@ mod tests { assert_eq!(rx.len(), 0); // Job should be rescheduled - let jobs = scheduler.storage.list_jobs( - Some(JobStatus::Pending as i32), None, None, 10, 0 - ).unwrap(); + let jobs = scheduler + .storage + .list_jobs(Some(JobStatus::Pending as i32), None, None, 10, 0) + .unwrap(); assert_eq!(jobs.len(), 1); assert!(jobs[0].scheduled_at > now_millis()); } @@ -618,11 +667,18 @@ mod tests { #[test] fn test_reap_stale_jobs() { let mut scheduler = test_scheduler(); - scheduler.register_task("stale_task".to_string(), TaskConfig { - retry_policy: RetryPolicy { max_retries: 3, base_delay_ms: 100, max_delay_ms: 1000 }, - rate_limit: None, - circuit_breaker: None, - }); + scheduler.register_task( + "stale_task".to_string(), + TaskConfig { + retry_policy: RetryPolicy { + max_retries: 3, + base_delay_ms: 100, + max_delay_ms: 1000, + }, + rate_limit: None, + circuit_breaker: None, + }, + ); // Create a job with a very short timeout let mut new_job = make_job("stale_task"); @@ -667,26 +723,35 @@ mod tests { scheduler.check_periodic().unwrap(); // A job should have been enqueued - let jobs = scheduler.storage.list_jobs( - Some(JobStatus::Pending as i32), None, Some("periodic_task"), 10, 0 - ).unwrap(); + let jobs = scheduler + .storage + .list_jobs( + Some(JobStatus::Pending as i32), + None, + Some("periodic_task"), + 10, + 0, + ) + .unwrap(); assert_eq!(jobs.len(), 1); } #[test] fn test_auto_cleanup() { - let storage = SqliteStorage::in_memory().unwrap(); + let storage = + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()); let config = SchedulerConfig { result_ttl_ms: Some(1), // 1ms TTL ..SchedulerConfig::default() }; - let scheduler = Scheduler::new( - storage, vec!["default".to_string()], config, - ); + let scheduler = Scheduler::new(storage, vec!["default".to_string()], config); // Enqueue, dequeue, and complete a job let job = scheduler.storage.enqueue(make_job("cleanup_task")).unwrap(); - scheduler.storage.dequeue("default", now_millis() + 1000).unwrap(); + scheduler + .storage + .dequeue("default", now_millis() + 1000) + .unwrap(); scheduler.storage.complete(&job.id, Some(vec![1])).unwrap(); // Wait for the TTL to expire @@ -701,16 +766,15 @@ mod tests { #[test] fn test_tick_dispatches_and_maintains() { - let storage = SqliteStorage::in_memory().unwrap(); + let storage = + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()); let config = SchedulerConfig { reap_interval: 1, periodic_check_interval: 1, cleanup_interval: 1, ..SchedulerConfig::default() }; - let scheduler = Scheduler::new( - storage, vec!["default".to_string()], config, - ); + let scheduler = Scheduler::new(storage, vec!["default".to_string()], config); scheduler.storage.enqueue(make_job("tick_task")).unwrap(); diff --git a/crates/taskito-core/src/storage/mod.rs b/crates/taskito-core/src/storage/mod.rs index b7434925..71f0ccbb 100644 --- a/crates/taskito-core/src/storage/mod.rs +++ b/crates/taskito-core/src/storage/mod.rs @@ -1,6 +1,299 @@ pub mod models; +#[cfg(feature = "postgres")] +pub mod postgres; pub mod schema; pub mod sqlite; pub mod traits; pub use traits::Storage; + +use crate::error::Result; +use crate::job::{Job, NewJob}; + +// ── Shared helper types ──────────────────────────────────────────────── + +#[derive(Debug, Clone, Default)] +pub struct QueueStats { + pub pending: i64, + pub running: i64, + pub completed: i64, + pub failed: i64, + pub dead: i64, + pub cancelled: i64, +} + +#[derive(Debug, Clone)] +pub struct DeadJob { + pub id: String, + pub original_job_id: String, + pub queue: String, + pub task_name: String, + pub payload: Vec, + pub error: Option, + pub retry_count: i32, + pub failed_at: i64, + pub metadata: Option, + pub priority: i32, + pub max_retries: i32, + pub timeout_ms: i64, + pub result_ttl_ms: Option, +} + +impl From for DeadJob { + fn from(row: models::DeadLetterRow) -> Self { + Self { + id: row.id, + original_job_id: row.original_job_id, + queue: row.queue, + task_name: row.task_name, + payload: row.payload, + error: row.error, + retry_count: row.retry_count, + failed_at: row.failed_at, + metadata: row.metadata, + priority: row.priority, + max_retries: row.max_retries, + timeout_ms: row.timeout_ms, + result_ttl_ms: row.result_ttl_ms, + } + } +} + +// ── Backend-agnostic storage wrapper ────────────────────────────────── + +/// Storage backend enum that dispatches to either SQLite or PostgreSQL. +#[derive(Clone)] +pub enum StorageBackend { + Sqlite(sqlite::SqliteStorage), + #[cfg(feature = "postgres")] + Postgres(postgres::PostgresStorage), +} + +macro_rules! delegate { + ($self:expr, $method:ident $(, $arg:expr)*) => { + match $self { + StorageBackend::Sqlite(s) => s.$method($($arg),*), + #[cfg(feature = "postgres")] + StorageBackend::Postgres(s) => s.$method($($arg),*), + } + }; +} + +impl Storage for StorageBackend { + fn enqueue(&self, new_job: NewJob) -> Result { + delegate!(self, enqueue, new_job) + } + fn enqueue_batch(&self, new_jobs: Vec) -> Result> { + delegate!(self, enqueue_batch, new_jobs) + } + fn enqueue_unique(&self, new_job: NewJob) -> Result { + delegate!(self, enqueue_unique, new_job) + } + fn dequeue(&self, queue_name: &str, now: i64) -> Result> { + delegate!(self, dequeue, queue_name, now) + } + fn dequeue_from(&self, queues: &[String], now: i64) -> Result> { + delegate!(self, dequeue_from, queues, now) + } + fn complete(&self, id: &str, result_bytes: Option>) -> Result<()> { + delegate!(self, complete, id, result_bytes) + } + fn fail(&self, id: &str, error: &str) -> Result<()> { + delegate!(self, fail, id, error) + } + fn retry(&self, id: &str, next_scheduled_at: i64) -> Result<()> { + delegate!(self, retry, id, next_scheduled_at) + } + fn cancel_job(&self, id: &str) -> Result { + delegate!(self, cancel_job, id) + } + fn request_cancel(&self, id: &str) -> Result { + delegate!(self, request_cancel, id) + } + fn is_cancel_requested(&self, id: &str) -> Result { + delegate!(self, is_cancel_requested, id) + } + fn mark_cancelled(&self, id: &str) -> Result<()> { + delegate!(self, mark_cancelled, id) + } + fn cascade_cancel(&self, failed_job_id: &str, reason: &str) -> Result<()> { + delegate!(self, cascade_cancel, failed_job_id, reason) + } + fn get_dependencies(&self, job_id: &str) -> Result> { + delegate!(self, get_dependencies, job_id) + } + fn get_dependents(&self, job_id: &str) -> Result> { + delegate!(self, get_dependents, job_id) + } + fn update_progress(&self, id: &str, progress: i32) -> Result<()> { + delegate!(self, update_progress, id, progress) + } + fn list_jobs( + &self, + status: Option, + queue_name: Option<&str>, + task_name: Option<&str>, + limit: i64, + offset: i64, + ) -> Result> { + delegate!(self, list_jobs, status, queue_name, task_name, limit, offset) + } + fn get_job(&self, id: &str) -> Result> { + delegate!(self, get_job, id) + } + fn stats(&self) -> Result { + delegate!(self, stats) + } + fn purge_completed(&self, older_than_ms: i64) -> Result { + delegate!(self, purge_completed, older_than_ms) + } + fn purge_completed_with_ttl(&self, global_cutoff_ms: i64) -> Result { + delegate!(self, purge_completed_with_ttl, global_cutoff_ms) + } + fn reap_stale_jobs(&self, now: i64) -> Result> { + delegate!(self, reap_stale_jobs, now) + } + fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()> { + delegate!(self, record_error, job_id, attempt, error) + } + fn get_job_errors(&self, job_id: &str) -> Result> { + delegate!(self, get_job_errors, job_id) + } + fn purge_job_errors(&self, older_than_ms: i64) -> Result { + delegate!(self, purge_job_errors, older_than_ms) + } + fn move_to_dlq(&self, job: &Job, error: &str, metadata: Option<&str>) -> Result<()> { + delegate!(self, move_to_dlq, job, error, metadata) + } + fn list_dead(&self, limit: i64, offset: i64) -> Result> { + delegate!(self, list_dead, limit, offset) + } + fn retry_dead(&self, dead_id: &str) -> Result { + delegate!(self, retry_dead, dead_id) + } + fn purge_dead(&self, older_than_ms: i64) -> Result { + delegate!(self, purge_dead, older_than_ms) + } + fn get_rate_limit(&self, key: &str) -> Result> { + delegate!(self, get_rate_limit, key) + } + fn upsert_rate_limit(&self, row: &models::RateLimitRow) -> Result<()> { + delegate!(self, upsert_rate_limit, row) + } + fn try_acquire_token(&self, key: &str, max_tokens: f64, refill_rate: f64) -> Result { + delegate!(self, try_acquire_token, key, max_tokens, refill_rate) + } + fn register_periodic(&self, task: &models::NewPeriodicTaskRow) -> Result<()> { + delegate!(self, register_periodic, task) + } + fn get_due_periodic(&self, now: i64) -> Result> { + delegate!(self, get_due_periodic, now) + } + fn update_periodic_schedule(&self, name: &str, last_run: i64, next_run: i64) -> Result<()> { + delegate!(self, update_periodic_schedule, name, last_run, next_run) + } + fn record_metric( + &self, + task_name: &str, + job_id: &str, + wall_time_ns: i64, + memory_bytes: i64, + succeeded: bool, + ) -> Result<()> { + delegate!( + self, + record_metric, + task_name, + job_id, + wall_time_ns, + memory_bytes, + succeeded + ) + } + fn get_metrics(&self, name: Option<&str>, since_ms: i64) -> Result> { + delegate!(self, get_metrics, name, since_ms) + } + fn purge_metrics(&self, older_than_ms: i64) -> Result { + delegate!(self, purge_metrics, older_than_ms) + } + fn record_replay( + &self, + original_job_id: &str, + replay_job_id: &str, + original_result: Option<&[u8]>, + replay_result: Option<&[u8]>, + original_error: Option<&str>, + replay_error: Option<&str>, + ) -> Result<()> { + delegate!( + self, + record_replay, + original_job_id, + replay_job_id, + original_result, + replay_result, + original_error, + replay_error + ) + } + fn get_replay_history(&self, original_job_id: &str) -> Result> { + delegate!(self, get_replay_history, original_job_id) + } + fn write_task_log( + &self, + job_id: &str, + task_name: &str, + level: &str, + message: &str, + extra: Option<&str>, + ) -> Result<()> { + delegate!( + self, + write_task_log, + job_id, + task_name, + level, + message, + extra + ) + } + fn get_task_logs(&self, job_id: &str) -> Result> { + delegate!(self, get_task_logs, job_id) + } + fn query_task_logs( + &self, + task_name: Option<&str>, + level: Option<&str>, + since_ms: i64, + limit: i64, + ) -> Result> { + delegate!(self, query_task_logs, task_name, level, since_ms, limit) + } + fn purge_task_logs(&self, older_than_ms: i64) -> Result { + delegate!(self, purge_task_logs, older_than_ms) + } + fn get_circuit_breaker(&self, task_name: &str) -> Result> { + delegate!(self, get_circuit_breaker, task_name) + } + fn upsert_circuit_breaker(&self, row: &models::CircuitBreakerRow) -> Result<()> { + delegate!(self, upsert_circuit_breaker, row) + } + fn list_circuit_breakers(&self) -> Result> { + delegate!(self, list_circuit_breakers) + } + fn register_worker(&self, worker_id: &str, queues: &str) -> Result<()> { + delegate!(self, register_worker, worker_id, queues) + } + fn heartbeat(&self, worker_id: &str) -> Result<()> { + delegate!(self, heartbeat, worker_id) + } + fn list_workers(&self) -> Result> { + delegate!(self, list_workers) + } + fn reap_dead_workers(&self) -> Result { + delegate!(self, reap_dead_workers) + } + fn unregister_worker(&self, worker_id: &str) -> Result<()> { + delegate!(self, unregister_worker, worker_id) + } +} diff --git a/crates/taskito-core/src/storage/models.rs b/crates/taskito-core/src/storage/models.rs index 47258205..3d2fb111 100644 --- a/crates/taskito-core/src/storage/models.rs +++ b/crates/taskito-core/src/storage/models.rs @@ -67,6 +67,10 @@ pub struct DeadLetterRow { pub retry_count: i32, pub failed_at: i64, pub metadata: Option, + pub priority: i32, + pub max_retries: i32, + pub timeout_ms: i64, + pub result_ttl_ms: Option, } /// Insertable struct for dead letter entries. @@ -82,6 +86,10 @@ pub struct NewDeadLetterRow<'a> { pub retry_count: i32, pub failed_at: i64, pub metadata: Option<&'a str>, + pub priority: i32, + pub max_retries: i32, + pub timeout_ms: i64, + pub result_ttl_ms: Option, } /// A row in the `rate_limits` table. diff --git a/crates/taskito-core/src/storage/postgres/circuit_breakers.rs b/crates/taskito-core/src/storage/postgres/circuit_breakers.rs new file mode 100644 index 00000000..4d53d7fe --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/circuit_breakers.rs @@ -0,0 +1,40 @@ +use diesel::prelude::*; + +use super::super::models::CircuitBreakerRow; +use super::super::schema::circuit_breakers; +use super::PostgresStorage; +use crate::error::Result; + +impl PostgresStorage { + /// Get or create a circuit breaker state for a task. + pub fn get_circuit_breaker(&self, task_name: &str) -> Result> { + let mut conn = self.conn()?; + let row = circuit_breakers::table + .find(task_name) + .select(CircuitBreakerRow::as_select()) + .first(&mut conn) + .optional()?; + Ok(row) + } + + /// Upsert circuit breaker state. + pub fn upsert_circuit_breaker(&self, row: &CircuitBreakerRow) -> Result<()> { + let mut conn = self.conn()?; + diesel::insert_into(circuit_breakers::table) + .values(row) + .on_conflict(circuit_breakers::task_name) + .do_update() + .set(row) + .execute(&mut conn)?; + Ok(()) + } + + /// Get all circuit breaker states. + pub fn list_circuit_breakers(&self) -> Result> { + let mut conn = self.conn()?; + let rows = circuit_breakers::table + .select(CircuitBreakerRow::as_select()) + .load(&mut conn)?; + Ok(rows) + } +} diff --git a/crates/taskito-core/src/storage/postgres/dead_letter.rs b/crates/taskito-core/src/storage/postgres/dead_letter.rs new file mode 100644 index 00000000..86596794 --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/dead_letter.rs @@ -0,0 +1,143 @@ +use diesel::prelude::*; + +use super::super::models::*; +use super::super::schema::{dead_letter, jobs}; +use super::PostgresStorage; +use crate::error::{QueueError, Result}; +use crate::job::{now_millis, Job, JobStatus, NewJob}; +use crate::storage::DeadJob; + +impl PostgresStorage { + /// Move a job to the dead letter queue and cascade-cancel dependents. + pub fn move_to_dlq(&self, job: &Job, error: &str, metadata: Option<&str>) -> Result<()> { + let now = now_millis(); + let dlq_id = uuid::Uuid::now_v7().to_string(); + let mut conn = self.conn()?; + let job_id = job.id.clone(); + + conn.transaction(|conn| { + let dlq_row = NewDeadLetterRow { + id: &dlq_id, + original_job_id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + error: Some(error), + retry_count: job.retry_count, + failed_at: now, + metadata, + priority: job.priority, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + result_ttl_ms: job.result_ttl_ms, + }; + + diesel::insert_into(dead_letter::table) + .values(&dlq_row) + .execute(conn)?; + + diesel::update(jobs::table) + .filter(jobs::id.eq(&job.id)) + .set(( + jobs::status.eq(JobStatus::Dead as i32), + jobs::error.eq(error), + jobs::completed_at.eq(now), + )) + .execute(conn)?; + + Ok::<(), diesel::result::Error>(()) + })?; + + // Drop connection before cascade (needed for single-connection pools) + drop(conn); + + // Cascade cancel dependents + self.cascade_cancel(&job_id, "dependency failed")?; + + Ok(()) + } + + /// List dead letter entries. + pub fn list_dead(&self, limit: i64, offset: i64) -> Result> { + let mut conn = self.conn()?; + + let rows: Vec = dead_letter::table + .order(dead_letter::failed_at.desc()) + .limit(limit) + .offset(offset) + .select(DeadLetterRow::as_select()) + .load(&mut conn)?; + + Ok(rows.into_iter().map(DeadJob::from).collect()) + } + + /// Re-enqueue a dead letter job. Returns the new job ID. + pub fn retry_dead(&self, dead_id: &str) -> Result { + let mut conn = self.conn()?; + + let dead_row: DeadLetterRow = dead_letter::table + .find(dead_id) + .select(DeadLetterRow::as_select()) + .first(&mut conn) + .map_err(|_| QueueError::JobNotFound(dead_id.to_string()))?; + + let new_job = NewJob { + queue: dead_row.queue, + task_name: dead_row.task_name, + payload: dead_row.payload, + priority: dead_row.priority, + scheduled_at: now_millis(), + max_retries: dead_row.max_retries, + timeout_ms: dead_row.timeout_ms, + unique_key: None, + metadata: dead_row.metadata, + depends_on: vec![], + expires_at: None, + result_ttl_ms: dead_row.result_ttl_ms, + }; + + let job = new_job.into_job(); + + conn.transaction(|conn| { + let row = super::super::models::NewJobRow { + id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + status: job.status as i32, + priority: job.priority, + created_at: job.created_at, + scheduled_at: job.scheduled_at, + retry_count: job.retry_count, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + unique_key: job.unique_key.as_deref(), + metadata: job.metadata.as_deref(), + cancel_requested: 0, + expires_at: job.expires_at, + result_ttl_ms: job.result_ttl_ms, + }; + + diesel::insert_into(jobs::table) + .values(&row) + .execute(conn)?; + + diesel::delete(dead_letter::table.find(dead_id)).execute(conn)?; + + Ok::<(), diesel::result::Error>(()) + })?; + + Ok(job.id) + } + + /// Purge dead letter entries older than the given timestamp. + pub fn purge_dead(&self, older_than_ms: i64) -> Result { + let mut conn = self.conn()?; + + let affected = + diesel::delete(dead_letter::table.filter(dead_letter::failed_at.lt(older_than_ms))) + .execute(&mut conn)?; + + Ok(affected as u64) + } +} diff --git a/crates/taskito-core/src/storage/postgres/jobs.rs b/crates/taskito-core/src/storage/postgres/jobs.rs new file mode 100644 index 00000000..79f1bf80 --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/jobs.rs @@ -0,0 +1,739 @@ +use diesel::pg::PgConnection; +use diesel::prelude::*; + +use super::super::models::*; +use super::super::schema::{ + job_dependencies, job_errors, jobs, replay_history, task_logs, task_metrics, +}; +use super::PostgresStorage; +use crate::error::{QueueError, Result}; +use crate::job::{now_millis, Job, JobStatus, NewJob}; +use crate::storage::QueueStats; + +fn delete_job_children( + conn: &mut PgConnection, + job_ids: &[String], +) -> diesel::result::QueryResult<()> { + if job_ids.is_empty() { + return Ok(()); + } + + diesel::delete(job_errors::table.filter(job_errors::job_id.eq_any(job_ids))).execute(conn)?; + diesel::delete(task_logs::table.filter(task_logs::job_id.eq_any(job_ids))).execute(conn)?; + diesel::delete(task_metrics::table.filter(task_metrics::job_id.eq_any(job_ids))) + .execute(conn)?; + diesel::delete( + job_dependencies::table.filter( + job_dependencies::job_id + .eq_any(job_ids) + .or(job_dependencies::depends_on_job_id.eq_any(job_ids)), + ), + ) + .execute(conn)?; + diesel::delete(replay_history::table.filter(replay_history::original_job_id.eq_any(job_ids))) + .execute(conn)?; + + Ok(()) +} + +impl PostgresStorage { + /// Insert a new job into the queue. Returns the job. + pub fn enqueue(&self, new_job: NewJob) -> Result { + let depends_on = new_job.depends_on.clone(); + let job = new_job.into_job(); + let mut conn = self.conn()?; + + conn.transaction(|conn| { + // Validate dependencies exist and aren't dead/cancelled + for dep_id in &depends_on { + let dep: Option = jobs::table + .find(dep_id) + .select(JobRow::as_select()) + .first(conn) + .optional()?; + + match dep { + None => return Err(diesel::result::Error::RollbackTransaction), + Some(d) + if d.status == JobStatus::Dead as i32 + || d.status == JobStatus::Cancelled as i32 => + { + return Err(diesel::result::Error::RollbackTransaction); + } + _ => {} + } + } + + let row = NewJobRow { + id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + status: job.status as i32, + priority: job.priority, + created_at: job.created_at, + scheduled_at: job.scheduled_at, + retry_count: job.retry_count, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + unique_key: job.unique_key.as_deref(), + metadata: job.metadata.as_deref(), + cancel_requested: 0, + expires_at: job.expires_at, + result_ttl_ms: job.result_ttl_ms, + }; + + diesel::insert_into(jobs::table) + .values(&row) + .execute(conn)?; + + // Insert dependency rows + for dep_id in &depends_on { + let dep_row = NewJobDependencyRow { + id: &uuid::Uuid::now_v7().to_string(), + job_id: &job.id, + depends_on_job_id: dep_id, + }; + diesel::insert_into(job_dependencies::table) + .values(&dep_row) + .execute(conn)?; + } + + Ok(()) + }) + .map_err(|e| match e { + diesel::result::Error::RollbackTransaction => QueueError::DependencyNotFound( + "dependency not found or already dead/cancelled".to_string(), + ), + other => QueueError::Storage(other), + })?; + + Ok(job) + } + + /// Enqueue multiple jobs in a single transaction. + pub fn enqueue_batch(&self, new_jobs: Vec) -> Result> { + let mut conn = self.conn()?; + let jobs: Vec = new_jobs.into_iter().map(|nj| nj.into_job()).collect(); + + conn.transaction(|conn| { + for job in &jobs { + let row = NewJobRow { + id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + status: job.status as i32, + priority: job.priority, + created_at: job.created_at, + scheduled_at: job.scheduled_at, + retry_count: job.retry_count, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + unique_key: job.unique_key.as_deref(), + metadata: job.metadata.as_deref(), + cancel_requested: 0, + expires_at: job.expires_at, + result_ttl_ms: job.result_ttl_ms, + }; + + diesel::insert_into(jobs::table) + .values(&row) + .execute(conn)?; + } + Ok(jobs) + }) + } + + /// Enqueue with unique_key deduplication. Returns existing job if duplicate. + pub fn enqueue_unique(&self, new_job: NewJob) -> Result { + let depends_on = new_job.depends_on.clone(); + let job = new_job.into_job(); + let mut conn = self.conn()?; + + let result = conn.transaction(|conn| { + // Check for existing active job with same unique_key + if let Some(ref uk) = job.unique_key { + let existing: Option = jobs::table + .filter(jobs::unique_key.eq(uk)) + .filter( + jobs::status.eq_any([JobStatus::Pending as i32, JobStatus::Running as i32]), + ) + .select(JobRow::as_select()) + .first(conn) + .optional()?; + + if let Some(row) = existing { + return Ok(Job::from(row)); + } + } + + let row = NewJobRow { + id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + status: job.status as i32, + priority: job.priority, + created_at: job.created_at, + scheduled_at: job.scheduled_at, + retry_count: job.retry_count, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + unique_key: job.unique_key.as_deref(), + metadata: job.metadata.as_deref(), + cancel_requested: 0, + expires_at: job.expires_at, + result_ttl_ms: job.result_ttl_ms, + }; + + diesel::insert_into(jobs::table) + .values(&row) + .execute(conn)?; + + // Insert dependency rows + for dep_id in &depends_on { + let dep_row = NewJobDependencyRow { + id: &uuid::Uuid::now_v7().to_string(), + job_id: &job.id, + depends_on_job_id: dep_id, + }; + diesel::insert_into(job_dependencies::table) + .values(&dep_row) + .execute(conn)?; + } + + Ok(job.clone()) + }); + + // Handle unique constraint violation by returning existing job + match result { + Ok(j) => Ok(j), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::UniqueViolation, + _, + )) => { + if let Some(ref uk) = job.unique_key { + let mut conn = self.conn()?; + let existing: Option = jobs::table + .filter(jobs::unique_key.eq(uk)) + .filter( + jobs::status + .eq_any([JobStatus::Pending as i32, JobStatus::Running as i32]), + ) + .select(JobRow::as_select()) + .first(&mut conn) + .optional()?; + if let Some(row) = existing { + return Ok(Job::from(row)); + } + } + Ok(job) + } + Err(e) => Err(QueueError::Storage(e)), + } + } + + /// Check if all dependencies for a job are complete. + fn deps_satisfied(conn: &mut PgConnection, job_id: &str) -> diesel::result::QueryResult { + let dep_job_ids: Vec = job_dependencies::table + .filter(job_dependencies::job_id.eq(job_id)) + .select(job_dependencies::depends_on_job_id) + .load(conn)?; + + if dep_job_ids.is_empty() { + return Ok(true); + } + + let incomplete: i64 = jobs::table + .filter(jobs::id.eq_any(&dep_job_ids)) + .filter(jobs::status.ne(JobStatus::Complete as i32)) + .count() + .get_result(conn)?; + + Ok(incomplete == 0) + } + + /// Atomically dequeue the highest-priority ready job from the given queue. + /// Skips expired jobs. + pub fn dequeue(&self, queue_name: &str, now: i64) -> Result> { + let mut conn = self.conn()?; + + conn.transaction(|conn| { + let candidates: Vec = jobs::table + .filter(jobs::queue.eq(queue_name)) + .filter(jobs::status.eq(JobStatus::Pending as i32)) + .filter(jobs::scheduled_at.le(now)) + .order((jobs::priority.desc(), jobs::scheduled_at.asc())) + .limit(100) + .select(JobRow::as_select()) + .load(conn)?; + + for row in candidates { + // Skip expired jobs + if let Some(expires_at) = row.expires_at { + if now > expires_at { + // Mark as cancelled (expired) + diesel::update(jobs::table) + .filter(jobs::id.eq(&row.id)) + .set(( + jobs::status.eq(JobStatus::Cancelled as i32), + jobs::completed_at.eq(now), + jobs::error.eq("expired before execution"), + )) + .execute(conn)?; + continue; + } + } + + if !Self::deps_satisfied(conn, &row.id)? { + continue; + } + + diesel::update(jobs::table) + .filter(jobs::id.eq(&row.id)) + .filter(jobs::status.eq(JobStatus::Pending as i32)) + .set(( + jobs::status.eq(JobStatus::Running as i32), + jobs::started_at.eq(now), + )) + .execute(conn)?; + + let updated: JobRow = jobs::table + .find(&row.id) + .select(JobRow::as_select()) + .first(conn)?; + + return Ok(Some(Job::from(updated))); + } + + Ok(None) + }) + } + + /// Dequeue from multiple queues, checking each in order. + pub fn dequeue_from(&self, queues: &[String], now: i64) -> Result> { + for queue_name in queues { + if let Some(job) = self.dequeue(queue_name, now)? { + return Ok(Some(job)); + } + } + Ok(None) + } + + /// Mark a job as complete with the given result. + pub fn complete(&self, id: &str, result_bytes: Option>) -> Result<()> { + let now = now_millis(); + let mut conn = self.conn()?; + + let affected = diesel::update(jobs::table) + .filter(jobs::id.eq(id)) + .filter(jobs::status.eq(JobStatus::Running as i32)) + .set(( + jobs::status.eq(JobStatus::Complete as i32), + jobs::completed_at.eq(now), + jobs::result.eq(result_bytes), + )) + .execute(&mut conn)?; + + if affected == 0 { + return Err(QueueError::JobNotFound(id.to_string())); + } + Ok(()) + } + + /// Mark a job as failed with the given error message. + pub fn fail(&self, id: &str, error: &str) -> Result<()> { + let now = now_millis(); + let mut conn = self.conn()?; + + let affected = diesel::update(jobs::table) + .filter(jobs::id.eq(id)) + .filter(jobs::status.eq(JobStatus::Running as i32)) + .set(( + jobs::status.eq(JobStatus::Failed as i32), + jobs::completed_at.eq(now), + jobs::error.eq(error), + )) + .execute(&mut conn)?; + + if affected == 0 { + return Err(QueueError::JobNotFound(id.to_string())); + } + Ok(()) + } + + /// Re-schedule a job for retry. + pub fn retry(&self, id: &str, next_scheduled_at: i64) -> Result<()> { + let mut conn = self.conn()?; + + let affected = diesel::update(jobs::table) + .filter(jobs::id.eq(id)) + .set(( + jobs::status.eq(JobStatus::Pending as i32), + jobs::scheduled_at.eq(next_scheduled_at), + jobs::retry_count.eq(jobs::retry_count + 1), + jobs::started_at.eq(None::), + jobs::completed_at.eq(None::), + jobs::error.eq(None::), + )) + .execute(&mut conn)?; + + if affected == 0 { + return Err(QueueError::JobNotFound(id.to_string())); + } + Ok(()) + } + + /// Cancel a pending job and cascade-cancel its dependents. + pub fn cancel_job(&self, id: &str) -> Result { + let now = now_millis(); + let mut conn = self.conn()?; + + let affected = diesel::update(jobs::table) + .filter(jobs::id.eq(id)) + .filter(jobs::status.eq(JobStatus::Pending as i32)) + .set(( + jobs::status.eq(JobStatus::Cancelled as i32), + jobs::completed_at.eq(now), + )) + .execute(&mut conn)?; + + drop(conn); + if affected > 0 { + self.cascade_cancel(id, "dependency cancelled")?; + } + + Ok(affected > 0) + } + + /// Request cancellation of a running job. The task must check for this. + pub fn request_cancel(&self, id: &str) -> Result { + let mut conn = self.conn()?; + + let affected = diesel::update(jobs::table) + .filter(jobs::id.eq(id)) + .filter(jobs::status.eq(JobStatus::Running as i32)) + .set(jobs::cancel_requested.eq(1)) + .execute(&mut conn)?; + + Ok(affected > 0) + } + + /// Check if cancellation has been requested for a job. + pub fn is_cancel_requested(&self, id: &str) -> Result { + let mut conn = self.conn()?; + + let val: Option = jobs::table + .find(id) + .select(jobs::cancel_requested) + .first(&mut conn) + .optional()?; + + Ok(val.unwrap_or(0) != 0) + } + + /// Mark a job as cancelled (used when a running job detects cancellation). + pub fn mark_cancelled(&self, id: &str) -> Result<()> { + let now = now_millis(); + let mut conn = self.conn()?; + + diesel::update(jobs::table) + .filter(jobs::id.eq(id)) + .set(( + jobs::status.eq(JobStatus::Cancelled as i32), + jobs::completed_at.eq(now), + jobs::error.eq("cancelled by request"), + )) + .execute(&mut conn)?; + + Ok(()) + } + + /// Cascade-cancel all pending jobs that depend (directly or transitively) + /// on the given job. Uses BFS to handle deep chains. + pub fn cascade_cancel(&self, failed_job_id: &str, reason: &str) -> Result<()> { + let mut conn = self.conn()?; + let now = now_millis(); + + let mut queue: Vec = vec![failed_job_id.to_string()]; + let mut visited = std::collections::HashSet::new(); + visited.insert(failed_job_id.to_string()); + let mut idx = 0; + + while idx < queue.len() { + let current_id = queue[idx].clone(); + idx += 1; + + let dependents: Vec = job_dependencies::table + .filter(job_dependencies::depends_on_job_id.eq(¤t_id)) + .select(job_dependencies::job_id) + .load(&mut conn)?; + + for dep_id in dependents { + if visited.insert(dep_id.clone()) { + queue.push(dep_id); + } + } + } + + // Remove the original job from the list (only cancel dependents) + if !queue.is_empty() { + queue.remove(0); + } + + if !queue.is_empty() { + let error_msg = format!("{reason}: {failed_job_id}"); + diesel::update(jobs::table) + .filter(jobs::id.eq_any(&queue)) + .filter(jobs::status.eq(JobStatus::Pending as i32)) + .set(( + jobs::status.eq(JobStatus::Cancelled as i32), + jobs::completed_at.eq(now), + jobs::error.eq(&error_msg), + )) + .execute(&mut conn)?; + } + + Ok(()) + } + + /// Get the IDs of jobs that a given job depends on. + pub fn get_dependencies(&self, job_id: &str) -> Result> { + let mut conn = self.conn()?; + let ids: Vec = job_dependencies::table + .filter(job_dependencies::job_id.eq(job_id)) + .select(job_dependencies::depends_on_job_id) + .load(&mut conn)?; + Ok(ids) + } + + /// Get the IDs of jobs that depend on a given job. + pub fn get_dependents(&self, job_id: &str) -> Result> { + let mut conn = self.conn()?; + let ids: Vec = job_dependencies::table + .filter(job_dependencies::depends_on_job_id.eq(job_id)) + .select(job_dependencies::job_id) + .load(&mut conn)?; + Ok(ids) + } + + /// Update progress for a running job (0-100). + pub fn update_progress(&self, id: &str, progress: i32) -> Result<()> { + let mut conn = self.conn()?; + + let affected = diesel::update(jobs::table) + .filter(jobs::id.eq(id)) + .set(jobs::progress.eq(progress)) + .execute(&mut conn)?; + + if affected == 0 { + return Err(QueueError::JobNotFound(id.to_string())); + } + Ok(()) + } + + /// List jobs with optional filters and pagination. + pub fn list_jobs( + &self, + status: Option, + queue_name: Option<&str>, + task_name: Option<&str>, + limit: i64, + offset: i64, + ) -> Result> { + let mut conn = self.conn()?; + + let mut query = jobs::table.into_boxed().order(jobs::created_at.desc()); + + if let Some(s) = status { + query = query.filter(jobs::status.eq(s)); + } + if let Some(q) = queue_name { + query = query.filter(jobs::queue.eq(q)); + } + if let Some(t) = task_name { + query = query.filter(jobs::task_name.eq(t)); + } + + let rows: Vec = query + .limit(limit) + .offset(offset) + .select(JobRow::as_select()) + .load(&mut conn)?; + + Ok(rows.into_iter().map(Job::from).collect()) + } + + /// Get a job by ID. + pub fn get_job(&self, id: &str) -> Result> { + let mut conn = self.conn()?; + + let row: Option = jobs::table + .find(id) + .select(JobRow::as_select()) + .first(&mut conn) + .optional()?; + + Ok(row.map(Job::from)) + } + + /// Get queue statistics. + pub fn stats(&self) -> Result { + let mut conn = self.conn()?; + + let rows: Vec<(i32, i64)> = jobs::table + .group_by(jobs::status) + .select((jobs::status, diesel::dsl::count(jobs::id))) + .load(&mut conn)?; + + let mut stats = QueueStats::default(); + for (status, count) in rows { + match JobStatus::from_i32(status) { + Some(JobStatus::Pending) => stats.pending = count, + Some(JobStatus::Running) => stats.running = count, + Some(JobStatus::Complete) => stats.completed = count, + Some(JobStatus::Failed) => stats.failed = count, + Some(JobStatus::Dead) => stats.dead = count, + Some(JobStatus::Cancelled) => stats.cancelled = count, + None => {} + } + } + + Ok(stats) + } + + /// Purge completed jobs older than the given timestamp. + pub fn purge_completed(&self, older_than_ms: i64) -> Result { + let mut conn = self.conn()?; + + conn.transaction(|conn| { + let job_ids: Vec = jobs::table + .filter(jobs::status.eq(JobStatus::Complete as i32)) + .filter(jobs::completed_at.lt(older_than_ms)) + .select(jobs::id) + .load(conn)?; + + delete_job_children(conn, &job_ids)?; + + let affected = + diesel::delete(jobs::table.filter(jobs::id.eq_any(&job_ids))).execute(conn)?; + + Ok(affected as u64) + }) + } + + /// Purge completed jobs respecting per-job result_ttl_ms. + /// Jobs with their own result_ttl_ms use that; others use the given global cutoff. + pub fn purge_completed_with_ttl(&self, global_cutoff_ms: i64) -> Result { + let now = now_millis(); + let mut conn = self.conn()?; + + conn.transaction(|conn| { + // For global TTL: collect IDs of completed jobs without per-job TTL + let global_ids: Vec = jobs::table + .filter(jobs::status.eq(JobStatus::Complete as i32)) + .filter(jobs::result_ttl_ms.is_null()) + .filter(jobs::completed_at.lt(global_cutoff_ms)) + .select(jobs::id) + .load(conn)?; + + // For per-job TTL: fetch candidates and collect expired IDs + let rows_with_ttl: Vec = jobs::table + .filter(jobs::status.eq(JobStatus::Complete as i32)) + .filter(jobs::result_ttl_ms.is_not_null()) + .select(JobRow::as_select()) + .load(conn)?; + + let per_job_ids: Vec = rows_with_ttl + .into_iter() + .filter(|row| { + matches!((row.completed_at, row.result_ttl_ms), (Some(completed), Some(ttl)) if completed + ttl < now) + }) + .map(|row| row.id) + .collect(); + + let all_ids: Vec = global_ids.into_iter().chain(per_job_ids).collect(); + + delete_job_children(conn, &all_ids)?; + + let affected = diesel::delete( + jobs::table.filter(jobs::id.eq_any(&all_ids)) + ).execute(conn)?; + + Ok(affected as u64) + }) + } + + /// Find stale running jobs that exceeded their timeout. + pub fn reap_stale_jobs(&self, now: i64) -> Result> { + let mut conn = self.conn()?; + + let rows: Vec = jobs::table + .filter(jobs::status.eq(JobStatus::Running as i32)) + .filter(jobs::started_at.is_not_null()) + .select(JobRow::as_select()) + .load(&mut conn)?; + + let stale: Vec = rows + .into_iter() + .filter(|r| { + if let Some(started) = r.started_at { + (started + r.timeout_ms) < now + } else { + false + } + }) + .map(Job::from) + .collect(); + + Ok(stale) + } + + /// Record an error for a job attempt. + pub fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()> { + let mut conn = self.conn()?; + let id = uuid::Uuid::now_v7().to_string(); + let now = now_millis(); + + let row = NewJobErrorRow { + id: &id, + job_id, + attempt, + error, + failed_at: now, + }; + + diesel::insert_into(super::super::schema::job_errors::table) + .values(&row) + .execute(&mut conn)?; + + Ok(()) + } + + /// Get all errors for a job, ordered by attempt. + pub fn get_job_errors(&self, job_id: &str) -> Result> { + use super::super::schema::job_errors; + let mut conn = self.conn()?; + + let rows = job_errors::table + .filter(job_errors::job_id.eq(job_id)) + .order(job_errors::attempt.asc()) + .select(JobErrorRow::as_select()) + .load(&mut conn)?; + + Ok(rows) + } + + /// Purge job errors older than the given timestamp. + pub fn purge_job_errors(&self, older_than_ms: i64) -> Result { + use super::super::schema::job_errors; + let mut conn = self.conn()?; + + let affected = + diesel::delete(job_errors::table.filter(job_errors::failed_at.lt(older_than_ms))) + .execute(&mut conn)?; + + Ok(affected as u64) + } +} diff --git a/crates/taskito-core/src/storage/postgres/logs.rs b/crates/taskito-core/src/storage/postgres/logs.rs new file mode 100644 index 00000000..d6216745 --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/logs.rs @@ -0,0 +1,89 @@ +use diesel::prelude::*; + +use super::super::models::*; +use super::super::schema::task_logs; +use super::PostgresStorage; +use crate::error::Result; +use crate::job::now_millis; + +impl PostgresStorage { + /// Write a structured log entry for a task. + pub fn write_task_log( + &self, + job_id: &str, + task_name: &str, + level: &str, + message: &str, + extra: Option<&str>, + ) -> Result<()> { + let mut conn = self.conn()?; + let id = uuid::Uuid::now_v7().to_string(); + let now = now_millis(); + + let row = NewTaskLogRow { + id: &id, + job_id, + task_name, + level, + message, + extra, + logged_at: now, + }; + + diesel::insert_into(task_logs::table) + .values(&row) + .execute(&mut conn)?; + + Ok(()) + } + + /// Get logs for a specific job. + pub fn get_task_logs(&self, job_id: &str) -> Result> { + let mut conn = self.conn()?; + let rows = task_logs::table + .filter(task_logs::job_id.eq(job_id)) + .order(task_logs::logged_at.asc()) + .select(TaskLogRow::as_select()) + .load(&mut conn)?; + Ok(rows) + } + + /// Query logs by task name, level, etc. + pub fn query_task_logs( + &self, + task_name: Option<&str>, + level: Option<&str>, + since_ms: i64, + limit: i64, + ) -> Result> { + let mut conn = self.conn()?; + + let mut query = task_logs::table + .filter(task_logs::logged_at.ge(since_ms)) + .into_boxed(); + + if let Some(n) = task_name { + query = query.filter(task_logs::task_name.eq(n)); + } + if let Some(l) = level { + query = query.filter(task_logs::level.eq(l)); + } + + let rows = query + .order(task_logs::logged_at.desc()) + .limit(limit) + .select(TaskLogRow::as_select()) + .load(&mut conn)?; + + Ok(rows) + } + + /// Purge old log records. + pub fn purge_task_logs(&self, older_than_ms: i64) -> Result { + let mut conn = self.conn()?; + let affected = + diesel::delete(task_logs::table.filter(task_logs::logged_at.lt(older_than_ms))) + .execute(&mut conn)?; + Ok(affected as u64) + } +} diff --git a/crates/taskito-core/src/storage/postgres/metrics.rs b/crates/taskito-core/src/storage/postgres/metrics.rs new file mode 100644 index 00000000..8546b291 --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/metrics.rs @@ -0,0 +1,115 @@ +use diesel::prelude::*; + +use super::super::models::*; +use super::super::schema::task_metrics; +use super::PostgresStorage; +use crate::error::Result; +use crate::job::now_millis; + +impl PostgresStorage { + /// Record a task execution metric. + pub fn record_metric( + &self, + task_name: &str, + job_id: &str, + wall_time_ns: i64, + memory_bytes: i64, + succeeded: bool, + ) -> Result<()> { + let mut conn = self.conn()?; + let id = uuid::Uuid::now_v7().to_string(); + let now = now_millis(); + + let row = NewTaskMetricRow { + id: &id, + task_name, + job_id, + wall_time_ns, + memory_bytes, + succeeded, + recorded_at: now, + }; + + diesel::insert_into(task_metrics::table) + .values(&row) + .execute(&mut conn)?; + + Ok(()) + } + + /// Get aggregated metrics for a task (or all tasks if name is None). + pub fn get_metrics(&self, name: Option<&str>, since_ms: i64) -> Result> { + let mut conn = self.conn()?; + + let mut query = task_metrics::table + .filter(task_metrics::recorded_at.ge(since_ms)) + .into_boxed(); + + if let Some(n) = name { + query = query.filter(task_metrics::task_name.eq(n)); + } + + let rows = query + .order(task_metrics::recorded_at.desc()) + .select(TaskMetricRow::as_select()) + .load(&mut conn)?; + + Ok(rows) + } + + /// Purge old metric records. + pub fn purge_metrics(&self, older_than_ms: i64) -> Result { + let mut conn = self.conn()?; + let affected = + diesel::delete(task_metrics::table.filter(task_metrics::recorded_at.lt(older_than_ms))) + .execute(&mut conn)?; + Ok(affected as u64) + } + + /// Record a replay event. + pub fn record_replay( + &self, + original_job_id: &str, + replay_job_id: &str, + original_result: Option<&[u8]>, + replay_result: Option<&[u8]>, + original_error: Option<&str>, + replay_error: Option<&str>, + ) -> Result<()> { + use super::super::schema::replay_history; + + let mut conn = self.conn()?; + let id = uuid::Uuid::now_v7().to_string(); + let now = now_millis(); + + let row = NewReplayHistoryRow { + id: &id, + original_job_id, + replay_job_id, + replayed_at: now, + original_result, + replay_result, + original_error, + replay_error, + }; + + diesel::insert_into(replay_history::table) + .values(&row) + .execute(&mut conn)?; + + Ok(()) + } + + /// Get replay history for a job. + pub fn get_replay_history(&self, original_job_id: &str) -> Result> { + use super::super::schema::replay_history; + + let mut conn = self.conn()?; + let rows = replay_history::table + .filter(replay_history::original_job_id.eq(original_job_id)) + .order(replay_history::replayed_at.desc()) + .select(ReplayHistoryRow::as_select()) + .load(&mut conn)?; + Ok(rows) + } +} diff --git a/crates/taskito-core/src/storage/postgres/mod.rs b/crates/taskito-core/src/storage/postgres/mod.rs new file mode 100644 index 00000000..e2be3251 --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/mod.rs @@ -0,0 +1,330 @@ +mod circuit_breakers; +mod dead_letter; +mod jobs; +mod logs; +mod metrics; +mod periodic; +mod rate_limits; +mod trait_impl; +mod workers; + +use diesel::pg::PgConnection; +use diesel::prelude::*; +use diesel::r2d2::{self, ConnectionManager, Pool}; + +use crate::error::Result; + +type PgPool = Pool>; + +/// Validate a PostgreSQL schema name (alphanumeric + underscores, non-empty). +fn validate_schema_name(schema: &str) -> Result<()> { + if schema.is_empty() { + return Err(crate::error::QueueError::Config( + "Schema name cannot be empty".into(), + )); + } + if !schema + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_') + { + return Err(crate::error::QueueError::Config( + format!("Invalid schema name '{schema}': only alphanumeric characters and underscores are allowed"), + )); + } + Ok(()) +} + +/// Sets `search_path` on every new pooled connection. +#[derive(Debug)] +struct SetSearchPath { + schema: String, +} + +impl r2d2::CustomizeConnection for SetSearchPath { + fn on_acquire(&self, conn: &mut PgConnection) -> std::result::Result<(), r2d2::Error> { + diesel::sql_query(format!("SET search_path TO {}", self.schema)) + .execute(conn) + .map_err(r2d2::Error::QueryError)?; + Ok(()) + } +} + +/// PostgreSQL-backed storage for the task queue, using Diesel ORM. +#[derive(Clone)] +pub struct PostgresStorage { + pool: PgPool, + schema: String, +} + +impl PostgresStorage { + /// Connect to a PostgreSQL database at the given URL. + /// Tables are created in the `taskito` schema by default. + pub fn new(database_url: &str) -> Result { + Self::with_schema(database_url, "taskito") + } + + /// Connect with a custom schema name. + pub fn with_schema(database_url: &str, schema: &str) -> Result { + Self::build(database_url, 10, schema) + } + + /// Connect with a custom connection pool size and schema. + pub fn with_pool_size(database_url: &str, pool_size: u32) -> Result { + Self::build(database_url, pool_size, "taskito") + } + + fn build(database_url: &str, pool_size: u32, schema: &str) -> Result { + validate_schema_name(schema)?; + + let manager = ConnectionManager::::new(database_url); + let customizer = SetSearchPath { + schema: schema.to_string(), + }; + let pool = Pool::builder() + .max_size(pool_size) + .connection_customizer(Box::new(customizer)) + .build(manager)?; + + let storage = Self { + pool, + schema: schema.to_string(), + }; + storage.run_migrations()?; + Ok(storage) + } + + pub(crate) fn conn( + &self, + ) -> Result>> { + Ok(self.pool.get()?) + } + + fn run_migrations(&self) -> Result<()> { + let mut conn = self.conn()?; + + // Ensure the schema exists before creating tables + diesel::sql_query(format!("CREATE SCHEMA IF NOT EXISTS {}", self.schema)) + .execute(&mut conn)?; + + // Use PG-native types: TEXT, BYTEA, BIGINT, INTEGER, BOOLEAN, DOUBLE PRECISION + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS jobs ( + id TEXT PRIMARY KEY, + queue TEXT NOT NULL DEFAULT 'default', + task_name TEXT NOT NULL, + payload BYTEA NOT NULL, + status INTEGER NOT NULL DEFAULT 0, + priority INTEGER NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + scheduled_at BIGINT NOT NULL, + started_at BIGINT, + completed_at BIGINT, + retry_count INTEGER NOT NULL DEFAULT 0, + max_retries INTEGER NOT NULL DEFAULT 3, + result BYTEA, + error TEXT, + timeout_ms BIGINT NOT NULL DEFAULT 300000, + unique_key TEXT, + progress INTEGER, + metadata TEXT, + cancel_requested INTEGER NOT NULL DEFAULT 0, + expires_at BIGINT, + result_ttl_ms BIGINT + )", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE INDEX IF NOT EXISTS idx_jobs_dequeue + ON jobs(queue, status, priority DESC, scheduled_at)", + ) + .execute(&mut conn)?; + + diesel::sql_query("CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status)") + .execute(&mut conn)?; + + // PG partial unique index + diesel::sql_query( + "CREATE UNIQUE INDEX IF NOT EXISTS idx_jobs_unique_key + ON jobs(unique_key) WHERE unique_key IS NOT NULL AND status IN (0, 1)", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS dead_letter ( + id TEXT PRIMARY KEY, + original_job_id TEXT NOT NULL, + queue TEXT NOT NULL, + task_name TEXT NOT NULL, + payload BYTEA NOT NULL, + error TEXT, + retry_count INTEGER NOT NULL, + failed_at BIGINT NOT NULL, + metadata TEXT, + priority INTEGER NOT NULL DEFAULT 0, + max_retries INTEGER NOT NULL DEFAULT 3, + timeout_ms BIGINT NOT NULL DEFAULT 300000, + result_ttl_ms BIGINT + )", + ) + .execute(&mut conn)?; + + // Migration: add columns if they don't exist (for existing databases) + for col in &[ + "ALTER TABLE dead_letter ADD COLUMN IF NOT EXISTS priority INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE dead_letter ADD COLUMN IF NOT EXISTS max_retries INTEGER NOT NULL DEFAULT 3", + "ALTER TABLE dead_letter ADD COLUMN IF NOT EXISTS timeout_ms BIGINT NOT NULL DEFAULT 300000", + "ALTER TABLE dead_letter ADD COLUMN IF NOT EXISTS result_ttl_ms BIGINT", + ] { + let _ = diesel::sql_query(*col).execute(&mut conn); + } + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS rate_limits ( + key TEXT PRIMARY KEY, + tokens DOUBLE PRECISION NOT NULL, + max_tokens DOUBLE PRECISION NOT NULL, + refill_rate DOUBLE PRECISION NOT NULL, + last_refill BIGINT NOT NULL + )", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS periodic_tasks ( + name TEXT PRIMARY KEY, + task_name TEXT NOT NULL, + cron_expr TEXT NOT NULL, + args BYTEA, + kwargs BYTEA, + queue TEXT NOT NULL DEFAULT 'default', + enabled BOOLEAN NOT NULL DEFAULT TRUE, + last_run BIGINT, + next_run BIGINT NOT NULL + )", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS job_errors ( + id TEXT PRIMARY KEY, + job_id TEXT NOT NULL, + attempt INTEGER NOT NULL, + error TEXT NOT NULL, + failed_at BIGINT NOT NULL + )", + ) + .execute(&mut conn)?; + + diesel::sql_query("CREATE INDEX IF NOT EXISTS idx_job_errors_job_id ON job_errors(job_id)") + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS job_dependencies ( + id TEXT PRIMARY KEY, + job_id TEXT NOT NULL, + depends_on_job_id TEXT NOT NULL + )", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE INDEX IF NOT EXISTS idx_job_deps_job_id ON job_dependencies(job_id)", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE INDEX IF NOT EXISTS idx_job_deps_depends_on ON job_dependencies(depends_on_job_id)" + ).execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS task_metrics ( + id TEXT PRIMARY KEY, + task_name TEXT NOT NULL, + job_id TEXT NOT NULL, + wall_time_ns BIGINT NOT NULL, + memory_bytes BIGINT NOT NULL DEFAULT 0, + succeeded BOOLEAN NOT NULL DEFAULT TRUE, + recorded_at BIGINT NOT NULL + )", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE INDEX IF NOT EXISTS idx_task_metrics_task_name ON task_metrics(task_name)", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE INDEX IF NOT EXISTS idx_task_metrics_recorded_at ON task_metrics(recorded_at)", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS replay_history ( + id TEXT PRIMARY KEY, + original_job_id TEXT NOT NULL, + replay_job_id TEXT NOT NULL, + replayed_at BIGINT NOT NULL, + original_result BYTEA, + replay_result BYTEA, + original_error TEXT, + replay_error TEXT + )", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE INDEX IF NOT EXISTS idx_replay_original ON replay_history(original_job_id)", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS task_logs ( + id TEXT PRIMARY KEY, + job_id TEXT NOT NULL, + task_name TEXT NOT NULL, + level TEXT NOT NULL DEFAULT 'info', + message TEXT NOT NULL, + extra TEXT, + logged_at BIGINT NOT NULL + )", + ) + .execute(&mut conn)?; + + diesel::sql_query("CREATE INDEX IF NOT EXISTS idx_task_logs_job_id ON task_logs(job_id)") + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE INDEX IF NOT EXISTS idx_task_logs_recorded ON task_logs(logged_at)", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS circuit_breakers ( + task_name TEXT PRIMARY KEY, + state INTEGER NOT NULL DEFAULT 0, + failure_count INTEGER NOT NULL DEFAULT 0, + last_failure_at BIGINT, + opened_at BIGINT, + half_open_at BIGINT, + threshold INTEGER NOT NULL DEFAULT 5, + window_ms BIGINT NOT NULL DEFAULT 60000, + cooldown_ms BIGINT NOT NULL DEFAULT 300000 + )", + ) + .execute(&mut conn)?; + + diesel::sql_query( + "CREATE TABLE IF NOT EXISTS workers ( + worker_id TEXT PRIMARY KEY, + last_heartbeat BIGINT NOT NULL, + queues TEXT NOT NULL DEFAULT 'default', + status TEXT NOT NULL DEFAULT 'active' + )", + ) + .execute(&mut conn)?; + + Ok(()) + } +} diff --git a/crates/taskito-core/src/storage/postgres/periodic.rs b/crates/taskito-core/src/storage/postgres/periodic.rs new file mode 100644 index 00000000..6a825283 --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/periodic.rs @@ -0,0 +1,49 @@ +use diesel::prelude::*; + +use super::super::models::*; +use super::super::schema::periodic_tasks; +use super::PostgresStorage; +use crate::error::Result; + +impl PostgresStorage { + /// Register or update a periodic task. + pub fn register_periodic(&self, task: &NewPeriodicTaskRow) -> Result<()> { + let mut conn = self.conn()?; + + diesel::insert_into(periodic_tasks::table) + .values(task) + .on_conflict(periodic_tasks::name) + .do_update() + .set(task) + .execute(&mut conn)?; + + Ok(()) + } + + /// Get all periodic tasks that are due to run. + pub fn get_due_periodic(&self, now: i64) -> Result> { + let mut conn = self.conn()?; + + let rows = periodic_tasks::table + .filter(periodic_tasks::enabled.eq(true)) + .filter(periodic_tasks::next_run.le(now)) + .select(PeriodicTaskRow::as_select()) + .load(&mut conn)?; + + Ok(rows) + } + + /// Update a periodic task's schedule after execution. + pub fn update_periodic_schedule(&self, name: &str, last_run: i64, next_run: i64) -> Result<()> { + let mut conn = self.conn()?; + + diesel::update(periodic_tasks::table.find(name)) + .set(( + periodic_tasks::last_run.eq(last_run), + periodic_tasks::next_run.eq(next_run), + )) + .execute(&mut conn)?; + + Ok(()) + } +} diff --git a/crates/taskito-core/src/storage/postgres/rate_limits.rs b/crates/taskito-core/src/storage/postgres/rate_limits.rs new file mode 100644 index 00000000..8e0abd9a --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/rate_limits.rs @@ -0,0 +1,85 @@ +use diesel::prelude::*; + +use super::super::models::RateLimitRow; +use super::super::schema::rate_limits; +use super::PostgresStorage; +use crate::error::Result; +use crate::job::now_millis; + +impl PostgresStorage { + pub fn get_rate_limit(&self, key: &str) -> Result> { + let mut conn = self.conn()?; + + let row: Option = rate_limits::table + .find(key) + .select(RateLimitRow::as_select()) + .first(&mut conn) + .optional()?; + + Ok(row) + } + + pub fn upsert_rate_limit(&self, row: &RateLimitRow) -> Result<()> { + let mut conn = self.conn()?; + + diesel::insert_into(rate_limits::table) + .values(row) + .on_conflict(rate_limits::key) + .do_update() + .set(row) + .execute(&mut conn)?; + + Ok(()) + } + + /// Atomically try to acquire a rate limit token. + /// Does the read-refill-consume-write in a single transaction to prevent + /// race conditions between concurrent workers. + pub fn try_acquire_token(&self, key: &str, max_tokens: f64, refill_rate: f64) -> Result { + let mut conn = self.conn()?; + let now = now_millis(); + + conn.transaction(|conn| { + let existing: Option = rate_limits::table + .find(key) + .select(RateLimitRow::as_select()) + .first(conn) + .optional()?; + + let mut row = match existing { + Some(r) => r, + None => RateLimitRow { + key: key.to_string(), + tokens: max_tokens, + max_tokens, + refill_rate, + last_refill: now, + }, + }; + + // Refill tokens based on elapsed time + let elapsed_ms = (now - row.last_refill).max(0) as f64; + let elapsed_sec = elapsed_ms / 1000.0; + let refilled = row.tokens + elapsed_sec * refill_rate; + row.tokens = refilled.min(max_tokens); + row.last_refill = now; + + // Try to consume one token + let acquired = if row.tokens >= 1.0 { + row.tokens -= 1.0; + true + } else { + false + }; + + diesel::insert_into(rate_limits::table) + .values(&row) + .on_conflict(rate_limits::key) + .do_update() + .set(&row) + .execute(conn)?; + + Ok(acquired) + }) + } +} diff --git a/crates/taskito-core/src/storage/postgres/trait_impl.rs b/crates/taskito-core/src/storage/postgres/trait_impl.rs new file mode 100644 index 00000000..698f10cd --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/trait_impl.rs @@ -0,0 +1,207 @@ +use super::PostgresStorage; +use crate::error::Result; +use crate::job::{Job, NewJob}; +use crate::storage::models::*; +use crate::storage::traits::Storage; +use crate::storage::{DeadJob, QueueStats}; + +impl Storage for PostgresStorage { + fn enqueue(&self, new_job: NewJob) -> Result { + self.enqueue(new_job) + } + fn enqueue_batch(&self, new_jobs: Vec) -> Result> { + self.enqueue_batch(new_jobs) + } + fn enqueue_unique(&self, new_job: NewJob) -> Result { + self.enqueue_unique(new_job) + } + fn dequeue(&self, queue_name: &str, now: i64) -> Result> { + self.dequeue(queue_name, now) + } + fn dequeue_from(&self, queues: &[String], now: i64) -> Result> { + self.dequeue_from(queues, now) + } + fn complete(&self, id: &str, result_bytes: Option>) -> Result<()> { + self.complete(id, result_bytes) + } + fn fail(&self, id: &str, error: &str) -> Result<()> { + self.fail(id, error) + } + fn retry(&self, id: &str, next_scheduled_at: i64) -> Result<()> { + self.retry(id, next_scheduled_at) + } + fn cancel_job(&self, id: &str) -> Result { + self.cancel_job(id) + } + fn request_cancel(&self, id: &str) -> Result { + self.request_cancel(id) + } + fn is_cancel_requested(&self, id: &str) -> Result { + self.is_cancel_requested(id) + } + fn mark_cancelled(&self, id: &str) -> Result<()> { + self.mark_cancelled(id) + } + fn cascade_cancel(&self, failed_job_id: &str, reason: &str) -> Result<()> { + self.cascade_cancel(failed_job_id, reason) + } + fn get_dependencies(&self, job_id: &str) -> Result> { + self.get_dependencies(job_id) + } + fn get_dependents(&self, job_id: &str) -> Result> { + self.get_dependents(job_id) + } + fn update_progress(&self, id: &str, progress: i32) -> Result<()> { + self.update_progress(id, progress) + } + fn list_jobs( + &self, + status: Option, + queue_name: Option<&str>, + task_name: Option<&str>, + limit: i64, + offset: i64, + ) -> Result> { + self.list_jobs(status, queue_name, task_name, limit, offset) + } + fn get_job(&self, id: &str) -> Result> { + self.get_job(id) + } + fn stats(&self) -> Result { + self.stats() + } + fn purge_completed(&self, older_than_ms: i64) -> Result { + self.purge_completed(older_than_ms) + } + fn purge_completed_with_ttl(&self, global_cutoff_ms: i64) -> Result { + self.purge_completed_with_ttl(global_cutoff_ms) + } + fn reap_stale_jobs(&self, now: i64) -> Result> { + self.reap_stale_jobs(now) + } + fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()> { + self.record_error(job_id, attempt, error) + } + fn get_job_errors(&self, job_id: &str) -> Result> { + self.get_job_errors(job_id) + } + fn purge_job_errors(&self, older_than_ms: i64) -> Result { + self.purge_job_errors(older_than_ms) + } + fn move_to_dlq(&self, job: &Job, error: &str, metadata: Option<&str>) -> Result<()> { + self.move_to_dlq(job, error, metadata) + } + fn list_dead(&self, limit: i64, offset: i64) -> Result> { + self.list_dead(limit, offset) + } + fn retry_dead(&self, dead_id: &str) -> Result { + self.retry_dead(dead_id) + } + fn purge_dead(&self, older_than_ms: i64) -> Result { + self.purge_dead(older_than_ms) + } + fn get_rate_limit(&self, key: &str) -> Result> { + self.get_rate_limit(key) + } + fn upsert_rate_limit(&self, row: &RateLimitRow) -> Result<()> { + self.upsert_rate_limit(row) + } + fn try_acquire_token(&self, key: &str, max_tokens: f64, refill_rate: f64) -> Result { + self.try_acquire_token(key, max_tokens, refill_rate) + } + fn register_periodic(&self, task: &NewPeriodicTaskRow) -> Result<()> { + self.register_periodic(task) + } + fn get_due_periodic(&self, now: i64) -> Result> { + self.get_due_periodic(now) + } + fn update_periodic_schedule(&self, name: &str, last_run: i64, next_run: i64) -> Result<()> { + self.update_periodic_schedule(name, last_run, next_run) + } + fn record_metric( + &self, + task_name: &str, + job_id: &str, + wall_time_ns: i64, + memory_bytes: i64, + succeeded: bool, + ) -> Result<()> { + self.record_metric(task_name, job_id, wall_time_ns, memory_bytes, succeeded) + } + fn get_metrics(&self, name: Option<&str>, since_ms: i64) -> Result> { + self.get_metrics(name, since_ms) + } + fn purge_metrics(&self, older_than_ms: i64) -> Result { + self.purge_metrics(older_than_ms) + } + fn record_replay( + &self, + original_job_id: &str, + replay_job_id: &str, + original_result: Option<&[u8]>, + replay_result: Option<&[u8]>, + original_error: Option<&str>, + replay_error: Option<&str>, + ) -> Result<()> { + self.record_replay( + original_job_id, + replay_job_id, + original_result, + replay_result, + original_error, + replay_error, + ) + } + fn get_replay_history(&self, original_job_id: &str) -> Result> { + self.get_replay_history(original_job_id) + } + fn write_task_log( + &self, + job_id: &str, + task_name: &str, + level: &str, + message: &str, + extra: Option<&str>, + ) -> Result<()> { + self.write_task_log(job_id, task_name, level, message, extra) + } + fn get_task_logs(&self, job_id: &str) -> Result> { + self.get_task_logs(job_id) + } + fn query_task_logs( + &self, + task_name: Option<&str>, + level: Option<&str>, + since_ms: i64, + limit: i64, + ) -> Result> { + self.query_task_logs(task_name, level, since_ms, limit) + } + fn purge_task_logs(&self, older_than_ms: i64) -> Result { + self.purge_task_logs(older_than_ms) + } + fn get_circuit_breaker(&self, task_name: &str) -> Result> { + self.get_circuit_breaker(task_name) + } + fn upsert_circuit_breaker(&self, row: &CircuitBreakerRow) -> Result<()> { + self.upsert_circuit_breaker(row) + } + fn list_circuit_breakers(&self) -> Result> { + self.list_circuit_breakers() + } + fn register_worker(&self, worker_id: &str, queues: &str) -> Result<()> { + self.register_worker(worker_id, queues) + } + fn heartbeat(&self, worker_id: &str) -> Result<()> { + self.heartbeat(worker_id) + } + fn list_workers(&self) -> Result> { + self.list_workers() + } + fn reap_dead_workers(&self) -> Result { + self.reap_dead_workers() + } + fn unregister_worker(&self, worker_id: &str) -> Result<()> { + self.unregister_worker(worker_id) + } +} diff --git a/crates/taskito-core/src/storage/postgres/workers.rs b/crates/taskito-core/src/storage/postgres/workers.rs new file mode 100644 index 00000000..63d771d9 --- /dev/null +++ b/crates/taskito-core/src/storage/postgres/workers.rs @@ -0,0 +1,79 @@ +use diesel::prelude::*; + +use super::super::models::*; +use super::super::schema::workers; +use super::PostgresStorage; +use crate::error::Result; +use crate::job::now_millis; + +/// Dead worker threshold: 30 seconds without heartbeat. +const DEAD_WORKER_THRESHOLD_MS: i64 = 30_000; + +impl PostgresStorage { + /// Register a new worker or update an existing one. + pub fn register_worker(&self, worker_id: &str, queues: &str) -> Result<()> { + let mut conn = self.conn()?; + let now = now_millis(); + + let row = NewWorkerRow { + worker_id, + last_heartbeat: now, + queues, + status: "active", + }; + + diesel::insert_into(workers::table) + .values(&row) + .on_conflict(workers::worker_id) + .do_update() + .set(&row) + .execute(&mut conn)?; + + Ok(()) + } + + /// Update the heartbeat timestamp for a worker. + pub fn heartbeat(&self, worker_id: &str) -> Result<()> { + let mut conn = self.conn()?; + let now = now_millis(); + + diesel::update(workers::table) + .filter(workers::worker_id.eq(worker_id)) + .set(workers::last_heartbeat.eq(now)) + .execute(&mut conn)?; + + Ok(()) + } + + /// List all workers with their heartbeat status. + pub fn list_workers(&self) -> Result> { + let mut conn = self.conn()?; + + let rows = workers::table + .select(WorkerRow::as_select()) + .load(&mut conn)?; + + Ok(rows) + } + + /// Remove workers that haven't sent a heartbeat within the threshold. + pub fn reap_dead_workers(&self) -> Result { + let mut conn = self.conn()?; + let cutoff = now_millis() - DEAD_WORKER_THRESHOLD_MS; + + let affected = diesel::delete(workers::table.filter(workers::last_heartbeat.lt(cutoff))) + .execute(&mut conn)?; + + Ok(affected as u64) + } + + /// Unregister a worker (called on shutdown). + pub fn unregister_worker(&self, worker_id: &str) -> Result<()> { + let mut conn = self.conn()?; + + diesel::delete(workers::table.filter(workers::worker_id.eq(worker_id))) + .execute(&mut conn)?; + + Ok(()) + } +} diff --git a/crates/taskito-core/src/storage/schema.rs b/crates/taskito-core/src/storage/schema.rs index cd245857..fe1c40e8 100644 --- a/crates/taskito-core/src/storage/schema.rs +++ b/crates/taskito-core/src/storage/schema.rs @@ -35,6 +35,10 @@ diesel::table! { retry_count -> Integer, failed_at -> BigInt, metadata -> Nullable, + priority -> Integer, + max_retries -> Integer, + timeout_ms -> BigInt, + result_ttl_ms -> Nullable, } } diff --git a/crates/taskito-core/src/storage/sqlite/circuit_breakers.rs b/crates/taskito-core/src/storage/sqlite/circuit_breakers.rs index 8847a8f4..9ac948f2 100644 --- a/crates/taskito-core/src/storage/sqlite/circuit_breakers.rs +++ b/crates/taskito-core/src/storage/sqlite/circuit_breakers.rs @@ -1,9 +1,9 @@ use diesel::prelude::*; -use crate::error::Result; use super::super::models::CircuitBreakerRow; use super::super::schema::circuit_breakers; use super::SqliteStorage; +use crate::error::Result; impl SqliteStorage { /// Get or create a circuit breaker state for a task. diff --git a/crates/taskito-core/src/storage/sqlite/dead_letter.rs b/crates/taskito-core/src/storage/sqlite/dead_letter.rs index 208bf356..90400669 100644 --- a/crates/taskito-core/src/storage/sqlite/dead_letter.rs +++ b/crates/taskito-core/src/storage/sqlite/dead_letter.rs @@ -1,10 +1,11 @@ use diesel::prelude::*; -use crate::error::{QueueError, Result}; -use crate::job::{NewJob, now_millis, Job, JobStatus}; use super::super::models::*; use super::super::schema::{dead_letter, jobs}; -use super::{DeadJob, SqliteStorage}; +use super::SqliteStorage; +use crate::error::{QueueError, Result}; +use crate::job::{now_millis, Job, JobStatus, NewJob}; +use crate::storage::DeadJob; impl SqliteStorage { /// Move a job to the dead letter queue and cascade-cancel dependents. @@ -25,6 +26,10 @@ impl SqliteStorage { retry_count: job.retry_count, failed_at: now, metadata, + priority: job.priority, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + result_ttl_ms: job.result_ttl_ms, }; diesel::insert_into(dead_letter::table) @@ -76,29 +81,51 @@ impl SqliteStorage { .first(&mut conn) .map_err(|_| QueueError::JobNotFound(dead_id.to_string()))?; - // Drop conn before calling enqueue (which gets its own conn) - drop(conn); - let new_job = NewJob { queue: dead_row.queue, task_name: dead_row.task_name, payload: dead_row.payload, - priority: 0, + priority: dead_row.priority, scheduled_at: now_millis(), - max_retries: 3, - timeout_ms: 300_000, + max_retries: dead_row.max_retries, + timeout_ms: dead_row.timeout_ms, unique_key: None, - metadata: None, + metadata: dead_row.metadata, depends_on: vec![], expires_at: None, - result_ttl_ms: None, + result_ttl_ms: dead_row.result_ttl_ms, }; - let job = self.enqueue(new_job)?; + let job = new_job.into_job(); - let mut conn = self.conn()?; - diesel::delete(dead_letter::table.find(dead_id)) - .execute(&mut conn)?; + conn.transaction(|conn| { + let row = super::super::models::NewJobRow { + id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + status: job.status as i32, + priority: job.priority, + created_at: job.created_at, + scheduled_at: job.scheduled_at, + retry_count: job.retry_count, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + unique_key: job.unique_key.as_deref(), + metadata: job.metadata.as_deref(), + cancel_requested: 0, + expires_at: job.expires_at, + result_ttl_ms: job.result_ttl_ms, + }; + + diesel::insert_into(jobs::table) + .values(&row) + .execute(conn)?; + + diesel::delete(dead_letter::table.find(dead_id)).execute(conn)?; + + Ok::<(), diesel::result::Error>(()) + })?; Ok(job.id) } @@ -107,9 +134,9 @@ impl SqliteStorage { pub fn purge_dead(&self, older_than_ms: i64) -> Result { let mut conn = self.conn()?; - let affected = diesel::delete( - dead_letter::table.filter(dead_letter::failed_at.lt(older_than_ms)) - ).execute(&mut conn)?; + let affected = + diesel::delete(dead_letter::table.filter(dead_letter::failed_at.lt(older_than_ms))) + .execute(&mut conn)?; Ok(affected as u64) } diff --git a/crates/taskito-core/src/storage/sqlite/jobs.rs b/crates/taskito-core/src/storage/sqlite/jobs.rs index 85c25907..64fffb09 100644 --- a/crates/taskito-core/src/storage/sqlite/jobs.rs +++ b/crates/taskito-core/src/storage/sqlite/jobs.rs @@ -1,11 +1,40 @@ use diesel::prelude::*; use diesel::sqlite::SqliteConnection; -use crate::error::{QueueError, Result}; -use crate::job::{Job, JobStatus, NewJob, now_millis}; use super::super::models::*; -use super::super::schema::{job_dependencies, jobs}; -use super::{QueueStats, SqliteStorage}; +use super::super::schema::{ + job_dependencies, job_errors, jobs, replay_history, task_logs, task_metrics, +}; +use super::SqliteStorage; +use crate::error::{QueueError, Result}; +use crate::job::{now_millis, Job, JobStatus, NewJob}; +use crate::storage::QueueStats; + +fn delete_job_children( + conn: &mut SqliteConnection, + job_ids: &[String], +) -> diesel::result::QueryResult<()> { + if job_ids.is_empty() { + return Ok(()); + } + + diesel::delete(job_errors::table.filter(job_errors::job_id.eq_any(job_ids))).execute(conn)?; + diesel::delete(task_logs::table.filter(task_logs::job_id.eq_any(job_ids))).execute(conn)?; + diesel::delete(task_metrics::table.filter(task_metrics::job_id.eq_any(job_ids))) + .execute(conn)?; + diesel::delete( + job_dependencies::table.filter( + job_dependencies::job_id + .eq_any(job_ids) + .or(job_dependencies::depends_on_job_id.eq_any(job_ids)), + ), + ) + .execute(conn)?; + diesel::delete(replay_history::table.filter(replay_history::original_job_id.eq_any(job_ids))) + .execute(conn)?; + + Ok(()) +} impl SqliteStorage { /// Insert a new job into the queue. Returns the job. @@ -25,8 +54,9 @@ impl SqliteStorage { match dep { None => return Err(diesel::result::Error::RollbackTransaction), - Some(d) if d.status == JobStatus::Dead as i32 - || d.status == JobStatus::Cancelled as i32 => + Some(d) + if d.status == JobStatus::Dead as i32 + || d.status == JobStatus::Cancelled as i32 => { return Err(diesel::result::Error::RollbackTransaction); } @@ -70,10 +100,11 @@ impl SqliteStorage { } Ok(()) - }).map_err(|e| match e { - diesel::result::Error::RollbackTransaction => { - QueueError::DependencyNotFound("dependency not found or already dead/cancelled".to_string()) - } + }) + .map_err(|e| match e { + diesel::result::Error::RollbackTransaction => QueueError::DependencyNotFound( + "dependency not found or already dead/cancelled".to_string(), + ), other => QueueError::Storage(other), })?; @@ -120,63 +151,94 @@ impl SqliteStorage { let job = new_job.into_job(); let mut conn = self.conn()?; - // Check for existing active job with same unique_key - if let Some(ref uk) = job.unique_key { - let existing: Option = jobs::table - .filter(jobs::unique_key.eq(uk)) - .filter(jobs::status.eq_any([ - JobStatus::Pending as i32, - JobStatus::Running as i32, - ])) - .select(JobRow::as_select()) - .first(&mut conn) - .optional()?; + let result = conn.transaction(|conn| { + // Check for existing active job with same unique_key + if let Some(ref uk) = job.unique_key { + let existing: Option = jobs::table + .filter(jobs::unique_key.eq(uk)) + .filter( + jobs::status.eq_any([JobStatus::Pending as i32, JobStatus::Running as i32]), + ) + .select(JobRow::as_select()) + .first(conn) + .optional()?; - if let Some(row) = existing { - return Ok(Job::from(row)); + if let Some(row) = existing { + return Ok(Job::from(row)); + } } - } - let row = NewJobRow { - id: &job.id, - queue: &job.queue, - task_name: &job.task_name, - payload: &job.payload, - status: job.status as i32, - priority: job.priority, - created_at: job.created_at, - scheduled_at: job.scheduled_at, - retry_count: job.retry_count, - max_retries: job.max_retries, - timeout_ms: job.timeout_ms, - unique_key: job.unique_key.as_deref(), - metadata: job.metadata.as_deref(), - cancel_requested: 0, - expires_at: job.expires_at, - result_ttl_ms: job.result_ttl_ms, - }; + let row = NewJobRow { + id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + status: job.status as i32, + priority: job.priority, + created_at: job.created_at, + scheduled_at: job.scheduled_at, + retry_count: job.retry_count, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + unique_key: job.unique_key.as_deref(), + metadata: job.metadata.as_deref(), + cancel_requested: 0, + expires_at: job.expires_at, + result_ttl_ms: job.result_ttl_ms, + }; - diesel::insert_into(jobs::table) - .values(&row) - .execute(&mut conn)?; + diesel::insert_into(jobs::table) + .values(&row) + .execute(conn)?; - // Insert dependency rows - for dep_id in &depends_on { - let dep_row = NewJobDependencyRow { - id: &uuid::Uuid::now_v7().to_string(), - job_id: &job.id, - depends_on_job_id: dep_id, - }; - diesel::insert_into(job_dependencies::table) - .values(&dep_row) - .execute(&mut conn)?; - } + // Insert dependency rows + for dep_id in &depends_on { + let dep_row = NewJobDependencyRow { + id: &uuid::Uuid::now_v7().to_string(), + job_id: &job.id, + depends_on_job_id: dep_id, + }; + diesel::insert_into(job_dependencies::table) + .values(&dep_row) + .execute(conn)?; + } - Ok(job) + Ok(job.clone()) + }); + + // Handle unique constraint violation by returning existing job + match result { + Ok(j) => Ok(j), + Err(diesel::result::Error::DatabaseError( + diesel::result::DatabaseErrorKind::UniqueViolation, + _, + )) => { + if let Some(ref uk) = job.unique_key { + let mut conn = self.conn()?; + let existing: Option = jobs::table + .filter(jobs::unique_key.eq(uk)) + .filter( + jobs::status + .eq_any([JobStatus::Pending as i32, JobStatus::Running as i32]), + ) + .select(JobRow::as_select()) + .first(&mut conn) + .optional()?; + if let Some(row) = existing { + return Ok(Job::from(row)); + } + } + Ok(job) + } + Err(e) => Err(QueueError::Storage(e)), + } } /// Check if all dependencies for a job are complete. - fn deps_satisfied(conn: &mut SqliteConnection, job_id: &str) -> diesel::result::QueryResult { + fn deps_satisfied( + conn: &mut SqliteConnection, + job_id: &str, + ) -> diesel::result::QueryResult { let dep_job_ids: Vec = job_dependencies::table .filter(job_dependencies::job_id.eq(job_id)) .select(job_dependencies::depends_on_job_id) @@ -206,7 +268,7 @@ impl SqliteStorage { .filter(jobs::status.eq(JobStatus::Pending as i32)) .filter(jobs::scheduled_at.le(now)) .order((jobs::priority.desc(), jobs::scheduled_at.asc())) - .limit(10) + .limit(100) .select(JobRow::as_select()) .load(conn)?; @@ -398,6 +460,8 @@ impl SqliteStorage { let now = now_millis(); let mut queue: Vec = vec![failed_job_id.to_string()]; + let mut visited = std::collections::HashSet::new(); + visited.insert(failed_job_id.to_string()); let mut idx = 0; while idx < queue.len() { @@ -410,7 +474,7 @@ impl SqliteStorage { .load(&mut conn)?; for dep_id in dependents { - if !queue.contains(&dep_id) { + if visited.insert(dep_id.clone()) { queue.push(dep_id); } } @@ -483,9 +547,7 @@ impl SqliteStorage { ) -> Result> { let mut conn = self.conn()?; - let mut query = jobs::table - .into_boxed() - .order(jobs::created_at.desc()); + let mut query = jobs::table.into_boxed().order(jobs::created_at.desc()); if let Some(s) = status { query = query.filter(jobs::status.eq(s)); @@ -548,11 +610,16 @@ impl SqliteStorage { pub fn purge_completed(&self, older_than_ms: i64) -> Result { let mut conn = self.conn()?; - let affected = diesel::delete( - jobs::table - .filter(jobs::status.eq(JobStatus::Complete as i32)) - .filter(jobs::completed_at.lt(older_than_ms)) - ).execute(&mut conn)?; + let job_ids: Vec = jobs::table + .filter(jobs::status.eq(JobStatus::Complete as i32)) + .filter(jobs::completed_at.lt(older_than_ms)) + .select(jobs::id) + .load(&mut conn)?; + + delete_job_children(&mut conn, &job_ids)?; + + let affected = + diesel::delete(jobs::table.filter(jobs::id.eq_any(&job_ids))).execute(&mut conn)?; Ok(affected as u64) } @@ -563,33 +630,40 @@ impl SqliteStorage { let now = now_millis(); let mut conn = self.conn()?; - // For global TTL: delete completed jobs without per-job TTL - let global = diesel::delete( - jobs::table + conn.transaction(|conn| { + // For global TTL: collect IDs of completed jobs without per-job TTL + let global_ids: Vec = jobs::table .filter(jobs::status.eq(JobStatus::Complete as i32)) .filter(jobs::result_ttl_ms.is_null()) .filter(jobs::completed_at.lt(global_cutoff_ms)) - ).execute(&mut conn)?; + .select(jobs::id) + .load(conn)?; - // For per-job TTL: fetch candidates and delete expired ones - let rows_with_ttl: Vec = jobs::table - .filter(jobs::status.eq(JobStatus::Complete as i32)) - .filter(jobs::result_ttl_ms.is_not_null()) - .select(JobRow::as_select()) - .load(&mut conn)?; + // For per-job TTL: fetch candidates and collect expired IDs + let rows_with_ttl: Vec = jobs::table + .filter(jobs::status.eq(JobStatus::Complete as i32)) + .filter(jobs::result_ttl_ms.is_not_null()) + .select(JobRow::as_select()) + .load(conn)?; - let mut per_job_count = 0u64; - for row in rows_with_ttl { - if let (Some(completed), Some(ttl)) = (row.completed_at, row.result_ttl_ms) { - if completed + ttl < now { - diesel::delete(jobs::table.filter(jobs::id.eq(&row.id))) - .execute(&mut conn)?; - per_job_count += 1; - } - } - } + let per_job_ids: Vec = rows_with_ttl + .into_iter() + .filter(|row| { + matches!((row.completed_at, row.result_ttl_ms), (Some(completed), Some(ttl)) if completed + ttl < now) + }) + .map(|row| row.id) + .collect(); + + let all_ids: Vec = global_ids.into_iter().chain(per_job_ids).collect(); - Ok(global as u64 + per_job_count) + delete_job_children(conn, &all_ids)?; + + let affected = diesel::delete( + jobs::table.filter(jobs::id.eq_any(&all_ids)) + ).execute(conn)?; + + Ok(affected as u64) + }) } /// Find stale running jobs that exceeded their timeout. @@ -657,9 +731,9 @@ impl SqliteStorage { use super::super::schema::job_errors; let mut conn = self.conn()?; - let affected = diesel::delete( - job_errors::table.filter(job_errors::failed_at.lt(older_than_ms)) - ).execute(&mut conn)?; + let affected = + diesel::delete(job_errors::table.filter(job_errors::failed_at.lt(older_than_ms))) + .execute(&mut conn)?; Ok(affected as u64) } diff --git a/crates/taskito-core/src/storage/sqlite/logs.rs b/crates/taskito-core/src/storage/sqlite/logs.rs index 9e9479b3..c6064ed6 100644 --- a/crates/taskito-core/src/storage/sqlite/logs.rs +++ b/crates/taskito-core/src/storage/sqlite/logs.rs @@ -1,10 +1,10 @@ use diesel::prelude::*; -use crate::error::Result; -use crate::job::now_millis; use super::super::models::*; use super::super::schema::task_logs; use super::SqliteStorage; +use crate::error::Result; +use crate::job::now_millis; impl SqliteStorage { /// Write a structured log entry for a task. @@ -81,9 +81,9 @@ impl SqliteStorage { /// Purge old log records. pub fn purge_task_logs(&self, older_than_ms: i64) -> Result { let mut conn = self.conn()?; - let affected = diesel::delete( - task_logs::table.filter(task_logs::logged_at.lt(older_than_ms)) - ).execute(&mut conn)?; + let affected = + diesel::delete(task_logs::table.filter(task_logs::logged_at.lt(older_than_ms))) + .execute(&mut conn)?; Ok(affected as u64) } } diff --git a/crates/taskito-core/src/storage/sqlite/metrics.rs b/crates/taskito-core/src/storage/sqlite/metrics.rs index 6c84a5d3..2356be4c 100644 --- a/crates/taskito-core/src/storage/sqlite/metrics.rs +++ b/crates/taskito-core/src/storage/sqlite/metrics.rs @@ -1,10 +1,10 @@ use diesel::prelude::*; -use crate::error::Result; -use crate::job::now_millis; use super::super::models::*; use super::super::schema::task_metrics; use super::SqliteStorage; +use crate::error::Result; +use crate::job::now_millis; impl SqliteStorage { /// Record a task execution metric. @@ -60,9 +60,9 @@ impl SqliteStorage { /// Purge old metric records. pub fn purge_metrics(&self, older_than_ms: i64) -> Result { let mut conn = self.conn()?; - let affected = diesel::delete( - task_metrics::table.filter(task_metrics::recorded_at.lt(older_than_ms)) - ).execute(&mut conn)?; + let affected = + diesel::delete(task_metrics::table.filter(task_metrics::recorded_at.lt(older_than_ms))) + .execute(&mut conn)?; Ok(affected as u64) } diff --git a/crates/taskito-core/src/storage/sqlite/mod.rs b/crates/taskito-core/src/storage/sqlite/mod.rs index bbd17207..e028fde1 100644 --- a/crates/taskito-core/src/storage/sqlite/mod.rs +++ b/crates/taskito-core/src/storage/sqlite/mod.rs @@ -1,10 +1,10 @@ -mod jobs; +mod circuit_breakers; mod dead_letter; -mod rate_limits; -mod periodic; -mod metrics; +mod jobs; mod logs; -mod circuit_breakers; +mod metrics; +mod periodic; +mod rate_limits; mod trait_impl; mod workers; @@ -21,7 +21,10 @@ type DbPool = Pool>; struct SqlitePragmaCustomizer; impl CustomizeConnection for SqlitePragmaCustomizer { - fn on_acquire(&self, conn: &mut SqliteConnection) -> std::result::Result<(), diesel::r2d2::Error> { + fn on_acquire( + &self, + conn: &mut SqliteConnection, + ) -> std::result::Result<(), diesel::r2d2::Error> { diesel::sql_query("PRAGMA journal_mode = WAL") .execute(conn) .map_err(diesel::r2d2::Error::QueryError)?; @@ -76,7 +79,9 @@ impl SqliteStorage { Ok(storage) } - pub(crate) fn conn(&self) -> Result>> { + pub(crate) fn conn( + &self, + ) -> Result>> { Ok(self.pool.get()?) } @@ -106,33 +111,34 @@ impl SqliteStorage { cancel_requested INTEGER NOT NULL DEFAULT 0, expires_at INTEGER, result_ttl_ms INTEGER - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; // Add new columns if they don't exist (migration for existing DBs) let _ = diesel::sql_query( - "ALTER TABLE jobs ADD COLUMN cancel_requested INTEGER NOT NULL DEFAULT 0" - ).execute(&mut conn); - let _ = diesel::sql_query( - "ALTER TABLE jobs ADD COLUMN expires_at INTEGER" - ).execute(&mut conn); - let _ = diesel::sql_query( - "ALTER TABLE jobs ADD COLUMN result_ttl_ms INTEGER" - ).execute(&mut conn); + "ALTER TABLE jobs ADD COLUMN cancel_requested INTEGER NOT NULL DEFAULT 0", + ) + .execute(&mut conn); + let _ = + diesel::sql_query("ALTER TABLE jobs ADD COLUMN expires_at INTEGER").execute(&mut conn); + let _ = diesel::sql_query("ALTER TABLE jobs ADD COLUMN result_ttl_ms INTEGER") + .execute(&mut conn); diesel::sql_query( "CREATE INDEX IF NOT EXISTS idx_jobs_dequeue - ON jobs(queue, status, priority DESC, scheduled_at)" - ).execute(&mut conn)?; + ON jobs(queue, status, priority DESC, scheduled_at)", + ) + .execute(&mut conn)?; - diesel::sql_query( - "CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status)" - ).execute(&mut conn)?; + diesel::sql_query("CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status)") + .execute(&mut conn)?; diesel::sql_query( "CREATE UNIQUE INDEX IF NOT EXISTS idx_jobs_unique_key - ON jobs(unique_key) WHERE unique_key IS NOT NULL AND status IN (0, 1)" - ).execute(&mut conn)?; + ON jobs(unique_key) WHERE unique_key IS NOT NULL AND status IN (0, 1)", + ) + .execute(&mut conn)?; diesel::sql_query( "CREATE TABLE IF NOT EXISTS dead_letter ( @@ -144,9 +150,24 @@ impl SqliteStorage { error TEXT, retry_count INTEGER NOT NULL, failed_at INTEGER NOT NULL, - metadata TEXT - )" - ).execute(&mut conn)?; + metadata TEXT, + priority INTEGER NOT NULL DEFAULT 0, + max_retries INTEGER NOT NULL DEFAULT 3, + timeout_ms INTEGER NOT NULL DEFAULT 300000, + result_ttl_ms INTEGER + )", + ) + .execute(&mut conn)?; + + // Migration: add columns if they don't exist (for existing databases) + for col in &[ + "ALTER TABLE dead_letter ADD COLUMN priority INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE dead_letter ADD COLUMN max_retries INTEGER NOT NULL DEFAULT 3", + "ALTER TABLE dead_letter ADD COLUMN timeout_ms INTEGER NOT NULL DEFAULT 300000", + "ALTER TABLE dead_letter ADD COLUMN result_ttl_ms INTEGER", + ] { + let _ = diesel::sql_query(*col).execute(&mut conn); + } diesel::sql_query( "CREATE TABLE IF NOT EXISTS rate_limits ( @@ -155,8 +176,9 @@ impl SqliteStorage { max_tokens REAL NOT NULL, refill_rate REAL NOT NULL, last_refill INTEGER NOT NULL - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; diesel::sql_query( "CREATE TABLE IF NOT EXISTS periodic_tasks ( @@ -169,8 +191,9 @@ impl SqliteStorage { enabled INTEGER NOT NULL DEFAULT 1, last_run INTEGER, next_run INTEGER NOT NULL - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; diesel::sql_query( "CREATE TABLE IF NOT EXISTS job_errors ( @@ -179,24 +202,26 @@ impl SqliteStorage { attempt INTEGER NOT NULL, error TEXT NOT NULL, failed_at INTEGER NOT NULL - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; - diesel::sql_query( - "CREATE INDEX IF NOT EXISTS idx_job_errors_job_id ON job_errors(job_id)" - ).execute(&mut conn)?; + diesel::sql_query("CREATE INDEX IF NOT EXISTS idx_job_errors_job_id ON job_errors(job_id)") + .execute(&mut conn)?; diesel::sql_query( "CREATE TABLE IF NOT EXISTS job_dependencies ( id TEXT PRIMARY KEY, job_id TEXT NOT NULL, depends_on_job_id TEXT NOT NULL - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; diesel::sql_query( - "CREATE INDEX IF NOT EXISTS idx_job_deps_job_id ON job_dependencies(job_id)" - ).execute(&mut conn)?; + "CREATE INDEX IF NOT EXISTS idx_job_deps_job_id ON job_dependencies(job_id)", + ) + .execute(&mut conn)?; diesel::sql_query( "CREATE INDEX IF NOT EXISTS idx_job_deps_depends_on ON job_dependencies(depends_on_job_id)" @@ -212,16 +237,19 @@ impl SqliteStorage { memory_bytes INTEGER NOT NULL DEFAULT 0, succeeded INTEGER NOT NULL DEFAULT 1, recorded_at INTEGER NOT NULL - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; diesel::sql_query( - "CREATE INDEX IF NOT EXISTS idx_task_metrics_task_name ON task_metrics(task_name)" - ).execute(&mut conn)?; + "CREATE INDEX IF NOT EXISTS idx_task_metrics_task_name ON task_metrics(task_name)", + ) + .execute(&mut conn)?; diesel::sql_query( - "CREATE INDEX IF NOT EXISTS idx_task_metrics_recorded_at ON task_metrics(recorded_at)" - ).execute(&mut conn)?; + "CREATE INDEX IF NOT EXISTS idx_task_metrics_recorded_at ON task_metrics(recorded_at)", + ) + .execute(&mut conn)?; // ── Replay History ──────────────────────────────── diesel::sql_query( @@ -234,12 +262,14 @@ impl SqliteStorage { replay_result BLOB, original_error TEXT, replay_error TEXT - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; diesel::sql_query( - "CREATE INDEX IF NOT EXISTS idx_replay_original ON replay_history(original_job_id)" - ).execute(&mut conn)?; + "CREATE INDEX IF NOT EXISTS idx_replay_original ON replay_history(original_job_id)", + ) + .execute(&mut conn)?; // ── Task Logs ───────────────────────────────────── diesel::sql_query( @@ -251,16 +281,17 @@ impl SqliteStorage { message TEXT NOT NULL, extra TEXT, logged_at INTEGER NOT NULL - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; - diesel::sql_query( - "CREATE INDEX IF NOT EXISTS idx_task_logs_job_id ON task_logs(job_id)" - ).execute(&mut conn)?; + diesel::sql_query("CREATE INDEX IF NOT EXISTS idx_task_logs_job_id ON task_logs(job_id)") + .execute(&mut conn)?; diesel::sql_query( - "CREATE INDEX IF NOT EXISTS idx_task_logs_recorded ON task_logs(logged_at)" - ).execute(&mut conn)?; + "CREATE INDEX IF NOT EXISTS idx_task_logs_recorded ON task_logs(logged_at)", + ) + .execute(&mut conn)?; // ── Circuit Breakers ────────────────────────────── diesel::sql_query( @@ -274,8 +305,9 @@ impl SqliteStorage { threshold INTEGER NOT NULL DEFAULT 5, window_ms INTEGER NOT NULL DEFAULT 60000, cooldown_ms INTEGER NOT NULL DEFAULT 300000 - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; // ── Workers ─────────────────────────────────────── diesel::sql_query( @@ -284,53 +316,15 @@ impl SqliteStorage { last_heartbeat INTEGER NOT NULL, queues TEXT NOT NULL DEFAULT 'default', status TEXT NOT NULL DEFAULT 'active' - )" - ).execute(&mut conn)?; + )", + ) + .execute(&mut conn)?; Ok(()) } } -// ── Helper types ─────────────────────────────────────────────────────── - -#[derive(Debug, Clone, Default)] -pub struct QueueStats { - pub pending: i64, - pub running: i64, - pub completed: i64, - pub failed: i64, - pub dead: i64, - pub cancelled: i64, -} - -#[derive(Debug, Clone)] -pub struct DeadJob { - pub id: String, - pub original_job_id: String, - pub queue: String, - pub task_name: String, - pub payload: Vec, - pub error: Option, - pub retry_count: i32, - pub failed_at: i64, - pub metadata: Option, -} - -impl From for DeadJob { - fn from(row: super::models::DeadLetterRow) -> Self { - Self { - id: row.id, - original_job_id: row.original_job_id, - queue: row.queue, - task_name: row.task_name, - payload: row.payload, - error: row.error, - retry_count: row.retry_count, - failed_at: row.failed_at, - metadata: row.metadata, - } - } -} +pub use crate::storage::{DeadJob, QueueStats}; #[cfg(test)] mod tests; diff --git a/crates/taskito-core/src/storage/sqlite/periodic.rs b/crates/taskito-core/src/storage/sqlite/periodic.rs index 696d3196..d6c5ff00 100644 --- a/crates/taskito-core/src/storage/sqlite/periodic.rs +++ b/crates/taskito-core/src/storage/sqlite/periodic.rs @@ -1,9 +1,9 @@ use diesel::prelude::*; -use crate::error::Result; use super::super::models::*; use super::super::schema::periodic_tasks; use super::SqliteStorage; +use crate::error::Result; impl SqliteStorage { /// Register or update a periodic task. @@ -31,12 +31,7 @@ impl SqliteStorage { } /// Update a periodic task's schedule after execution. - pub fn update_periodic_schedule( - &self, - name: &str, - last_run: i64, - next_run: i64, - ) -> Result<()> { + pub fn update_periodic_schedule(&self, name: &str, last_run: i64, next_run: i64) -> Result<()> { let mut conn = self.conn()?; diesel::update(periodic_tasks::table.find(name)) diff --git a/crates/taskito-core/src/storage/sqlite/rate_limits.rs b/crates/taskito-core/src/storage/sqlite/rate_limits.rs index abd64554..6b38c480 100644 --- a/crates/taskito-core/src/storage/sqlite/rate_limits.rs +++ b/crates/taskito-core/src/storage/sqlite/rate_limits.rs @@ -1,10 +1,10 @@ use diesel::prelude::*; -use crate::error::Result; -use crate::job::now_millis; use super::super::models::RateLimitRow; use super::super::schema::rate_limits; use super::SqliteStorage; +use crate::error::Result; +use crate::job::now_millis; impl SqliteStorage { pub fn get_rate_limit(&self, key: &str) -> Result> { @@ -32,12 +32,7 @@ impl SqliteStorage { /// Atomically try to acquire a rate limit token. /// Does the read-refill-consume-write in a single transaction to prevent /// race conditions between concurrent workers. - pub fn try_acquire_token( - &self, - key: &str, - max_tokens: f64, - refill_rate: f64, - ) -> Result { + pub fn try_acquire_token(&self, key: &str, max_tokens: f64, refill_rate: f64) -> Result { let mut conn = self.conn()?; let now = now_millis(); diff --git a/crates/taskito-core/src/storage/sqlite/tests.rs b/crates/taskito-core/src/storage/sqlite/tests.rs index 08adc05e..147362e2 100644 --- a/crates/taskito-core/src/storage/sqlite/tests.rs +++ b/crates/taskito-core/src/storage/sqlite/tests.rs @@ -1,5 +1,5 @@ use super::*; -use crate::job::{JobStatus, NewJob, now_millis}; +use crate::job::{now_millis, JobStatus, NewJob}; fn test_storage() -> SqliteStorage { SqliteStorage::in_memory().unwrap() @@ -37,7 +37,10 @@ fn test_dequeue() { let storage = test_storage(); let job = storage.enqueue(make_job("dequeue_task")).unwrap(); - let dequeued = storage.dequeue("default", now_millis() + 1000).unwrap().unwrap(); + let dequeued = storage + .dequeue("default", now_millis() + 1000) + .unwrap() + .unwrap(); assert_eq!(dequeued.id, job.id); assert_eq!(dequeued.status, JobStatus::Running); @@ -127,7 +130,13 @@ fn test_dead_letter_queue() { let job = storage.enqueue(make_job("dlq_task")).unwrap(); storage.dequeue("default", now_millis() + 1000).unwrap(); - storage.move_to_dlq(&storage.get_job(&job.id).unwrap().unwrap(), "max retries exceeded", None).unwrap(); + storage + .move_to_dlq( + &storage.get_job(&job.id).unwrap().unwrap(), + "max retries exceeded", + None, + ) + .unwrap(); let fetched = storage.get_job(&job.id).unwrap().unwrap(); assert_eq!(fetched.status, JobStatus::Dead); @@ -144,7 +153,9 @@ fn test_retry_dead() { storage.dequeue("default", now_millis() + 1000).unwrap(); let running_job = storage.get_job(&job.id).unwrap().unwrap(); - storage.move_to_dlq(&running_job, "fatal error", None).unwrap(); + storage + .move_to_dlq(&running_job, "fatal error", None) + .unwrap(); let dead = storage.list_dead(10, 0).unwrap(); let new_id = storage.retry_dead(&dead[0].id).unwrap(); @@ -201,11 +212,13 @@ fn test_unique_key_dedup() { #[test] fn test_enqueue_batch() { let storage = test_storage(); - let jobs: Vec = (0..5).map(|i| { - let mut j = make_job(&format!("batch_task_{i}")); - j.priority = i; - j - }).collect(); + let jobs: Vec = (0..5) + .map(|i| { + let mut j = make_job(&format!("batch_task_{i}")); + j.priority = i; + j + }) + .collect(); let result = storage.enqueue_batch(jobs).unwrap(); assert_eq!(result.len(), 5); diff --git a/crates/taskito-core/src/storage/sqlite/trait_impl.rs b/crates/taskito-core/src/storage/sqlite/trait_impl.rs index 9c89655b..e7dbd09e 100644 --- a/crates/taskito-core/src/storage/sqlite/trait_impl.rs +++ b/crates/taskito-core/src/storage/sqlite/trait_impl.rs @@ -1,68 +1,207 @@ +use super::SqliteStorage; use crate::error::Result; use crate::job::{Job, NewJob}; use crate::storage::models::*; use crate::storage::traits::Storage; -use super::{DeadJob, QueueStats, SqliteStorage}; +use crate::storage::{DeadJob, QueueStats}; impl Storage for SqliteStorage { - fn enqueue(&self, new_job: NewJob) -> Result { self.enqueue(new_job) } - fn enqueue_batch(&self, new_jobs: Vec) -> Result> { self.enqueue_batch(new_jobs) } - fn enqueue_unique(&self, new_job: NewJob) -> Result { self.enqueue_unique(new_job) } - fn dequeue(&self, queue_name: &str, now: i64) -> Result> { self.dequeue(queue_name, now) } - fn dequeue_from(&self, queues: &[String], now: i64) -> Result> { self.dequeue_from(queues, now) } - fn complete(&self, id: &str, result_bytes: Option>) -> Result<()> { self.complete(id, result_bytes) } - fn fail(&self, id: &str, error: &str) -> Result<()> { self.fail(id, error) } - fn retry(&self, id: &str, next_scheduled_at: i64) -> Result<()> { self.retry(id, next_scheduled_at) } - fn cancel_job(&self, id: &str) -> Result { self.cancel_job(id) } - fn cascade_cancel(&self, failed_job_id: &str, reason: &str) -> Result<()> { self.cascade_cancel(failed_job_id, reason) } - fn get_dependencies(&self, job_id: &str) -> Result> { self.get_dependencies(job_id) } - fn get_dependents(&self, job_id: &str) -> Result> { self.get_dependents(job_id) } - fn update_progress(&self, id: &str, progress: i32) -> Result<()> { self.update_progress(id, progress) } - fn list_jobs(&self, status: Option, queue_name: Option<&str>, task_name: Option<&str>, limit: i64, offset: i64) -> Result> { + fn enqueue(&self, new_job: NewJob) -> Result { + self.enqueue(new_job) + } + fn enqueue_batch(&self, new_jobs: Vec) -> Result> { + self.enqueue_batch(new_jobs) + } + fn enqueue_unique(&self, new_job: NewJob) -> Result { + self.enqueue_unique(new_job) + } + fn dequeue(&self, queue_name: &str, now: i64) -> Result> { + self.dequeue(queue_name, now) + } + fn dequeue_from(&self, queues: &[String], now: i64) -> Result> { + self.dequeue_from(queues, now) + } + fn complete(&self, id: &str, result_bytes: Option>) -> Result<()> { + self.complete(id, result_bytes) + } + fn fail(&self, id: &str, error: &str) -> Result<()> { + self.fail(id, error) + } + fn retry(&self, id: &str, next_scheduled_at: i64) -> Result<()> { + self.retry(id, next_scheduled_at) + } + fn cancel_job(&self, id: &str) -> Result { + self.cancel_job(id) + } + fn cascade_cancel(&self, failed_job_id: &str, reason: &str) -> Result<()> { + self.cascade_cancel(failed_job_id, reason) + } + fn get_dependencies(&self, job_id: &str) -> Result> { + self.get_dependencies(job_id) + } + fn get_dependents(&self, job_id: &str) -> Result> { + self.get_dependents(job_id) + } + fn update_progress(&self, id: &str, progress: i32) -> Result<()> { + self.update_progress(id, progress) + } + fn list_jobs( + &self, + status: Option, + queue_name: Option<&str>, + task_name: Option<&str>, + limit: i64, + offset: i64, + ) -> Result> { self.list_jobs(status, queue_name, task_name, limit, offset) } - fn get_job(&self, id: &str) -> Result> { self.get_job(id) } - fn stats(&self) -> Result { self.stats() } - fn purge_completed(&self, older_than_ms: i64) -> Result { self.purge_completed(older_than_ms) } - fn reap_stale_jobs(&self, now: i64) -> Result> { self.reap_stale_jobs(now) } - fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()> { self.record_error(job_id, attempt, error) } - fn get_job_errors(&self, job_id: &str) -> Result> { self.get_job_errors(job_id) } - fn purge_job_errors(&self, older_than_ms: i64) -> Result { self.purge_job_errors(older_than_ms) } - fn move_to_dlq(&self, job: &Job, error: &str, metadata: Option<&str>) -> Result<()> { self.move_to_dlq(job, error, metadata) } - fn list_dead(&self, limit: i64, offset: i64) -> Result> { self.list_dead(limit, offset) } - fn retry_dead(&self, dead_id: &str) -> Result { self.retry_dead(dead_id) } - fn purge_dead(&self, older_than_ms: i64) -> Result { self.purge_dead(older_than_ms) } - fn get_rate_limit(&self, key: &str) -> Result> { self.get_rate_limit(key) } - fn upsert_rate_limit(&self, row: &RateLimitRow) -> Result<()> { self.upsert_rate_limit(row) } - fn register_periodic(&self, task: &NewPeriodicTaskRow) -> Result<()> { self.register_periodic(task) } - fn get_due_periodic(&self, now: i64) -> Result> { self.get_due_periodic(now) } - fn update_periodic_schedule(&self, name: &str, last_run: i64, next_run: i64) -> Result<()> { self.update_periodic_schedule(name, last_run, next_run) } - fn record_metric(&self, task_name: &str, job_id: &str, wall_time_ns: i64, memory_bytes: i64, succeeded: bool) -> Result<()> { + fn get_job(&self, id: &str) -> Result> { + self.get_job(id) + } + fn stats(&self) -> Result { + self.stats() + } + fn purge_completed(&self, older_than_ms: i64) -> Result { + self.purge_completed(older_than_ms) + } + fn purge_completed_with_ttl(&self, global_cutoff_ms: i64) -> Result { + self.purge_completed_with_ttl(global_cutoff_ms) + } + fn reap_stale_jobs(&self, now: i64) -> Result> { + self.reap_stale_jobs(now) + } + fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()> { + self.record_error(job_id, attempt, error) + } + fn get_job_errors(&self, job_id: &str) -> Result> { + self.get_job_errors(job_id) + } + fn purge_job_errors(&self, older_than_ms: i64) -> Result { + self.purge_job_errors(older_than_ms) + } + fn move_to_dlq(&self, job: &Job, error: &str, metadata: Option<&str>) -> Result<()> { + self.move_to_dlq(job, error, metadata) + } + fn list_dead(&self, limit: i64, offset: i64) -> Result> { + self.list_dead(limit, offset) + } + fn retry_dead(&self, dead_id: &str) -> Result { + self.retry_dead(dead_id) + } + fn purge_dead(&self, older_than_ms: i64) -> Result { + self.purge_dead(older_than_ms) + } + fn get_rate_limit(&self, key: &str) -> Result> { + self.get_rate_limit(key) + } + fn upsert_rate_limit(&self, row: &RateLimitRow) -> Result<()> { + self.upsert_rate_limit(row) + } + fn try_acquire_token(&self, key: &str, max_tokens: f64, refill_rate: f64) -> Result { + self.try_acquire_token(key, max_tokens, refill_rate) + } + fn register_periodic(&self, task: &NewPeriodicTaskRow) -> Result<()> { + self.register_periodic(task) + } + fn get_due_periodic(&self, now: i64) -> Result> { + self.get_due_periodic(now) + } + fn update_periodic_schedule(&self, name: &str, last_run: i64, next_run: i64) -> Result<()> { + self.update_periodic_schedule(name, last_run, next_run) + } + fn record_metric( + &self, + task_name: &str, + job_id: &str, + wall_time_ns: i64, + memory_bytes: i64, + succeeded: bool, + ) -> Result<()> { self.record_metric(task_name, job_id, wall_time_ns, memory_bytes, succeeded) } - fn get_metrics(&self, name: Option<&str>, since_ms: i64) -> Result> { self.get_metrics(name, since_ms) } - fn purge_metrics(&self, older_than_ms: i64) -> Result { self.purge_metrics(older_than_ms) } - fn record_replay(&self, original_job_id: &str, replay_job_id: &str, original_result: Option<&[u8]>, replay_result: Option<&[u8]>, original_error: Option<&str>, replay_error: Option<&str>) -> Result<()> { - self.record_replay(original_job_id, replay_job_id, original_result, replay_result, original_error, replay_error) + fn get_metrics(&self, name: Option<&str>, since_ms: i64) -> Result> { + self.get_metrics(name, since_ms) + } + fn purge_metrics(&self, older_than_ms: i64) -> Result { + self.purge_metrics(older_than_ms) + } + fn record_replay( + &self, + original_job_id: &str, + replay_job_id: &str, + original_result: Option<&[u8]>, + replay_result: Option<&[u8]>, + original_error: Option<&str>, + replay_error: Option<&str>, + ) -> Result<()> { + self.record_replay( + original_job_id, + replay_job_id, + original_result, + replay_result, + original_error, + replay_error, + ) } - fn get_replay_history(&self, original_job_id: &str) -> Result> { self.get_replay_history(original_job_id) } - fn write_task_log(&self, job_id: &str, task_name: &str, level: &str, message: &str, extra: Option<&str>) -> Result<()> { + fn get_replay_history(&self, original_job_id: &str) -> Result> { + self.get_replay_history(original_job_id) + } + fn write_task_log( + &self, + job_id: &str, + task_name: &str, + level: &str, + message: &str, + extra: Option<&str>, + ) -> Result<()> { self.write_task_log(job_id, task_name, level, message, extra) } - fn get_task_logs(&self, job_id: &str) -> Result> { self.get_task_logs(job_id) } - fn query_task_logs(&self, task_name: Option<&str>, level: Option<&str>, since_ms: i64, limit: i64) -> Result> { + fn get_task_logs(&self, job_id: &str) -> Result> { + self.get_task_logs(job_id) + } + fn query_task_logs( + &self, + task_name: Option<&str>, + level: Option<&str>, + since_ms: i64, + limit: i64, + ) -> Result> { self.query_task_logs(task_name, level, since_ms, limit) } - fn purge_task_logs(&self, older_than_ms: i64) -> Result { self.purge_task_logs(older_than_ms) } - fn get_circuit_breaker(&self, task_name: &str) -> Result> { self.get_circuit_breaker(task_name) } - fn upsert_circuit_breaker(&self, row: &CircuitBreakerRow) -> Result<()> { self.upsert_circuit_breaker(row) } - fn list_circuit_breakers(&self) -> Result> { self.list_circuit_breakers() } - fn request_cancel(&self, id: &str) -> Result { self.request_cancel(id) } - fn is_cancel_requested(&self, id: &str) -> Result { self.is_cancel_requested(id) } - fn mark_cancelled(&self, id: &str) -> Result<()> { self.mark_cancelled(id) } - fn register_worker(&self, worker_id: &str, queues: &str) -> Result<()> { self.register_worker(worker_id, queues) } - fn heartbeat(&self, worker_id: &str) -> Result<()> { self.heartbeat(worker_id) } - fn list_workers(&self) -> Result> { self.list_workers() } - fn reap_dead_workers(&self) -> Result { self.reap_dead_workers() } - fn unregister_worker(&self, worker_id: &str) -> Result<()> { self.unregister_worker(worker_id) } + fn purge_task_logs(&self, older_than_ms: i64) -> Result { + self.purge_task_logs(older_than_ms) + } + fn get_circuit_breaker(&self, task_name: &str) -> Result> { + self.get_circuit_breaker(task_name) + } + fn upsert_circuit_breaker(&self, row: &CircuitBreakerRow) -> Result<()> { + self.upsert_circuit_breaker(row) + } + fn list_circuit_breakers(&self) -> Result> { + self.list_circuit_breakers() + } + fn request_cancel(&self, id: &str) -> Result { + self.request_cancel(id) + } + fn is_cancel_requested(&self, id: &str) -> Result { + self.is_cancel_requested(id) + } + fn mark_cancelled(&self, id: &str) -> Result<()> { + self.mark_cancelled(id) + } + fn register_worker(&self, worker_id: &str, queues: &str) -> Result<()> { + self.register_worker(worker_id, queues) + } + fn heartbeat(&self, worker_id: &str) -> Result<()> { + self.heartbeat(worker_id) + } + fn list_workers(&self) -> Result> { + self.list_workers() + } + fn reap_dead_workers(&self) -> Result { + self.reap_dead_workers() + } + fn unregister_worker(&self, worker_id: &str) -> Result<()> { + self.unregister_worker(worker_id) + } } diff --git a/crates/taskito-core/src/storage/sqlite/workers.rs b/crates/taskito-core/src/storage/sqlite/workers.rs index b10a2c26..21464735 100644 --- a/crates/taskito-core/src/storage/sqlite/workers.rs +++ b/crates/taskito-core/src/storage/sqlite/workers.rs @@ -1,10 +1,10 @@ use diesel::prelude::*; -use crate::error::Result; -use crate::job::now_millis; use super::super::models::*; use super::super::schema::workers; use super::SqliteStorage; +use crate::error::Result; +use crate::job::now_millis; /// Dead worker threshold: 30 seconds without heartbeat. const DEAD_WORKER_THRESHOLD_MS: i64 = 30_000; @@ -58,9 +58,8 @@ impl SqliteStorage { let mut conn = self.conn()?; let cutoff = now_millis() - DEAD_WORKER_THRESHOLD_MS; - let affected = diesel::delete( - workers::table.filter(workers::last_heartbeat.lt(cutoff)) - ).execute(&mut conn)?; + let affected = diesel::delete(workers::table.filter(workers::last_heartbeat.lt(cutoff))) + .execute(&mut conn)?; Ok(affected as u64) } diff --git a/crates/taskito-core/src/storage/traits.rs b/crates/taskito-core/src/storage/traits.rs index 99de484d..94be8119 100644 --- a/crates/taskito-core/src/storage/traits.rs +++ b/crates/taskito-core/src/storage/traits.rs @@ -1,12 +1,12 @@ use crate::error::Result; use crate::job::{Job, NewJob}; use crate::storage::models::*; -use crate::storage::sqlite::{DeadJob, QueueStats}; +use crate::storage::{DeadJob, QueueStats}; /// Trait abstracting the storage backend for the task queue. /// -/// `SqliteStorage` is the primary implementation. This trait enables -/// alternative backends and simplifies testing with mock storage. +/// Implementations include `SqliteStorage` and `PostgresStorage`. This trait +/// enables alternative backends and simplifies testing with mock storage. pub trait Storage: Send + Sync + Clone { // ── Job operations ────────────────────────────────────────────── @@ -37,6 +37,7 @@ pub trait Storage: Send + Sync + Clone { fn get_job(&self, id: &str) -> Result>; fn stats(&self) -> Result; fn purge_completed(&self, older_than_ms: i64) -> Result; + fn purge_completed_with_ttl(&self, global_cutoff_ms: i64) -> Result; fn reap_stale_jobs(&self, now: i64) -> Result>; fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()>; fn get_job_errors(&self, job_id: &str) -> Result>; @@ -53,6 +54,7 @@ pub trait Storage: Send + Sync + Clone { fn get_rate_limit(&self, key: &str) -> Result>; fn upsert_rate_limit(&self, row: &RateLimitRow) -> Result<()>; + fn try_acquire_token(&self, key: &str, max_tokens: f64, refill_rate: f64) -> Result; // ── Periodic task operations ──────────────────────────────────── diff --git a/crates/taskito-python/Cargo.toml b/crates/taskito-python/Cargo.toml index 1332c725..7f8eecaf 100644 --- a/crates/taskito-python/Cargo.toml +++ b/crates/taskito-python/Cargo.toml @@ -3,6 +3,10 @@ name = "taskito-python" version = "0.2.2" edition = "2021" +[features] +default = [] +postgres = ["taskito-core/postgres"] + [lib] name = "_taskito" crate-type = ["cdylib"] diff --git a/crates/taskito-python/src/lib.rs b/crates/taskito-python/src/lib.rs index 0c8056a8..8f73e5cb 100644 --- a/crates/taskito-python/src/lib.rs +++ b/crates/taskito-python/src/lib.rs @@ -1,13 +1,13 @@ use pyo3::prelude::*; -mod py_queue; +mod py_config; mod py_job; +mod py_queue; mod py_worker; -mod py_config; -use py_queue::PyQueue; -use py_job::PyJob; use py_config::PyTaskConfig; +use py_job::PyJob; +use py_queue::PyQueue; #[pymodule] fn _taskito(m: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/crates/taskito-python/src/py_config.rs b/crates/taskito-python/src/py_config.rs index 2ec61cf8..ea3e91a6 100644 --- a/crates/taskito-python/src/py_config.rs +++ b/crates/taskito-python/src/py_config.rs @@ -27,6 +27,7 @@ pub struct PyTaskConfig { } #[pymethods] +#[allow(clippy::too_many_arguments)] impl PyTaskConfig { #[new] #[pyo3(signature = (name, max_retries=3, retry_backoff=1.0, timeout=300, priority=0, rate_limit=None, queue="default".to_string(), circuit_breaker_threshold=None, circuit_breaker_window=None, circuit_breaker_cooldown=None))] diff --git a/crates/taskito-python/src/py_queue.rs b/crates/taskito-python/src/py_queue.rs index d4e1d830..783fe387 100644 --- a/crates/taskito-python/src/py_queue.rs +++ b/crates/taskito-python/src/py_queue.rs @@ -1,17 +1,23 @@ -use std::sync::Arc; +// pyo3's #[pymethods] macro generates Into conversions that trigger this lint +#![allow(clippy::useless_conversion)] + use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; use pyo3::prelude::*; use pyo3::types::PyDict; -use taskito_core::job::{NewJob, now_millis}; +use taskito_core::job::{now_millis, NewJob}; use taskito_core::periodic::next_cron_time; use taskito_core::resilience::circuit_breaker::CircuitBreakerConfig; use taskito_core::resilience::rate_limiter::RateLimitConfig; use taskito_core::resilience::retry::RetryPolicy; use taskito_core::scheduler::{Scheduler, SchedulerConfig, TaskConfig}; use taskito_core::storage::models::NewPeriodicTaskRow; +#[cfg(feature = "postgres")] +use taskito_core::storage::postgres::PostgresStorage; use taskito_core::storage::sqlite::SqliteStorage; +use taskito_core::storage::{Storage, StorageBackend}; use crate::py_config::PyTaskConfig; use crate::py_job::PyJob; @@ -20,7 +26,7 @@ use crate::py_worker::WorkerPool; /// The core queue engine exposed to Python. #[pyclass] pub struct PyQueue { - storage: SqliteStorage, + storage: StorageBackend, db_path: String, num_workers: usize, default_retry: i32, @@ -31,9 +37,14 @@ pub struct PyQueue { } #[pymethods] +#[allow( + clippy::too_many_arguments, + clippy::useless_conversion, + unused_variables +)] impl PyQueue { #[new] - #[pyo3(signature = (db_path=".taskito/taskito.db", workers=0, default_retry=3, default_timeout=300, default_priority=0, result_ttl=None))] + #[pyo3(signature = (db_path=".taskito/taskito.db", workers=0, default_retry=3, default_timeout=300, default_priority=0, result_ttl=None, backend="sqlite", db_url=None, schema="taskito"))] pub fn new( db_path: &str, workers: usize, @@ -41,9 +52,33 @@ impl PyQueue { default_timeout: i64, default_priority: i32, result_ttl: Option, + backend: &str, + db_url: Option<&str>, + schema: &str, ) -> PyResult { - let storage = SqliteStorage::new(db_path) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + let storage = match backend { + "sqlite" => { + let s = SqliteStorage::new(db_path) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + StorageBackend::Sqlite(s) + } + #[cfg(feature = "postgres")] + "postgres" | "postgresql" => { + let url = db_url.ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err( + "db_url is required for postgres backend", + ) + })?; + let s = PostgresStorage::with_schema(url, schema) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + StorageBackend::Postgres(s) + } + _ => { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "Unknown backend: '{backend}'. Use 'sqlite' or 'postgres'." + ))); + } + }; let num_workers = if workers == 0 { std::thread::available_parallelism() @@ -195,8 +230,8 @@ impl PyQueue { "cancelled" => 5, _ => { return Err(pyo3::exceptions::PyValueError::new_err(format!( - "Invalid status: {s}. Use: pending, running, complete, failed, dead, cancelled" - ))) + "Invalid status: {s}. Use: pending, running, complete, failed, dead, cancelled" + ))) } }), None => None, @@ -426,14 +461,17 @@ impl PyQueue { base_delay_ms: (tc.retry_backoff * 1000.0) as i64, max_delay_ms: 300_000, }; - let rate_limit = tc.rate_limit.as_ref().and_then(|s| RateLimitConfig::parse(s)); - let circuit_breaker = tc.circuit_breaker_threshold.map(|threshold| { - CircuitBreakerConfig { - threshold, - window_ms: tc.circuit_breaker_window.unwrap_or(60) * 1000, - cooldown_ms: tc.circuit_breaker_cooldown.unwrap_or(300) * 1000, - } - }); + let rate_limit = tc + .rate_limit + .as_ref() + .and_then(|s| RateLimitConfig::parse(s)); + let circuit_breaker = + tc.circuit_breaker_threshold + .map(|threshold| CircuitBreakerConfig { + threshold, + window_ms: tc.circuit_breaker_window.unwrap_or(60) * 1000, + cooldown_ms: tc.circuit_breaker_cooldown.unwrap_or(300) * 1000, + }); scheduler.register_task( tc.name.clone(), TaskConfig { @@ -451,8 +489,13 @@ impl PyQueue { let registry_arc = Arc::new(task_registry); let filters_arc = Arc::new(retry_filters.into()); - let worker_pool = - WorkerPool::start(self.num_workers, job_rx, result_tx, registry_arc, filters_arc); + let worker_pool = WorkerPool::start( + self.num_workers, + job_rx, + result_tx, + registry_arc, + filters_arc, + ); let scheduler_arc = Arc::new(scheduler); let scheduler_for_dispatch = scheduler_arc.clone(); @@ -474,10 +517,16 @@ impl PyQueue { }); let scheduler_handle = std::thread::spawn(move || { - let rt = tokio::runtime::Builder::new_current_thread() + let rt = match tokio::runtime::Builder::new_current_thread() .enable_all() .build() - .expect("failed to build tokio runtime"); + { + Ok(rt) => rt, + Err(e) => { + eprintln!("taskito: failed to build tokio runtime: {e}"); + return; + } + }; rt.block_on(scheduler_for_dispatch.run(job_tx)); }); @@ -543,7 +592,11 @@ impl PyQueue { /// Get raw metric records for a task (or all tasks). #[pyo3(signature = (task_name=None, since_seconds=3600))] - pub fn get_metrics(&self, task_name: Option<&str>, since_seconds: i64) -> PyResult> { + pub fn get_metrics( + &self, + task_name: Option<&str>, + since_seconds: i64, + ) -> PyResult> { let since_ms = now_millis() - (since_seconds * 1000); let rows = self .storage @@ -575,7 +628,9 @@ impl PyQueue { .storage .get_job(job_id) .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))? - .ok_or_else(|| pyo3::exceptions::PyValueError::new_err(format!("Job not found: {job_id}")))?; + .ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err(format!("Job not found: {job_id}")) + })?; let new_job = NewJob { queue: original.queue, diff --git a/crates/taskito-python/src/py_worker.rs b/crates/taskito-python/src/py_worker.rs index f779c3a6..1a83f6cf 100644 --- a/crates/taskito-python/src/py_worker.rs +++ b/crates/taskito-python/src/py_worker.rs @@ -70,7 +70,7 @@ fn worker_loop( execute_task(py, &task_registry, &job) }); - let wall_time_ns = start.elapsed().as_nanos() as i64; + let wall_time_ns: i64 = start.elapsed().as_nanos().try_into().unwrap_or(i64::MAX); let job_result = match result { Ok(result_bytes) => JobResult::Success { @@ -118,24 +118,15 @@ fn worker_loop( } } -fn execute_task( - py: Python<'_>, - task_registry: &PyObject, - job: &Job, -) -> PyResult>> { +fn execute_task(py: Python<'_>, task_registry: &PyObject, job: &Job) -> PyResult>> { let cloudpickle = py.import_bound("cloudpickle")?; let registry = task_registry.bind(py); // Look up the task function let registry_dict: &Bound<'_, PyDict> = registry.downcast()?; - let task_fn = registry_dict - .get_item(&job.task_name)? - .ok_or_else(|| { - pyo3::exceptions::PyKeyError::new_err(format!( - "task '{}' not registered", - job.task_name - )) - })?; + let task_fn = registry_dict.get_item(&job.task_name)?.ok_or_else(|| { + pyo3::exceptions::PyKeyError::new_err(format!("task '{}' not registered", job.task_name)) + })?; // Set job context before execution let context_mod = py.import_bound("taskito.context")?; @@ -144,23 +135,26 @@ fn execute_task( (&job.id, &job.task_name, job.retry_count, &job.queue), )?; - // Deserialize arguments: (args, kwargs) - let payload_bytes = PyBytes::new_bound(py, &job.payload); - let unpickled = cloudpickle.call_method1("loads", (payload_bytes,))?; - let args_tuple: Bound<'_, PyTuple> = unpickled.downcast_into()?; - - let args = args_tuple.get_item(0)?; - let kwargs = args_tuple.get_item(1)?; - - // Call the task function - let result = if kwargs.is_none() { - let args_tuple_inner: Bound<'_, PyTuple> = args.downcast_into()?; - task_fn.call(args_tuple_inner, None) - } else { - let kwargs_dict: Bound<'_, PyDict> = kwargs.downcast_into()?; - let args_tuple_inner: Bound<'_, PyTuple> = args.downcast_into()?; - task_fn.call(args_tuple_inner, Some(&kwargs_dict)) - }; + // Wrap deserialization + call so _clear_context is always called + let result = (|| -> PyResult> { + // Deserialize arguments: (args, kwargs) + let payload_bytes = PyBytes::new_bound(py, &job.payload); + let unpickled = cloudpickle.call_method1("loads", (payload_bytes,))?; + let args_tuple: Bound<'_, PyTuple> = unpickled.downcast_into()?; + + let args = args_tuple.get_item(0)?; + let kwargs = args_tuple.get_item(1)?; + + // Call the task function + if kwargs.is_none() { + let args_tuple_inner: Bound<'_, PyTuple> = args.downcast_into()?; + task_fn.call(args_tuple_inner, None) + } else { + let kwargs_dict: Bound<'_, PyDict> = kwargs.downcast_into()?; + let args_tuple_inner: Bound<'_, PyTuple> = args.downcast_into()?; + task_fn.call(args_tuple_inner, Some(&kwargs_dict)) + } + })(); // Clear context after execution (whether success or failure) let _ = context_mod.call_method0("_clear_context"); @@ -200,7 +194,10 @@ fn format_python_error(py: Python<'_>, e: &PyErr) -> String { fn is_cancelled_error(py: Python<'_>, e: &PyErr) -> bool { if let Ok(exceptions_mod) = py.import_bound("taskito.exceptions") { if let Ok(cancelled_cls) = exceptions_mod.getattr("TaskCancelledError") { - return e.get_type_bound(py).is_subclass(&cancelled_cls).unwrap_or(false); + return e + .get_type_bound(py) + .is_subclass(&cancelled_cls) + .unwrap_or(false); } } false @@ -209,10 +206,12 @@ fn is_cancelled_error(py: Python<'_>, e: &PyErr) -> bool { /// Get the fully-qualified class name of a Python exception. fn get_exception_class_name(py: Python<'_>, e: &PyErr) -> String { let type_obj = e.get_type_bound(py); - let module = type_obj.getattr("__module__") + let module = type_obj + .getattr("__module__") .and_then(|m| m.extract::()) .unwrap_or_default(); - let qualname = type_obj.getattr("__qualname__") + let qualname = type_obj + .getattr("__qualname__") .and_then(|q| q.extract::()) .unwrap_or_else(|_| "Exception".to_string()); diff --git a/docs/architecture.md b/docs/architecture.md index 04f0ca61..1a484ac4 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -125,7 +125,8 @@ jobs (id, queue, task_name, payload, status, priority, -- Dead letter queue dead_letter (id, original_job_id, queue, task_name, - payload, error, retry_count, failed_at, metadata) + payload, error, retry_count, failed_at, metadata, + priority, max_retries, timeout_ms, result_ttl_ms) -- Token bucket rate limiting rate_limits (key, tokens, max_tokens, refill_rate, last_refill) diff --git a/docs/changelog.md b/docs/changelog.md index f2821721..027b73bf 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -2,6 +2,47 @@ All notable changes to taskito are documented here. +## 0.2.3 + +### Build & Tooling + +- **Pre-commit hooks** — Added `.pre-commit-config.yaml` with local hooks for `cargo fmt`, `cargo clippy`, `ruff check`, `ruff format`, and `mypy` + +### Critical Fixes + +- **Dashboard dead routes** — Moved `/logs` and `/replay-history` handlers above the generic catch-all in `dashboard.py`, fixing 404s on these endpoints +- **Stale `__version__`** — Replaced hardcoded version with `importlib.metadata.version()` with fallback +- **`retry_dead` non-atomic** — Wrapped enqueue + delete in a single transaction (SQLite & Postgres), preventing ghost dead letters on partial failure +- **`retry_dead` hardcoded defaults** — Added `priority`, `max_retries`, `timeout_ms`, `result_ttl_ms` columns to `dead_letter` table; replayed jobs now preserve their original configuration +- **`enqueue_unique` race condition** — Wrapped check + insert in a transaction; catches unique constraint violations to return the existing job instead of erroring +- **`now_millis()` panic** — Replaced `.expect()` with `.unwrap_or(Duration::ZERO)` to prevent scheduler panic on clock issues +- **`reap_stale` double error records** — Removed redundant `storage.fail()` call; `handle_result` already records the failure +- **README cron format** — Updated example to correct 6-field format: `"0 0 */6 * * *"` + +### Important Fixes + +- **`result.py` hardcoded cloudpickle** — `job.result()` now uses the queue's configured serializer for deserialization +- **Context leak on deserialization failure** — Wrapped deserialization + call in closure; `_clear_context` always runs via `finally` +- **OTel spans not thread-safe** — Added `threading.Lock` around all `_spans` dict access in `OpenTelemetryMiddleware` +- **`build_periodic_payload` misleading `_kwargs` param** — Removed unused parameter, added explanatory comment +- **Tokio runtime panic** — Replaced `.expect()` with graceful error handling on runtime creation +- **`dequeue` LIMIT 10** — Increased to 100 for better throughput under load (both SQLite & Postgres) +- **`check_periodic` not atomic** — Uses `enqueue_unique` with deterministic key to prevent duplicate periodic jobs +- **SQLite `purge_completed_with_ttl` no transaction** — Wrapped in transaction for consistency +- **Django admin status validation** — Added try/except around `queue.list_jobs()` to handle connection errors gracefully +- **Silent job loss on `get_job` None** — Added `warn!` logging when a dequeued job ID returns None +- **Cascade cleanup on job purge** — `purge_completed()` and `purge_completed_with_ttl()` now automatically delete orphaned child records (`job_errors`, `task_logs`, `task_metrics`, `job_dependencies`, `replay_history`) when removing completed jobs + +### Minor Fixes + +- **`cascade_cancel` O(n²)** — Replaced `Vec::contains` with `HashSet` for dependency lookups (both backends) +- **`chain.apply()` hardcoded 300s timeout** — Now derives timeout from `sig.options.get("timeout", 300)` +- **`_FakeJobResult` missing `refresh()`** — Added no-op method for test mode compatibility +- **Storage trait doc outdated** — Updated to mention both SQLite and Postgres backends +- **`wall_time_ns` truncation** — Uses `.try_into().unwrap_or(i64::MAX)` to prevent silent overflow + +--- + ## 0.2.2 - Added `readme` field to `pyproject.toml` so PyPI displays the project description. diff --git a/docs/guide/advanced.md b/docs/guide/advanced.md index ce9d5a26..4a83f2a1 100644 --- a/docs/guide/advanced.md +++ b/docs/guide/advanced.md @@ -13,7 +13,7 @@ assert job1.id == job2.id # Same job, not duplicated Once the original job completes (or fails to DLQ), the key is released and a new job can be created with the same key. !!! info "Implementation" - Deduplication uses a partial unique index: `CREATE UNIQUE INDEX ... ON jobs(unique_key) WHERE unique_key IS NOT NULL AND status IN (0, 1)`. Only pending and running jobs participate. + Deduplication uses a partial unique index: `CREATE UNIQUE INDEX ... ON jobs(unique_key) WHERE unique_key IS NOT NULL AND status IN (0, 1)`. Only pending and running jobs participate. The check-and-insert is atomic (transaction-protected), so concurrent calls with the same `unique_key` are handled gracefully without race conditions. ## Job Cancellation @@ -59,6 +59,35 @@ The scheduler checks every ~60 seconds and purges: Set to `None` (default) to disable auto-cleanup. +### Cascade Cleanup + +When jobs are purged — either manually via `purge_completed()` or automatically via `result_ttl` — their related child records are also deleted: + +- Error history (`job_errors`) +- Task logs (`task_logs`) +- Task metrics (`task_metrics`) +- Job dependencies (`job_dependencies`) +- Replay history (`replay_history`) + +This prevents orphaned records from accumulating when parent jobs are removed. + +```python +# Manual purge — child records are cleaned up automatically +deleted = queue.purge_completed(older_than=3600) +print(f"Purged {deleted} jobs and their related records") + +# With per-job TTL — cascade cleanup still applies +job = resize_image.apply_async( + args=("photo.jpg",), + result_ttl=600, # This job's results expire after 10 minutes +) +# When this job is purged (after 10 min), its errors, logs, +# metrics, dependencies, and replay history are also removed. +``` + +!!! note + Dead letter entries are **not** cascade-deleted — they have their own lifecycle managed by `purge_dead()`. Timestamp-based cleanup (`result_ttl`) of error history, logs, and metrics also continues to run independently, catching old records regardless of whether the parent job still exists. + ## Async Support All inspection methods have async variants that run in a thread pool: diff --git a/docs/guide/dashboard.md b/docs/guide/dashboard.md index dae28d9a..4ed236e6 100644 --- a/docs/guide/dashboard.md +++ b/docs/guide/dashboard.md @@ -162,6 +162,51 @@ Delete a single dead letter entry. Cancel a pending job. Returns `204 No Content` on success or `409 Conflict` if the job is not in a cancellable state. +### `GET /api/jobs/{id}/logs` + +Task execution logs for a specific job. + +```bash +curl http://localhost:8080/api/jobs/01H5K6X.../logs +``` + +### `GET /api/jobs/{id}/replay-history` + +Replay history for a job that has been retried from the DLQ. + +```bash +curl http://localhost:8080/api/jobs/01H5K6X.../replay-history +``` + +### `GET /api/logs` + +Query task execution logs across all jobs. + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `limit` | `int` | `50` | Page size | +| `offset` | `int` | `0` | Pagination offset | + +```bash +curl http://localhost:8080/api/logs?limit=20 +``` + +### `GET /api/metrics` + +Task execution metrics (timing, throughput). + +```bash +curl http://localhost:8080/api/metrics +``` + +### `GET /api/circuit-breakers` + +Current state of all circuit breakers. + +```bash +curl http://localhost:8080/api/circuit-breakers +``` + ### `GET /api/workers` List registered workers with heartbeat status. diff --git a/docs/guide/error-handling.md b/docs/guide/error-handling.md index 59aa4ae9..c60afd25 100644 --- a/docs/guide/error-handling.md +++ b/docs/guide/error-handling.md @@ -71,6 +71,9 @@ for d in dead: queue.retry_dead(d["id"]) ``` +!!! note "Config preservation" + Replayed jobs preserve the original job's `priority`, `max_retries`, `timeout`, and `result_ttl` settings — no need to re-specify them. + **Error message mentions serialization** — See [Serialization Errors](#serialization-errors) below. ## Common Error Scenarios diff --git a/docs/guide/otel.md b/docs/guide/otel.md index 11adab24..29634d64 100644 --- a/docs/guide/otel.md +++ b/docs/guide/otel.md @@ -77,4 +77,7 @@ queue = Queue(middleware=[ ]) ``` +!!! note "Thread safety" + `OpenTelemetryMiddleware` is thread-safe and can be used with multi-worker configurations. Internal span tracking is protected by a lock. + See the [Middleware guide](middleware.md) for more on combining middleware. diff --git a/docs/guide/retries.md b/docs/guide/retries.md index df679815..8de16e4a 100644 --- a/docs/guide/retries.md +++ b/docs/guide/retries.md @@ -107,6 +107,9 @@ for d in dead: new_job_id = queue.retry_dead(dead[0]["id"]) ``` +!!! note "Config preservation" + Replayed jobs preserve the original job's `priority`, `max_retries`, `timeout`, and `result_ttl` settings. You don't need to re-specify them — the DLQ stores the full configuration. + ### Purge Old Dead Letters ```python diff --git a/docs/guide/serializers.md b/docs/guide/serializers.md index fbe8873f..6d3f0940 100644 --- a/docs/guide/serializers.md +++ b/docs/guide/serializers.md @@ -64,6 +64,9 @@ The protocol requires: The serializer is used for both task arguments (`(args, kwargs)` tuple) and return values. +!!! note "Result deserialization" + `job.result()` uses the queue's configured serializer for deserialization. If you're using `JsonSerializer` or a custom serializer, results are correctly deserialized with that serializer — not hardcoded cloudpickle. + ## Configuration Pass the serializer to the `Queue` constructor: diff --git a/py_src/taskito/__init__.py b/py_src/taskito/__init__.py index 028b04d8..0e35f4ff 100644 --- a/py_src/taskito/__init__.py +++ b/py_src/taskito/__init__.py @@ -50,4 +50,9 @@ "group", "starmap", ] -__version__ = "0.1.0" +try: + from importlib.metadata import version as _get_version + + __version__ = _get_version("taskito") +except Exception: + __version__ = "0.2.2" diff --git a/py_src/taskito/_taskito.pyi b/py_src/taskito/_taskito.pyi index fb20afea..13324232 100644 --- a/py_src/taskito/_taskito.pyi +++ b/py_src/taskito/_taskito.pyi @@ -68,6 +68,9 @@ class PyQueue: default_timeout: int = 300, default_priority: int = 0, result_ttl: int | None = None, + backend: str = "sqlite", + db_url: str | None = None, + schema: str = "taskito", ) -> None: ... def request_shutdown(self) -> None: ... def enqueue( diff --git a/py_src/taskito/app.py b/py_src/taskito/app.py index 3beb87ba..d19873cc 100644 --- a/py_src/taskito/app.py +++ b/py_src/taskito/app.py @@ -8,6 +8,7 @@ import logging import os import signal +import threading from collections.abc import Sequence from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, Callable @@ -65,13 +66,16 @@ def __init__( result_ttl: int | None = None, serializer: Serializer | None = None, middleware: list[TaskMiddleware] | None = None, + backend: str = "sqlite", + db_url: str | None = None, + schema: str = "taskito", ): - """Initialize a new task queue backed by SQLite. + """Initialize a new task queue. Args: db_path: Path to the SQLite database file. Defaults to ``.taskito/taskito.db``. Parent directories are created - automatically. + automatically. Ignored when backend is ``"postgres"``. workers: Number of worker threads (0 = auto-detect CPU count). default_retry: Default max retry attempts for tasks. default_timeout: Default task timeout in seconds. @@ -80,11 +84,17 @@ def __init__( seconds. None disables auto-cleanup. serializer: Serializer for task payloads. Defaults to CloudpickleSerializer. middleware: List of global middleware instances applied to all tasks. + backend: Storage backend — ``"sqlite"`` (default) or ``"postgres"``. + db_url: PostgreSQL connection URL (required when backend is ``"postgres"``). + Example: ``"postgresql://user:pass@localhost/taskito"``. + schema: PostgreSQL schema name for all taskito tables. Defaults to + ``"taskito"``. Ignored when backend is ``"sqlite"``. """ - # Ensure parent directory exists - db_dir = os.path.dirname(db_path) - if db_dir: - os.makedirs(db_dir, exist_ok=True) + if backend == "sqlite": + # Ensure parent directory exists for SQLite + db_dir = os.path.dirname(db_path) + if db_dir: + os.makedirs(db_dir, exist_ok=True) self._inner = PyQueue( db_path=db_path, @@ -93,7 +103,13 @@ def __init__( default_timeout=default_timeout, default_priority=default_priority, result_ttl=result_ttl, + backend=backend, + db_url=db_url, + schema=schema, ) + self._backend = backend + self._db_url = db_url + self._schema = schema self._db_path = db_path self._workers = workers or os.cpu_count() or 1 self._task_registry: dict[str, Callable] = {} @@ -321,7 +337,12 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: try: ret = fn(*args, **kwargs) if asyncio.iscoroutine(ret): - ret = asyncio.run(ret) + loop = getattr(queue_ref, "_async_loop", None) + if loop is not None and loop.is_running(): + future = asyncio.run_coroutine_threadsafe(ret, loop) + ret = future.result() + else: + ret = asyncio.run(ret) result = ret except Exception as exc: error = exc @@ -525,7 +546,19 @@ def _print_banner(self, queues: list[str]) -> None: \__\__,_|___/_|\_\_|\__\___/ v{__version__} """ lines = [banner] - lines.append(f"> DB: {self._db_path}") + lines.append(f"> Backend: {self._backend}") + if self._backend == "sqlite": + lines.append(f"> DB: {self._db_path}") + else: + # Mask password in connection URL for display + url = self._db_url or "" + if "@" in url: + pre, post = url.split("@", 1) + if ":" in pre: + scheme_user = pre.rsplit(":", 1)[0] + url = f"{scheme_user}:****@{post}" + lines.append(f"> DB: {url}") + lines.append(f"> Schema: {self._schema}") lines.append(f"> Concurrency: {self._workers} (threads)") lines.append(f"> Queues: {', '.join(queues)}") lines.append("") @@ -545,6 +578,23 @@ def _print_banner(self, queues: list[str]) -> None: print("\n".join(lines)) + def _start_async_loop(self) -> None: + """Start a shared event loop in a background thread for async tasks.""" + self._async_loop = asyncio.new_event_loop() + self._async_thread = threading.Thread( + target=self._async_loop.run_forever, + daemon=True, + name="taskito-async-loop", + ) + self._async_thread.start() + + def _stop_async_loop(self) -> None: + """Stop the shared async event loop.""" + loop = getattr(self, "_async_loop", None) + if loop is not None and loop.is_running(): + loop.call_soon_threadsafe(loop.stop) + self._async_thread.join(timeout=5) + def run_worker(self, queues: Sequence[str] | None = None) -> None: """Start the worker loop. Blocks until interrupted. @@ -572,9 +622,10 @@ def run_worker(self, queues: Sequence[str] | None = None) -> None: worker_queues = queue_list or ["default"] self._print_banner(worker_queues) - # Set up signal handler for graceful shutdown (only in main thread) - import threading + # Start shared async event loop for async tasks + self._start_async_loop() + # Set up signal handler for graceful shutdown (only in main thread) is_main = threading.current_thread() is threading.main_thread() original_sigint = None @@ -598,6 +649,7 @@ def shutdown_handler(signum: int, frame: Any) -> None: except KeyboardInterrupt: logger.info("Worker force-stopped.") finally: + self._stop_async_loop() logger.info("Worker stopped.") if is_main and original_sigint is not None: signal.signal(signal.SIGINT, original_sigint) diff --git a/py_src/taskito/canvas.py b/py_src/taskito/canvas.py index 6db36e27..4a161cdd 100644 --- a/py_src/taskito/canvas.py +++ b/py_src/taskito/canvas.py @@ -65,7 +65,7 @@ def apply(self, queue: Queue | None = None) -> JobResult: kwargs=sig.kwargs if sig.kwargs else None, **sig.options, ) - prev_result = last_job.result(timeout=300) + prev_result = last_job.result(timeout=sig.options.get("timeout", 300)) return last_job # type: ignore[return-value] diff --git a/py_src/taskito/contrib/django/__init__.py b/py_src/taskito/contrib/django/__init__.py new file mode 100644 index 00000000..efa1b06c --- /dev/null +++ b/py_src/taskito/contrib/django/__init__.py @@ -0,0 +1,25 @@ +"""Django integration for taskito. + +Add ``"taskito.contrib.django"`` to your ``INSTALLED_APPS`` and configure +via Django settings:: + + INSTALLED_APPS = [ + ... + "taskito.contrib.django", + ] + + TASKITO_DB_PATH = ".taskito/taskito.db" + TASKITO_BACKEND = "sqlite" # or "postgres" + TASKITO_DB_URL = None # required for postgres + TASKITO_WORKERS = 0 # 0 = auto-detect + TASKITO_DEFAULT_RETRY = 3 + TASKITO_DEFAULT_TIMEOUT = 300 + TASKITO_DEFAULT_PRIORITY = 0 + TASKITO_RESULT_TTL = None # seconds, or None to disable + +Requires the ``django`` optional dependency:: + + pip install taskito[django] +""" + +default_app_config = "taskito.contrib.django.apps.TaskitoConfig" diff --git a/py_src/taskito/contrib/django/admin.py b/py_src/taskito/contrib/django/admin.py new file mode 100644 index 00000000..0f94158f --- /dev/null +++ b/py_src/taskito/contrib/django/admin.py @@ -0,0 +1,176 @@ +"""Django admin views for taskito queue inspection. + +Registers custom admin views for browsing jobs, dead letters, and queue stats. +Uses taskito's Python API directly — no Django ORM models needed. +""" + +from __future__ import annotations + +from typing import Any + +try: + from django.contrib import admin + from django.http import HttpRequest, HttpResponse + from django.template.response import TemplateResponse + from django.urls import path +except ImportError as e: + raise ImportError( + "Django integration requires 'django'. Install with: pip install taskito[django]" + ) from e + + +class TaskitoAdminSite(admin.AdminSite): + """Custom admin site with taskito queue views.""" + + site_header = "Taskito Admin" + site_title = "Taskito" + + def get_urls(self) -> list: + urls = super().get_urls() + custom = [ + path("taskito/", self.admin_view(self.dashboard_view), name="taskito_dashboard"), + path("taskito/jobs/", self.admin_view(self.jobs_view), name="taskito_jobs"), + path( + "taskito/jobs//", + self.admin_view(self.job_detail_view), + name="taskito_job_detail", + ), + path( + "taskito/dead-letters/", + self.admin_view(self.dead_letters_view), + name="taskito_dead_letters", + ), + ] + return custom + urls # type: ignore[no-any-return] + + def dashboard_view(self, request: HttpRequest) -> HttpResponse: + from taskito.contrib.django.settings import get_queue + + queue = get_queue() + stats = queue.stats() + context = {**self.each_context(request), "stats": stats, "title": "Taskito Dashboard"} + return TemplateResponse(request, "taskito/admin/dashboard.html", context) + + def jobs_view(self, request: HttpRequest) -> HttpResponse: + from taskito.contrib.django.settings import get_queue + + queue = get_queue() + status = request.GET.get("status") + queue_name = request.GET.get("queue") + task_name = request.GET.get("task_name") + page = int(request.GET.get("page", "1")) + per_page = 50 + + try: + jobs = queue.list_jobs( + status=status, + queue=queue_name, + task_name=task_name, + limit=per_page, + offset=(page - 1) * per_page, + ) + except (ValueError, Exception): + jobs = [] + context = { + **self.each_context(request), + "jobs": [j.to_dict() for j in jobs], + "filters": {"status": status, "queue": queue_name, "task_name": task_name}, + "page": page, + "title": "Taskito Jobs", + } + return TemplateResponse(request, "taskito/admin/jobs.html", context) + + def job_detail_view(self, request: HttpRequest, job_id: str) -> HttpResponse: + from taskito.contrib.django.settings import get_queue + + queue = get_queue() + job = queue.get_job(job_id) + errors = queue.job_errors(job_id) if job else [] + context = { + **self.each_context(request), + "job": job.to_dict() if job else None, + "errors": errors, + "title": f"Job {job_id}", + } + return TemplateResponse(request, "taskito/admin/job_detail.html", context) + + def dead_letters_view(self, request: HttpRequest) -> HttpResponse: + from taskito.contrib.django.settings import get_queue + + queue = get_queue() + + if request.method == "POST": + action = request.POST.get("action") + dead_id = request.POST.get("dead_id") + if action == "retry" and dead_id: + queue.retry_dead(dead_id) + + page = int(request.GET.get("page", "1")) + per_page = 50 + dead = queue.dead_letters(limit=per_page, offset=(page - 1) * per_page) + context = { + **self.each_context(request), + "dead_letters": dead, + "page": page, + "title": "Taskito Dead Letters", + } + return TemplateResponse(request, "taskito/admin/dead_letters.html", context) + + +def register_taskito_admin(site: Any = None) -> None: + """Register taskito views on an existing admin site. + + Call this in your project's ``admin.py`` or ``urls.py``:: + + from taskito.contrib.django.admin import register_taskito_admin + register_taskito_admin() + """ + target = site or admin.site + + def dashboard_view(request: HttpRequest) -> HttpResponse: + from taskito.contrib.django.settings import get_queue + + queue = get_queue() + stats = queue.stats() + context = {**target.each_context(request), "stats": stats, "title": "Taskito Dashboard"} + return TemplateResponse(request, "taskito/admin/dashboard.html", context) + + def jobs_view(request: HttpRequest) -> HttpResponse: + from taskito.contrib.django.settings import get_queue + + queue = get_queue() + status = request.GET.get("status") + queue_name = request.GET.get("queue") + task_name = request.GET.get("task_name") + page = int(request.GET.get("page", "1")) + per_page = 50 + try: + jobs = queue.list_jobs( + status=status, + queue=queue_name, + task_name=task_name, + limit=per_page, + offset=(page - 1) * per_page, + ) + except (ValueError, Exception): + jobs = [] + context = { + **target.each_context(request), + "jobs": [j.to_dict() for j in jobs], + "filters": {"status": status, "queue": queue_name, "task_name": task_name}, + "page": page, + "title": "Taskito Jobs", + } + return TemplateResponse(request, "taskito/admin/jobs.html", context) + + original_get_urls = target.get_urls + + def patched_get_urls() -> list: + urls = original_get_urls() + custom = [ + path("taskito/", target.admin_view(dashboard_view), name="taskito_dashboard"), + path("taskito/jobs/", target.admin_view(jobs_view), name="taskito_jobs"), + ] + return custom + urls # type: ignore[no-any-return] + + target.get_urls = patched_get_urls diff --git a/py_src/taskito/contrib/django/apps.py b/py_src/taskito/contrib/django/apps.py new file mode 100644 index 00000000..0de7b161 --- /dev/null +++ b/py_src/taskito/contrib/django/apps.py @@ -0,0 +1,24 @@ +"""Django AppConfig for taskito.""" + +from __future__ import annotations + +try: + from django.apps import AppConfig +except ImportError as e: + raise ImportError( + "Django integration requires 'django'. Install with: pip install taskito[django]" + ) from e + + +class TaskitoConfig(AppConfig): + """Django application configuration for taskito.""" + + name = "taskito.contrib.django" + verbose_name = "Taskito" + default_auto_field = "django.db.models.BigAutoField" + + def ready(self) -> None: + """Auto-discover ``tasks.py`` modules in all installed apps.""" + from django.utils.module_loading import autodiscover_modules + + autodiscover_modules("tasks") diff --git a/py_src/taskito/contrib/django/management/__init__.py b/py_src/taskito/contrib/django/management/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/py_src/taskito/contrib/django/management/commands/__init__.py b/py_src/taskito/contrib/django/management/commands/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/py_src/taskito/contrib/django/management/commands/taskito_dashboard.py b/py_src/taskito/contrib/django/management/commands/taskito_dashboard.py new file mode 100644 index 00000000..ab70da21 --- /dev/null +++ b/py_src/taskito/contrib/django/management/commands/taskito_dashboard.py @@ -0,0 +1,27 @@ +"""Management command to start the taskito web dashboard.""" + +from __future__ import annotations + +try: + from django.core.management.base import BaseCommand +except ImportError as e: + raise ImportError( + "Django integration requires 'django'. Install with: pip install taskito[django]" + ) from e + + +class Command(BaseCommand): + help = "Start the taskito web dashboard" + + def add_arguments(self, parser): # type: ignore[no-untyped-def] + parser.add_argument( + "--host", default="127.0.0.1", help="Bind address (default: 127.0.0.1)" + ) + parser.add_argument("--port", type=int, default=8080, help="Bind port (default: 8080)") + + def handle(self, **options): # type: ignore[no-untyped-def] + from taskito.contrib.django.settings import get_queue + from taskito.dashboard import serve_dashboard + + queue = get_queue() + serve_dashboard(queue, host=options["host"], port=options["port"]) diff --git a/py_src/taskito/contrib/django/management/commands/taskito_info.py b/py_src/taskito/contrib/django/management/commands/taskito_info.py new file mode 100644 index 00000000..4f28f19a --- /dev/null +++ b/py_src/taskito/contrib/django/management/commands/taskito_info.py @@ -0,0 +1,62 @@ +"""Management command to show taskito queue statistics.""" + +from __future__ import annotations + +try: + from django.core.management.base import BaseCommand +except ImportError as e: + raise ImportError( + "Django integration requires 'django'. Install with: pip install taskito[django]" + ) from e + + +class Command(BaseCommand): + help = "Show taskito queue statistics and registered tasks" + + def add_arguments(self, parser): # type: ignore[no-untyped-def] + parser.add_argument( + "--watch", + action="store_true", + default=False, + help="Continuously refresh stats every 2 seconds", + ) + + def handle(self, **options): # type: ignore[no-untyped-def] + from taskito.contrib.django.settings import get_queue + + queue = get_queue() + + if options["watch"]: + self._watch(queue) + else: + self._print(queue) + + def _print(self, queue): # type: ignore[no-untyped-def] + stats = queue.stats() + self.stdout.write("taskito queue statistics") + self.stdout.write("-" * 30) + for key in ("pending", "running", "completed", "failed", "dead", "cancelled"): + self.stdout.write(f" {key:<12} {stats.get(key, 0)}") + total = sum(stats.values()) + self.stdout.write("-" * 30) + self.stdout.write(f" {'total':<12} {total}") + + def _watch(self, queue): # type: ignore[no-untyped-def] + import time + + prev_completed = 0 + try: + while True: + self.stdout.write("\033[2J\033[H", ending="") + stats = queue.stats() + completed = stats.get("completed", 0) + throughput = (completed - prev_completed) / 2.0 + prev_completed = completed + + self._print(queue) + if throughput > 0: + self.stdout.write(f"\n throughput {throughput:.1f} jobs/s") + self.stdout.write("\nRefreshing every 2s... (Ctrl+C to stop)") + time.sleep(2) + except KeyboardInterrupt: + pass diff --git a/py_src/taskito/contrib/django/management/commands/taskito_worker.py b/py_src/taskito/contrib/django/management/commands/taskito_worker.py new file mode 100644 index 00000000..b65541f6 --- /dev/null +++ b/py_src/taskito/contrib/django/management/commands/taskito_worker.py @@ -0,0 +1,29 @@ +"""Management command to start the taskito worker.""" + +from __future__ import annotations + +try: + from django.core.management.base import BaseCommand +except ImportError as e: + raise ImportError( + "Django integration requires 'django'. Install with: pip install taskito[django]" + ) from e + + +class Command(BaseCommand): + help = "Start the taskito worker process" + + def add_arguments(self, parser): # type: ignore[no-untyped-def] + parser.add_argument( + "--queues", + type=str, + default=None, + help="Comma-separated list of queues to process (default: all)", + ) + + def handle(self, **options): # type: ignore[no-untyped-def] + from taskito.contrib.django.settings import get_queue + + queue = get_queue() + queues = options["queues"].split(",") if options["queues"] else None + queue.run_worker(queues=queues) diff --git a/py_src/taskito/contrib/django/settings.py b/py_src/taskito/contrib/django/settings.py new file mode 100644 index 00000000..ebc49e84 --- /dev/null +++ b/py_src/taskito/contrib/django/settings.py @@ -0,0 +1,38 @@ +"""Settings helper — reads TASKITO_* from Django settings and provides a singleton Queue.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from taskito.app import Queue + +_queue_instance: Queue | None = None + + +def _get_setting(name: str, default: Any = None) -> Any: + from django.conf import settings + + return getattr(settings, name, default) + + +def get_queue() -> Queue: + """Return a singleton Queue instance configured from Django settings.""" + global _queue_instance + if _queue_instance is not None: + return _queue_instance + + from taskito.app import Queue + + _queue_instance = Queue( + db_path=_get_setting("TASKITO_DB_PATH", ".taskito/taskito.db"), + workers=_get_setting("TASKITO_WORKERS", 0), + default_retry=_get_setting("TASKITO_DEFAULT_RETRY", 3), + default_timeout=_get_setting("TASKITO_DEFAULT_TIMEOUT", 300), + default_priority=_get_setting("TASKITO_DEFAULT_PRIORITY", 0), + result_ttl=_get_setting("TASKITO_RESULT_TTL", None), + backend=_get_setting("TASKITO_BACKEND", "sqlite"), + db_url=_get_setting("TASKITO_DB_URL", None), + schema=_get_setting("TASKITO_SCHEMA", "taskito"), + ) + return _queue_instance diff --git a/py_src/taskito/contrib/otel.py b/py_src/taskito/contrib/otel.py index dc9d363b..71895d00 100644 --- a/py_src/taskito/contrib/otel.py +++ b/py_src/taskito/contrib/otel.py @@ -46,8 +46,11 @@ def __init__(self, tracer_name: str = _TRACER_NAME): "opentelemetry-api is required for OpenTelemetryMiddleware. " "Install it with: pip install taskito[otel]" ) + import threading + self._tracer = trace.get_tracer(tracer_name) self._spans: dict[str, Any] = {} + self._lock = threading.Lock() def before(self, ctx: JobContext) -> None: span = self._tracer.start_span( @@ -59,10 +62,12 @@ def before(self, ctx: JobContext) -> None: "taskito.retry_count": ctx.retry_count, }, ) - self._spans[ctx.id] = span + with self._lock: + self._spans[ctx.id] = span def after(self, ctx: JobContext, result: Any, error: Exception | None) -> None: - span = self._spans.pop(ctx.id, None) + with self._lock: + span = self._spans.pop(ctx.id, None) if span is None: return @@ -75,7 +80,8 @@ def after(self, ctx: JobContext, result: Any, error: Exception | None) -> None: span.end() def on_retry(self, ctx: JobContext, error: Exception, retry_count: int) -> None: - span = self._spans.get(ctx.id) + with self._lock: + span = self._spans.get(ctx.id) if span is not None: span.add_event( "retry", diff --git a/py_src/taskito/dashboard.py b/py_src/taskito/dashboard.py index cda8fee2..12e48727 100644 --- a/py_src/taskito/dashboard.py +++ b/py_src/taskito/dashboard.py @@ -82,6 +82,12 @@ def do_GET(self) -> None: elif path.startswith("/api/jobs/") and path.endswith("/errors"): job_id = path[len("/api/jobs/") : -len("/errors")] self._json_response(queue.job_errors(job_id)) + elif path.startswith("/api/jobs/") and path.endswith("/logs"): + job_id = path[len("/api/jobs/") : -len("/logs")] + self._json_response(queue.task_logs(job_id)) + elif path.startswith("/api/jobs/") and path.endswith("/replay-history"): + job_id = path[len("/api/jobs/") : -len("/replay-history")] + self._json_response(queue.replay_history(job_id)) elif path.startswith("/api/jobs/"): job_id = path[len("/api/jobs/") :] job = queue.get_job(job_id) @@ -97,12 +103,6 @@ def do_GET(self) -> None: task = qs.get("task", [None])[0] since = int(qs.get("since", ["3600"])[0]) self._json_response(queue.metrics(task_name=task, since=since)) - elif path.startswith("/api/jobs/") and path.endswith("/logs"): - job_id = path[len("/api/jobs/") : -len("/logs")] - self._json_response(queue.task_logs(job_id)) - elif path.startswith("/api/jobs/") and path.endswith("/replay-history"): - job_id = path[len("/api/jobs/") : -len("/replay-history")] - self._json_response(queue.replay_history(job_id)) elif path == "/api/logs": task = qs.get("task", [None])[0] level = qs.get("level", [None])[0] diff --git a/py_src/taskito/result.py b/py_src/taskito/result.py index 41f378ce..9f6a732e 100644 --- a/py_src/taskito/result.py +++ b/py_src/taskito/result.py @@ -6,8 +6,6 @@ import time from typing import TYPE_CHECKING, Any -import cloudpickle - if TYPE_CHECKING: from taskito._taskito import PyJob from taskito.app import Queue @@ -87,7 +85,7 @@ def _poll_once(self) -> tuple[str, Any]: result_bytes = self._py_job.result_bytes if result_bytes is None: return status, None - return status, cloudpickle.loads(result_bytes) + return status, self._queue._serializer.loads(result_bytes) if status in ("failed", "dead"): raise RuntimeError(f"Job {self.id} {status}: {self._py_job.error or 'unknown error'}") diff --git a/py_src/taskito/testing.py b/py_src/taskito/testing.py index 915fac8a..50090007 100644 --- a/py_src/taskito/testing.py +++ b/py_src/taskito/testing.py @@ -187,6 +187,9 @@ def result(self, timeout: float = 30.0, **kwargs: Any) -> Any: async def aresult(self, timeout: float = 30.0, **kwargs: Any) -> Any: return self.result(timeout=timeout) + def refresh(self) -> None: + pass + def to_dict(self) -> dict[str, Any]: return { "id": self._tr.job_id, diff --git a/pyproject.toml b/pyproject.toml index d1b248d0..7b27e0ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,8 @@ dependencies = ["cloudpickle>=3.0"] [project.optional-dependencies] dev = ["pytest>=7.0", "pytest-asyncio>=0.21", "ruff>=0.8", "mypy>=1.13"] fastapi = ["fastapi>=0.100.0", "pydantic>=2.0"] +django = ["django>=3.2"] +postgres = [] # PG driver is bundled in the Rust binary; extra exists for discoverability otel = ["opentelemetry-api", "opentelemetry-sdk"] docs = ["zensical"] @@ -95,3 +97,7 @@ ignore_missing_imports = true [[tool.mypy.overrides]] module = "pydantic" ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = ["django", "django.*"] +ignore_missing_imports = true diff --git a/tests/python/test_dashboard.py b/tests/python/test_dashboard.py index 242e8d03..79586a58 100644 --- a/tests/python/test_dashboard.py +++ b/tests/python/test_dashboard.py @@ -115,6 +115,7 @@ def test_list_jobs_invalid_status(queue): def test_to_dict_fields(queue): """to_dict() returns all expected fields.""" + @queue.task() def dummy(): pass @@ -123,10 +124,22 @@ def dummy(): d = job.to_dict() expected_keys = { - "id", "queue", "task_name", "status", "priority", "progress", - "retry_count", "max_retries", "created_at", "scheduled_at", - "started_at", "completed_at", "error", "timeout_ms", - "unique_key", "metadata", + "id", + "queue", + "task_name", + "status", + "priority", + "progress", + "retry_count", + "max_retries", + "created_at", + "scheduled_at", + "started_at", + "completed_at", + "error", + "timeout_ms", + "unique_key", + "metadata", } assert set(d.keys()) == expected_keys assert d["status"] == "pending" @@ -135,6 +148,7 @@ def dummy(): def test_to_dict_is_json_serializable(queue): """to_dict() output can be serialized to JSON.""" + @queue.task() def dummy(): pass diff --git a/tests/python/test_dlq.py b/tests/python/test_dlq.py index 6b247033..5faa4c4f 100644 --- a/tests/python/test_dlq.py +++ b/tests/python/test_dlq.py @@ -35,9 +35,11 @@ def instant_fail(): ) worker_thread.start() - time.sleep(2) - - dead = queue.dead_letters() + for _ in range(20): + time.sleep(0.5) + dead = queue.dead_letters() + if len(dead) >= 1: + break assert len(dead) >= 1 # Purge entries older than 0 seconds (purge everything) diff --git a/tests/python/test_fastapi.py b/tests/python/test_fastapi.py index 0db93c38..d00888a6 100644 --- a/tests/python/test_fastapi.py +++ b/tests/python/test_fastapi.py @@ -168,6 +168,7 @@ def test_progress_stream(populated): assert len(lines) >= 1 # The first (and only) event should show complete status import json + data = json.loads(lines[0].replace("data: ", "")) assert data["status"] == "complete" diff --git a/uv.lock b/uv.lock index 082289eb..a50983cb 100644 --- a/uv.lock +++ b/uv.lock @@ -2,7 +2,8 @@ version = 1 revision = 3 requires-python = ">=3.9" resolution-markers = [ - "python_full_version >= '3.10'", + "python_full_version >= '3.12'", + "python_full_version >= '3.10' and python_full_version < '3.12'", "python_full_version < '3.10'", ] @@ -38,6 +39,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/0e/27be9fdef66e72d64c0cdc3cc2823101b80585f8119b5c112c2e8f5f7dab/anyio-4.12.1-py3-none-any.whl", hash = "sha256:d405828884fc140aa80a3c667b8beed277f1dfedec42ba031bd6ac3db606ab6c", size = 113592, upload-time = "2026-01-06T11:45:19.497Z" }, ] +[[package]] +name = "asgiref" +version = "3.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/63/40/f03da1264ae8f7cfdbf9146542e5e7e8100a4c66ab48e791df9a03d3f6c0/asgiref-3.11.1.tar.gz", hash = "sha256:5f184dc43b7e763efe848065441eac62229c9f7b0475f41f80e207a114eda4ce", size = 38550, upload-time = "2026-02-03T13:30:14.33Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/0a/a72d10ed65068e115044937873362e6e32fab1b7dce0046aeb224682c989/asgiref-3.11.1-py3-none-any.whl", hash = "sha256:e8667a091e69529631969fd45dc268fa79b99c92c5fcdda727757e52146ec133", size = 24345, upload-time = "2026-02-03T13:30:13.039Z" }, +] + [[package]] name = "backports-asyncio-runner" version = "1.2.0" @@ -86,6 +99,57 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2d/82/e5d2c1c67d19841e9edc74954c827444ae826978499bde3dfc1d007c8c11/deepmerge-2.0-py3-none-any.whl", hash = "sha256:6de9ce507115cff0bed95ff0ce9ecc31088ef50cbdf09bc90a09349a318b3d00", size = 13475, upload-time = "2024-08-30T05:31:48.659Z" }, ] +[[package]] +name = "django" +version = "4.2.29" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.10'", +] +dependencies = [ + { name = "asgiref", marker = "python_full_version < '3.10'" }, + { name = "sqlparse", marker = "python_full_version < '3.10'" }, + { name = "tzdata", marker = "python_full_version < '3.10' and sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1e/7d/7306757cf2ac016d718d8a5dbae66de630addaa73dca2c341fc388458e71/django-4.2.29.tar.gz", hash = "sha256:86d91bc8086569c8d08f9c55888b583a921ac1f95ed3bdc7d5659d4709542014", size = 10438980, upload-time = "2026-03-03T13:56:42.083Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/2b/bd0a1d1846d5580e9f209b9e0128b4859e381ec1b39d6d175c61294bb530/django-4.2.29-py3-none-any.whl", hash = "sha256:074d7c4d2808050e528388bda442bd491f06def4df4fe863f27066851bba010c", size = 7996481, upload-time = "2026-03-03T13:56:36.495Z" }, +] + +[[package]] +name = "django" +version = "5.2.12" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.10' and python_full_version < '3.12'", +] +dependencies = [ + { name = "asgiref", marker = "python_full_version >= '3.10' and python_full_version < '3.12'" }, + { name = "sqlparse", marker = "python_full_version >= '3.10' and python_full_version < '3.12'" }, + { name = "tzdata", marker = "python_full_version >= '3.10' and python_full_version < '3.12' and sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bd/55/b9445fc0695b03746f355c05b2eecc54c34e05198c686f4fc4406b722b52/django-5.2.12.tar.gz", hash = "sha256:6b809af7165c73eff5ce1c87fdae75d4da6520d6667f86401ecf55b681eb1eeb", size = 10860574, upload-time = "2026-03-03T13:56:05.509Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/32/4b144e125678efccf5d5b61581de1c4088d6b0286e46096e3b8de0d556c8/django-5.2.12-py3-none-any.whl", hash = "sha256:4853482f395c3a151937f6991272540fcbf531464f254a347bf7c89f53c8cff7", size = 8310245, upload-time = "2026-03-03T13:56:01.174Z" }, +] + +[[package]] +name = "django" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.12'", +] +dependencies = [ + { name = "asgiref", marker = "python_full_version >= '3.12'" }, + { name = "sqlparse", marker = "python_full_version >= '3.12'" }, + { name = "tzdata", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/80/e1/894115c6bd70e2c8b66b0c40a3c367d83a5a48c034a4d904d31b62f7c53a/django-6.0.3.tar.gz", hash = "sha256:90be765ee756af8a6cbd6693e56452404b5ad15294f4d5e40c0a55a0f4870fe1", size = 10872701, upload-time = "2026-03-03T13:55:15.026Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/72/b1/23f2556967c45e34d3d3cf032eb1bd3ef925ee458667fb99052a0b3ea3a6/django-6.0.3-py3-none-any.whl", hash = "sha256:2e5974441491ddb34c3f13d5e7a9f97b07ba03bf70234c0a9c68b79bbb235bc3", size = 8358527, upload-time = "2026-03-03T13:55:10.552Z" }, +] + [[package]] name = "exceptiongroup" version = "1.3.1" @@ -122,7 +186,8 @@ name = "fastapi" version = "0.135.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.10'", + "python_full_version >= '3.12'", + "python_full_version >= '3.10' and python_full_version < '3.12'", ] dependencies = [ { name = "annotated-doc", marker = "python_full_version >= '3.10'" }, @@ -174,7 +239,8 @@ name = "iniconfig" version = "2.3.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.10'", + "python_full_version >= '3.12'", + "python_full_version >= '3.10' and python_full_version < '3.12'", ] sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } wheels = [ @@ -609,7 +675,8 @@ name = "pytest" version = "9.0.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.10'", + "python_full_version >= '3.12'", + "python_full_version >= '3.10' and python_full_version < '3.12'", ] dependencies = [ { name = "colorama", marker = "python_full_version >= '3.10' and sys_platform == 'win32'" }, @@ -647,7 +714,8 @@ name = "pytest-asyncio" version = "1.3.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.10'", + "python_full_version >= '3.12'", + "python_full_version >= '3.10' and python_full_version < '3.12'", ] dependencies = [ { name = "backports-asyncio-runner", marker = "python_full_version == '3.10.*'" }, @@ -757,6 +825,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fe/4e/cd76eca6db6115604b7626668e891c9dd03330384082e33662fb0f113614/ruff-0.15.5-py3-none-win_arm64.whl", hash = "sha256:b498d1c60d2fe5c10c45ec3f698901065772730b411f164ae270bb6bfcc4740b", size = 10965572, upload-time = "2026-03-05T20:06:16.984Z" }, ] +[[package]] +name = "sqlparse" +version = "0.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/90/76/437d71068094df0726366574cf3432a4ed754217b436eb7429415cf2d480/sqlparse-0.5.5.tar.gz", hash = "sha256:e20d4a9b0b8585fdf63b10d30066c7c94c5d7a7ec47c889a2d83a3caa93ff28e", size = 120815, upload-time = "2025-12-19T07:17:45.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/4b/359f28a903c13438ef59ebeee215fb25da53066db67b305c125f1c6d2a25/sqlparse-0.5.5-py3-none-any.whl", hash = "sha256:12a08b3bf3eec877c519589833aed092e2444e68240a3577e8e26148acc7b1ba", size = 46138, upload-time = "2025-12-19T07:17:46.573Z" }, +] + [[package]] name = "starlette" version = "0.49.3" @@ -778,7 +855,8 @@ name = "starlette" version = "0.52.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.10'", + "python_full_version >= '3.12'", + "python_full_version >= '3.10' and python_full_version < '3.12'", ] dependencies = [ { name = "anyio", marker = "python_full_version >= '3.10'" }, @@ -791,7 +869,7 @@ wheels = [ [[package]] name = "taskito" -version = "0.2.0" +version = "0.2.2" source = { editable = "." } dependencies = [ { name = "cloudpickle" }, @@ -806,6 +884,11 @@ dev = [ { name = "pytest-asyncio", version = "1.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, { name = "ruff" }, ] +django = [ + { name = "django", version = "4.2.29", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, + { name = "django", version = "5.2.12", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10' and python_full_version < '3.12'" }, + { name = "django", version = "6.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, +] docs = [ { name = "zensical", version = "0.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" }, { name = "zensical", version = "0.0.24", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, @@ -823,6 +906,7 @@ otel = [ [package.metadata] requires-dist = [ { name = "cloudpickle", specifier = ">=3.0" }, + { name = "django", marker = "extra == 'django'", specifier = ">=3.2" }, { name = "fastapi", marker = "extra == 'fastapi'", specifier = ">=0.100.0" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.13" }, { name = "opentelemetry-api", marker = "extra == 'otel'" }, @@ -833,7 +917,7 @@ requires-dist = [ { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.8" }, { name = "zensical", marker = "extra == 'docs'" }, ] -provides-extras = ["dev", "fastapi", "otel", "docs"] +provides-extras = ["dev", "fastapi", "django", "postgres", "otel", "docs"] [[package]] name = "tomli" @@ -910,6 +994,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] +[[package]] +name = "tzdata" +version = "2025.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5e/a7/c202b344c5ca7daf398f3b8a477eeb205cf3b6f32e7ec3a6bac0629ca975/tzdata-2025.3.tar.gz", hash = "sha256:de39c2ca5dc7b0344f2eba86f49d614019d29f060fc4ebc8a417896a620b56a7", size = 196772, upload-time = "2025-12-13T17:45:35.667Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/b0/003792df09decd6849a5e39c28b513c06e84436a54440380862b5aeff25d/tzdata-2025.3-py2.py3-none-any.whl", hash = "sha256:06a47e5700f3081aab02b2e513160914ff0694bce9947d6b76ebd6bf57cfc5d1", size = 348521, upload-time = "2025-12-13T17:45:33.889Z" }, +] + [[package]] name = "zensical" version = "0.0.2" @@ -927,7 +1020,8 @@ name = "zensical" version = "0.0.24" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.10'", + "python_full_version >= '3.12'", + "python_full_version >= '3.10' and python_full_version < '3.12'", ] dependencies = [ { name = "click", marker = "python_full_version >= '3.10'" }, diff --git a/zensical.toml b/zensical.toml index f34f2327..39a3c2a1 100644 --- a/zensical.toml +++ b/zensical.toml @@ -1,7 +1,7 @@ [project] site_name = "taskito" site_description = "Rust-powered task queue for Python. No broker required." -site_url = "https://taskito-sepia.vercel.app" +site_url = "https://taskito.grigori.in" repo_url = "https://github.com/pratyush618/taskito" repo_name = "pratyush618/taskito"