From e43a82b78a81769f0354d944dce675391230afc1 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 30 May 2026 02:30:17 +0530 Subject: [PATCH 1/2] feat(scheduler): batch dequeue to claim N jobs per round-trip Adds dequeue_batch/dequeue_batch_from to the Storage trait (all three backends) and a scheduler_batch_size knob (default 1, behavior unchanged). The poller claims up to N eligible jobs in one transaction while still enforcing per-task and per-queue concurrency caps per job. --- crates/taskito-core/src/scheduler/mod.rs | 12 +- crates/taskito-core/src/scheduler/poller.rs | 57 ++++++- .../src/storage/diesel_common/jobs.rs | 122 +++++++++++++++ crates/taskito-core/src/storage/mod.rs | 36 +++++ .../src/storage/redis_backend/jobs/dequeue.rs | 141 ++++++++++++++++++ .../taskito-core/src/storage/sqlite/tests.rs | 114 ++++++++++++++ crates/taskito-core/src/storage/traits.rs | 19 +++ .../taskito-core/tests/rust/storage_tests.rs | 39 +++++ crates/taskito-python/src/py_queue/mod.rs | 5 +- crates/taskito-python/src/py_queue/worker.rs | 1 + .../src/py_queue/workflow_ops/test_helpers.rs | 1 + py_src/taskito/_taskito.pyi | 1 + py_src/taskito/app.py | 5 + tests/worker/test_batch_dequeue.py | 62 ++++++++ 14 files changed, 612 insertions(+), 3 deletions(-) create mode 100644 tests/worker/test_batch_dequeue.py diff --git a/crates/taskito-core/src/scheduler/mod.rs b/crates/taskito-core/src/scheduler/mod.rs index 6ebee2b7..e0b64e9a 100644 --- a/crates/taskito-core/src/scheduler/mod.rs +++ b/crates/taskito-core/src/scheduler/mod.rs @@ -33,6 +33,10 @@ pub struct SchedulerConfig { pub cleanup_interval: u32, /// TTL for job results in milliseconds. None means no auto-cleanup. pub result_ttl_ms: Option, + /// Maximum number of jobs claimed per dispatch round. `1` (the default) + /// preserves the original one-job-per-round-trip behavior; values above + /// `1` enable batch claiming for higher throughput. + pub batch_size: usize, } impl Default for SchedulerConfig { @@ -44,6 +48,7 @@ impl Default for SchedulerConfig { periodic_check_interval: 60, cleanup_interval: 1200, result_ttl_ms: None, + batch_size: 1, } } } @@ -225,7 +230,12 @@ impl Scheduler { /// Returns true if any work was done (job dispatched or periodic task enqueued), /// which resets the adaptive poll interval. fn tick(&self, job_tx: &tokio::sync::mpsc::Sender, counters: &mut TickCounters) -> bool { - let dispatched = match self.try_dispatch(job_tx) { + let dispatch_result = if self.config.batch_size > 1 { + self.try_dispatch_batch(job_tx) + } else { + self.try_dispatch(job_tx) + }; + let dispatched = match dispatch_result { Ok(d) => d, Err(e) => { error!("scheduler error: {e}"); diff --git a/crates/taskito-core/src/scheduler/poller.rs b/crates/taskito-core/src/scheduler/poller.rs index 0023665d..f1191e09 100644 --- a/crates/taskito-core/src/scheduler/poller.rs +++ b/crates/taskito-core/src/scheduler/poller.rs @@ -43,6 +43,61 @@ impl Scheduler { None => return Ok(false), }; + self.gate_and_dispatch(job, now, job_tx) + } + + /// Batch variant of [`try_dispatch`]. Claims up to `batch_size` jobs in a + /// single round-trip, then runs the exact same per-job gate/claim/ + /// concurrency/dispatch logic the single path uses. The pre-sized batch is + /// advisory only — the hard per-task and per-queue caps are still enforced + /// per job after the claim, rolling back any job that exceeds a limit. + pub(super) fn try_dispatch_batch( + &self, + job_tx: &tokio::sync::mpsc::Sender, + ) -> Result { + let now = now_millis(); + + let active_queues = self.active_queues(); + if active_queues.is_empty() { + return Ok(false); + } + + // Size the batch to the worker pool's free capacity so we never claim + // more jobs than the channel can immediately accept. `try_send` still + // guards each hand-off, but pre-sizing avoids needless claim/rollback + // churn. Always claim at least one. + let budget = self.config.batch_size.min(job_tx.capacity().max(1)); + + let jobs = self.storage.dequeue_batch_from( + &active_queues, + now, + self.namespace.as_deref(), + budget, + )?; + if jobs.is_empty() { + return Ok(false); + } + + let mut dispatched_any = false; + for job in jobs { + if self.gate_and_dispatch(job, now, job_tx)? { + dispatched_any = true; + } + } + Ok(dispatched_any) + } + + /// Run the post-dequeue pipeline for a single already-claimed (Running) + /// job: soft pre-claim gates, exactly-once claim, hard concurrency caps, + /// and hand-off to the worker pool. Returns `Ok(true)` if the job was + /// dispatched, `Ok(false)` if it was gated/rolled back. Shared by both the + /// single and batch dispatch paths so limit enforcement never drifts. + fn gate_and_dispatch( + &self, + job: Job, + now: i64, + job_tx: &tokio::sync::mpsc::Sender, + ) -> Result { // Pre-claim soft gates: rate limits and circuit breaker. // // These don't need to be atomic with the claim — if two schedulers @@ -60,7 +115,7 @@ impl Scheduler { // Post-claim hard gate: concurrency caps must be checked AFTER the // claim so two schedulers cannot both pass the cap. Status was - // already transitioned to `Running` by `dequeue_from`, so the + // already transitioned to `Running` by the dequeue, so the // running-count includes this job — use strict `>` to allow exactly // `max_concurrent` jobs. // diff --git a/crates/taskito-core/src/storage/diesel_common/jobs.rs b/crates/taskito-core/src/storage/diesel_common/jobs.rs index cf6bce14..6f175f3b 100644 --- a/crates/taskito-core/src/storage/diesel_common/jobs.rs +++ b/crates/taskito-core/src/storage/diesel_common/jobs.rs @@ -363,6 +363,128 @@ macro_rules! impl_diesel_job_ops { Ok(None) } + /// Atomically claim up to `max` ready jobs from a single queue in + /// one transaction. Generalizes `dequeue` to the batch case: scans + /// a bounded candidate set and claims eligible rows until the + /// budget is met or candidates are exhausted. + /// + /// Each claim uses an `UPDATE ... WHERE status = Pending` guarded + /// by the affected-row count: if another worker already moved the + /// row out of `Pending`, the update affects zero rows and the + /// candidate is skipped — avoiding a double-claim race. + pub fn dequeue_batch( + &self, + queue_name: &str, + now: i64, + namespace: Option<&str>, + max: usize, + ) -> Result> { + if max == 0 { + return Ok(Vec::new()); + } + + // Scan more candidates than `max` so dependency/expiry skips + // still leave enough eligible rows to fill the batch, bounded + // to keep the loaded set small. + let scan_limit = (max.saturating_mul(4)).min(400) as i64; + + let mut conn = self.conn()?; + + conn.transaction(|conn| { + let mut query = 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(scan_limit) + .into_boxed(); + + if let Some(ns) = namespace { + query = query.filter(jobs::namespace.eq(ns)); + } else { + query = query.filter(jobs::namespace.is_null()); + } + + let candidates: Vec = query.select(JobRow::as_select()).load(conn)?; + + let mut claimed: Vec = Vec::with_capacity(max.min(candidates.len())); + + for row in candidates { + if claimed.len() == max { + break; + } + + // Skip expired jobs (and mark them cancelled). + if let Some(expires_at) = row.expires_at { + if now > expires_at { + 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; + } + } + + // Common case: jobs with no dependencies skip the + // job_dependencies lookup entirely. + if row.has_deps && !Self::deps_satisfied(conn, &row.id)? { + continue; + } + + // Claim guarded by the affected-row count: if another + // worker already moved this row out of `Pending`, the + // update touches zero rows — skip it rather than + // claiming a job we don't own. + let affected = 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)?; + + if affected == 0 { + continue; + } + + let updated: JobRow = jobs::table + .find(&row.id) + .select(JobRow::as_select()) + .first(conn)?; + + claimed.push(Job::from(updated)); + } + + Ok(claimed) + }) + } + + /// Claim up to `max` ready jobs across the given queues, checking + /// each in order until the budget is exhausted. + pub fn dequeue_batch_from( + &self, + queues: &[String], + now: i64, + namespace: Option<&str>, + max: usize, + ) -> Result> { + let mut claimed: Vec = Vec::new(); + for queue_name in queues { + if claimed.len() >= max { + break; + } + let remaining = max - claimed.len(); + let mut batch = self.dequeue_batch(queue_name, now, namespace, remaining)?; + claimed.append(&mut batch); + } + Ok(claimed) + } + /// Mark a job as complete with the given result. pub fn complete(&self, id: &str, result_bytes: Option>) -> Result<()> { let now = now_millis(); diff --git a/crates/taskito-core/src/storage/mod.rs b/crates/taskito-core/src/storage/mod.rs index d37d5e7d..b532755e 100644 --- a/crates/taskito-core/src/storage/mod.rs +++ b/crates/taskito-core/src/storage/mod.rs @@ -117,6 +117,24 @@ macro_rules! impl_storage { ) -> $crate::error::Result> { self.dequeue_from(queues, now, namespace) } + fn dequeue_batch( + &self, + queue_name: &str, + now: i64, + namespace: Option<&str>, + max: usize, + ) -> $crate::error::Result> { + self.dequeue_batch(queue_name, now, namespace, max) + } + fn dequeue_batch_from( + &self, + queues: &[String], + now: i64, + namespace: Option<&str>, + max: usize, + ) -> $crate::error::Result> { + self.dequeue_batch_from(queues, now, namespace, max) + } fn complete( &self, id: &str, @@ -581,6 +599,24 @@ impl Storage for StorageBackend { ) -> Result> { delegate!(self, dequeue_from, queues, now, namespace) } + fn dequeue_batch( + &self, + queue_name: &str, + now: i64, + namespace: Option<&str>, + max: usize, + ) -> Result> { + delegate!(self, dequeue_batch, queue_name, now, namespace, max) + } + fn dequeue_batch_from( + &self, + queues: &[String], + now: i64, + namespace: Option<&str>, + max: usize, + ) -> Result> { + delegate!(self, dequeue_batch_from, queues, now, namespace, max) + } fn complete(&self, id: &str, result_bytes: Option>) -> Result<()> { delegate!(self, complete, id, result_bytes) } diff --git a/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs b/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs index 86ec6acc..daa56712 100644 --- a/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs +++ b/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs @@ -122,4 +122,145 @@ impl RedisStorage { } Ok(None) } + + /// Claim up to `max` ready jobs from a single queue. Mirrors `dequeue` + /// (ZRANGEBYSCORE candidates + MGET + claim loop) but accumulates up to + /// `max` jobs. Shares `dequeue`'s TOCTOU window — `claim_execution` is the + /// external exactly-once guard. + pub fn dequeue_batch( + &self, + queue_name: &str, + now: i64, + namespace: Option<&str>, + max: usize, + ) -> Result> { + if max == 0 { + return Ok(Vec::new()); + } + + let mut conn = self.conn()?; + let queue_key = self.key(&["queue", queue_name, "pending"]); + + // Scan more candidates than `max` so dependency/expiry skips still + // leave enough eligible rows to fill the batch, bounded to keep the + // loaded set small. + let scan_limit = (max.saturating_mul(4)).min(400) as isize; + + // Get candidates ordered by score (lowest first = highest priority) + let candidates: Vec = conn + .zrangebyscore_limit(&queue_key, "-inf", "+inf", 0, scan_limit) + .map_err(map_err)?; + + if candidates.is_empty() { + return Ok(Vec::new()); + } + + // Batch-load every candidate's JSON in one MGET instead of one GET per + // candidate. + let job_keys: Vec = candidates.iter().map(|id| self.key(&["job", id])).collect(); + let blobs: Vec> = conn.mget(&job_keys).map_err(map_err)?; + + let mut claimed: Vec = Vec::with_capacity(max.min(candidates.len())); + + for (job_id, data) in candidates.into_iter().zip(blobs) { + if claimed.len() == max { + break; + } + + let data = match data { + Some(d) => d, + None => { + // Stale entry — remove from queue + conn.zrem::<_, _, ()>(&queue_key, &job_id) + .map_err(map_err)?; + continue; + } + }; + + let mut job: Job = + serde_json::from_str(&data).map_err(|e| QueueError::Other(e.to_string()))?; + + // Must be pending and scheduled_at <= now + if job.status != JobStatus::Pending || job.scheduled_at > now { + continue; + } + + // Filter by namespace: Some(ns) matches that namespace, None matches only jobs without a namespace + if let Some(ns) = namespace { + if job.namespace.as_deref() != Some(ns) { + continue; + } + } else if job.namespace.is_some() { + continue; + } + + // Skip expired jobs + if let Some(expires_at) = job.expires_at { + if now > expires_at { + job.status = JobStatus::Cancelled; + job.completed_at = Some(now); + job.error = Some("expired before execution".to_string()); + self.save_job_and_move_status(&mut conn, &job, JobStatus::Pending)?; + conn.zrem::<_, _, ()>(&queue_key, &job_id) + .map_err(map_err)?; + continue; + } + } + + // Check dependencies — only for jobs that actually have them. + if job.has_deps { + let deps_key = self.key(&["job", &job_id, "depends_on"]); + let dep_ids: Vec = conn.smembers(&deps_key).map_err(map_err)?; + if !dep_ids.is_empty() { + let mut all_complete = true; + for dep_id in &dep_ids { + if let Some(dep_job) = self.load_job(&mut conn, dep_id)? { + if dep_job.status != JobStatus::Complete { + all_complete = false; + break; + } + } else { + all_complete = false; + break; + } + } + if !all_complete { + continue; + } + } + } + + // Claim the job + job.status = JobStatus::Running; + job.started_at = Some(now); + self.save_job_and_move_status(&mut conn, &job, JobStatus::Pending)?; + conn.zrem::<_, _, ()>(&queue_key, &job_id) + .map_err(map_err)?; + + claimed.push(job); + } + + Ok(claimed) + } + + /// Claim up to `max` ready jobs across the given queues, checking each in + /// order until the budget is exhausted. + pub fn dequeue_batch_from( + &self, + queues: &[String], + now: i64, + namespace: Option<&str>, + max: usize, + ) -> Result> { + let mut claimed: Vec = Vec::new(); + for queue_name in queues { + if claimed.len() >= max { + break; + } + let remaining = max - claimed.len(); + let mut batch = self.dequeue_batch(queue_name, now, namespace, remaining)?; + claimed.append(&mut batch); + } + Ok(claimed) + } } diff --git a/crates/taskito-core/src/storage/sqlite/tests.rs b/crates/taskito-core/src/storage/sqlite/tests.rs index f004f0ef..738178ae 100644 --- a/crates/taskito-core/src/storage/sqlite/tests.rs +++ b/crates/taskito-core/src/storage/sqlite/tests.rs @@ -99,6 +99,120 @@ fn test_dequeue() { assert!(none.is_none()); } +#[test] +fn test_dequeue_batch_claims_n() { + let storage = test_storage(); + for _ in 0..5 { + storage.enqueue(make_job("batch_task")).unwrap(); + } + + let claimed = storage + .dequeue_batch("default", now_millis() + 1000, None, 3) + .unwrap(); + assert_eq!(claimed.len(), 3); + for job in &claimed { + assert_eq!(job.status, JobStatus::Running); + } + + let running = storage + .list_jobs(Some(JobStatus::Running as i32), None, None, 100, 0, None) + .unwrap(); + assert_eq!(running.len(), 3); +} + +#[test] +fn test_dequeue_batch_respects_available() { + let storage = test_storage(); + storage.enqueue(make_job("batch_task")).unwrap(); + storage.enqueue(make_job("batch_task")).unwrap(); + + let claimed = storage + .dequeue_batch("default", now_millis() + 1000, None, 10) + .unwrap(); + assert_eq!(claimed.len(), 2, "only claims what's available"); +} + +#[test] +fn test_dequeue_batch_empty_and_zero_max() { + let storage = test_storage(); + + // Empty queue → empty batch. + let empty = storage + .dequeue_batch("default", now_millis() + 1000, None, 5) + .unwrap(); + assert!(empty.is_empty()); + + // max == 0 claims nothing even when jobs exist. + storage.enqueue(make_job("batch_task")).unwrap(); + let zero = storage + .dequeue_batch("default", now_millis() + 1000, None, 0) + .unwrap(); + assert!(zero.is_empty()); + + // The job must still be pending after a zero-max batch. + let pending = storage + .list_jobs(Some(JobStatus::Pending as i32), None, None, 100, 0, None) + .unwrap(); + assert_eq!(pending.len(), 1); +} + +#[test] +fn test_dequeue_batch_no_double_claim() { + let storage = test_storage(); + for _ in 0..4 { + storage.enqueue(make_job("batch_task")).unwrap(); + } + + let now = now_millis() + 1000; + let first = storage.dequeue_batch("default", now, None, 2).unwrap(); + let second = storage.dequeue_batch("default", now, None, 2).unwrap(); + assert_eq!(first.len(), 2); + assert_eq!(second.len(), 2); + + let mut ids: Vec = first + .iter() + .chain(second.iter()) + .map(|j| j.id.clone()) + .collect(); + ids.sort(); + ids.dedup(); + assert_eq!(ids.len(), 4, "two batches must claim disjoint jobs"); +} + +#[test] +fn test_dequeue_batch_from_across_queues() { + let storage = test_storage(); + + let mut a = make_job("batch_task"); + a.queue = "qa".to_string(); + storage.enqueue(a).unwrap(); + storage + .enqueue({ + let mut j = make_job("batch_task"); + j.queue = "qa".to_string(); + j + }) + .unwrap(); + storage + .enqueue({ + let mut j = make_job("batch_task"); + j.queue = "qb".to_string(); + j + }) + .unwrap(); + + let queues = vec!["qa".to_string(), "qb".to_string()]; + let claimed = storage + .dequeue_batch_from(&queues, now_millis() + 1000, None, 10) + .unwrap(); + assert_eq!(claimed.len(), 3, "claims across both queues"); + + let queue_names: std::collections::HashSet<&str> = + claimed.iter().map(|j| j.queue.as_str()).collect(); + assert!(queue_names.contains("qa")); + assert!(queue_names.contains("qb")); +} + #[test] fn test_dequeue_respects_schedule() { let storage = test_storage(); diff --git a/crates/taskito-core/src/storage/traits.rs b/crates/taskito-core/src/storage/traits.rs index 89d653b5..134e1f3c 100644 --- a/crates/taskito-core/src/storage/traits.rs +++ b/crates/taskito-core/src/storage/traits.rs @@ -21,6 +21,25 @@ pub trait Storage: Send + Sync + Clone { now: i64, namespace: Option<&str>, ) -> Result>; + /// Atomically claim up to `max` ready jobs from a single queue in one + /// transaction. Returns the claimed jobs (now in `Running` state). May + /// return fewer than `max` if the queue lacks enough eligible jobs. + fn dequeue_batch( + &self, + queue_name: &str, + now: i64, + namespace: Option<&str>, + max: usize, + ) -> Result>; + /// Claim up to `max` ready jobs across the given queues, checking each in + /// order until the budget is exhausted. + fn dequeue_batch_from( + &self, + queues: &[String], + now: i64, + namespace: Option<&str>, + max: usize, + ) -> Result>; fn complete(&self, id: &str, result_bytes: Option>) -> Result<()>; fn fail(&self, id: &str, error: &str) -> Result<()>; fn retry(&self, id: &str, next_scheduled_at: i64) -> Result<()>; diff --git a/crates/taskito-core/tests/rust/storage_tests.rs b/crates/taskito-core/tests/rust/storage_tests.rs index 4be533c0..e03f5978 100644 --- a/crates/taskito-core/tests/rust/storage_tests.rs +++ b/crates/taskito-core/tests/rust/storage_tests.rs @@ -50,6 +50,44 @@ fn test_dequeue(s: &impl Storage) { assert!(none.is_none()); } +fn test_dequeue_batch(s: &impl Storage) { + let q = "q-dequeue-batch"; + let mut ids = Vec::new(); + for _ in 0..5 { + ids.push(s.enqueue(make_job(q, "batch_task")).unwrap().id); + } + + // Claim 3 of the 5 in one round-trip. + let now = now_millis() + 1000; + let first = s.dequeue_batch(q, now, None, 3).unwrap(); + assert_eq!(first.len(), 3); + for job in &first { + assert_eq!(job.status, JobStatus::Running); + } + + // A second batch of 10 returns only the 2 remaining — and no id overlaps. + let second = s.dequeue_batch(q, now, None, 10).unwrap(); + assert_eq!(second.len(), 2); + + let mut all: Vec = first + .iter() + .chain(second.iter()) + .map(|j| j.id.clone()) + .collect(); + all.sort(); + all.dedup(); + assert_eq!(all.len(), 5, "batches must claim disjoint jobs"); + + // Queue is now empty. + let empty = s.dequeue_batch(q, now, None, 4).unwrap(); + assert!(empty.is_empty()); + + // max == 0 claims nothing even when jobs exist. + s.enqueue(make_job(q, "batch_task")).unwrap(); + let zero = s.dequeue_batch(q, now, None, 0).unwrap(); + assert!(zero.is_empty()); +} + fn test_complete(s: &impl Storage) { let q = "q-complete"; let job = s.enqueue(make_job(q, "complete_task")).unwrap(); @@ -340,6 +378,7 @@ fn test_circuit_breakers(s: &impl Storage) { fn run_storage_tests(s: &impl Storage) { test_enqueue_and_get(s); test_dequeue(s); + test_dequeue_batch(s); test_complete(s); test_fail(s); test_retry(s); diff --git a/crates/taskito-python/src/py_queue/mod.rs b/crates/taskito-python/src/py_queue/mod.rs index 8da355e4..2e212df5 100644 --- a/crates/taskito-python/src/py_queue/mod.rs +++ b/crates/taskito-python/src/py_queue/mod.rs @@ -39,6 +39,7 @@ pub struct PyQueue { pub(crate) scheduler_poll_interval_ms: u64, pub(crate) scheduler_reap_interval: u32, pub(crate) scheduler_cleanup_interval: u32, + pub(crate) scheduler_batch_size: usize, pub(crate) namespace: Option, /// Active worker dispatcher, set while `run_worker` is executing. Used by /// `request_cancel` to deliver a side-channel signal to pools that run @@ -60,7 +61,7 @@ pub struct PyQueue { )] impl PyQueue { #[new] - #[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", pool_size=None, scheduler_poll_interval_ms=50, scheduler_reap_interval=100, scheduler_cleanup_interval=1200, namespace=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", pool_size=None, scheduler_poll_interval_ms=50, scheduler_reap_interval=100, scheduler_cleanup_interval=1200, scheduler_batch_size=1, namespace=None))] #[allow(clippy::too_many_arguments)] pub fn new( py: Python<'_>, @@ -77,6 +78,7 @@ impl PyQueue { scheduler_poll_interval_ms: u64, scheduler_reap_interval: u32, scheduler_cleanup_interval: u32, + scheduler_batch_size: usize, namespace: Option, ) -> PyResult { // Storage init blocks on connection-pool builders that may emit @@ -151,6 +153,7 @@ impl PyQueue { scheduler_poll_interval_ms, scheduler_reap_interval, scheduler_cleanup_interval, + scheduler_batch_size: scheduler_batch_size.max(1), namespace, dispatcher: Arc::new(Mutex::new(None)), #[cfg(feature = "workflows")] diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index df286584..bdf30f12 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -214,6 +214,7 @@ impl PyQueue { reap_interval: self.scheduler_reap_interval, cleanup_interval: self.scheduler_cleanup_interval, result_ttl_ms: self.result_ttl_ms, + batch_size: self.scheduler_batch_size, ..SchedulerConfig::default() }; let mut scheduler = Scheduler::new( diff --git a/crates/taskito-python/src/py_queue/workflow_ops/test_helpers.rs b/crates/taskito-python/src/py_queue/workflow_ops/test_helpers.rs index e056ba18..156a129c 100644 --- a/crates/taskito-python/src/py_queue/workflow_ops/test_helpers.rs +++ b/crates/taskito-python/src/py_queue/workflow_ops/test_helpers.rs @@ -54,6 +54,7 @@ pub(crate) fn make_test_pyqueue() -> PyQueue { scheduler_poll_interval_ms: 50, scheduler_reap_interval: 100, scheduler_cleanup_interval: 1200, + scheduler_batch_size: 1, namespace: None, dispatcher: Arc::new(Mutex::new(None)), workflow_storage, diff --git a/py_src/taskito/_taskito.pyi b/py_src/taskito/_taskito.pyi index 079f365f..5b217ca9 100644 --- a/py_src/taskito/_taskito.pyi +++ b/py_src/taskito/_taskito.pyi @@ -91,6 +91,7 @@ class PyQueue: scheduler_poll_interval_ms: int = 50, scheduler_reap_interval: int = 100, scheduler_cleanup_interval: int = 1200, + scheduler_batch_size: int = 1, namespace: str | None = None, ) -> None: ... def request_shutdown(self) -> None: ... diff --git a/py_src/taskito/app.py b/py_src/taskito/app.py index c25b0150..0ffa9fa8 100644 --- a/py_src/taskito/app.py +++ b/py_src/taskito/app.py @@ -133,6 +133,7 @@ def __init__( scheduler_poll_interval_ms: int = 50, scheduler_reap_interval: int = 100, scheduler_cleanup_interval: int = 1200, + scheduler_batch_size: int = 1, namespace: str | None = None, ): """Initialize a new task queue. @@ -180,6 +181,9 @@ def __init__( (default 100). scheduler_cleanup_interval: Cleanup old jobs every N poll iterations (default 1200). + scheduler_batch_size: Maximum number of jobs the scheduler claims + per dispatch round. ``1`` (default) preserves one-job-per-round + behavior; higher values batch-claim for greater throughput. """ if backend == "sqlite": # Ensure parent directory exists for SQLite @@ -201,6 +205,7 @@ def __init__( scheduler_poll_interval_ms=scheduler_poll_interval_ms, scheduler_reap_interval=scheduler_reap_interval, scheduler_cleanup_interval=scheduler_cleanup_interval, + scheduler_batch_size=scheduler_batch_size, namespace=namespace, ) self._backend = backend diff --git a/tests/worker/test_batch_dequeue.py b/tests/worker/test_batch_dequeue.py new file mode 100644 index 00000000..1fd98ace --- /dev/null +++ b/tests/worker/test_batch_dequeue.py @@ -0,0 +1,62 @@ +"""Tests for batch dequeue (claiming several jobs per scheduler round).""" + +import threading +from pathlib import Path +from typing import Any + +from taskito import Queue + +PollUntil = Any # the conftest fixture's runtime type + + +def test_batch_dequeue_processes_burst(tmp_path: Path, poll_until: PollUntil) -> None: + """A queue with scheduler_batch_size=4 drains a burst of jobs. + + Default behavior (batch_size=1) is exercised by the rest of the suite; + here we assert the batch-claiming path completes every job correctly. + """ + db_path = str(tmp_path / "batch.db") + queue = Queue(db_path=db_path, workers=4, scheduler_batch_size=4) + + results: dict[int, int] = {} + lock = threading.Lock() + + @queue.task() + def square(n: int) -> int: + value = n * n + with lock: + results[n] = value + return value + + jobs = [square.delay(i) for i in range(20)] + + worker = threading.Thread(target=queue.run_worker, daemon=True) + worker.start() + try: + for job in jobs: + assert job.result(timeout=15) is not None + poll_until(lambda: len(results) == 20, timeout=15) + assert results == {i: i * i for i in range(20)} + finally: + queue._inner.request_shutdown() + worker.join(timeout=5) + + +def test_batch_dequeue_default_is_one(tmp_path: Path) -> None: + """scheduler_batch_size defaults to 1 (back-compat preserved).""" + db_path = str(tmp_path / "default.db") + queue = Queue(db_path=db_path, workers=2) + + @queue.task() + def echo(x: str) -> str: + return x + + job = echo.delay("hello") + + worker = threading.Thread(target=queue.run_worker, daemon=True) + worker.start() + try: + assert job.result(timeout=10) == "hello" + finally: + queue._inner.request_shutdown() + worker.join(timeout=5) From 46442fc14e7398e7f9b8445206d95b9304a2ccf7 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 30 May 2026 10:13:48 +0530 Subject: [PATCH 2/2] fix(scheduler): isolate per-job failures in batch dispatch --- crates/taskito-core/src/scheduler/poller.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/taskito-core/src/scheduler/poller.rs b/crates/taskito-core/src/scheduler/poller.rs index f1191e09..46568d0d 100644 --- a/crates/taskito-core/src/scheduler/poller.rs +++ b/crates/taskito-core/src/scheduler/poller.rs @@ -78,10 +78,15 @@ impl Scheduler { return Ok(false); } + // Every job in `jobs` is already claimed (Running). Isolate per-job + // failures so one error doesn't strand the rest of the batch in + // Running until the reaper times them out. let mut dispatched_any = false; for job in jobs { - if self.gate_and_dispatch(job, now, job_tx)? { - dispatched_any = true; + match self.gate_and_dispatch(job, now, job_tx) { + Ok(true) => dispatched_any = true, + Ok(false) => {} + Err(e) => log::warn!("batch dispatch failed for a claimed job: {e}"), } } Ok(dispatched_any)